u64 elem_map_flags, map_flags;
        struct hlist_nulls_head *head;
        struct hlist_nulls_node *n;
-       unsigned long flags;
+       unsigned long flags = 0;
+       bool locked = false;
        struct htab_elem *l;
        struct bucket *b;
        int ret = 0;
        dst_val = values;
        b = &htab->buckets[batch];
        head = &b->head;
-       raw_spin_lock_irqsave(&b->lock, flags);
+       /* do not grab the lock unless need it (bucket_cnt > 0). */
+       if (locked)
+               raw_spin_lock_irqsave(&b->lock, flags);
 
        bucket_cnt = 0;
        hlist_nulls_for_each_entry_rcu(l, n, head, hash_node)
                bucket_cnt++;
 
+       if (bucket_cnt && !locked) {
+               locked = true;
+               goto again_nocopy;
+       }
+
        if (bucket_cnt > (max_count - total)) {
                if (total == 0)
                        ret = -ENOSPC;
+               /* Note that since bucket_cnt > 0 here, it is implicit
+                * that the locked was grabbed, so release it.
+                */
                raw_spin_unlock_irqrestore(&b->lock, flags);
                rcu_read_unlock();
                this_cpu_dec(bpf_prog_active);
 
        if (bucket_cnt > bucket_size) {
                bucket_size = bucket_cnt;
+               /* Note that since bucket_cnt > 0 here, it is implicit
+                * that the locked was grabbed, so release it.
+                */
                raw_spin_unlock_irqrestore(&b->lock, flags);
                rcu_read_unlock();
                this_cpu_dec(bpf_prog_active);
                goto alloc;
        }
 
+       /* Next block is only safe to run if you have grabbed the lock */
+       if (!locked)
+               goto next_batch;
+
        hlist_nulls_for_each_entry_safe(l, n, head, hash_node) {
                memcpy(dst_key, l->key, key_size);
 
        }
 
        raw_spin_unlock_irqrestore(&b->lock, flags);
+       locked = false;
+next_batch:
        /* If we are not copying data, we can go to next bucket and avoid
         * unlocking the rcu.
         */