cherry picked nexus (#10004)

Signed-off-by: Lazar Petrovic <lpetrovic05@gmail.com>

platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/signed/ReservedSignedState.java

@@ -76,8 +76,34 @@ private ReservedSignedState() {
     ReservedSignedState(@NonNull final SignedState signedState, @NonNull final String reason) {
         this.signedState = Objects.requireNonNull(signedState);
         this.reason = Objects.requireNonNull(reason);
+    }
 
-        signedState.incrementReservationCount(reason, reservationId);
+    /**
+     * Create a new reserved signed state and increment the reservation count on the underlying signed state.
+     *
+     * @param signedState the signed state to reserve
+     * @param reason      a short description of why this SignedState is being reserved. Each location where a
+     *                    SignedState is reserved should attempt to use a unique reason, as this makes debugging
+     *                    reservation bugs easier.
+     */
+    static @NonNull ReservedSignedState createAndReserve(
+            @NonNull final SignedState signedState, @NonNull final String reason) {
+        final ReservedSignedState reservedSignedState = new ReservedSignedState(signedState, reason);
+        signedState.incrementReservationCount(reason, reservedSignedState.getReservationId());
+        return reservedSignedState;
+    }
+
+    /**
+     * Create a new reserved signed state. This method assumes that the reservation count will be incremented by the
+     * caller.
+     *
+     * @param signedState the signed state to reserve
+     * @param reason      a short description of why this SignedState is being reserved. Each location where a
+     *                    SignedState is reserved should attempt to use a unique reason, as this makes debugging
+     *                    reservation bugs easier.
+     */
+    static @NonNull ReservedSignedState create(@NonNull final SignedState signedState, @NonNull final String reason) {
+        return new ReservedSignedState(signedState, reason);
     }
 
     /**
@@ -113,7 +139,25 @@ public boolean isNotNull() {
         if (signedState == null) {
             return new ReservedSignedState();
         }
-        return new ReservedSignedState(signedState, reason);
+        return createAndReserve(signedState, reason);
+    }
+
+    /**
+     * Try to get another reservation on the signed state. If the signed state is not closed, then a new reservation
+     * will be returned. If the signed state is closed, then null will be returned.
+     *
+     * @param reason a short description of why this SignedState is being reserved. Each location where a SignedState is
+     *               reserved should attempt to use a unique reason, as this makes debugging reservation bugs easier.
+     * @return a new wrapper around the state that holds a new reservation, or null if the signed state is closed
+     */
+    public @Nullable ReservedSignedState tryGetAndReserve(@NonNull final String reason) {
+        if (signedState == null) {
+            return new ReservedSignedState();
+        }
+        if (!signedState.tryIncrementReservationCount(reason, reservationId)) {
+            return null;
+        }
+        return create(signedState, reason);
     }
 
     /**

platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/signed/SignedState.java

@@ -296,7 +296,7 @@ public void markAsRecoveryState() {
      * @return a wrapper that holds the state and the reservation
      */
     public @NonNull ReservedSignedState reserve(@NonNull final String reason) {
-        return new ReservedSignedState(this, reason);
+        return ReservedSignedState.createAndReserve(this, reason);
     }
 
     /**
@@ -309,6 +309,19 @@ void incrementReservationCount(@NonNull final String reason, final long reservat
         reservations.reserve();
     }
 
+    /**
+     * Try to increment the reservation count.
+     */
+    boolean tryIncrementReservationCount(@NonNull final String reason, final long reservationId) {
+        if (!reservations.tryReserve()) {
+            return false;
+        }
+        if (history != null) {
+            history.recordAction(RESERVE, getReservationCount(), reason, reservationId);
+        }
+        return true;
+    }
+
     /**
      * Decrement reservation count.
      */

platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/signed/SignedStateMap.java

@@ -128,7 +128,7 @@ public void put(@NonNull final SignedState signedState, @NonNull final String re
 
         try (final Locked l = lock.lock()) {
             final ReservedSignedState previousState =
-                    map.put(signedState.getRound(), new ReservedSignedState(signedState, reason));
+                    map.put(signedState.getRound(), ReservedSignedState.createAndReserve(signedState, reason));
             if (previousState != null) {
                 previousState.close();
             }

platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/signed/SignedStateNexus.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2023 Hedera Hashgraph, LLC
+ *
+ * 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 com.swirlds.platform.state.signed;
+
+import com.swirlds.common.utility.Clearable;
+import com.swirlds.platform.consensus.ConsensusConstants;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+/**
+ * A thread-safe container that also manages reservations for a single signed state.
+ */
+public class SignedStateNexus implements Consumer<ReservedSignedState>, Clearable {
+    private final AtomicReference<ReservedSignedState> currentState = new AtomicReference<>();
+    private final AtomicLong currentStateRound = new AtomicLong(ConsensusConstants.ROUND_UNDEFINED);
+
+    /**
+     * Returns the current signed state and reserves it. If the current signed state is null, or cannot be reserved,
+     * then null is returned.
+     *
+     * @param reason a short description of why this SignedState is being reserved
+     * @return the current signed state, or null if there is no state set or if it cannot be reserved
+     */
+    public @Nullable ReservedSignedState getState(@NonNull final String reason) {
+        ReservedSignedState state;
+        do {
+            state = currentState.get();
+            if (state == null) {
+                return null;
+            }
+
+            // between the get method on the atomic reference and the tryGetAndReserve method on the state, the state
+            // could have been closed. If this happens, tryGetAndReserve will return null
+            final ReservedSignedState newReservation = state.tryGetAndReserve(reason);
+            if (newReservation != null) {
+                return newReservation;
+            }
+            // if tryGetAndReserve returned null, then we should check if set() was called in the meantime
+            // if yes, then we should try again and reserve the new state
+        } while (state != currentState.get());
+        // this means we cannot reserve the state we are holding, this is probably an error, since we should hold a
+        // reservation on it
+        return null;
+    }
+
+    /**
+     * Sets the current signed state to the given state, and releases the previous state if it exists.
+     *
+     * @param reservedSignedState the new signed state
+     */
+    public void setState(@Nullable final ReservedSignedState reservedSignedState) {
+        final ReservedSignedState oldState = currentState.getAndSet(reservedSignedState);
+        currentStateRound.set(
+                reservedSignedState == null
+                        ? ConsensusConstants.ROUND_UNDEFINED
+                        : reservedSignedState.get().getRound());
+        if (oldState != null) {
+            oldState.close();
+        }
+    }
+
+    /**
+     * Returns the round of the current signed state
+     *
+     * @return the round of the current signed state, or {@link ConsensusConstants#ROUND_UNDEFINED} if there is no
+     * current signed state
+     */
+    public long getRound() {
+        return currentStateRound.get();
+    }
+
+    /**
+     * Same as {@link #setState(ReservedSignedState)} with a null argument
+     */
+    @Override
+    public void clear() {
+        setState(null);
+    }
+
+    /**
+     * Same as {@link #setState(ReservedSignedState)}
+     */
+    @Override
+    public void accept(@Nullable final ReservedSignedState reservedSignedState) {
+        setState(reservedSignedState);
+    }
+}

platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/signed/SignedStateReference.java

@@ -66,8 +66,9 @@ public void set(@Nullable final SignedState signedState, @NonNull final String r
             }
 
             reservedSignedState.close();
-            reservedSignedState =
-                    signedState == null ? createNullReservation() : new ReservedSignedState(signedState, reason);
+            reservedSignedState = signedState == null
+                    ? createNullReservation()
+                    : ReservedSignedState.createAndReserve(signedState, reason);
         }
     }
 

platform-sdk/swirlds-platform-core/src/test/java/com/swirlds/platform/state/signed/ReservedSignedStateTests.java

@@ -20,6 +20,8 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -63,7 +65,8 @@ void NonNullStateTest() {
         final SignedState signedState = new RandomSignedStateGenerator().build();
         assertEquals(0, signedState.getReservationCount());
 
-        try (final ReservedSignedState reservedSignedState = new ReservedSignedState(signedState, "reason")) {
+        try (final ReservedSignedState reservedSignedState =
+                ReservedSignedState.createAndReserve(signedState, "reason")) {
 
             assertSame(signedState, reservedSignedState.get());
             assertSame(signedState, reservedSignedState.getNullable());
@@ -105,7 +108,7 @@ void nullBadLifecycleTest() {
     @DisplayName("Non-Null Bad Lifecycle Test")
     void nonNullBadLifecycleTest() {
         final ReservedSignedState reservedSignedState =
-                new ReservedSignedState(new RandomSignedStateGenerator().build(), "reason");
+                ReservedSignedState.createAndReserve(new RandomSignedStateGenerator().build(), "reason");
         reservedSignedState.close();
 
         assertThrows(ReferenceCountException.class, reservedSignedState::get);
@@ -115,4 +118,38 @@ void nonNullBadLifecycleTest() {
         assertThrows(ReferenceCountException.class, reservedSignedState::isNotNull);
         assertThrows(ReferenceCountException.class, reservedSignedState::close);
     }
+
+    @Test
+    @DisplayName("Try-reserve Paradigm Test")
+    void tryReserveTest() {
+        final SignedState signedState = new RandomSignedStateGenerator().build();
+        assertEquals(0, signedState.getReservationCount());
+
+        final ReservedSignedState reservedSignedState = ReservedSignedState.createAndReserve(signedState, "reason");
+        try {
+            assertEquals(1, signedState.getReservationCount());
+
+            try (final ReservedSignedState reservedSignedState2 =
+                    reservedSignedState.tryGetAndReserve("successful try")) {
+                assertNotNull(reservedSignedState2);
+                assertNotSame(reservedSignedState, reservedSignedState2);
+                assertSame(signedState, reservedSignedState2.get());
+                assertSame(signedState, reservedSignedState2.getNullable());
+                assertEquals("successful try", reservedSignedState2.getReason());
+                assertFalse(reservedSignedState2.isNull());
+                assertTrue(reservedSignedState2.isNotNull());
+                assertEquals(2, signedState.getReservationCount());
+
+                assertNotEquals(reservedSignedState.getReservationId(), reservedSignedState2.getReservationId());
+            }
+            assertEquals(1, signedState.getReservationCount());
+        } finally {
+            reservedSignedState.close();
+        }
+        assertEquals(-1, signedState.getReservationCount());
+        try (final ReservedSignedState reservedSignedState2 = reservedSignedState.tryGetAndReserve("failed try")) {
+            assertNull(reservedSignedState2);
+        }
+        assertEquals(-1, signedState.getReservationCount());
+    }
 }