From 7ee652c980ee9132b585b24383176c6b262980be Mon Sep 17 00:00:00 2001
From: Omri Mor <omrimor2@illinois.edu>
Date: Wed, 10 Jun 2020 16:56:17 -0500
Subject: [PATCH] pool: change steal API for better design

---
 include/lc/pool.h | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/include/lc/pool.h b/include/lc/pool.h
index 45886cf0..04eccda7 100644
--- a/include/lc/pool.h
+++ b/include/lc/pool.h
@@ -90,29 +90,29 @@ LC_INLINE int32_t lc_pool_get_local_id(struct lc_pool* pool)
   return pid;
 }
 
-LC_INLINE int32_t lc_pool_get_steal_id(struct lc_pool* pool, int32_t pid)
+LC_INLINE int32_t lc_pool_get_steal_id(int32_t npools, int32_t pid)
 {
-  int32_t npools = pool->npools;
   if (npools == 1)
-    return pid; /* if only one pool, no one else to steal from */
+    return -1; /* if only one pool, no one else to steal from */
   int32_t r = rand() % (npools - 1);
   return (r + pid + 1) % npools;
 }
 
-LC_INLINE void* lc_pool_steal(struct lc_pool* pool, int32_t pid)
+LC_INLINE void* lc_pool_steal_from(struct lc_pool* pool, int32_t pid)
 {
   void* elm = NULL;
-  int32_t target = lc_pool_get_steal_id(pool, pid);
-  if (target != pid && likely(pool->lpools[target] != NULL))
-    elm = dq_pop_bot(pool->lpools[target]);
+  if (likely(pool->lpools[pid] != NULL))
+    elm = dq_pop_bot(pool->lpools[pid]);
   return elm;
 }
 
-LC_INLINE void lc_pool_put(struct lc_pool* pool, void* elm)
+LC_INLINE void* lc_pool_steal(struct lc_pool* pool, int32_t pid)
 {
-  int32_t pid = lc_pool_get_local_id(pool);
-  struct dequeue* lpool = pool->lpools[pid];
-  dq_push_top(lpool, elm);
+  void* elm = NULL;
+  int32_t target = lc_pool_get_steal_id(pool->npools, pid);
+  if (target != -1)
+    elm = lc_pool_steal_from(pool, pid);
+  return elm;
 }
 
 LC_INLINE void lc_pool_put_to(struct lc_pool* pool, void* elm, int32_t pid)
@@ -121,6 +121,12 @@ LC_INLINE void lc_pool_put_to(struct lc_pool* pool, void* elm, int32_t pid)
   dq_push_top(lpool, elm);
 }
 
+LC_INLINE void lc_pool_put(struct lc_pool* pool, void* elm)
+{
+  int32_t pid = lc_pool_get_local_id(pool);
+  lc_pool_put_to(pool, elm, pid);
+}
+
 LC_INLINE void* lc_pool_get_nb(struct lc_pool* pool)
 {
   int32_t pid = lc_pool_get_local_id(pool);
@@ -136,12 +142,12 @@ LC_INLINE void* lc_pool_get(struct lc_pool* pool)
   int32_t pid = lc_pool_get_local_id(pool);
   struct dequeue* lpool = pool->lpools[pid];
   void* elm = NULL;
-  while (elm == NULL) {
+  do {
     /* must try self every iteration since we never steal from self */
     elm = dq_pop_top(lpool);
     if (elm == NULL)
       elm = lc_pool_steal(pool, pid);
-  }
+  } while (elm == NULL);
   return elm;
 }