IGNITE-26349 Refactoring setting socket timeout for TcpDiscoverySpi#writeToSocket (#12305)

Author: timoninmaxim

modules/core/src/main/java/org/apache/ignite/spi/discovery/tcp/TcpDiscoverySpi.java

@@ -1674,18 +1674,11 @@ Socket createSocket() throws IOException {
      * @param timeout Socket write timeout.
      * @throws IOException If IO failed or write timed out.
      */
-    @SuppressWarnings("ThrowFromFinallyBlock")
     protected void writeToSocket(Socket sock, TcpDiscoveryAbstractMessage msg, byte[] data, long timeout) throws IOException {
         assert sock != null;
         assert data != null;
 
-        SocketTimeoutObject obj = new SocketTimeoutObject(sock, U.currentTimeMillis() + timeout);
-
-        addTimeoutObject(obj);
-
-        IOException err = null;
-
-        try {
+        try (SocketTimeoutObject ignored = startTimer(sock, timeout)) {
             OutputStream out = sock.getOutputStream();
 
             out.write(data);
@@ -1695,20 +1688,7 @@ protected void writeToSocket(Socket sock, TcpDiscoveryAbstractMessage msg, byte[
         catch (IOException e) {
             SSLException sslEx = checkSslException(sock, e);
 
-            err = sslEx == null ? e : sslEx;
-        }
-        finally {
-            boolean cancelled = obj.cancel();
-
-            if (cancelled)
-                removeTimeoutObject(obj);
-
-            // Throw original exception.
-            if (err != null)
-                throw err;
-
-            if (!cancelled)
-                throw new SocketTimeoutException("Write timed out (socket was concurrently closed).");
+            throw sslEx == null ? e : sslEx;
         }
     }
 
@@ -1744,41 +1724,20 @@ protected void startMessageProcess(TcpDiscoveryAbstractMessage msg) {
      * @throws IOException If IO failed or write timed out.
      * @throws IgniteCheckedException If marshalling failed.
      */
-    @SuppressWarnings("ThrowFromFinallyBlock")
     protected void writeToSocket(Socket sock,
         OutputStream out,
         TcpDiscoveryAbstractMessage msg,
         long timeout) throws IOException, IgniteCheckedException {
         assert sock != null;
         assert msg != null;
-        assert out != null;
 
-        SocketTimeoutObject obj = new SocketTimeoutObject(sock, U.currentTimeMillis() + timeout);
-
-        addTimeoutObject(obj);
-
-        IgniteCheckedException err = null;
-
-        try {
+        try (SocketTimeoutObject ignored = startTimer(sock, timeout)) {
             U.marshal(marshaller(), msg, out);
         }
         catch (IgniteCheckedException e) {
             SSLException sslEx = checkSslException(sock, e);
 
-            err = sslEx == null ? e : new IgniteCheckedException(sslEx);
-        }
-        finally {
-            boolean cancelled = obj.cancel();
-
-            if (cancelled)
-                removeTimeoutObject(obj);
-
-            // Throw original exception.
-            if (err != null)
-                throw err;
-
-            if (!cancelled)
-                throw new SocketTimeoutException("Write timed out (socket was concurrently closed).");
+            throw sslEx == null ? e : new IgniteCheckedException(sslEx);
         }
     }
 
@@ -1791,41 +1750,21 @@ protected void writeToSocket(Socket sock,
      * @param timeout Socket timeout.
      * @throws IOException If IO failed or write timed out.
      */
-    @SuppressWarnings("ThrowFromFinallyBlock")
     protected void writeToSocket(TcpDiscoveryAbstractMessage msg, Socket sock, int res, long timeout)
         throws IOException {
         assert sock != null;
 
-        SocketTimeoutObject obj = new SocketTimeoutObject(sock, U.currentTimeMillis() + timeout);
-
-        addTimeoutObject(obj);
-
-        OutputStream out = sock.getOutputStream();
-
-        IOException err = null;
+        try (SocketTimeoutObject ignored = startTimer(sock, timeout)) {
+            OutputStream out = sock.getOutputStream();
 
-        try {
             out.write(res);
 
             out.flush();
         }
         catch (IOException e) {
             SSLException sslEx = checkSslException(sock, e);
 
-            err = sslEx == null ? e : sslEx;
-        }
-        finally {
-            boolean cancelled = obj.cancel();
-
-            if (cancelled)
-                removeTimeoutObject(obj);
-
-            // Throw original exception.
-            if (err != null)
-                throw err;
-
-            if (!cancelled)
-                throw new SocketTimeoutException("Write timed out (socket was concurrently closed).");
+            throw (sslEx == null) ? e : sslEx;
         }
     }
 
@@ -2489,10 +2428,19 @@ protected Marshaller marshaller() {
         return S.toString(TcpDiscoverySpi.class, this);
     }
 
+    /** Starts a timer for a socket operation. */
+    private SocketTimeoutObject startTimer(Socket sock, long timeout) {
+        SocketTimeoutObject obj = new SocketTimeoutObject(sock, U.currentTimeMillis() + timeout);
+
+        addTimeoutObject(obj);
+
+        return obj;
+    }
+
     /**
      * Socket timeout object.
      */
-    private class SocketTimeoutObject implements IgniteSpiTimeoutObject {
+    private class SocketTimeoutObject implements IgniteSpiTimeoutObject, AutoCloseable {
         /** */
         private final IgniteUuid id = IgniteUuid.randomUuid();
 
@@ -2550,6 +2498,14 @@ boolean cancel() {
             return id;
         }
 
+        /** */
+        @Override public void close() throws SocketTimeoutException {
+            if (cancel())
+                removeTimeoutObject(this);
+            else
+                throw new SocketTimeoutException("Write timed out (socket was concurrently closed).");
+        }
+
         /** {@inheritDoc} */
         @Override public String toString() {
             return S.toString(SocketTimeoutObject.class, this);