From 0cd1d39db2cd158c1aead68b26b55d597a8727fb Mon Sep 17 00:00:00 2001
From: baranowb <baranowb@gmail.com>
Date: Wed, 13 Nov 2024 09:01:31 +0100
Subject: [PATCH] [UNDERTOW-2356] AbstractFramedStreamSinkConduit

---
 .../conduits/AbstractFramedStreamSinkConduit.java  | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/core/src/main/java/io/undertow/conduits/AbstractFramedStreamSinkConduit.java b/core/src/main/java/io/undertow/conduits/AbstractFramedStreamSinkConduit.java
index 8ca6ed87af..34ae8cb727 100644
--- a/core/src/main/java/io/undertow/conduits/AbstractFramedStreamSinkConduit.java
+++ b/core/src/main/java/io/undertow/conduits/AbstractFramedStreamSinkConduit.java
@@ -19,6 +19,7 @@
 package io.undertow.conduits;
 
 import io.undertow.UndertowMessages;
+
 import org.xnio.Buffers;
 import org.xnio.IoUtils;
 import io.undertow.connector.PooledByteBuffer;
@@ -33,6 +34,7 @@
 import java.nio.channels.FileChannel;
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 
 import static org.xnio.Bits.allAreClear;
 import static org.xnio.Bits.anyAreSet;
@@ -58,8 +60,9 @@ public class AbstractFramedStreamSinkConduit extends AbstractStreamSinkConduit<S
      */
     private int bufferCount = 0;
 
-    private int state;
-
+    private volatile int state;
+    private static final AtomicIntegerFieldUpdater<AbstractFramedStreamSinkConduit> stateUpdater = AtomicIntegerFieldUpdater.newUpdater(
+            AbstractFramedStreamSinkConduit.class, "state");
     private static final int FLAG_WRITES_TERMINATED = 1;
     private static final int FLAG_DELEGATE_SHUTDOWN = 2;
 
@@ -191,11 +194,12 @@ public void terminateWrites() throws IOException {
             return;
         }
         queueCloseFrames();
-        state |= FLAG_WRITES_TERMINATED;
         if (queuedData == 0) {
-            state |= FLAG_DELEGATE_SHUTDOWN;
+            stateUpdater.accumulateAndGet(this, FLAG_WRITES_TERMINATED | FLAG_DELEGATE_SHUTDOWN, (current, flag) -> current | flag);
             doTerminateWrites();
             finished();
+        } else {
+            stateUpdater.accumulateAndGet(this, FLAG_WRITES_TERMINATED, (current, flag) -> current | flag);
         }
     }
 
@@ -212,7 +216,7 @@ protected boolean flushQueuedData() throws IOException {
         }
         if (anyAreSet(state, FLAG_WRITES_TERMINATED) && allAreClear(state, FLAG_DELEGATE_SHUTDOWN)) {
             doTerminateWrites();
-            state |= FLAG_DELEGATE_SHUTDOWN;
+            stateUpdater.accumulateAndGet(this, FLAG_DELEGATE_SHUTDOWN, (current, flag) -> current | flag);
             finished();
         }
         return next.flush();