/*
 * Copyright (c) 2010-2018. Axon Framework
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.axonframework.modelling.command;

import org.axonframework.common.lock.Lock;
import org.axonframework.messaging.Message;

import java.util.function.Consumer;
import java.util.function.Function;

/**
 * Aggregate implementation that provides access to the lock held by the aggregate while a command is handled.
 *
 * @param <AR> The aggregate root type
 * @param <A>  The {@link Aggregate} implementation type
 */
public class LockAwareAggregate<AR, A extends Aggregate<AR>> implements Aggregate<AR> {

    private final A wrappedAggregate;
    private final Lock lock;

    /**
     * Initializes a new {@link LockAwareAggregate} for given {@code wrappedAggregate} and {@code lock}.
     *
     * @param wrappedAggregate the aggregate instance to which the LockAwareAggregate will delegate
     * @param lock             the lock held by the aggregate
     */
    public LockAwareAggregate(A wrappedAggregate, Lock lock) {
        this.wrappedAggregate = wrappedAggregate;
        this.lock = lock;
    }

    /**
     * Get the delegate aggregate wrapped by this instance.
     *
     * @return the wrapped aggregate
     */
    public A getWrappedAggregate() {
        return wrappedAggregate;
    }

    /**
     * Check if the aggregate currently holds a lock.
     *
     * @return {@code true} if the lock is held, {@code false} otherwise
     */
    public boolean isLockHeld() {
        return lock.isHeld();
    }

    @Override
    public String type() {
        return wrappedAggregate.type();
    }

    @Override
    public Object identifier() {
        return wrappedAggregate.identifier();
    }

    @Override
    public Long version() {
        return wrappedAggregate.version();
    }

    @Override
    public Object handle(Message<?> message) throws Exception {
        return wrappedAggregate.handle(message);
    }

    @Override
    public <R> R invoke(Function<AR, R> invocation) {
        return wrappedAggregate.invoke(invocation);
    }

    @Override
    public void execute(Consumer<AR> invocation) {
        wrappedAggregate.execute(invocation);
    }

    @Override
    public boolean isDeleted() {
        return wrappedAggregate.isDeleted();
    }

    @Override
    public Class<? extends AR> rootType() {
        return wrappedAggregate.rootType();
    }
}

@@ -24,10 +24,10 @@
 0-20
+2
 1
-8
 . Ax
@@ -806,16 +806,52 @@
 unction;
+%0Aimport java.util.function.Supplier;
 %0A%0A/**%0A *
@@ -977,17 +977,17 @@
 am %3CAR%3E 
-T
+t
 he aggre
@@ -1016,17 +1016,17 @@
 am %3CA%3E  
-T
+t
 he %7B@lin
@@ -1058,16 +1058,55 @@
 on type%0A
+ * @author Allard Buijze%0A * @since 3.0%0A
  */%0Apubl
@@ -1246,25 +1246,43 @@
 e final 
+Supplier%3C
 Lock
+%3E
  lock
+Supplier
 ;%0A%0A    /
@@ -1696,16 +1696,523 @@
 his.lock
+Supplier = () -%3E lock;%0A    %7D%0A%0A    /**%0A     * Initializes a new %7B@link LockAwareAggregate%7D for given %7B@code wrappedAggregate%7D and a supplier of the %7B@code%0A     * lock%7D.%0A     *%0A     * @param wrappedAggregate the aggregate instance to which the LockAwareAggregate will delegate%0A     * @param lock             a supplier of the lock held by the aggregate%0A     */%0A    public LockAwareAggregate(A wrappedAggregate, Supplier%3CLock%3E lock) %7B%0A        this.wrappedAggregate = wrappedAggregate;%0A        this.lockSupplier
  = lock;
@@ -2615,20 +2615,39 @@
  return 
-lock
+this.lockSupplier.get()
 .isHeld(

Allow usage of a Lock supplier