Merge pull request #7018 from yossigo/fix-accept-issues

Fix issues with failed/rejected accepts.

5 files changed, 58 insertions, 40 deletions

diff --git a/src/cluster.c b/src/cluster.c
index 72755823a..a2e9ff5b6 100644
--- a/src/cluster.c
+++ b/src/cluster.c
@@ -681,9 +681,10 @@ void clusterAcceptHandler(aeEventLoop *el, int fd, void *privdata, int mask) {
          * or schedule it for later depending on connection implementation.
          */
         if (connAccept(conn, clusterConnAcceptHandler) == C_ERR) {
-            serverLog(LL_VERBOSE,
-                    "Error accepting cluster node connection: %s",
-                    connGetLastError(conn));
+            if (connGetState(conn) == CONN_STATE_ERROR)
+                serverLog(LL_VERBOSE,
+                        "Error accepting cluster node connection: %s",
+                        connGetLastError(conn));
             connClose(conn);
             return;
         }
diff --git a/src/connection.c b/src/connection.c
index 58d86c31b..2015c9195 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -152,7 +152,7 @@ static void connSocketClose(connection *conn) {
     /* If called from within a handler, schedule the close but
      * keep the connection until the handler returns.
      */
-    if (conn->flags & CONN_FLAG_IN_HANDLER) {
+    if (connHasRefs(conn)) {
         conn->flags |= CONN_FLAG_CLOSE_SCHEDULED;
         return;
     }
@@ -183,10 +183,16 @@ static int connSocketRead(connection *conn, void *buf, size_t buf_len) {
 }
 
 static int connSocketAccept(connection *conn, ConnectionCallbackFunc accept_handler) {
+    int ret = C_OK;
+
     if (conn->state != CONN_STATE_ACCEPTING) return C_ERR;
     conn->state = CONN_STATE_CONNECTED;
-    if (!callHandler(conn, accept_handler)) return C_ERR;
-    return C_OK;
+
+    connIncrRefs(conn);
+    if (!callHandler(conn, accept_handler)) ret = C_ERR;
+    connDecrRefs(conn);
+
+    return ret;
 }
 
 /* Register a write handler, to be called when the connection is writable.
diff --git a/src/connection.h b/src/connection.h
index 97622f8d6..db09dfd83 100644
--- a/src/connection.h
+++ b/src/connection.h
@@ -45,9 +45,8 @@ typedef enum {
     CONN_STATE_ERROR
 } ConnectionState;
 
-#define CONN_FLAG_IN_HANDLER        (1<<0)      /* A handler execution is in progress */
-#define CONN_FLAG_CLOSE_SCHEDULED   (1<<1)      /* Closed scheduled by a handler */
-#define CONN_FLAG_WRITE_BARRIER     (1<<2)      /* Write barrier requested */
+#define CONN_FLAG_CLOSE_SCHEDULED   (1<<0)      /* Closed scheduled by a handler */
+#define CONN_FLAG_WRITE_BARRIER     (1<<1)      /* Write barrier requested */
 
 typedef void (*ConnectionCallbackFunc)(struct connection *conn);
 
@@ -70,7 +69,8 @@ typedef struct ConnectionType {
 struct connection {
     ConnectionType *type;
     ConnectionState state;
-    int flags;
+    short int flags;
+    short int refs;
     int last_errno;
     void *private_data;
     ConnectionCallbackFunc conn_handler;
@@ -88,6 +88,13 @@ struct connection {
  * connAccept() may directly call accept_handler(), or return and call it
  * at a later time. This behavior is a bit awkward but aims to reduce the need
  * to wait for the next event loop, if no additional handshake is required.
+ *
+ * IMPORTANT: accept_handler may decide to close the connection, calling connClose().
+ * To make this safe, the connection is only marked with CONN_FLAG_CLOSE_SCHEDULED
+ * in this case, and connAccept() returns with an error.
+ *
+ * connAccept() callers must always check the return value and on error (C_ERR)
+ * a connClose() must be called.
  */
 
 static inline int connAccept(connection *conn, ConnectionCallbackFunc accept_handler) {
diff --git a/src/connhelpers.h b/src/connhelpers.h
index f237c9b1d..86250d09e 100644
--- a/src/connhelpers.h
+++ b/src/connhelpers.h
@@ -37,46 +37,49 @@
  * implementations (currently sockets in connection.c and TLS in tls.c).
  *
  * Currently helpers implement the mechanisms for invoking connection
- * handlers, tracking in-handler states and dealing with deferred
- * destruction (if invoked by a handler).
+ * handlers and tracking connection references, to allow safe destruction
+ * of connections from within a handler.
  */
 
-/* Called whenever a handler is invoked on a connection and sets the
- * CONN_FLAG_IN_HANDLER flag to indicate we're in a handler context.
+/* Incremenet connection references.
  *
- * An attempt to close a connection while CONN_FLAG_IN_HANDLER is
- * set will result with deferred close, i.e. setting the CONN_FLAG_CLOSE_SCHEDULED
- * instead of destructing it.
+ * Inside a connection handler, we guarantee refs >= 1 so it is always
+ * safe to connClose().
+ *
+ * In other cases where we don't want to prematurely lose the connection,
+ * it can go beyond 1 as well; currently it is only done by connAccept().
  */
-static inline void enterHandler(connection *conn) {
-    conn->flags |= CONN_FLAG_IN_HANDLER;
+static inline void connIncrRefs(connection *conn) {
+    conn->refs++;
 }
 
-/* Called whenever a handler returns. This unsets the CONN_FLAG_IN_HANDLER
- * flag and performs actual close/destruction if a deferred close was
- * scheduled by the handler.
+/* Decrement connection references.
+ *
+ * Note that this is not intended to provide any automatic free logic!
+ * callHandler() takes care of that for the common flows, and anywhere an
+ * explicit connIncrRefs() is used, the caller is expected to take care of
+ * that.
  */
-static inline int exitHandler(connection *conn) {
-    conn->flags &= ~CONN_FLAG_IN_HANDLER;
-    if (conn->flags & CONN_FLAG_CLOSE_SCHEDULED) {
-        connClose(conn);
-        return 0;
-    }
-    return 1;
+
+static inline void connDecrRefs(connection *conn) {
+    conn->refs--;
+}
+
+static inline int connHasRefs(connection *conn) {
+    return conn->refs;
 }
 
 /* Helper for connection implementations to call handlers:
- * 1. Mark the handler in use.
+ * 1. Increment refs to protect the connection.
  * 2. Execute the handler (if set).
- * 3. Mark the handler as NOT in use and perform deferred close if was
- * requested by the handler at any time.
+ * 3. Decrement refs and perform deferred close, if refs==0.
  */
 static inline int callHandler(connection *conn, ConnectionCallbackFunc handler) {
-    conn->flags |= CONN_FLAG_IN_HANDLER;
+    connIncrRefs(conn);
     if (handler) handler(conn);
-    conn->flags &= ~CONN_FLAG_IN_HANDLER;
+    connDecrRefs(conn);
     if (conn->flags & CONN_FLAG_CLOSE_SCHEDULED) {
-        connClose(conn);
+        if (!connHasRefs(conn)) connClose(conn);
         return 0;
     }
     return 1;
diff --git a/src/networking.c b/src/networking.c
index 69d59a59b..a550e4040 100644
--- a/src/networking.c
+++ b/src/networking.c
@@ -786,7 +786,7 @@ void clientAcceptHandler(connection *conn) {
         serverLog(LL_WARNING,
                 "Error accepting a client connection: %s",
                 connGetLastError(conn));
-        freeClient(c);
+        freeClientAsync(c);
         return;
     }
 
@@ -828,7 +828,7 @@ void clientAcceptHandler(connection *conn) {
                 /* Nothing to do, Just to avoid the warning... */
             }
             server.stat_rejected_conn++;
-            freeClient(c);
+            freeClientAsync(c);
             return;
         }
     }
@@ -887,9 +887,10 @@ static void acceptCommonHandler(connection *conn, int flags, char *ip) {
      */
     if (connAccept(conn, clientAcceptHandler) == C_ERR) {
         char conninfo[100];
-        serverLog(LL_WARNING,
-                "Error accepting a client connection: %s (conn: %s)",
-                connGetLastError(conn), connGetInfo(conn, conninfo, sizeof(conninfo)));
+        if (connGetState(conn) == CONN_STATE_ERROR)
+            serverLog(LL_WARNING,
+                    "Error accepting a client connection: %s (conn: %s)",
+                    connGetLastError(conn), connGetInfo(conn, conninfo, sizeof(conninfo)));
         freeClient(connGetPrivateData(conn));
         return;
     }