1 files changed, 15 insertions, 16 deletions

diff --git a/net/sctp/associola.c b/net/sctp/associola.c
index 2f95f5a5145d..43cd0dd9149d 100644
--- a/net/sctp/associola.c
+++ b/net/sctp/associola.c
@@ -1591,32 +1591,31 @@ int sctp_assoc_lookup_laddr(struct sctp_association *asoc,
 /* Set an association id for a given association */
 int sctp_assoc_set_id(struct sctp_association *asoc, gfp_t gfp)
 {
-        int assoc_id;
-        int error = 0;
+        bool preload = gfp & __GFP_WAIT;
+        int ret;
 
         /* If the id is already assigned, keep it. */
         if (asoc->assoc_id)
-                return error;
-retry:
-        if (unlikely(!idr_pre_get(&sctp_assocs_id, gfp)))
-                return -ENOMEM;
+                return 0;
 
+        if (preload)
+                idr_preload(gfp);
         spin_lock_bh(&sctp_assocs_id_lock);
-        error = idr_get_new_above(&sctp_assocs_id, (void *)asoc,
-                                    idr_low, &assoc_id);
-        if (!error) {
-                idr_low = assoc_id + 1;
+        /* 0 is not a valid id, idr_low is always >= 1 */
+        ret = idr_alloc(&sctp_assocs_id, asoc, idr_low, 0, GFP_NOWAIT);
+        if (ret >= 0) {
+                idr_low = ret + 1;
                 if (idr_low == INT_MAX)
                         idr_low = 1;
         }
         spin_unlock_bh(&sctp_assocs_id_lock);
-        if (error == -EAGAIN)
-                goto retry;
-        else if (error)
-                return error;
+        if (preload)
+                idr_preload_end();
+        if (ret < 0)
+                return ret;
 
-        asoc->assoc_id = (sctp_assoc_t) assoc_id;
-        return error;
+        asoc->assoc_id = (sctp_assoc_t)ret;
+        return 0;
 }
 
 /* Free the ASCONF queue */