#include <linux/sched.h>
 #include <linux/workqueue.h>
 #include <linux/kthread.h>
+#include <linux/completion.h>
 #include <trace/events/xdp.h>
 #include <linux/btf_ids.h>
 
        struct rcu_head rcu;
 
        struct work_struct kthread_stop_wq;
+       struct completion kthread_running;
 };
 
 struct bpf_cpu_map {
 static void cpu_map_kthread_stop(struct work_struct *work)
 {
        struct bpf_cpu_map_entry *rcpu;
-       int err;
 
        rcpu = container_of(work, struct bpf_cpu_map_entry, kthread_stop_wq);
 
        rcu_barrier();
 
        /* kthread_stop will wake_up_process and wait for it to complete */
-       err = kthread_stop(rcpu->kthread);
-       if (err) {
-               /* kthread_stop may be called before cpu_map_kthread_run
-                * is executed, so we need to release the memory related
-                * to rcpu.
-                */
-               put_cpu_map_entry(rcpu);
-       }
+       kthread_stop(rcpu->kthread);
 }
 
 static void cpu_map_bpf_prog_run_skb(struct bpf_cpu_map_entry *rcpu,
        return nframes;
 }
 
-
 static int cpu_map_kthread_run(void *data)
 {
        struct bpf_cpu_map_entry *rcpu = data;
 
+       complete(&rcpu->kthread_running);
        set_current_state(TASK_INTERRUPTIBLE);
 
        /* When kthread gives stop order, then rcpu have been disconnected
                goto free_ptr_ring;
 
        /* Setup kthread */
+       init_completion(&rcpu->kthread_running);
        rcpu->kthread = kthread_create_on_node(cpu_map_kthread_run, rcpu, numa,
                                               "cpumap/%d/map:%d", cpu,
                                               map->id);
        kthread_bind(rcpu->kthread, cpu);
        wake_up_process(rcpu->kthread);
 
+       /* Make sure kthread has been running, so kthread_stop() will not
+        * stop the kthread prematurely and all pending frames or skbs
+        * will be handled by the kthread before kthread_stop() returns.
+        */
+       wait_for_completion(&rcpu->kthread_running);
+
        return rcpu;
 
 free_prog: