quic: dynamically disable keep_alive

Author: akhinvasara-jumptrading

src/waltz/quic/fd_quic.c

@@ -2957,7 +2957,7 @@ fd_quic_svc_poll( fd_quic_t *      quic,
         fd_quic_set_conn_state( conn, FD_QUIC_CONN_STATE_DEAD );
         quic->metrics.conn_timeout_cnt++;
       }
-    } else if( quic->config.keep_alive ) {
+    } else if( quic->config.keep_alive & !!(conn->let_die_ticks > now) ) {
       /* send PING */
       if( !( conn->flags & FD_QUIC_CONN_FLAGS_PING ) ) {
         conn->flags         |= FD_QUIC_CONN_FLAGS_PING;
@@ -4446,6 +4446,7 @@ fd_quic_conn_create( fd_quic_t *               quic,
   /* idle timeout */
   conn->idle_timeout_ticks  = config->idle_timeout;
   conn->last_activity       = state->now;
+  conn->let_die_ticks       = ULONG_MAX;
 
   /* update metrics */
   quic->metrics.conn_active_cnt++;
@@ -5645,3 +5646,10 @@ fd_quic_conn_close( fd_quic_conn_t * conn,
   /* set connection to be serviced ASAP */
   fd_quic_svc_schedule1( conn, FD_QUIC_SVC_INSTANT );
 }
+
+void
+fd_quic_conn_let_die( fd_quic_conn_t * conn,
+                      ulong            keep_alive_duration_ticks ) {
+  ulong const now = fd_quic_get_state( conn->quic )->now;
+  conn->let_die_ticks = fd_ulong_sat_add( now, keep_alive_duration_ticks );
+}

src/waltz/quic/fd_quic.h

@@ -555,6 +555,23 @@ FD_QUIC_API void
 fd_quic_conn_close( fd_quic_conn_t * conn,
                     uint             reason );
 
+/* fd_quic_conn_let_die stops keeping a conn alive after
+   'keep_alive_duration_ticks'. No-op if keep-alive is not configured.
+   Safe to call on a connection in any state.
+
+   If called multiple times on the same connection, only the latest
+   call will stay in effect. However, it may not take effect if we
+   already skipped a keep-alive due to a previous call. 'Undoing' a
+   previous call can be done by passing ULONG_MAX.
+
+   This function does NOT guarantee that the connection will be closed
+   immediately after the given duration. Rather, it just disables keep-alive
+   behavior after the given duration. */
+
+FD_QUIC_API void
+fd_quic_conn_let_die( fd_quic_conn_t * conn,
+                      ulong            keep_alive_duration_ticks );
+
 /* Service API ********************************************************/
 
 /* fd_quic_get_next_wakeup returns the next requested service time.

src/waltz/quic/fd_quic_conn.h

@@ -212,6 +212,7 @@ struct fd_quic_conn {
   ulong                idle_timeout_ticks;
   ulong                last_activity;
   ulong                last_ack;
+  ulong                let_die_ticks; /* stop keep-alive after this time */
 
   /* round trip time related members */
   fd_rtt_estimate_t rtt[1];

src/waltz/quic/tests/test_quic_keep_alive.c

@@ -59,17 +59,24 @@ test_quic_keep_alive( fd_quic_t * client_quic, fd_quic_t * server_quic, int keep
   ulong const idle_timeout = client_quic->config.idle_timeout;
   ulong const timestep     = idle_timeout>>3;
 
-  for( int i=0; i<10; ++i ) {
-    now+=timestep;
-    fd_quic_service( client_quic );
-    fd_quic_service( server_quic );
-  }
-  if( keep_alive ) {
-    FD_TEST( client_conn->state == FD_QUIC_CONN_STATE_ACTIVE );
-  } else {
-    FD_TEST( client_conn->state == FD_QUIC_CONN_STATE_DEAD ||
-            client_conn->state == FD_QUIC_CONN_STATE_INVALID );
+
+  for( int let_die=0; let_die<2; ++let_die ) {
+    for( int i=0; i<10; ++i ) {
+      now+=timestep;
+      fd_quic_service( client_quic );
+      fd_quic_service( server_quic );
+    }
+
+    if( keep_alive & !let_die ) {
+      FD_TEST( client_conn->state == FD_QUIC_CONN_STATE_ACTIVE );
+    } else {
+      FD_TEST( client_conn->state == FD_QUIC_CONN_STATE_DEAD ||
+              client_conn->state == FD_QUIC_CONN_STATE_INVALID );
+    }
+
+    fd_quic_conn_let_die( client_conn, timestep );
   }
+
 }