quic: fix pkt_meta remove and abandon_enc_level

Author: akhinvasara-jumptrading

src/waltz/quic/fd_quic.c

@@ -1201,36 +1201,59 @@ fd_quic_conn_set_rx_max_data( fd_quic_conn_t * conn, ulong rx_max_data ) {
 /* packet processing */
 
 /* fd_quic_conn_free_pkt_meta frees all pkt_meta associated with
-   encryption levels less or equal to enc_level. Returns the number
-   of freed pkt_meta. */
+   enc_level. Returns the number of freed pkt_meta. */
 static ulong
 fd_quic_conn_free_pkt_meta( fd_quic_conn_t * conn,
                             uint             enc_level ) {
-  ulong                        freed   = 0UL;
   fd_quic_pkt_meta_tracker_t * tracker = &conn->pkt_meta_tracker;
   fd_quic_pkt_meta_t         * pool    = tracker->pool;
+  fd_quic_pkt_meta_ds_t      * sent = &tracker->sent_pkt_metas[enc_level];
 
-  for( uint j=0; j<=enc_level; ++j ) {
-    fd_quic_pkt_meta_ds_t * sent = &tracker->sent_pkt_metas[j];
-    fd_quic_pkt_meta_t    * prev = NULL;
-    for( fd_quic_pkt_meta_ds_fwd_iter_t iter = fd_quic_pkt_meta_ds_fwd_iter_init( sent, pool );
-                                                !fd_quic_pkt_meta_ds_fwd_iter_done( iter );
-                                                iter = fd_quic_pkt_meta_ds_fwd_iter_next( iter, pool ) ) {
-      fd_quic_pkt_meta_t * e = fd_quic_pkt_meta_ds_fwd_iter_ele( iter, pool );
-      if( FD_LIKELY( prev ) ) {
-        fd_quic_pkt_meta_pool_ele_release( pool, prev );
-      }
-      prev = e;
-    }
+  fd_quic_pkt_meta_t * prev = NULL;
+  for( fd_quic_pkt_meta_ds_fwd_iter_t iter = fd_quic_pkt_meta_ds_fwd_iter_init( sent, pool );
+                                             !fd_quic_pkt_meta_ds_fwd_iter_done( iter );
+                                             iter = fd_quic_pkt_meta_ds_fwd_iter_next( iter, pool ) ) {
+    fd_quic_pkt_meta_t * e = fd_quic_pkt_meta_ds_fwd_iter_ele( iter, pool );
     if( FD_LIKELY( prev ) ) {
       fd_quic_pkt_meta_pool_ele_release( pool, prev );
     }
+    prev = e;
+  }
+  if( FD_LIKELY( prev ) ) {
+    fd_quic_pkt_meta_pool_ele_release( pool, prev );
+  }
+
 
-    conn->used_pkt_meta -= fd_quic_pkt_meta_ds_ele_cnt( sent );
-    freed               += fd_quic_pkt_meta_ds_ele_cnt( sent );
-    fd_quic_pkt_meta_ds_clear( tracker, j );
+  /* Instead of paying log(n) to maintain the treap structure during ele_remove
+     on each iteration, we've returned all pkt_meta to the pool, but left them
+     in the treap. Then, we reinitialize the treap, in constant time deleting
+     all elements that were in the 'old treap'. This function now runs in linear
+     time, instead of O(n*lg(n)). */
+
+  ulong pre_ele_cnt = fd_quic_pkt_meta_ds_ele_cnt( sent );
+
+  fd_quic_pkt_meta_ds_clear( tracker, enc_level );
+  conn->used_pkt_meta -= pre_ele_cnt;
+
+  return pre_ele_cnt;
+}
+
+/* reclaim and free all pkt_meta for a given encryption level,
+   and return the number affected. */
+static ulong
+fd_quic_reclaim_pkt_meta_level( fd_quic_conn_t * conn, uint enc_level ) {
+  fd_quic_pkt_meta_tracker_t * tracker  = &conn->pkt_meta_tracker;
+  fd_quic_pkt_meta_t         * pool     = tracker->pool;
+  fd_quic_pkt_meta_ds_t      * sent     = &tracker->sent_pkt_metas[enc_level];
+
+  for( fd_quic_pkt_meta_ds_fwd_iter_t iter = fd_quic_pkt_meta_ds_fwd_iter_init( sent, pool );
+                                             !fd_quic_pkt_meta_ds_fwd_iter_done( iter );
+                                             iter = fd_quic_pkt_meta_ds_fwd_iter_next( iter, pool ) ) {
+    fd_quic_pkt_meta_t * e = fd_quic_pkt_meta_ds_fwd_iter_ele( iter, pool );
+    fd_quic_reclaim_pkt_meta( conn, e, enc_level );
   }
-  return freed;
+
+  return fd_quic_conn_free_pkt_meta( conn, enc_level );
 }
 
 
@@ -1246,23 +1269,14 @@ fd_quic_abandon_enc_level( fd_quic_conn_t * conn,
 
   fd_quic_ack_gen_abandon_enc_level( conn->ack_gen, enc_level );
 
-  fd_quic_pkt_meta_tracker_t * tracker = &conn->pkt_meta_tracker;
-  fd_quic_pkt_meta_t         * pool    = tracker->pool;
-
+  ulong freed = 0UL;
   for( uint j = 0; j <= enc_level; ++j ) {
     conn->keys_avail = fd_uint_clear_bit( conn->keys_avail, (int)j );
     /* treat all packets as ACKed (freeing handshake data, etc.) */
-    fd_quic_pkt_meta_ds_t * sent  =  &tracker->sent_pkt_metas[j];
-
-    for( fd_quic_pkt_meta_ds_fwd_iter_t iter = fd_quic_pkt_meta_ds_fwd_iter_init( sent, pool );
-                                               !fd_quic_pkt_meta_ds_fwd_iter_done( iter );
-                                               iter = fd_quic_pkt_meta_ds_fwd_iter_next( iter, pool ) ) {
-      fd_quic_pkt_meta_t * e = fd_quic_pkt_meta_ds_fwd_iter_ele( iter, pool );
-      fd_quic_reclaim_pkt_meta( conn, e, j );
-    }
+    freed += fd_quic_reclaim_pkt_meta_level( conn, j );
   }
 
-  return fd_quic_conn_free_pkt_meta( conn, enc_level );
+  return freed;
 }
 
 static void
@@ -3968,7 +3982,7 @@ fd_quic_conn_free( fd_quic_t *      quic,
        enqueued once */
 
   /* free pkt_meta */
-  fd_quic_conn_free_pkt_meta( conn, fd_quic_enc_level_appdata_id );
+  for( uint j=0; j<=fd_quic_enc_level_appdata_id; ++j ) fd_quic_conn_free_pkt_meta( conn, j );
 
   /* remove all stream ids from map, and free stream */
 
@@ -4427,7 +4441,7 @@ fd_quic_pkt_meta_retry( fd_quic_t *          quic,
 
     /* already moved to another enc_level */
     if( enc_level < peer_enc_level ) {
-      cnt_freed += fd_quic_abandon_enc_level( conn, peer_enc_level );
+      cnt_freed += fd_quic_reclaim_pkt_meta_level( conn, enc_level );
       continue;
     }
 
@@ -4529,10 +4543,9 @@ fd_quic_pkt_meta_retry( fd_quic_t *          quic,
     fd_quic_svc_prep_schedule_now( conn );
 
     /* free pkt_meta */
-    fd_quic_pkt_meta_remove_range( &tracker->sent_pkt_metas[enc_level],
-                                    pool,
-                                    pkt_meta->key.pkt_num,
-                                    pkt_meta->key.pkt_num );
+    fd_quic_pkt_meta_remove( &tracker->sent_pkt_metas[enc_level],
+                             pool,
+                             pkt_meta );
 
     conn->used_pkt_meta -= 1;
     cnt_freed++;
@@ -4776,6 +4789,7 @@ fd_quic_reclaim_pkt_meta( fd_quic_conn_t *     conn,
       break;
   }
 }
+
 /* process lost packets
  * These packets will be declared lost and relevant data potentially resent */
 void

src/waltz/quic/fd_quic_pkt_meta.c

@@ -59,6 +59,14 @@ fd_quic_pkt_meta_remove_range( fd_quic_pkt_meta_ds_t * ds,
   return cnt_removed;
 }
 
+void
+fd_quic_pkt_meta_remove( fd_quic_pkt_meta_ds_t * ds,
+                         fd_quic_pkt_meta_t    * pool,
+                         fd_quic_pkt_meta_t    * pkt_meta ) {
+  fd_quic_pkt_meta_treap_ele_remove( ds, pkt_meta, pool );
+  fd_quic_pkt_meta_pool_ele_release( pool, pkt_meta );
+}
+
 fd_quic_pkt_meta_t *
 fd_quic_pkt_meta_min( fd_quic_pkt_meta_ds_t * ds,
                       fd_quic_pkt_meta_t    * pool ) {

src/waltz/quic/fd_quic_pkt_meta.h

@@ -258,6 +258,16 @@ fd_quic_pkt_meta_remove_range( fd_quic_pkt_meta_ds_t * ds,
                                ulong                   pkt_number_lo,
                                ulong                   pkt_number_hi );
 
+/* fd_quic_pkt_meta_remove removes a single pkt_meta from the ds and returns it to the pool.
+@arguments:
+- ds: pointer to the ds
+- pool: pointer to the backing pool
+- pkt_meta: pointer to the pkt_meta to remove. Assumed non-null */
+void
+fd_quic_pkt_meta_remove( fd_quic_pkt_meta_ds_t * ds,
+                       fd_quic_pkt_meta_t    * pool,
+                       fd_quic_pkt_meta_t    * pkt_meta );
+
 /* fd_quic_pkt_meta_min returns pointer to pkt_meta with smallest pkt_number in the ds
   @arguments:
   - ds: pointer to the ds
@@ -268,10 +278,10 @@ fd_quic_pkt_meta_t *
 fd_quic_pkt_meta_min( fd_quic_pkt_meta_ds_t * ds,
                       fd_quic_pkt_meta_t    * pool );
 
-/* fd_quic_pkt_meta_ds_clear clears all pkt_meta tracking for a given encoding level
+/* fd_quic_pkt_meta_ds_clear clears all pkt_meta tracking for a given encryption level
   @arguments:
   - tracker: pointer to the pkt_meta tracker
-  - enc_level: encoding level to clear */
+  - enc_level: encryption level to clear */
 void
 fd_quic_pkt_meta_ds_clear( fd_quic_pkt_meta_tracker_t * tracker,
                            uint                         enc_level );