static DEFINE_MUTEX(reading_mutex);
 static int data_avail;
 static u8 *rng_buffer, *rng_fillbuf;
-static DECLARE_WAIT_QUEUE_HEAD(rng_done);
 static unsigned short current_quality;
 static unsigned short default_quality; /* = 0; default to "off" */
 
        if (rng->cleanup)
                rng->cleanup(rng);
 
-       /* cleanup_done should be updated after cleanup finishes */
-       smp_wmb();
-       rng->cleanup_done = true;
-       wake_up_all(&rng_done);
+       complete(&rng->cleanup_done);
 }
 
 static void set_current_rng(struct hwrng *rng)
                add_early_randomness(rng);
        }
 
-       rng->cleanup_done = false;
+       init_completion(&rng->cleanup_done);
 
 out_unlock:
        mutex_unlock(&rng_mutex);
        } else
                mutex_unlock(&rng_mutex);
 
-       /* Just in case rng is reading right now, wait. */
-       wait_event(rng_done, rng->cleanup_done &&
-                  atomic_read(&rng->ref.refcount) == 0);
+       wait_for_completion(&rng->cleanup_done);
 }
 EXPORT_SYMBOL_GPL(hwrng_unregister);
 
 
 #ifndef LINUX_HWRANDOM_H_
 #define LINUX_HWRANDOM_H_
 
+#include <linux/completion.h>
 #include <linux/types.h>
 #include <linux/list.h>
 #include <linux/kref.h>
        /* internal. */
        struct list_head list;
        struct kref ref;
-       bool cleanup_done;
+       struct completion cleanup_done;
 };
 
 /** Register a new Hardware Random Number Generator driver. */