void *handle, u32 id, int lock)
 {
        int ret;
-       int newid;
 
-       do {
-               if (!idr_pre_get(idr, lock ? GFP_KERNEL : GFP_ATOMIC))
-                       return -ENOMEM;
-               if (lock)
-                       spin_lock_irq(&rhp->lock);
-               ret = idr_get_new_above(idr, handle, id, &newid);
-               BUG_ON(!ret && newid != id);
-               if (lock)
-                       spin_unlock_irq(&rhp->lock);
-       } while (ret == -EAGAIN);
-
-       return ret;
+       if (lock) {
+               idr_preload(GFP_KERNEL);
+               spin_lock_irq(&rhp->lock);
+       }
+
+       ret = idr_alloc(idr, handle, id, id + 1, GFP_ATOMIC);
+
+       if (lock) {
+               spin_unlock_irq(&rhp->lock);
+               idr_preload_end();
+       }
+
+       BUG_ON(ret == -ENOSPC);
+       return ret < 0 ? ret : 0;
 }
 
 static inline int insert_handle(struct c4iw_dev *rhp, struct idr *idr,