1717
1818import java .io .IOException ;
1919import java .util .concurrent .CompletableFuture ;
20- import java .util .concurrent .CountDownLatch ;
20+ import java .util .concurrent .CompletionException ;
21+ import java .util .concurrent .atomic .AtomicBoolean ;
2122import software .amazon .awssdk .annotations .SdkInternalApi ;
2223import software .amazon .awssdk .crt .http .HttpStreamBase ;
24+ import software .amazon .awssdk .http .async .SdkAsyncHttpResponseHandler ;
25+ import software .amazon .awssdk .utils .CompletableFutureUtils ;
26+ import software .amazon .awssdk .utils .Logger ;
2327
2428/**
2529 * Manages the lifecycle of a CRT HTTP stream, providing thread-safe access to stream operations.
2630 * Shared between the request executor (for writing body data) and the response handler (for
2731 * incrementing the window and releasing/closing the connection).
32+ *
33+ * <p>The handler is constructed with a {@link CompletableFuture} representing stream acquisition.
34+ * The caller (request executor) completes that future once the underlying CRT stream manager has
35+ * either acquired the stream or failed. All operations on this handler chain off that future, so
36+ * writes issued before acquisition completes are queued.
2837 */
2938@ SdkInternalApi
3039public final class CrtStreamHandler {
3140
41+ private static final Logger log = Logger .loggerFor (CrtStreamHandler .class );
42+
3243 private final Object streamLock = new Object ();
33- private final CountDownLatch streamLatch = new CountDownLatch ( 1 ) ;
34- private HttpStreamBase stream ;
44+ private final CompletableFuture < HttpStreamBase > streamFuture ;
45+ private final AtomicBoolean responseHandlerNotified = new AtomicBoolean ( false ) ;
3546 private boolean streamClosed ;
3647
37- /**
38- * Sets the stream. Called once when the stream is acquired from the connection pool.
39- */
40- public void setStream (HttpStreamBase stream ) {
41- this .stream = stream ;
42- streamLatch .countDown ();
48+ public CrtStreamHandler (CompletableFuture <HttpStreamBase > streamFuture ) {
49+ this .streamFuture = streamFuture ;
4350 }
4451
4552 /**
46- * Blocks until the stream has been acquired.
53+ * Atomically notifies the {@link SdkAsyncHttpResponseHandler#onError(Throwable)} callback at most
54+ * once across all callers sharing this handler. Returns {@code true} if this caller delivered the
55+ * notification, {@code false} if another caller already did. Exceptions thrown by the handler are
56+ * caught and logged so callers can proceed with their own cleanup.
4757 */
48- public void waitForStream () {
58+ public boolean tryNotifyResponseHandlerError (SdkAsyncHttpResponseHandler handler , Throwable t ) {
59+ if (!responseHandlerNotified .compareAndSet (false , true )) {
60+ return false ;
61+ }
4962 try {
50- streamLatch . await ( );
51- } catch (InterruptedException e ) {
52- Thread . currentThread (). interrupt ();
53- throw new RuntimeException ( "Interrupted while waiting for stream "e );
63+ handler . onError ( t );
64+ } catch (Exception e ) {
65+ log . error (() -> "SdkAsyncHttpResponseHandler " + handler + " threw an exception in onError. It will be "
66+ + "ignored. "e );
5467 }
68+ return true ;
69+ }
70+
71+ /**
72+ * Blocks until the stream has been acquired or acquisition has failed. Returns the acquired
73+ * stream on success. If acquisition failed, the failure cause is rethrown wrapped in a
74+ * {@link CompletionException} so callers can use the same handling as for response futures.
75+ */
76+ public HttpStreamBase waitForStream () {
77+ return CompletableFutureUtils .joinInterruptibly (streamFuture );
5578 }
5679
5780 /**
58- * Write data to the stream. The caller must ensure the stream is ready (via {@link #waitForStream()})
59- * before calling this method.
81+ * Write data to the stream. The returned future chains on stream acquisition: if the stream
82+ * is not yet ready, the write is queued until the {@code streamFuture} passed to the
83+ * constructor completes. Failures from either stream acquisition or the underlying CRT write
84+ * are propagated as the original cause (not wrapped in {@link CompletionException}) so callers
85+ * see the same exception type whether the failure happens before or after {@code thenCompose}-
86+ * style chaining.
6087 */
6188 public CompletableFuture <Void > writeData (byte [] data , boolean endStream ) {
62- if (streamLatch .getCount () != 0 ) {
63- CompletableFuture <Void > future = new CompletableFuture <>();
64- future .completeExceptionally (
65- new IllegalStateException ("writeData called before stream is ready. Call waitForStream() first." ));
66- return future ;
67- }
68- synchronized (streamLock ) {
69- if (streamClosed ) {
70- CompletableFuture <Void > future = new CompletableFuture <>();
71- future .completeExceptionally (
72- new IOException ("Stream is already closed, cannot write data." ));
73- return future ;
89+ CompletableFuture <Void > result = new CompletableFuture <>();
90+ streamFuture .whenComplete ((s , t ) -> {
91+ if (t != null ) {
92+ result .completeExceptionally (unwrap (t ));
93+ return ;
7494 }
75- return stream .writeData (data , endStream );
76- }
95+ CompletableFuture <Void > writeFuture ;
96+ synchronized (streamLock ) {
97+ if (streamClosed ) {
98+ result .completeExceptionally (new IOException ("Stream is already closed, cannot write data." ));
99+ return ;
100+ }
101+ writeFuture = s .writeData (data , endStream );
102+ }
103+ writeFuture .whenComplete ((v , err ) -> {
104+ if (err != null ) {
105+ result .completeExceptionally (unwrap (err ));
106+ } else {
107+ result .complete (null );
108+ }
109+ });
110+ });
111+ return result ;
112+ }
113+
114+ private static Throwable unwrap (Throwable t ) {
115+ return t instanceof CompletionException && t .getCause () != null ? t .getCause () : t ;
77116 }
78117
79118 public void incrementWindow (int windowSize ) {
80- if (streamLatch .getCount () != 0 ) {
81- throw new IllegalStateException ("incrementWindow called before stream is ready." );
82- }
83119 synchronized (streamLock ) {
84- if (!streamClosed ) {
85- stream .incrementWindow (windowSize );
120+ HttpStreamBase s = streamFuture .getNow (null );
121+ if (!streamClosed && s != null ) {
122+ s .incrementWindow (windowSize );
86123 }
87124 }
88125 }
@@ -93,9 +130,10 @@ public void incrementWindow(int windowSize) {
93130 */
94131 public void releaseConnection () {
95132 synchronized (streamLock ) {
96- if (!streamClosed && stream != null ) {
133+ HttpStreamBase s = streamFuture .getNow (null );
134+ if (!streamClosed && s != null ) {
97135 streamClosed = true ;
98- stream .close ();
136+ s .close ();
99137 }
100138 }
101139 }
@@ -107,10 +145,11 @@ public void releaseConnection() {
107145 */
108146 public void closeConnection () {
109147 synchronized (streamLock ) {
110- if (!streamClosed && stream != null ) {
148+ HttpStreamBase s = streamFuture .getNow (null );
149+ if (!streamClosed && s != null ) {
111150 streamClosed = true ;
112- stream .cancel ();
113- stream .close ();
151+ s .cancel ();
152+ s .close ();
114153 }
115154 }
116155 }
0 commit comments