RATIS-2094. Avoid corruptions from TransactionContext's stateMachineL…

…ogEntry and stateMachineContext. (#1106)
apache · Jun 3, 2024 · a2bdd10 · a2bdd10
1 parent 924a0cd
commit a2bdd10
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 6 deletions.
diff --git a/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java b/ratis-common/src/main/java/org/apache/ratis/util/ReferenceCountedObject.java
@@ -23,6 +23,7 @@
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * A reference-counted object can be retained for later use
@@ -152,6 +153,13 @@ public boolean release() {
     };
   }
 
+  /**
+   * @return a {@link ReferenceCountedObject} by apply the given function to this object.
+   */
+  default <V> ReferenceCountedObject<V> apply(Function<T, V> function) {
+    return delegate(function.apply(get()));
+  }
+
   /**
    * Wrap the given value as a {@link ReferenceCountedObject}.
    *

diff --git a/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java b/ratis-examples/src/main/java/org/apache/ratis/examples/filestore/FileStoreStateMachine.java
@@ -115,11 +115,12 @@ public TransactionContext startTransaction(RaftClientRequest request) throws IOE
 
   @Override
   public TransactionContext startTransaction(LogEntryProto entry, RaftProtos.RaftPeerRole role) {
+    ByteString copied = ByteString.copyFrom(entry.getStateMachineLogEntry().getLogData().asReadOnlyByteBuffer());
     return TransactionContext.newBuilder()
         .setStateMachine(this)
         .setLogEntry(entry)
         .setServerRole(role)
-        .setStateMachineContext(getProto(entry))
+        .setStateMachineContext(getProto(copied))
         .build();
   }
 
@@ -147,14 +148,14 @@ static FileStoreRequestProto getProto(TransactionContext context, LogEntryProto
         return proto;
       }
     }
-    return getProto(entry);
+    return getProto(entry.getStateMachineLogEntry().getLogData());
   }
 
-  static FileStoreRequestProto getProto(LogEntryProto entry) {
+  static FileStoreRequestProto getProto(ByteString bytes) {
     try {
-      return FileStoreRequestProto.parseFrom(entry.getStateMachineLogEntry().getLogData());
+      return FileStoreRequestProto.parseFrom(bytes);
     } catch (InvalidProtocolBufferException e) {
-      throw new IllegalArgumentException("Failed to parse data, entry=" + entry, e);
+      throw new IllegalArgumentException("Failed to parse data", e);
     }
   }
 

diff --git a/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java b/ratis-server-api/src/main/java/org/apache/ratis/statemachine/TransactionContext.java
@@ -60,7 +60,10 @@ public interface TransactionContext {
   /**
    * Returns the data from the {@link StateMachine}
    * @return the data from the {@link StateMachine}
+   * @deprecated access StateMachineLogEntry via {@link TransactionContext#getLogEntryRef()} or
+   * {@link TransactionContext#getLogEntryUnsafe()}
    */
+  @Deprecated
   StateMachineLogEntryProto getStateMachineLogEntry();
 
   /** Set exception in case of failure. */

diff --git a/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java b/ratis-test/src/test/java/org/apache/ratis/statemachine/TestStateMachine.java
@@ -92,7 +92,7 @@ public TransactionContext startTransaction(RaftClientRequest request) {
     public CompletableFuture<Message> applyTransaction(TransactionContext trx) {
       try {
         assertNotNull(trx.getLogEntryUnsafe());
-        assertNotNull(trx.getStateMachineLogEntry());
+        assertNotNull(trx.getLogEntryUnsafe().getStateMachineLogEntry());
         Object context = trx.getStateMachineContext();
         if (isLeader.get()) {
           assertNotNull(trx.getClientRequest());