Linpool state save and restore refactoring

7 files changed, 47 insertions, 64 deletions

diff --git a/filter/tree_test.c b/filter/tree_test.c
index 655205c2..8accdf58 100644
--- a/filter/tree_test.c
+++ b/filter/tree_test.c
@@ -182,8 +182,7 @@ t_balancing_random(void)
     uint i;
     for(i = 0; i < 10; i++)
     {
-      struct lp_state lps;
-      lp_save(tmp_linpool, &lps);
+      struct lp_state *lps = lp_save(tmp_linpool);
 
       struct f_tree *random_degenerated_tree = get_random_degenerated_left_tree(nodes_count);
       show_tree(random_degenerated_tree);
@@ -195,7 +194,7 @@ t_balancing_random(void)
 
       bt_assert(same_tree(balanced_tree_from_random, expected_balanced_tree));
 
-      lp_restore(tmp_linpool, &lps);
+      lp_restore(tmp_linpool, lps);
     }
 
     tmp_flush();
diff --git a/lib/mempool.c b/lib/mempool.c
index 578750c1..da1dc857 100644
--- a/lib/mempool.c
+++ b/lib/mempool.c
@@ -39,6 +39,7 @@ struct linpool {
   byte *ptr, *end;
   struct lp_chunk *first, *current;		/* Normal (reusable) chunks */
   struct lp_chunk *first_large;			/* Large chunks */
+  struct lp_state *initial;			/* Initial state to restore to */
   uint total, total_large;
 };
 
@@ -66,7 +67,9 @@ static struct resclass lp_class = {
 linpool
 *lp_new(pool *p)
 {
-  return ralloc(p, &lp_class);
+  linpool *m = ralloc(p, &lp_class);
+  m->initial = lp_save(m);
+  return m;
 }
 
 /**
@@ -185,27 +188,8 @@ lp_allocz(linpool *m, uint size)
 void
 lp_flush(linpool *m)
 {
-  ASSERT_DIE(DG_IS_LOCKED(resource_parent(&m->r)->domain));
-  struct lp_chunk *c;
-
-  /* Move ptr to the first chunk and free all other chunks */
-  m->current = c = m->first;
-  m->ptr = c ? c->data : NULL;
-  m->end = c ? c->data + LP_DATA_SIZE : NULL;
-
-  while (c && c->next)
-  {
-    struct lp_chunk *d = c->next;
-    c->next = d->next;
-    free_page(d);
-  }
-
-  while (c = m->first_large)
-    {
-      m->first_large = c->next;
-      xfree(c);
-    }
-  m->total_large = 0;
+  lp_restore(m, m->initial);
+  m->initial = lp_save(m);
 }
 
 /**
@@ -216,14 +200,19 @@ lp_flush(linpool *m)
  * This function saves the state of a linear memory pool. Saved state can be
  * used later to restore the pool (to free memory allocated since).
  */
-void
-lp_save(linpool *m, lp_state *p)
+struct lp_state *
+lp_save(linpool *m)
 {
   ASSERT_DIE(DG_IS_LOCKED(resource_parent(&m->r)->domain));
-  p->current = m->current;
-  p->large = m->first_large;
-  p->total_large = m->total_large;
-  p->ptr = m->ptr;
+  struct lp_state *p = lp_alloc(m, sizeof(struct lp_state));
+  ASSERT_DIE(m->current);
+  *p = (struct lp_state) {
+    .current = m->current,
+    .large = m->first_large,
+    .total_large = m->total_large,
+  };
+
+  return p;
 }
 
 /**
@@ -243,9 +232,10 @@ lp_restore(linpool *m, lp_state *p)
   ASSERT_DIE(DG_IS_LOCKED(resource_parent(&m->r)->domain));
 
   /* Move ptr to the saved pos and free all newer large chunks */
-  m->current = c = p->current ?: m->first;
-  m->ptr = p->ptr ?: (c ? c->data : NULL);
-  m->end = c ? (c->data + LP_DATA_SIZE) : NULL;
+  ASSERT_DIE(p->current);
+  m->current = c = p->current;
+  m->ptr = (byte *) p;
+  m->end = c->data + LP_DATA_SIZE;
   m->total_large = p->total_large;
 
   while ((c = m->first_large) && (c != p->large))
@@ -254,7 +244,7 @@ lp_restore(linpool *m, lp_state *p)
       xfree(c);
     }
 
-  while (m->current && (c = m->current->next))
+  while (c = m->current->next)
     {
       m->current->next = c->next;
       free_page(c);
diff --git a/lib/resource.h b/lib/resource.h
index 06af4289..abcf46fa 100644
--- a/lib/resource.h
+++ b/lib/resource.h
@@ -92,7 +92,6 @@ typedef struct linpool linpool;
 
 typedef struct lp_state {
   void *current, *large;
-  byte *ptr;
   uint total_large;
 } lp_state;
 
@@ -101,9 +100,12 @@ void *lp_alloc(linpool *, unsigned size);	/* Aligned */
 void *lp_allocu(linpool *, unsigned size);	/* Unaligned */
 void *lp_allocz(linpool *, unsigned size);	/* With clear */
 void lp_flush(linpool *);			/* Free everything, but leave linpool */
-void lp_save(linpool *m, lp_state *p);		/* Save state */
+lp_state *lp_save(linpool *m);			/* Save state */
 void lp_restore(linpool *m, lp_state *p);	/* Restore state */
 
+#define LP_SAVED(m)	for (struct lp_state *_lp_state = lp_save(m); _lp_state; lp_restore(m, _lp_state), _lp_state = NULL)
+#define TMP_SAVED	LP_SAVED(tmp_linpool)
+
 struct tmp_resources {
   pool *pool, *parent;
   linpool *lp;
diff --git a/nest/rt-table.c b/nest/rt-table.c
index f6d2ca0f..132057fe 100644
--- a/nest/rt-table.c
+++ b/nest/rt-table.c
@@ -3946,10 +3946,8 @@ rt_next_hop_update(struct rtable_private *tab)
 	  birdloop_flag(tab->loop, RTF_NHU);
 	  return;
 	}
-      lp_state lps;
-      lp_save(tmp_linpool, &lps);
-      max_feed -= rt_next_hop_update_net(tab, n);
-      lp_restore(tmp_linpool, &lps);
+      TMP_SAVED
+	max_feed -= rt_next_hop_update_net(tab, n);
     }
   FIB_ITERATE_END;
 
diff --git a/proto/bgp/attrs.c b/proto/bgp/attrs.c
index 82971c01..88639442 100644
--- a/proto/bgp/attrs.c
+++ b/proto/bgp/attrs.c
@@ -2691,8 +2691,7 @@ bgp_rte_modify_stale(struct rt_export_request *req, const net_addr *n,
       continue;
 
     /* Store the tmp_linpool state to aggresively save memory */
-    struct lp_state tmpp;
-    lp_save(tmp_linpool, &tmpp);
+    struct lp_state *tmpp = lp_save(tmp_linpool);
 
     /* Mark the route as LLGR */
     rte e0 = *r;
@@ -2705,7 +2704,7 @@ bgp_rte_modify_stale(struct rt_export_request *req, const net_addr *n,
     irh->stale_set++;
 
     /* Restore the memory state */
-    lp_restore(tmp_linpool, &tmpp);
+    lp_restore(tmp_linpool, tmpp);
   }
 
   rpe_mark_seen_all(req->hook, first, last, NULL);
diff --git a/proto/bgp/packets.c b/proto/bgp/packets.c
index 978a1891..b9fc38ce 100644
--- a/proto/bgp/packets.c
+++ b/proto/bgp/packets.c
@@ -2387,11 +2387,13 @@ bgp_create_update(struct bgp_channel *c, byte *buf)
   struct bgp_bucket *buck;
   byte *end = buf + (bgp_max_packet_length(p->conn) - BGP_HEADER_LENGTH);
   byte *res = NULL;
+  struct lp_state *tmpp = NULL;
 
-again: ;
+again:
+  if (tmpp)
+    lp_restore(tmp_linpool, tmpp);
 
-  struct lp_state tmpp;
-  lp_save(tmp_linpool, &tmpp);
+  tmpp = lp_save(tmp_linpool);
 
   /* Initialize write state */
   struct bgp_write_state s = {
@@ -2421,10 +2423,7 @@ again: ;
 
     /* Cleanup empty buckets */
     if (bgp_done_bucket(c, buck))
-    {
-      lp_restore(tmp_linpool, &tmpp);
       goto again;
-    }
 
     res = !s.mp_reach ?
       bgp_create_ip_reach(&s, buck, buf, end):
@@ -2433,21 +2432,19 @@ again: ;
     bgp_done_bucket(c, buck);
 
     if (!res)
-    {
-      lp_restore(tmp_linpool, &tmpp);
       goto again;
-    }
 
     goto done;
   }
 
   /* No more prefixes to send */
+  lp_restore(tmp_linpool, tmpp);
   return NULL;
 
 done:
   BGP_TRACE_RL(&rl_snd_update, D_PACKETS, "Sending UPDATE");
   p->stats.tx_updates++;
-  lp_restore(tmp_linpool, &tmpp);
+  lp_restore(tmp_linpool, tmpp);
 
   return res;
 }
@@ -2574,8 +2571,7 @@ bgp_rx_update(struct bgp_conn *conn, byte *pkt, uint len)
 
   bgp_start_timer(p, conn->hold_timer, conn->hold_time);
 
-  struct lp_state tmpp;
-  lp_save(tmp_linpool, &tmpp);
+  struct lp_state *tmpp = lp_save(tmp_linpool);
 
   /* Initialize parse state */
   struct bgp_parse_state s = {
@@ -2593,7 +2589,10 @@ bgp_rx_update(struct bgp_conn *conn, byte *pkt, uint len)
 
   /* Check minimal length */
   if (len < 23)
-  { bgp_error(conn, 1, 2, pkt+16, 2); return; }
+  {
+    bgp_error(conn, 1, 2, pkt+16, 2);
+    goto done;
+  }
 
   /* Skip fixed header */
   uint pos = 19;
@@ -2658,7 +2657,7 @@ bgp_rx_update(struct bgp_conn *conn, byte *pkt, uint len)
 
 done:
   rta_free(s.cached_ea);
-  lp_restore(tmp_linpool, &tmpp);
+  lp_restore(tmp_linpool, tmpp);
   return;
 }
 
diff --git a/proto/static/static.c b/proto/static/static.c
index 6bae827b..74603a93 100644
--- a/proto/static/static.c
+++ b/proto/static/static.c
@@ -494,12 +494,8 @@ static_start(struct proto *P)
   proto_notify_state(P, PS_UP);
 
   WALK_LIST(r, cf->routes)
-  {
-    struct lp_state lps;
-    lp_save(tmp_linpool, &lps);
-    static_add_rte(p, r);
-    lp_restore(tmp_linpool, &lps);
-  }
+    TMP_SAVED
+      static_add_rte(p, r);
 
   return PS_UP;
 }