Address ikhoon comments

core/src/main/java/com/linecorp/armeria/common/HttpRequest.java

@@ -288,6 +288,7 @@ static HttpRequest of(RequestHeaders headers, Publisher<? extends HttpObject> pu
      * {@link Subscription#cancel()} or {@link #abort()} is called. You should add a hook in order to
      * release the elements. See {@link PublisherBasedStreamMessage} for more information.
      */
+    @UnstableApi
     static HttpRequest of(RequestHeaders headers,
                           Publisher<? extends HttpData> publisher,
                           HttpHeaders trailers) {

core/src/main/java/com/linecorp/armeria/internal/common/stream/SurroundingPublisher.java

@@ -21,7 +21,6 @@
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
-import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
@@ -33,6 +32,7 @@
 import com.linecorp.armeria.common.stream.AbortedStreamException;
 import com.linecorp.armeria.common.stream.CancelledSubscriptionException;
 import com.linecorp.armeria.common.stream.NoopSubscriber;
+import com.linecorp.armeria.common.stream.PublisherBasedStreamMessage;
 import com.linecorp.armeria.common.stream.StreamMessage;
 import com.linecorp.armeria.common.stream.SubscriptionOption;
 import com.linecorp.armeria.common.util.EventLoopCheckingFuture;
@@ -48,7 +48,7 @@ public final class SurroundingPublisher<T> implements StreamMessage<T> {
 
     @Nullable
     private final T head;
-    private final Publisher<? extends T> publisher;
+    private final StreamMessage<T> publisher;
     @Nullable
     private final T tail;
 
@@ -58,10 +58,15 @@ public final class SurroundingPublisher<T> implements StreamMessage<T> {
     @Nullable
     private volatile SurroundingSubscriber<T> surroundingSubscriber;
 
+    @SuppressWarnings("unchecked")
     public SurroundingPublisher(@Nullable T head, Publisher<? extends T> publisher, @Nullable T tail) {
         requireNonNull(publisher, "publisher");
         this.head = head;
-        this.publisher = publisher;
+        if (publisher instanceof StreamMessage) {
+            this.publisher = (StreamMessage<T>) publisher;
+        } else {
+            this.publisher = new PublisherBasedStreamMessage<>(publisher);
+        }
         this.tail = tail;
     }
 
@@ -118,8 +123,7 @@ private void subscribe0(Subscriber<? super T> subscriber, EventExecutor executor
                             SubscriptionOption... options) {
 
         final SurroundingSubscriber<T> surroundingSubscriber = new SurroundingSubscriber<>(
-                head, publisher, tail, subscriber, executor,
-                completionFuture, containsNotifyCancellation(options));
+                head, publisher, tail, subscriber, executor, completionFuture, options);
         this.surroundingSubscriber = surroundingSubscriber;
         subscriber.onSubscribe(surroundingSubscriber);
 
@@ -154,10 +158,6 @@ public void abort(Throwable cause) {
 
     static final class SurroundingSubscriber<T> implements Subscriber<T>, Subscription {
 
-        @SuppressWarnings("rawtypes")
-        private static final AtomicReferenceFieldUpdater<SurroundingSubscriber, State> stateUpdater =
-                AtomicReferenceFieldUpdater.newUpdater(SurroundingSubscriber.class, State.class, "state");
-
         enum State {
             REQUIRE_HEAD,
             REQUIRE_BODY,
@@ -170,7 +170,7 @@ enum State {
 
         @Nullable
         private final T head;
-        private final Publisher<? extends T> publisher;
+        private final StreamMessage<T> publisher;
         @Nullable
         private final T tail;
 
@@ -185,11 +185,11 @@ enum State {
         private volatile boolean closed;
 
         private final CompletableFuture<Void> completionFuture;
-        private final boolean notifyCancellation;
+        private final SubscriptionOption[] options;
 
-        SurroundingSubscriber(@Nullable T head, Publisher<? extends T> publisher, @Nullable T tail,
+        SurroundingSubscriber(@Nullable T head, StreamMessage<T> publisher, @Nullable T tail,
                               Subscriber<? super T> downstream, EventExecutor executor,
-                              CompletableFuture<Void> completionFuture, boolean notifyCancellation) {
+                              CompletableFuture<Void> completionFuture, SubscriptionOption... options) {
             requireNonNull(publisher, "publisher");
             requireNonNull(downstream, "downstream");
             requireNonNull(executor, "executor");
@@ -200,7 +200,7 @@ enum State {
             this.downstream = downstream;
             this.executor = executor;
             this.completionFuture = completionFuture;
-            this.notifyCancellation = notifyCancellation;
+            this.options = options;
         }
 
         @Override
@@ -253,7 +253,7 @@ private void publish() {
                 case REQUIRE_BODY: {
                     if (!subscribed) {
                         subscribed = true;
-                        publisher.subscribe(this);
+                        publisher.subscribe(this, executor, options);
                         return;
                     }
                     if (upstream != null) {
@@ -362,7 +362,7 @@ private void cancel0() {
             closed = true;
 
             final CancelledSubscriptionException cause = CancelledSubscriptionException.get();
-            if (notifyCancellation) {
+            if (containsNotifyCancellation(options)) {
                 downstream.onError(cause);
             }
             downstream = NoopSubscriber.get();
@@ -407,9 +407,11 @@ private void release() {
             }
         }
 
-        private boolean setState(State oldState, State newState) {
+        private void setState(State oldState, State newState) {
+            assert state == oldState :
+                    "curState: " + state + ", oldState: " + oldState + ", newState: " + newState;
             assert newState != State.REQUIRE_HEAD : "oldState: " + oldState + ", newState: " + newState;
-            return stateUpdater.compareAndSet(this, oldState, newState);
+            state = newState;
         }
     }
 }

core/src/test/java/com/linecorp/armeria/internal/common/stream/SurroundingPublisherTckTest.java

@@ -50,20 +50,12 @@ public StreamMessage<Object> createPublisher(long elements) {
             return publisher;
         }
         if (elements == 1) {
-            // return new SurroundingPublisher<>("head", Mono.empty(), null);
-            // return new SurroundingPublisher<>(null, Mono.just(1), null);
-            return new SurroundingPublisher<>(null, Mono.empty(), "tail");
+            return new SurroundingPublisher<>("head", Mono.empty(), null);
         }
         if (elements == 2) {
-            // return new SurroundingPublisher<>("head", Mono.just(1), null);
-            // return new SurroundingPublisher<>("head", Mono.empty(), "tail");
-            // return new SurroundingPublisher<>(null, Mono.just(1), "tail");
-            return new SurroundingPublisher<>(null, Flux.just(1, 2), null);
+            return new SurroundingPublisher<>(null, Mono.just(1), "tail");
         }
-        // return new SurroundingPublisher<>("head", Flux.fromStream(LongStream.range(0, elements - 1).boxed()), null);
-        // return new SurroundingPublisher<>("head", Flux.fromStream(LongStream.range(0, elements - 2).boxed()), "tail");
-        // return new SurroundingPublisher<>(null, Flux.fromStream(LongStream.range(0, elements - 1).boxed()), "tail");
-        return new SurroundingPublisher<>(null, Flux.fromStream(LongStream.range(0, elements).boxed()), null);
+        return new SurroundingPublisher<>("head", Flux.fromStream(LongStream.range(0, elements - 2).boxed()), "tail");
     }
 
     /**