Make ProtoOutputStream thread-safe

p2p/src/main/java/bisq/network/p2p/network/Connection.java

@@ -142,7 +142,7 @@ public static int getShutdownTimeout() {
     private final ConnectionStatistics connectionStatistics;
 
     // set in init
-    private SynchronizedProtoOutputStream protoOutputStream;
+    private ProtoOutputStream protoOutputStream;
 
     // mutable data, set from other threads but not changed internally.
     @Getter
@@ -198,7 +198,7 @@ private void init(@Nullable NodeAddress peersNodeAddress) {
             // When you construct an ObjectInputStream, in the constructor the class attempts to read a header that
             // the associated ObjectOutputStream on the other end of the connection has written.
             // It will not return until that header has been read.
-            protoOutputStream = new SynchronizedProtoOutputStream(socket.getOutputStream(), statistic);
+            protoOutputStream = new ProtoOutputStream(socket.getOutputStream(), statistic);
             protoInputStream = socket.getInputStream();
             // We create a thread for handling inputStream data
             singleThreadExecutor.submit(this);

p2p/src/main/java/bisq/network/p2p/network/ProtoOutputStream.java

@@ -24,37 +24,65 @@
 import java.io.IOException;
 import java.io.OutputStream;
 
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.annotation.concurrent.NotThreadSafe;
+import javax.annotation.concurrent.ThreadSafe;
 
-@NotThreadSafe
+@ThreadSafe
 class ProtoOutputStream {
     private static final Logger log = LoggerFactory.getLogger(ProtoOutputStream.class);
 
     private final OutputStream outputStream;
     private final Statistic statistic;
 
+    private final AtomicBoolean isConnectionActive = new AtomicBoolean(true);
+    private final Lock lock = new ReentrantLock();
+
     ProtoOutputStream(OutputStream outputStream, Statistic statistic) {
         this.outputStream = outputStream;
         this.statistic = statistic;
     }
 
     void writeEnvelope(NetworkEnvelope envelope) {
+        lock.lock();
+
         try {
             writeEnvelopeOrThrow(envelope);
         } catch (IOException e) {
+            if (!isConnectionActive.get()) {
+                // Connection was closed by us.
+                return;
+            }
+
             log.error("Failed to write envelope", e);
             throw new BisqRuntimeException("Failed to write envelope", e);
+
+        } finally {
+            lock.unlock();
         }
     }
 
     void onConnectionShutdown() {
+        isConnectionActive.set(false);
+
+        boolean acquiredLock = tryToAcquireLock();
+        if (!acquiredLock) {
+            return;
+        }
+
         try {
             outputStream.close();
         } catch (Throwable t) {
             log.error("Failed to close connection", t);
+
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -74,4 +102,13 @@ private void writeEnvelopeOrThrow(NetworkEnvelope envelope) throws IOException {
             statistic.updateLastActivityTimestamp();
         }
     }
+
+    private boolean tryToAcquireLock() {
+        long shutdownTimeout = Connection.getShutdownTimeout();
+        try {
+            return lock.tryLock(shutdownTimeout, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            return false;
+        }
+    }
 }