From 9809f958fa7b7204e07e9b798dbdaab37fe75175 Mon Sep 17 00:00:00 2001
From: Omri Mor <omrimor2@illinois.edu>
Date: Mon, 8 Jun 2020 19:05:45 -0500
Subject: [PATCH] pool: prevent div by 0 error with npools == 1

If pool->npools == 1, there is no valid steal target: return self pid.
Also includes formatting fixes and changes discussed in #20.
---
 include/lc/pool.h | 49 +++++++++++++++++++++++++++++------------------
 1 file changed, 30 insertions(+), 19 deletions(-)

diff --git a/include/lc/pool.h b/include/lc/pool.h
index 6e4560ea..27767be2 100644
--- a/include/lc/pool.h
+++ b/include/lc/pool.h
@@ -26,7 +26,7 @@ struct dequeue;
 extern "C" {
 #endif
 
-LC_INLINE int lc_worker_id()
+LC_INLINE int lc_worker_id(void)
 {
   if (unlikely(lcg_core_id == -1)) {
     lcg_core_id = sched_getcpu();
@@ -90,47 +90,58 @@ LC_INLINE int32_t lc_pool_get_local(struct lc_pool* pool)
   return pid;
 }
 
-LC_INLINE int32_t lc_pool_get_steal(struct lc_pool* pool)
+LC_INLINE int32_t lc_pool_get_steal_id(struct lc_pool* pool, int32_t pid)
 {
-    int32_t pid = lc_pool_get_local(pool);
-    int32_t npools = pool->npools;
-    int32_t r = rand() % (npools - 1);
-    return (r + pid + 1) % npools;
+  int32_t npools = pool->npools;
+  if (npools == 1)
+    return pid; /* 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)
+LC_INLINE void* lc_pool_steal(struct lc_pool* pool, int32_t pid)
 {
-    void* elm = NULL;
-    int32_t steal = lc_pool_get_steal(pool);
-    if (likely(pool->lpools[steal] != NULL))
-      elm = dq_pop_bot(pool->lpools[steal]);
-    return elm;
+  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]);
+  return elm;
 }
 
-LC_INLINE void lc_pool_put(struct lc_pool* pool, void* elm) {
+LC_INLINE void lc_pool_put(struct lc_pool* pool, void* elm)
+{
   int32_t pid = lc_pool_get_local(pool);
   struct dequeue* lpool = pool->lpools[pid];
   dq_push_top(lpool, elm);
 }
 
-LC_INLINE void lc_pool_put_to(struct lc_pool* pool, void* elm, int32_t pid) {
+LC_INLINE void lc_pool_put_to(struct lc_pool* pool, void* elm, int32_t pid)
+{
   struct dequeue* lpool = pool->lpools[pid];
   dq_push_top(lpool, elm);
 }
 
-LC_INLINE void* lc_pool_get_nb(struct lc_pool* pool) {
+LC_INLINE void* lc_pool_get_nb(struct lc_pool* pool)
+{
   int32_t pid = lc_pool_get_local(pool);
   struct dequeue* lpool = pool->lpools[pid];
   void* elm = dq_pop_top(lpool);
   if (elm == NULL)
-    elm = lc_pool_steal(pool);
+    elm = lc_pool_steal(pool, pid);
   return elm;
 }
 
-LC_INLINE void* lc_pool_get(struct lc_pool* pool) {
+LC_INLINE void* lc_pool_get(struct lc_pool* pool)
+{
+  int32_t pid = lc_pool_get_local(pool);
+  struct dequeue* lpool = pool->lpools[pid];
   void* elm = NULL;
-  while (elm == NULL)
-    elm = lc_pool_get_nb(pool);
+  while (elm == NULL) {
+    /* 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);
+  }
   return elm;
 }