#include <uapi/linux/sock_diag.h>
 #include <uapi/linux/btf.h>
 
-static atomic_t cache_idx;
-
 #define SK_STORAGE_CREATE_FLAG_MASK                                    \
        (BPF_F_NO_PREALLOC | BPF_F_CLONE)
 
 #define SDATA(_SELEM) (&(_SELEM)->sdata)
 #define BPF_SK_STORAGE_CACHE_SIZE      16
 
+static DEFINE_SPINLOCK(cache_idx_lock);
+static u64 cache_idx_usage_counts[BPF_SK_STORAGE_CACHE_SIZE];
+
 struct bpf_sk_storage {
        struct bpf_sk_storage_data __rcu *cache[BPF_SK_STORAGE_CACHE_SIZE];
        struct hlist_head list; /* List of bpf_sk_storage_elem */
        return 0;
 }
 
+static u16 cache_idx_get(void)
+{
+       u64 min_usage = U64_MAX;
+       u16 i, res = 0;
+
+       spin_lock(&cache_idx_lock);
+
+       for (i = 0; i < BPF_SK_STORAGE_CACHE_SIZE; i++) {
+               if (cache_idx_usage_counts[i] < min_usage) {
+                       min_usage = cache_idx_usage_counts[i];
+                       res = i;
+
+                       /* Found a free cache_idx */
+                       if (!min_usage)
+                               break;
+               }
+       }
+       cache_idx_usage_counts[res]++;
+
+       spin_unlock(&cache_idx_lock);
+
+       return res;
+}
+
+static void cache_idx_free(u16 idx)
+{
+       spin_lock(&cache_idx_lock);
+       cache_idx_usage_counts[idx]--;
+       spin_unlock(&cache_idx_lock);
+}
+
 /* Called by __sk_destruct() & bpf_sk_storage_clone() */
 void bpf_sk_storage_free(struct sock *sk)
 {
 
        smap = (struct bpf_sk_storage_map *)map;
 
+       cache_idx_free(smap->cache_idx);
+
        /* Note that this map might be concurrently cloned from
         * bpf_sk_storage_clone. Wait for any existing bpf_sk_storage_clone
         * RCU read section to finish before proceeding. New RCU
        }
 
        smap->elem_size = sizeof(struct bpf_sk_storage_elem) + attr->value_size;
-       smap->cache_idx = (unsigned int)atomic_inc_return(&cache_idx) %
-               BPF_SK_STORAGE_CACHE_SIZE;
+       smap->cache_idx = cache_idx_get();
 
        return &smap->map;
 }