From 6623037b2dc933c2d224518da9f1ed8e50424e4b Mon Sep 17 00:00:00 2001 From: Thomas Draier Date: Thu, 30 Jan 2025 16:38:55 +0100 Subject: [PATCH 1/3] Improve log with multiple memberships, use concurrent execution --- front/admin/relocate_users.ts | 47 +++++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/front/admin/relocate_users.ts b/front/admin/relocate_users.ts index bf3912583069..110f0d286687 100644 --- a/front/admin/relocate_users.ts +++ b/front/admin/relocate_users.ts @@ -1,4 +1,4 @@ -import { removeNulls } from "@dust-tt/types"; +import { concurrentExecutor, removeNulls } from "@dust-tt/types"; import type { ApiResponse } from "auth0"; import { getAuth0ManagemementClient } from "@app/lib/api/auth0"; @@ -70,8 +70,13 @@ makeScript( (m) => m.workspaceId !== workspace.id ); if (externalMemberships.length > 0) { + const userIds = externalMemberships.map((m) => m.userId); + const users = await UserResource.fetchByModelIds(userIds); logger.error( - { users: externalMemberships.map((m) => m.user) }, + { + externalMemberships, + users: users.map((u) => ({ sId: u.sId, email: u.email })), + }, "Some users have mutiple memberships" ); process.exit(1); @@ -84,23 +89,29 @@ makeScript( let count = 0; - for (const auth0Id of auth0Ids) { - count++; - logger.info({ user: auth0Id, count }, "Setting region"); - if (execute) { - await throttleAuth0(() => - managementClient.users.update( - { - id: auth0Id, - }, - { - app_metadata: { - region: destinationRegion, + await concurrentExecutor( + auth0Ids, + async (auth0Id) => { + count++; + logger.info({ user: auth0Id, count }, "Setting region"); + if (execute) { + await throttleAuth0(() => + managementClient.users.update( + { + id: auth0Id, }, - } - ) - ); + { + app_metadata: { + region: destinationRegion, + }, + } + ) + ); + } + }, + { + concurrency: 10, } - } + ); } ); From b35c27dd3a6c48acb72694c2eb4087f33411ee28 Mon Sep 17 00:00:00 2001 From: Thomas Draier Date: Thu, 30 Jan 2025 16:52:07 +0100 Subject: [PATCH 2/3] share throttle --- front/admin/relocate_users.ts | 51 +++++++++-------------------------- front/lib/api/auth0.ts | 31 +++++++++++++++++++++ 2 files changed, 44 insertions(+), 38 deletions(-) diff --git a/front/admin/relocate_users.ts b/front/admin/relocate_users.ts index 110f0d286687..a6a7b7a3c90a 100644 --- a/front/admin/relocate_users.ts +++ b/front/admin/relocate_users.ts @@ -1,7 +1,6 @@ import { concurrentExecutor, removeNulls } from "@dust-tt/types"; -import type { ApiResponse } from "auth0"; -import { getAuth0ManagemementClient } from "@app/lib/api/auth0"; +import { getAuth0ManagemementClient, throttleAuth0 } from "@app/lib/api/auth0"; import { SUPPORTED_REGIONS } from "@app/lib/api/regions/config"; import { Authenticator } from "@app/lib/auth"; import { MembershipResource } from "@app/lib/resources/membership_resource"; @@ -34,32 +33,6 @@ makeScript( const auth = await Authenticator.internalAdminForWorkspace(workspaceId); const workspace = auth.getNonNullableWorkspace(); - let remaining = 10; - let resetTime = Date.now(); - - const throttleAuth0 = async (fn: () => Promise>) => { - if (remaining < rateLimitThreshold) { - const now = Date.now(); - const waitTime = resetTime * 1000 - now; - logger.info({ waitTime }, "Waiting"); - await new Promise((resolve) => setTimeout(resolve, waitTime)); - } - - const res = await fn(); - if (res.status !== 200) { - logger.error({ res }, "When calling Auth0"); - process.exit(1); - } - - remaining = Number(res.headers.get("x-ratelimit-remaining")); - resetTime = Number(res.headers.get("x-ratelimit-reset")); - - const limit = Number(res.headers.get("x-ratelimit-limit")); - logger.info({ limit, remaining, resetTime }, "Rate limit"); - - return res.data; - }; - const members = await MembershipResource.getMembershipsForWorkspace({ workspace, }); @@ -95,17 +68,19 @@ makeScript( count++; logger.info({ user: auth0Id, count }, "Setting region"); if (execute) { - await throttleAuth0(() => - managementClient.users.update( - { - id: auth0Id, - }, - { - app_metadata: { - region: destinationRegion, + await throttleAuth0( + () => + managementClient.users.update( + { + id: auth0Id, }, - } - ) + { + app_metadata: { + region: destinationRegion, + }, + } + ), + { rateLimitThreshold } ); } }, diff --git a/front/lib/api/auth0.ts b/front/lib/api/auth0.ts index 31c874916d55..d47d8ff1bbb9 100644 --- a/front/lib/api/auth0.ts +++ b/front/lib/api/auth0.ts @@ -1,6 +1,7 @@ import type { Session } from "@auth0/nextjs-auth0"; import type { Result } from "@dust-tt/types"; import { Err, Ok } from "@dust-tt/types"; +import type { ApiResponse } from "auth0"; import { ManagementClient } from "auth0"; import { isLeft } from "fp-ts/lib/Either"; import * as t from "io-ts"; @@ -266,3 +267,33 @@ export async function getAuth0UsersFromEmail(emails: string[]) { return auth0Users; } + +const throttleAuth0State = { + remaining: -1, + resetTime: 0, +}; + +export const throttleAuth0 = async ( + fn: () => Promise>, + { rateLimitThreshold }: { rateLimitThreshold: number } +) => { + const { remaining, resetTime } = throttleAuth0State; + if (remaining === -1 || remaining < rateLimitThreshold) { + const now = Date.now(); + const waitTime = resetTime * 1000 - now; + await new Promise((resolve) => setTimeout(resolve, waitTime)); + } + + const res = await fn(); + if (res.status !== 200) { + logger.error({ res }, "When calling Auth0"); + process.exit(1); + } + + throttleAuth0State.remaining = Number( + res.headers.get("x-ratelimit-remaining") + ); + throttleAuth0State.resetTime = Number(res.headers.get("x-ratelimit-reset")); + + return res.data; +}; From cacb951db917451167485cf1290fab855a3131d0 Mon Sep 17 00:00:00 2001 From: Thomas Draier Date: Thu, 30 Jan 2025 17:04:24 +0100 Subject: [PATCH 3/3] fix --- front/lib/api/auth0.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/front/lib/api/auth0.ts b/front/lib/api/auth0.ts index d47d8ff1bbb9..97dc6fb9a89a 100644 --- a/front/lib/api/auth0.ts +++ b/front/lib/api/auth0.ts @@ -278,7 +278,7 @@ export const throttleAuth0 = async ( { rateLimitThreshold }: { rateLimitThreshold: number } ) => { const { remaining, resetTime } = throttleAuth0State; - if (remaining === -1 || remaining < rateLimitThreshold) { + if (remaining !== -1 && remaining < rateLimitThreshold) { const now = Date.now(); const waitTime = resetTime * 1000 - now; await new Promise((resolve) => setTimeout(resolve, waitTime));