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+ try {
96+ CompletableFuture <Void > writeFuture ;
97+ synchronized (streamLock ) {
98+ if (streamClosed ) {
99+ result .completeExceptionally (new IOException ("Stream is already closed, cannot write data." ));
100+ return ;
101+ }
102+ writeFuture = s .writeData (data , endStream );
103+ }
104+ writeFuture .whenComplete ((v , err ) -> {
105+ if (err != null ) {
106+ result .completeExceptionally (unwrap (err ));
107+ } else {
108+ result .complete (null );
109+ }
110+ });
111+ } catch (Throwable th ) {
112+ // s.writeData can throw synchronously (e.g., AWS_ERROR_HTTP_STREAM_NOT_ACTIVATED).
113+ // Tear down the stream so the caller's connection is not leaked back to the pool.
114+ closeConnection ();
115+ result .completeExceptionally (th );
116+ }
117+ });
118+ return result ;
119+ }
120+
121+ private static Throwable unwrap (Throwable t ) {
122+ return t instanceof CompletionException && t .getCause () != null ? t .getCause () : t ;
77123 }
78124
79125 public void incrementWindow (int windowSize ) {
80- if (streamLatch .getCount () != 0 ) {
81- throw new IllegalStateException ("incrementWindow called before stream is ready." );
82- }
83126 synchronized (streamLock ) {
84- if (!streamClosed ) {
85- stream .incrementWindow (windowSize );
127+ HttpStreamBase s = streamFuture .getNow (null );
128+ if (!streamClosed && s != null ) {
129+ s .incrementWindow (windowSize );
86130 }
87131 }
88132 }
@@ -93,9 +137,10 @@ public void incrementWindow(int windowSize) {
93137 */
94138 public void releaseConnection () {
95139 synchronized (streamLock ) {
96- if (!streamClosed && stream != null ) {
140+ HttpStreamBase s = streamFuture .getNow (null );
141+ if (!streamClosed && s != null ) {
97142 streamClosed = true ;
98- stream .close ();
143+ s .close ();
99144 }
100145 }
101146 }
@@ -107,10 +152,11 @@ public void releaseConnection() {
107152 */
108153 public void closeConnection () {
109154 synchronized (streamLock ) {
110- if (!streamClosed && stream != null ) {
155+ HttpStreamBase s = streamFuture .getNow (null );
156+ if (!streamClosed && s != null ) {
111157 streamClosed = true ;
112- stream .cancel ();
113- stream .close ();
158+ s .cancel ();
159+ s .close ();
114160 }
115161 }
116162 }
0 commit comments