Revert removing writerLock in favor of setting a flag to call flushBuffer should be PR'ed

Author: scottf

src/main/java/io/nats/client/impl/NatsConnectionWriter.java

@@ -37,6 +37,7 @@ class NatsConnectionWriter implements Runnable {
 
     private final NatsConnection connection;
 
+    private final ReentrantLock writerLock;
     private Future<Boolean> stopped;
     private Future<DataPort> dataPortFuture;
     private DataPort dataPort;
@@ -50,10 +51,10 @@ class NatsConnectionWriter implements Runnable {
     private final MessageQueue outgoing;
     private final MessageQueue reconnectOutgoing;
     private final long reconnectBufferSize;
-    private final AtomicBoolean flushBuffer;
 
     NatsConnectionWriter(NatsConnection connection, NatsConnectionWriter sourceWriter) {
         this.connection = connection;
+        writerLock = new ReentrantLock();
 
         this.running = new AtomicBoolean(false);
         this.reconnectMode = new AtomicBoolean(sourceWriter != null);
@@ -76,8 +77,6 @@ class NatsConnectionWriter implements Runnable {
         reconnectOutgoing = new MessageQueue(true, options.getRequestCleanupInterval(),
             sourceWriter == null ? null : sourceWriter.reconnectOutgoing);
         reconnectBufferSize = options.getReconnectBufferSize();
-
-        flushBuffer = new AtomicBoolean(false);
     }
 
     // Should only be called if the current thread has exited.
@@ -123,65 +122,60 @@ boolean isRunning() {
     }
 
     void sendMessageBatch(NatsMessage msg, DataPort dataPort, StatisticsCollector stats) throws IOException {
-        int sendPosition = 0;
-        int sbl = sendBufferLength.get();
-
-        while (msg != null) {
-            long size = msg.getSizeInBytes();
-
-            if (sendPosition + size > sbl) {
-                if (sendPosition > 0) {
-                    dataPort.write(sendBuffer, sendPosition);
-                    connection.getNatsStatistics().registerWrite(sendPosition);
-                    sendPosition = 0;
-                }
-                if (size > sbl) { // have to resize b/c can't fit 1 message
-                    sbl = bufferAllocSize((int) size, BUFFER_BLOCK_SIZE);
-                    sendBufferLength.set(sbl);
-                    sendBuffer = new byte[sbl];
+        writerLock.lock();
+        try {
+            int sendPosition = 0;
+            int sbl = sendBufferLength.get();
+
+            while (msg != null) {
+                long size = msg.getSizeInBytes();
+
+                if (sendPosition + size > sbl) {
+                    if (sendPosition > 0) {
+                        dataPort.write(sendBuffer, sendPosition);
+                        connection.getNatsStatistics().registerWrite(sendPosition);
+                        sendPosition = 0;
+                    }
+                    if (size > sbl) { // have to resize b/c can't fit 1 message
+                        sbl = bufferAllocSize((int) size, BUFFER_BLOCK_SIZE);
+                        sendBufferLength.set(sbl);
+                        sendBuffer = new byte[sbl];
+                    }
                 }
-            }
-
-            ByteArrayBuilder bab = msg.getProtocolBab();
-            int babLen = bab.length();
-            System.arraycopy(bab.internalArray(), 0, sendBuffer, sendPosition, babLen);
-            sendPosition += babLen;
 
-            sendBuffer[sendPosition++] = CR;
-            sendBuffer[sendPosition++] = LF;
-
-            if (!msg.isProtocol()) {
-                sendPosition += msg.copyNotEmptyHeaders(sendPosition, sendBuffer);
-
-                byte[] bytes = msg.getData(); // guaranteed to not be null
-                if (bytes.length > 0) {
-                    System.arraycopy(bytes, 0, sendBuffer, sendPosition, bytes.length);
-                    sendPosition += bytes.length;
-                }
+                ByteArrayBuilder bab = msg.getProtocolBab();
+                int babLen = bab.length();
+                System.arraycopy(bab.internalArray(), 0, sendBuffer, sendPosition, babLen);
+                sendPosition += babLen;
 
                 sendBuffer[sendPosition++] = CR;
                 sendBuffer[sendPosition++] = LF;
-            }
 
-            stats.incrementOutMsgs();
-            stats.incrementOutBytes(size);
+                if (!msg.isProtocol()) {
+                    sendPosition += msg.copyNotEmptyHeaders(sendPosition, sendBuffer);
 
-            msg = msg.next;
-        }
+                    byte[] bytes = msg.getData(); // guaranteed to not be null
+                    if (bytes.length > 0) {
+                        System.arraycopy(bytes, 0, sendBuffer, sendPosition, bytes.length);
+                        sendPosition += bytes.length;
+                    }
 
-        // no need to write if there are no bytes
-        if (sendPosition > 0) {
-            dataPort.write(sendBuffer, sendPosition);
-        }
+                    sendBuffer[sendPosition++] = CR;
+                    sendBuffer[sendPosition++] = LF;
+                }
 
-        try {
-            if (flushBuffer.get()) {
-                dataPort.flush();
-                flushBuffer.set(false);
+                stats.incrementOutMsgs();
+                stats.incrementOutBytes(size);
+
+                msg = msg.next;
             }
-        } catch (Exception ignore) {}
 
-        connection.getNatsStatistics().registerWrite(sendPosition);
+            dataPort.write(sendBuffer, sendPosition);
+            connection.getNatsStatistics().registerWrite(sendPosition);
+        }
+        finally {
+            writerLock.unlock();
+        }
     }
 
     @Override
@@ -243,8 +237,18 @@ void queueInternalMessage(NatsMessage msg) {
     }
 
     void flushBuffer() {
-        if (running.get()) {
-            flushBuffer.set(true);
+        // Since there is no connection level locking, we rely on synchronization
+        // of the APIs here.
+        writerLock.lock();
+        try {
+            if (this.running.get()) {
+                dataPort.flush();
+            }
+        } catch (Exception e) {
+            // NOOP;
+        }
+        finally {
+            writerLock.unlock();
         }
     }
 }