diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index a0e58a2..5300126 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1679,7 +1679,15 @@
         return status;
     }
 
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_lock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     status = psa_reserve_free_key_slot(&volatile_key_id, p_slot);
+#if defined(MBEDTLS_THREADING_C)
+    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
+                              &mbedtls_threading_key_slot_mutex));
+#endif
     if (status != PSA_SUCCESS) {
         return status;
     }
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 07d7f35..dc38662 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -160,13 +160,9 @@
     size_t slot_idx;
     psa_key_slot_t *selected_slot, *unused_persistent_key_slot;
 
-#if defined(MBEDTLS_THREADING_C)
-    PSA_THREADING_CHK_GOTO_EXIT(mbedtls_mutex_lock(
-                                    &mbedtls_threading_key_slot_mutex));
-#endif
     if (!global_data.key_slots_initialized) {
         status = PSA_ERROR_BAD_STATE;
-        goto exit;
+        goto error;
     }
 
     selected_slot = unused_persistent_key_slot = NULL;
@@ -198,7 +194,7 @@
         psa_register_read(selected_slot);
         status = psa_wipe_key_slot(selected_slot);
         if (status != PSA_SUCCESS) {
-            goto exit;
+            goto error;
         }
     }
 
@@ -206,27 +202,21 @@
         status = psa_key_slot_state_transition(selected_slot, PSA_SLOT_EMPTY,
                                                PSA_SLOT_FILLING);
         if (status != PSA_SUCCESS) {
-            goto exit;
+            goto error;
         }
 
         *volatile_key_id = PSA_KEY_ID_VOLATILE_MIN +
                            ((psa_key_id_t) (selected_slot - global_data.key_slots));
         *p_slot = selected_slot;
 
-        goto exit;
+        return PSA_SUCCESS;
     }
     status = PSA_ERROR_INSUFFICIENT_MEMORY;
 
-exit:
-    if (status != PSA_SUCCESS) {
-        *p_slot = NULL;
-        *volatile_key_id = 0;
-    }
+error:
+    *p_slot = NULL;
+    *volatile_key_id = 0;
 
-#if defined(MBEDTLS_THREADING_C)
-    PSA_THREADING_CHK_RET(mbedtls_mutex_unlock(
-                              &mbedtls_threading_key_slot_mutex));
-#endif
     return status;
 }
 
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 18a9144..585de13 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -107,6 +107,9 @@
  * It is the responsibility of the caller to change the slot's state to
  * PSA_SLOT_EMPTY/FULL once key creation has finished.
  *
+ * If multi-threading is enabled, the caller must hold the
+ * global key slot mutex.
+ *
  * \param[out] volatile_key_id   On success, volatile key identifier
  *                               associated to the returned slot.
  * \param[out] p_slot            On success, a pointer to the slot.