From 7357a1a379ed79c6754a8093eb108cd82062880a Mon Sep 17 00:00:00 2001
From: Daniel Fuchs <dfuchs@openjdk.org>
Date: Tue, 15 Nov 2022 18:36:45 +0000
Subject: [PATCH] 8296889: Race condition when cancelling a request

Reviewed-by: jpai
---
 .../jdk/internal/net/http/Http1Exchange.java        |  6 +++++-
 .../share/classes/jdk/internal/net/http/Stream.java | 10 +++++++---
 .../net/http/common/HttpBodySubscriberWrapper.java  | 13 ++++++++++++-
 test/jdk/java/net/httpclient/CancelRequestTest.java |  2 +-
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/java.net.http/share/classes/jdk/internal/net/http/Http1Exchange.java b/src/java.net.http/share/classes/jdk/internal/net/http/Http1Exchange.java
index 947511002f9..f03b174943d 100644
--- a/src/java.net.http/share/classes/jdk/internal/net/http/Http1Exchange.java
+++ b/src/java.net.http/share/classes/jdk/internal/net/http/Http1Exchange.java
@@ -208,6 +208,11 @@ class Http1Exchange<T> extends ExchangeImpl<T> {
             this.exchange = exchange;
         }
 
+        @Override
+        public void onSubscribed() {
+            exchange.registerResponseSubscriber(this);
+        }
+
         @Override
         protected void complete(Throwable t) {
             try {
@@ -459,7 +464,6 @@ class Http1Exchange<T> extends ExchangeImpl<T> {
         BodySubscriber<T> subscriber = handler.apply(response);
         Http1ResponseBodySubscriber<T> bs =
                 new Http1ResponseBodySubscriber<T>(subscriber, this);
-        registerResponseSubscriber(bs);
         return bs;
     }
 
diff --git a/src/java.net.http/share/classes/jdk/internal/net/http/Stream.java b/src/java.net.http/share/classes/jdk/internal/net/http/Stream.java
index 273201b7218..fc97f2b413c 100644
--- a/src/java.net.http/share/classes/jdk/internal/net/http/Stream.java
+++ b/src/java.net.http/share/classes/jdk/internal/net/http/Stream.java
@@ -344,7 +344,6 @@ class Stream<T> extends ExchangeImpl<T> {
     Http2StreamResponseSubscriber<T> createResponseSubscriber(BodyHandler<T> handler, ResponseInfo response) {
         Http2StreamResponseSubscriber<T> subscriber =
                 new Http2StreamResponseSubscriber<>(handler.apply(response));
-        registerResponseSubscriber(subscriber);
         return subscriber;
     }
 
@@ -1543,17 +1542,22 @@ class Stream<T> extends ExchangeImpl<T> {
             super(subscriber);
         }
 
+        @Override
+        public void onSubscribed() {
+            registerResponseSubscriber(this);
+        }
+
         @Override
         protected void complete(Throwable t) {
             try {
-                Stream.this.unregisterResponseSubscriber(this);
+                unregisterResponseSubscriber(this);
             } finally {
                 super.complete(t);
             }
         }
         @Override
         protected void onCancel() {
-            Stream.this.unregisterResponseSubscriber(this);
+            unregisterResponseSubscriber(this);
         }
     }
 
diff --git a/src/java.net.http/share/classes/jdk/internal/net/http/common/HttpBodySubscriberWrapper.java b/src/java.net.http/share/classes/jdk/internal/net/http/common/HttpBodySubscriberWrapper.java
index 1a0457e263b..2e1af8bf3b6 100644
--- a/src/java.net.http/share/classes/jdk/internal/net/http/common/HttpBodySubscriberWrapper.java
+++ b/src/java.net.http/share/classes/jdk/internal/net/http/common/HttpBodySubscriberWrapper.java
@@ -127,6 +127,15 @@ public class HttpBodySubscriberWrapper<T> implements TrustedSubscriber<T> {
      */
     protected void onCancel() { }
 
+    /**
+     * Called right after the userSubscriber::onSubscribe is called.
+     * @apiNote
+     * This method may be used by subclasses to perform cleanup
+     * related actions after a subscription has been succesfully
+     * accepted.
+     */
+    protected void onSubscribed() { }
+
     /**
      * Complete the subscriber, either normally or exceptionally
      * ensure that the subscriber is completed only once.
@@ -169,8 +178,9 @@ public class HttpBodySubscriberWrapper<T> implements TrustedSubscriber<T> {
     public void onSubscribe(Flow.Subscription subscription) {
         // race condition with propagateError: we need to wait until
         // subscription is finished before calling onError;
+        boolean onSubscribed;
         synchronized (this) {
-            if (subscribed.compareAndSet(false, true)) {
+            if ((onSubscribed = subscribed.compareAndSet(false, true))) {
                 SubscriptionWrapper wrapped = new SubscriptionWrapper(subscription);
                 userSubscriber.onSubscribe(this.subscription = wrapped);
             } else {
@@ -181,6 +191,7 @@ public class HttpBodySubscriberWrapper<T> implements TrustedSubscriber<T> {
                 assert completed.get();
             }
         }
+        if (onSubscribed) onSubscribed();
     }
 
     @Override
diff --git a/test/jdk/java/net/httpclient/CancelRequestTest.java b/test/jdk/java/net/httpclient/CancelRequestTest.java
index d80968a4783..57e388cc85c 100644
--- a/test/jdk/java/net/httpclient/CancelRequestTest.java
+++ b/test/jdk/java/net/httpclient/CancelRequestTest.java
@@ -23,7 +23,7 @@
 
 /*
  * @test
- * @bug 8245462 8229822 8254786
+ * @bug 8245462 8229822 8254786 8296889
  * @summary Tests cancelling the request.
  * @library /test/lib http2/server
  * @key randomness