(ratis) branch master updated: RATIS-2094. Avoid corruptions from TransactionContext's stateMachineLogEntry and stateMachineContext. (#1106)

[szetszwo]

This is an automated email from the ASF dual-hosted git repository.

szetszwo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/ratis.git


The following commit(s) were added to refs/heads/master by this push:
     new a2bdd10fc RATIS-2094. Avoid corruptions from TransactionContext's 
stateMachineLogEntry and stateMachineContext. (#1106)
a2bdd10fc is described below

commit a2bdd10fccb346d02f12d3961d6a5b2d281c5d56
Author: Duong Nguyen <duongnt...@gmail.com>
AuthorDate: Mon Jun 3 10:35:07 2024 -0700

    RATIS-2094. Avoid corruptions from TransactionContext's 
stateMachineLogEntry and stateMachineContext. (#1106)
---
 .../java/org/apache/ratis/util/ReferenceCountedObject.java    |  8 ++++++++
 .../ratis/examples/filestore/FileStoreStateMachine.java       | 11 ++++++-----
 .../org/apache/ratis/statemachine/TransactionContext.java     |  3 +++
 .../java/org/apache/ratis/statemachine/TestStateMachine.java  |  2 +-
 4 files changed, 18 insertions(+), 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
index eb5ff30ae..b2c53182d 100644
--- 
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.Collection;
 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 interface ReferenceCountedObject<T> {
     };
   }
 
+  /**
+   * @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
index f870cbacf..2fb9cace6 100644
--- 
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 class FileStoreStateMachine extends 
BaseStateMachine {
 
   @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 @@ public class FileStoreStateMachine extends 
BaseStateMachine {
         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
index 2ec87e37a..6f08c4a55 100644
--- 
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
index 07ea4edbc..0c49b31af 100644
--- 
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 class TestStateMachine extends BaseTest implements 
MiniRaftClusterWithSim
     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());