2 files changed, 10 insertions, 14 deletions

diff --git a/src/vnet/util/throttle.c b/src/vnet/util/throttle.c
index 0985b4a81a3..8b8e030bf53 100644
--- a/src/vnet/util/throttle.c
+++ b/src/vnet/util/throttle.c
@@ -16,17 +16,18 @@
 #include <vnet/util/throttle.h>
 
 void
-throttle_init (throttle_t * t, u32 n_threads, f64 time)
+throttle_init (throttle_t *t, u32 n_threads, u32 buckets, f64 time)
 {
   u32 i;
 
   t->time = time;
+  t->buckets = 1 << max_log2 (buckets);
   vec_validate (t->bitmaps, n_threads);
   vec_validate (t->seeds, n_threads);
   vec_validate (t->last_seed_change_time, n_threads);
 
   for (i = 0; i < n_threads; i++)
-    vec_validate (t->bitmaps[i], (THROTTLE_BITS / BITS (uword)) - 1);
+    clib_bitmap_alloc (t->bitmaps[i], t->buckets);
 }
 
 /*
diff --git a/src/vnet/util/throttle.h b/src/vnet/util/throttle.h
index 38ace280131..53435c4a359 100644
--- a/src/vnet/util/throttle.h
+++ b/src/vnet/util/throttle.h
@@ -31,11 +31,13 @@ typedef struct throttle_t_
   uword **bitmaps;
   u64 *seeds;
   f64 *last_seed_change_time;
+  u32 buckets;
 } throttle_t;
 
 #define THROTTLE_BITS	(512)
 
-extern void throttle_init (throttle_t * t, u32 n_threads, f64 time);
+extern void throttle_init (throttle_t *t, u32 n_threads, u32 buckets,
+			   f64 time);
 
 always_inline u64
 throttle_seed (throttle_t * t, u32 thread_index, f64 time_now)
@@ -43,7 +45,7 @@ throttle_seed (throttle_t * t, u32 thread_index, f64 time_now)
   if (time_now - t->last_seed_change_time[thread_index] > t->time)
     {
       (void) random_u64 (&t->seeds[thread_index]);
-      clib_memset (t->bitmaps[thread_index], 0, THROTTLE_BITS / BITS (u8));
+      clib_bitmap_zero (t->bitmaps[thread_index]);
 
       t->last_seed_change_time[thread_index] = time_now;
     }
@@ -53,21 +55,14 @@ throttle_seed (throttle_t * t, u32 thread_index, f64 time_now)
 always_inline int
 throttle_check (throttle_t * t, u32 thread_index, u64 hash, u64 seed)
 {
-  int drop;
-  uword m;
-  u32 w;
+  ASSERT (is_pow2 (t->buckets));
 
   hash = clib_xxhash (hash ^ seed);
 
   /* Select bit number */
-  hash &= THROTTLE_BITS - 1;
-  w = hash / BITS (uword);
-  m = (uword) 1 << (hash % BITS (uword));
+  hash &= t->buckets - 1;
 
-  drop = (t->bitmaps[thread_index][w] & m) != 0;
-  t->bitmaps[thread_index][w] |= m;
-
-  return (drop);
+  return clib_bitmap_set_no_check (t->bitmaps[thread_index], hash, 1);
 }
 
 #endif