feat: add user locking to isolation reservation table

src/main/java/io/supertokens/storage/postgresql/LockedUserImpl.java

@@ -0,0 +1,104 @@
+/*
+ *    Copyright (c) 2024, VRAI Labs and/or its affiliates. All rights reserved.
+ *
+ *    This software is licensed under the Apache License, Version 2.0 (the
+ *    "License") as published by the Apache Software Foundation.
+ *
+ *    You may not use this file except in compliance with the License. You may
+ *    obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ *    License for the specific language governing permissions and limitations
+ *    under the License.
+ */
+
+package io.supertokens.storage.postgresql;
+
+import io.supertokens.pluginInterface.useridmapping.LockedUser;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.sql.Connection;
+import java.lang.ref.WeakReference;
+
+/**
+ * PostgreSQL implementation of LockedUser.
+ * Tracks the connection to validate the lock is still active.
+ */
+public class LockedUserImpl implements LockedUser {
+
+    @Nonnull
+    private final String recipeUserId;
+
+    @Nonnull
+    private final String recipeId;
+
+    @Nullable
+    private final String primaryUserId;
+
+    // WeakReference so we don't prevent connection from being garbage collected
+    private final WeakReference<Connection> connectionRef;
+
+    public LockedUserImpl(@Nonnull String recipeUserId, @Nonnull String recipeId,
+                          @Nullable String primaryUserId, @Nonnull Connection connection) {
+        this.recipeUserId = recipeUserId;
+        this.recipeId = recipeId;
+        this.primaryUserId = primaryUserId;
+        this.connectionRef = new WeakReference<>(connection);
+    }
+
+    @Override
+    @Nonnull
+    public String getRecipeUserId() {
+        return recipeUserId;
+    }
+
+    @Override
+    @Nonnull
+    public String getRecipeId() {
+        return recipeId;
+    }
+
+    @Override
+    @Nullable
+    public String getPrimaryUserId() {
+        return primaryUserId;
+    }
+
+    @Override
+    public boolean isValidForConnection(Object connection) {
+        Connection originalCon = connectionRef.get();
+        if (originalCon == null) {
+            return false;
+        }
+        // Check that the provided connection is the same instance as the one used to acquire the lock
+        if (originalCon != connection) {
+            return false;
+        }
+        try {
+            return !originalCon.isClosed();
+        } catch (Exception e) {
+            return false;
+        }
+    }
+
+    @Override
+    public String toString() {
+        Connection con = connectionRef.get();
+        boolean connectionAlive = false;
+        try {
+            connectionAlive = con != null && !con.isClosed();
+        } catch (Exception ignored) {
+        }
+        return "LockedUser{" +
+               "recipeUserId='" + recipeUserId + '\'' +
+               ", recipeId='" + recipeId + '\'' +
+               ", primaryUserId='" + primaryUserId + '\'' +
+               ", isLinked=" + isLinked() +
+               ", isPrimary=" + isPrimary() +
+               ", connectionAlive=" + connectionAlive +
+               '}';
+    }
+}

src/main/java/io/supertokens/storage/postgresql/queries/GeneralQueries.java

@@ -42,6 +42,7 @@
 import io.supertokens.pluginInterface.exceptions.StorageTransactionLogicException;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifier;
 import io.supertokens.pluginInterface.multitenancy.TenantIdentifier;
+import io.supertokens.pluginInterface.useridmapping.LockedUser;
 import io.supertokens.pluginInterface.opentelemetry.WithinOtelSpan;
 import io.supertokens.storage.postgresql.ConnectionPool;
 import io.supertokens.storage.postgresql.PreparedStatementValueSetter;
@@ -739,6 +740,7 @@ public static void createTablesIfNotExists(Start start, Connection con) throws S
                     // indexes
                     update(con, AccountInfoQueries.getQueryToCreateTenantIndexForRecipeUserTenantsTable(start), NO_OP_SETTER);
                     update(con, AccountInfoQueries.getQueryToCreateRecipeUserIdIndexForRecipeUserTenantsTable(start), NO_OP_SETTER);
+                    update(con, AccountInfoQueries.getQueryToCreateRecipeUserIdIndexForRecipeUserAccountInfoTable(start), NO_OP_SETTER);
                     update(con, AccountInfoQueries.getQueryToCreateAccountInfoIndexForRecipeUserTenantsTable(start), NO_OP_SETTER);
                 }
 
@@ -1574,15 +1576,8 @@ public static AuthRecipeUserInfo[] listPrimaryUsersByThirdPartyInfo_Transaction(
                                                                                     String thirdPartyId,
                                                                                     String thirdPartyUserId)
             throws SQLException, StorageQueryException {
-        // we first lock on the table based on thirdparty info and tenant - this will ensure that any other
-        // query happening related to the account linking on this third party info / tenant will wait for this to
-        // finish,
-        // and vice versa.
-
-        ThirdPartyQueries.lockThirdPartyInfoAndTenant_Transaction(start, sqlCon, appIdentifier, thirdPartyId,
-                thirdPartyUserId);
-
-        // now that we have locks on all the relevant tables, we can read from them safely
+        // Note: Locking is now done at the core level via UserLockingStorage.lockUser()
+        // This method just queries the users without acquiring locks.
         List<String> userIds = ThirdPartyQueries.listUserIdsByThirdPartyInfo_Transaction(start, sqlCon, appIdentifier,
                 thirdPartyId, thirdPartyUserId);
         List<AuthRecipeUserInfo> result = getPrimaryUserInfoForUserIds_Transaction(start, sqlCon, appIdentifier,
@@ -1923,21 +1918,16 @@ private static List<AuthRecipeUserInfo> getPrimaryUserInfoForUserIds_Transaction
                 .collect(Collectors.toList());
     }
 
-    public static String getRecipeIdForUser_Transaction(Start start, Connection sqlCon,
-                                                        TenantIdentifier tenantIdentifier, String userId)
-            throws SQLException, StorageQueryException {
-
-        String QUERY = "SELECT recipe_id FROM " + getConfig(start).getAppIdToUserIdTable()
-                + " WHERE app_id = ? AND user_id = ? FOR UPDATE";
-        return execute(sqlCon, QUERY, pst -> {
-            pst.setString(1, tenantIdentifier.getAppId());
-            pst.setString(2, userId);
-        }, result -> {
-            if (result.next()) {
-                return result.getString("recipe_id");
-            }
-            return null;
-        });
+    /**
+     * Gets the recipe ID for a user that is already locked.
+     * The recipe ID is stored in the LockedUser object, which was fetched from
+     * app_id_to_user_id during lock acquisition.
+     *
+     * @param lockedUser The locked user (lock must be held)
+     * @return The recipe ID string
+     */
+    public static String getRecipeIdForUser_Transaction(LockedUser lockedUser) {
+        return lockedUser.getRecipeId();
     }
 
     public static Map<String, List<String>> getTenantIdsForUserIds_transaction(Start start, Connection sqlCon,

src/main/java/io/supertokens/storage/postgresql/queries/TOTPQueries.java

@@ -319,8 +319,9 @@ public static Map<String, List<TOTPDevice>> getDevicesForMultipleUsers(Start sta
     public static TOTPDevice[] getDevices_Transaction(Start start, Connection con, AppIdentifier appIdentifier,
                                                       String userId)
             throws StorageQueryException, SQLException {
+        // Note: FOR UPDATE removed - caller should obtain user lock via UserLockingStorage before calling this method
         String QUERY = "SELECT * FROM " + Config.getConfig(start).getTotpUserDevicesTable()
-                + " WHERE app_id = ? AND user_id = ? FOR UPDATE;";
+                + " WHERE app_id = ? AND user_id = ?;";
 
         return execute(con, QUERY, pst -> {
             pst.setString(1, appIdentifier.getAppId());

src/main/java/io/supertokens/storage/postgresql/queries/ThirdPartyQueries.java

@@ -226,59 +226,6 @@ public static void deleteUser_Transaction(Connection sqlCon, Start start, AppIde
         }
     }
 
-    public static List<String> lockThirdPartyInfoAndTenant_Transaction(Start start, Connection con,
-                                                                       AppIdentifier appIdentifier,
-                                                                       String thirdPartyId, String thirdPartyUserId)
-            throws SQLException, StorageQueryException {
-        String QUERY = "SELECT user_id " +
-                " FROM " + getConfig(start).getThirdPartyUsersTable() +
-                " WHERE app_id = ? AND third_party_id = ? AND third_party_user_id = ? FOR UPDATE";
-
-        return execute(con, QUERY, pst -> {
-            pst.setString(1, appIdentifier.getAppId());
-            pst.setString(2, thirdPartyId);
-            pst.setString(3, thirdPartyUserId);
-        }, result -> {
-            List<String> finalResult = new ArrayList<>();
-            while (result.next()) {
-                finalResult.add(result.getString("user_id"));
-            }
-            return finalResult;
-        });
-    }
-
-    public static List<String> lockThirdPartyInfoAndTenant_Transaction(Start start, Connection con,
-                                                                       AppIdentifier appIdentifier,
-                                                                       Map<String, String> thirdPartyUserIdToThirdPartyId)
-            throws SQLException, StorageQueryException {
-        if(thirdPartyUserIdToThirdPartyId == null || thirdPartyUserIdToThirdPartyId.isEmpty()) {
-            return new ArrayList<>();
-        }
-
-        String QUERY = "SELECT user_id " +
-                " FROM " + getConfig(start).getThirdPartyUsersTable() +
-                " WHERE app_id = ? AND third_party_id IN ("+Utils.generateCommaSeperatedQuestionMarks(
-                thirdPartyUserIdToThirdPartyId.size())+") AND third_party_user_id IN ("+
-                Utils.generateCommaSeperatedQuestionMarks(thirdPartyUserIdToThirdPartyId.size())+") FOR UPDATE";
-
-        return execute(con, QUERY, pst -> {
-            pst.setString(1, appIdentifier.getAppId());
-            int counter = 2;
-            for (String thirdPartyId : thirdPartyUserIdToThirdPartyId.values()){
-                pst.setString(counter++, thirdPartyId);
-            }
-            for (String thirdPartyUserId : thirdPartyUserIdToThirdPartyId.keySet()) {
-                pst.setString(counter++, thirdPartyUserId);
-            }
-        }, result -> {
-            List<String> finalResult = new ArrayList<>();
-            while (result.next()) {
-                finalResult.add(result.getString("user_id"));
-            }
-            return finalResult;
-        });
-    }
-
     public static List<LoginMethod> getUsersInfoUsingIdList(Start start, Set<String> ids,
                                                             AppIdentifier appIdentifier)
             throws SQLException, StorageQueryException {