/*
 * Copyright (c) 2010-2021. 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;
import java.util.function.Supplier;

/**
 * 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
 * @author Allard Buijze
 * @since 3.0
 */
public class LockAwareAggregate<AR, A extends Aggregate<AR>> implements Aggregate<AR> {

    private final A wrappedAggregate;
    private final Supplier<Lock> lockSupplier;

    /**
     * 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.lockSupplier = () -> lock;
    }

    /**
     * Initializes a new {@link LockAwareAggregate} for given {@code wrappedAggregate} and a supplier of the {@code
     * lock}.
     *
     * @param wrappedAggregate the aggregate instance to which the LockAwareAggregate will delegate
     * @param lock             a supplier of the lock held by the aggregate
     */
    public LockAwareAggregate(A wrappedAggregate, Supplier<Lock> lock) {
        this.wrappedAggregate = wrappedAggregate;
        this.lockSupplier = 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 this.lockSupplier.get().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 {
        Object result = wrappedAggregate.handle(message);
        // we need to ensure the lock is acquired, as this may not have happened earlier
        lockSupplier.get();
        return result;
    }

    @Override
    public <R> R invoke(Function<AR, R> invocation) {
        R result = wrappedAggregate.invoke(invocation);
        // we need to ensure the lock is acquired, as this may not have happened earlier
        lockSupplier.get();

        return result;
    }

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

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

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

@@ -1246,24 +1246,28 @@
 e final 
+Lock
 Supplier
 %3CLock%3E l
@@ -1262,27 +1262,13 @@
 lier
-%3CLock%3E lockSupplier
+ lock
 ;%0A%0A 
@@ -1686,24 +1686,16 @@
 his.lock
-Supplier
  = () -%3E
@@ -2181,23 +2181,20 @@
 lock
-Supplier
  = lock
+::get
 ;%0A  
@@ -2603,28 +2603,24 @@
 his.lock
-Supplier.get
+.acquire
 ().isHel
@@ -3144,36 +3144,32 @@
         lock
-Supplier.get
+.acquire
 ();%0A        
@@ -3411,36 +3411,32 @@
         lock
-Supplier.get
+.acquire
 ();%0A%0A       
@@ -3627,20 +3627,16 @@
 lock
-Supplier.get
+.acquire
 ();%0A
@@ -3801,32 +3801,32 @@
 R%3E rootType() %7B%0A
-
         return w
@@ -3858,11 +3858,178 @@
 );%0A    %7D
+%0A%0A    @FunctionalInterface%0A    private interface LockSupplier extends Supplier%3CLock%3E %7B%0A%0A        default Lock acquire() %7B%0A            return this.get();%0A        %7D%0A    %7D
 %0A%7D%0A

Add private functional interface for clearer naming