Trivial optimizations on copies

Author: slinkydeveloper

sdk-core/src/main/java/dev/restate/sdk/core/HandlerContextImpl.java

@@ -530,8 +530,8 @@ public void proposeFailure(Throwable toWrite, @Nullable RetryPolicy retryPolicy)
   }
 
   private void pumpOutput() {
-    byte[] chunk = stateMachine.takeOutput();
-    if (chunk.length > 0) outputSink.accept(Slice.wrap(chunk));
+    Slice chunk = stateMachine.takeOutput();
+    if (chunk.readableBytes() > 0) outputSink.accept(chunk);
   }
 
   @Override

sdk-core/src/main/java/dev/restate/sdk/core/RequestProcessorImpl.java

@@ -131,7 +131,7 @@ public void onNext(Slice slice) {
     if (state == State.CLOSED) return;
 
     try {
-      stateMachine.notifyInput(slice.toByteArray());
+      stateMachine.notifyInput(slice);
       onInputEvent();
     } catch (Throwable e) {
       onError(e);
@@ -199,19 +199,19 @@ private void onClose() {
     cancelInputSubscription();
 
     // Pump remaining output
-    byte[] chunk;
+    Slice chunk;
     if (outputSubscriber != null) {
       chunk = stateMachine.takeOutput();
     } else {
-      chunk = new byte[0];
+      chunk = Slice.EMPTY;
     }
 
     // Close state machine
     this.state = State.CLOSED;
     stateMachine.close();
 
     // Send final bits and close output subscriber
-    if (chunk.length > 0) outputSubscriber.onNext(Slice.wrap(chunk));
+    if (chunk.readableBytes() > 0) outputSubscriber.onNext(chunk);
     outputSubscriber.onComplete();
     outputSubscriber = null;
   }

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/JavaStateMachine.java

@@ -41,7 +41,7 @@
  *
  * <p>This is a port of the legacy {@code dev.restate.sdk.core.statemachine.StateMachineImpl}, which
  * was a {@link Flow.Processor}. The canonical interface is imperative, so the legacy {@code
- * waitForReady()} / {@code onNextEvent()} signalling is replaced by {@link #notifyInput(byte[])} /
+ * waitForReady()} / {@code onNextEvent()} signalling is replaced by {@link #notifyInput(Slice)} /
  * {@link #isReadyToExecute()} and synchronous {@link #doAwait(UnresolvedFuture)}, while the Flow
  * output side is replaced by a buffer drained via {@link #takeOutput()}.
  */
@@ -101,9 +101,12 @@ public JavaStateMachine(HeadersAccessor headersAccessor) {
   // -------------------------------------------------------------------------
 
   @Override
-  public void notifyInput(byte[] bytes) {
+  public void notifyInput(Slice bytes) {
     LOG.trace("Received input slice");
-    this.messageDecoder.offer(Slice.wrap(bytes));
+    // Zero-copy: hand the slice straight to the decoder (which retains it via unsafeWrap). Safe
+    // because the transport hands out a retainable buffer (Vert.x wraps non-heap byte bufs as
+    // unreleasable), matching the legacy/main input path.
+    this.messageDecoder.offer(bytes);
 
     InvocationInput invocationInput = this.messageDecoder.next();
     while (invocationInput != null) {
@@ -132,8 +135,8 @@ public void notifyError(Throwable throwable) {
   }
 
   @Override
-  public byte[] takeOutput() {
-    return this.outputSink.take();
+  public Slice takeOutput() {
+    return Slice.wrap(this.outputSink.take());
   }
 
   @Override

sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachine.java

@@ -49,13 +49,18 @@ public interface StateMachine extends AutoCloseable {
   // Imperative replacement for main's Flow.Processor + waitForReady/onNextEvent: bytes are fed in
   // with notifyInput, drained out with takeOutput, and isReadyToExecute gates handler start.
 
-  void notifyInput(byte[] bytes);
+  /** Feed the next chunk of wire input. */
+  void notifyInput(Slice bytes);
 
   void notifyInputClosed();
 
   void notifyError(Throwable throwable);
 
-  byte[] takeOutput();
+  /**
+   * Drain the next chunk of serialized output ready for the wire, or {@link Slice#EMPTY} when
+   * nothing is buffered.
+   */
+  Slice takeOutput();
 
   boolean isReadyToExecute();
 

sdk-core/src/main/java23/dev/restate/sdk/core/statemachine/ffm/FfmEncoding.java

@@ -19,6 +19,7 @@
 import java.io.ByteArrayOutputStream;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SegmentAllocator;
 import java.lang.foreign.ValueLayout;
 import java.nio.charset.StandardCharsets;
 import java.util.Collection;
@@ -40,31 +41,45 @@ final class FfmEncoding {
   private FfmEncoding() {}
 
   // -------------------------------------------------------------------------
-  // Native segment allocation (inputs: copy-for-now into the call arena)
+  // Native segment allocation (inputs: copy-for-now into the call scratch)
   // -------------------------------------------------------------------------
 
   /**
    * Allocate a native segment holding the given bytes, or {@link MemorySegment#NULL} when {@code
    * bytes} is null/empty. The Rust side treats a null ptr / zero len as the empty slice.
    */
-  static MemorySegment allocateBytes(Arena arena, byte @Nullable [] bytes) {
+  static MemorySegment allocateBytes(SegmentAllocator alloc, byte @Nullable [] bytes) {
     if (bytes == null || bytes.length == 0) {
       return MemorySegment.NULL;
     }
-    MemorySegment seg = arena.allocate(bytes.length);
+    MemorySegment seg = alloc.allocate(bytes.length);
     MemorySegment.copy(bytes, 0, seg, ValueLayout.JAVA_BYTE, 0, bytes.length);
     return seg;
   }
 
+  /**
+   * Allocate a native segment holding the bytes of {@code slice}, copied straight from the slice's
+   * buffer (no intermediate heap {@code byte[]}), or {@link MemorySegment#NULL} when empty.
+   */
+  static MemorySegment allocateSlice(SegmentAllocator alloc, dev.restate.common.Slice slice) {
+    int len = slice.readableBytes();
+    if (len == 0) {
+      return MemorySegment.NULL;
+    }
+    MemorySegment seg = alloc.allocate(len);
+    slice.copyTo(seg.asByteBuffer());
+    return seg;
+  }
+
   /**
    * Allocate a native segment holding the UTF-8 bytes of {@code s}, or {@link MemorySegment#NULL}
    * when {@code s} is null. An empty (non-null) string yields a zero-length, non-null segment.
    */
-  static MemorySegment allocateUtf8(Arena arena, @Nullable String s) {
+  static MemorySegment allocateUtf8(SegmentAllocator alloc, @Nullable String s) {
     if (s == null) {
       return MemorySegment.NULL;
     }
-    return allocateBytes(arena, s.getBytes(StandardCharsets.UTF_8));
+    return allocateBytes(alloc, s.getBytes(StandardCharsets.UTF_8));
   }
 
   static long len(MemorySegment seg) {
@@ -102,6 +117,37 @@ static String takeSliceString(MemorySegment sliceStruct) {
     return new String(bytes, StandardCharsets.UTF_8);
   }
 
+  /**
+   * Wraps an owned result {@link Slice} (Rust-allocated; must be released via {@code free_buffer})
+   * in a {@link MemorySegmentSlice} read zero-copy — no copy-out into a heap {@code byte[]}. The
+   * native buffer is freed when the returned slice becomes unreachable (GC-tied via {@link
+   * Arena#ofAuto()}), so it is safe to read later and from another thread (deserialization on a
+   * different dispatcher, or an async socket write). Returns {@link dev.restate.common.Slice#EMPTY}
+   * for the null/empty slice.
+   */
+  static dev.restate.common.Slice wrapOwnedSlice(MemorySegment sliceStruct) {
+    MemorySegment ptr = Slice.ptr(sliceStruct);
+    long len = Slice.len(sliceStruct);
+    if (ptr.address() == 0 || len == 0) {
+      return dev.restate.common.Slice.EMPTY;
+    }
+    MemorySegment seg =
+        ptr.reinterpret(len, Arena.ofAuto(), s -> SharedCoreNative.free_buffer(s, len));
+    return new MemorySegmentSlice(seg);
+  }
+
+  /**
+   * Releases an owned result {@link Slice} we are not consuming (e.g. the unused {@code value}/
+   * {@code extra} arm of a notification). No-op for the null/empty slice.
+   */
+  static void freeSlice(MemorySegment sliceStruct) {
+    MemorySegment ptr = Slice.ptr(sliceStruct);
+    long len = Slice.len(sliceStruct);
+    if (ptr.address() != 0 && len != 0) {
+      SharedCoreNative.free_buffer(ptr, len);
+    }
+  }
+
   private static final byte[] EMPTY_BYTES = new byte[0];
 
   // -------------------------------------------------------------------------
@@ -114,54 +160,54 @@ static String takeSliceString(MemorySegment sliceStruct) {
    * header-list blob.
    */
   static MemorySegment buildTarget(
-      Arena arena,
+      SegmentAllocator alloc,
       Target target,
       @Nullable String idempotencyKey,
       @Nullable Collection<Map.Entry<String, String>> headers) {
-    MemorySegment t = TargetAbi.allocate(arena);
+    MemorySegment t = TargetAbi.allocate(alloc);
 
-    MemorySegment service = allocateUtf8(arena, target.getService());
+    MemorySegment service = allocateUtf8(alloc, target.getService());
     TargetAbi.service_ptr(t, service);
     TargetAbi.service_len(t, len(service));
 
-    MemorySegment handler = allocateUtf8(arena, target.getHandler());
+    MemorySegment handler = allocateUtf8(alloc, target.getHandler());
     TargetAbi.handler_ptr(t, handler);
     TargetAbi.handler_len(t, len(handler));
 
-    MemorySegment key = allocateUtf8(arena, target.getKey());
+    MemorySegment key = allocateUtf8(alloc, target.getKey());
     TargetAbi.key_ptr(t, key);
     TargetAbi.key_len(t, len(key));
 
-    MemorySegment idem = allocateUtf8(arena, idempotencyKey);
+    MemorySegment idem = allocateUtf8(alloc, idempotencyKey);
     TargetAbi.idempotency_key_ptr(t, idem);
     TargetAbi.idempotency_key_len(t, len(idem));
 
-    MemorySegment headersBlob = allocateBytes(arena, encodeHeaderList(headers));
+    MemorySegment headersBlob = allocateBytes(alloc, encodeHeaderList(headers));
     TargetAbi.headers_ptr(t, headersBlob);
     TargetAbi.headers_len(t, len(headersBlob));
 
     return t;
   }
 
-  /** Builds a success-valued {@link NonEmptyValueAbi} struct in {@code arena}. */
-  static MemorySegment buildSuccessValue(Arena arena, byte[] value) {
-    MemorySegment v = NonEmptyValueAbi.allocate(arena);
+  /** Builds a success-valued {@link NonEmptyValueAbi} struct, copying {@code value} zero-hop. */
+  static MemorySegment buildSuccessValue(SegmentAllocator alloc, dev.restate.common.Slice value) {
+    MemorySegment v = NonEmptyValueAbi.allocate(alloc);
     NonEmptyValueAbi.is_failure(v, 0);
-    MemorySegment val = allocateBytes(arena, value);
+    MemorySegment val = allocateSlice(alloc, value);
     NonEmptyValueAbi.value_ptr(v, val);
     NonEmptyValueAbi.value_len(v, len(val));
     NonEmptyValueAbi.failure_ptr(v, MemorySegment.NULL);
     NonEmptyValueAbi.failure_len(v, 0L);
     return v;
   }
 
-  /** Builds a failure-valued {@link NonEmptyValueAbi} struct in {@code arena}. */
-  static MemorySegment buildFailureValue(Arena arena, TerminalException failure) {
-    MemorySegment v = NonEmptyValueAbi.allocate(arena);
+  /** Builds a failure-valued {@link NonEmptyValueAbi} struct. */
+  static MemorySegment buildFailureValue(SegmentAllocator alloc, TerminalException failure) {
+    MemorySegment v = NonEmptyValueAbi.allocate(alloc);
     NonEmptyValueAbi.is_failure(v, 1);
     NonEmptyValueAbi.value_ptr(v, MemorySegment.NULL);
     NonEmptyValueAbi.value_len(v, 0L);
-    MemorySegment fail = allocateBytes(arena, encodeFailure(failure));
+    MemorySegment fail = allocateBytes(alloc, encodeFailure(failure));
     NonEmptyValueAbi.failure_ptr(v, fail);
     NonEmptyValueAbi.failure_len(v, len(fail));
     return v;