Merge pull request #2709 from fosrl/pool-update

Update pool and disable idp

server/db/pg/driver.ts

@@ -1,7 +1,7 @@
 import { drizzle as DrizzlePostgres } from "drizzle-orm/node-postgres";
-import { Pool } from "pg";
 import { readConfigFile } from "@server/lib/readConfigFile";
 import { withReplicas } from "drizzle-orm/pg-core";
+import { createPool } from "./poolConfig";
 
 function createDb() {
     const config = readConfigFile();
@@ -39,12 +39,17 @@ function createDb() {
 
     // Create connection pools instead of individual connections
     const poolConfig = config.postgres.pool;
-    const primaryPool = new Pool({
+    const maxConnections = poolConfig?.max_connections || 20;
+    const idleTimeoutMs = poolConfig?.idle_timeout_ms || 30000;
+    const connectionTimeoutMs = poolConfig?.connection_timeout_ms || 5000;
+
+    const primaryPool = createPool(
         connectionString,
-        max: poolConfig?.max_connections || 20,
-        idleTimeoutMillis: poolConfig?.idle_timeout_ms || 30000,
-        connectionTimeoutMillis: poolConfig?.connection_timeout_ms || 5000
-    });
+        maxConnections,
+        idleTimeoutMs,
+        connectionTimeoutMs,
+        "primary"
+    );
 
     const replicas = [];
 
@@ -55,14 +60,16 @@ function createDb() {
             })
         );
     } else {
+        const maxReplicaConnections =
+            poolConfig?.max_replica_connections || 20;
         for (const conn of replicaConnections) {
-            const replicaPool = new Pool({
-                connectionString: conn.connection_string,
-                max: poolConfig?.max_replica_connections || 20,
-                idleTimeoutMillis: poolConfig?.idle_timeout_ms || 30000,
-                connectionTimeoutMillis:
-                    poolConfig?.connection_timeout_ms || 5000
-            });
+            const replicaPool = createPool(
+                conn.connection_string,
+                maxReplicaConnections,
+                idleTimeoutMs,
+                connectionTimeoutMs,
+                "replica"
+            );
             replicas.push(
                 DrizzlePostgres(replicaPool, {
                     logger: process.env.QUERY_LOGGING == "true"

server/db/pg/logsDriver.ts

@@ -1,9 +1,9 @@
 import { drizzle as DrizzlePostgres } from "drizzle-orm/node-postgres";
-import { Pool } from "pg";
 import { readConfigFile } from "@server/lib/readConfigFile";
 import { withReplicas } from "drizzle-orm/pg-core";
 import { build } from "@server/build";
 import { db as mainDb, primaryDb as mainPrimaryDb } from "./driver";
+import { createPool } from "./poolConfig";
 
 function createLogsDb() {
     // Only use separate logs database in SaaS builds
@@ -42,12 +42,17 @@ function createLogsDb() {
 
     // Create separate connection pool for logs database
     const poolConfig = logsConfig?.pool || config.postgres?.pool;
-    const primaryPool = new Pool({
+    const maxConnections = poolConfig?.max_connections || 20;
+    const idleTimeoutMs = poolConfig?.idle_timeout_ms || 30000;
+    const connectionTimeoutMs = poolConfig?.connection_timeout_ms || 5000;
+
+    const primaryPool = createPool(
         connectionString,
-        max: poolConfig?.max_connections || 20,
-        idleTimeoutMillis: poolConfig?.idle_timeout_ms || 30000,
-        connectionTimeoutMillis: poolConfig?.connection_timeout_ms || 5000
-    });
+        maxConnections,
+        idleTimeoutMs,
+        connectionTimeoutMs,
+        "logs-primary"
+    );
 
     const replicas = [];
 
@@ -58,14 +63,16 @@ function createLogsDb() {
             })
         );
     } else {
+        const maxReplicaConnections =
+            poolConfig?.max_replica_connections || 20;
         for (const conn of replicaConnections) {
-            const replicaPool = new Pool({
-                connectionString: conn.connection_string,
-                max: poolConfig?.max_replica_connections || 20,
-                idleTimeoutMillis: poolConfig?.idle_timeout_ms || 30000,
-                connectionTimeoutMillis:
-                    poolConfig?.connection_timeout_ms || 5000
-            });
+            const replicaPool = createPool(
+                conn.connection_string,
+                maxReplicaConnections,
+                idleTimeoutMs,
+                connectionTimeoutMs,
+                "logs-replica"
+            );
             replicas.push(
                 DrizzlePostgres(replicaPool, {
                     logger: process.env.QUERY_LOGGING == "true"

server/db/pg/poolConfig.ts

@@ -0,0 +1,63 @@
+import { Pool, PoolConfig } from "pg";
+import logger from "@server/logger";
+
+export function createPoolConfig(
+    connectionString: string,
+    maxConnections: number,
+    idleTimeoutMs: number,
+    connectionTimeoutMs: number
+): PoolConfig {
+    return {
+        connectionString,
+        max: maxConnections,
+        idleTimeoutMillis: idleTimeoutMs,
+        connectionTimeoutMillis: connectionTimeoutMs,
+        // TCP keepalive to prevent silent connection drops by NAT gateways,
+        // load balancers, and other intermediate network devices (e.g. AWS
+        // NAT Gateway drops idle TCP connections after ~350s)
+        keepAlive: true,
+        keepAliveInitialDelayMillis: 10000, // send first keepalive after 10s of idle
+        // Allow connections to be released and recreated more aggressively
+        // to avoid stale connections building up
+        allowExitOnIdle: false
+    };
+}
+
+export function attachPoolErrorHandlers(pool: Pool, label: string): void {
+    pool.on("error", (err) => {
+        // This catches errors on idle clients in the pool. Without this
+        // handler an unexpected disconnect would crash the process.
+        logger.error(
+            `Unexpected error on idle ${label} database client: ${err.message}`
+        );
+    });
+
+    pool.on("connect", (client) => {
+        // Set a statement timeout on every new connection so a single slow
+        // query can't block the pool forever
+        client.query("SET statement_timeout = '30s'").catch((err: Error) => {
+            logger.warn(
+                `Failed to set statement_timeout on ${label} client: ${err.message}`
+            );
+        });
+    });
+}
+
+export function createPool(
+    connectionString: string,
+    maxConnections: number,
+    idleTimeoutMs: number,
+    connectionTimeoutMs: number,
+    label: string
+): Pool {
+    const pool = new Pool(
+        createPoolConfig(
+            connectionString,
+            maxConnections,
+            idleTimeoutMs,
+            connectionTimeoutMs
+        )
+    );
+    attachPoolErrorHandlers(pool, label);
+    return pool;
+}

server/private/lib/cache.ts

@@ -24,23 +24,31 @@ setInterval(() => {
  */
 class AdaptiveCache {
     private useRedis(): boolean {
-        return redisManager.isRedisEnabled() && redisManager.getHealthStatus().isHealthy;
+        return (
+            redisManager.isRedisEnabled() &&
+            redisManager.getHealthStatus().isHealthy
+        );
     }
 
     /**
      * Set a value in the cache
      * @param key - Cache key
      * @param value - Value to cache (will be JSON stringified for Redis)
-     * @param ttl - Time to live in seconds (0 = no expiration)
+     * @param ttl - Time to live in seconds (0 = no expiration; omit = 3600s for Redis)
      * @returns boolean indicating success
      */
     async set(key: string, value: any, ttl?: number): Promise<boolean> {
         const effectiveTtl = ttl === 0 ? undefined : ttl;
+        const redisTtl = ttl === 0 ? undefined : (ttl ?? 3600);
 
         if (this.useRedis()) {
             try {
                 const serialized = JSON.stringify(value);
-                const success = await redisManager.set(key, serialized, effectiveTtl);
+                const success = await redisManager.set(
+                    key,
+                    serialized,
+                    redisTtl
+                );
 
                 if (success) {
                     logger.debug(`Set key in Redis: ${key}`);
@@ -48,7 +56,9 @@ class AdaptiveCache {
                 }
 
                 // Redis failed, fall through to local cache
-                logger.debug(`Redis set failed for key ${key}, falling back to local cache`);
+                logger.debug(
+                    `Redis set failed for key ${key}, falling back to local cache`
+                );
             } catch (error) {
                 logger.error(`Redis set error for key ${key}:`, error);
                 // Fall through to local cache
@@ -120,9 +130,14 @@ class AdaptiveCache {
                 }
 
                 // Some Redis deletes failed, fall through to local cache
-                logger.debug(`Some Redis deletes failed, falling back to local cache`);
+                logger.debug(
+                    `Some Redis deletes failed, falling back to local cache`
+                );
             } catch (error) {
-                logger.error(`Redis del error for keys ${keys.join(", ")}:`, error);
+                logger.error(
+                    `Redis del error for keys ${keys.join(", ")}:`,
+                    error
+                );
                 // Fall through to local cache
                 deletedCount = 0;
             }
@@ -195,7 +210,9 @@ class AdaptiveCache {
      */
     async flushAll(): Promise<void> {
         if (this.useRedis()) {
-            logger.warn("Adaptive cache flushAll called - Redis flush not implemented, only local cache will be flushed");
+            logger.warn(
+                "Adaptive cache flushAll called - Redis flush not implemented, only local cache will be flushed"
+            );
         }
 
         localCache.flushAll();
@@ -239,7 +256,9 @@ class AdaptiveCache {
     getTtl(key: string): number {
         // Note: This only works for local cache, Redis TTL is not supported
         if (this.useRedis()) {
-            logger.warn(`getTtl called for key ${key} but Redis TTL lookup is not implemented`);
+            logger.warn(
+                `getTtl called for key ${key} but Redis TTL lookup is not implemented`
+            );
         }
 
         const ttl = localCache.getTtl(key);
@@ -255,7 +274,9 @@ class AdaptiveCache {
      */
     keys(): string[] {
         if (this.useRedis()) {
-            logger.warn("keys() called but Redis keys are not included, only local cache keys returned");
+            logger.warn(
+                "keys() called but Redis keys are not included, only local cache keys returned"
+            );
         }
         return localCache.keys();
     }

server/private/routers/hybrid.ts

@@ -15,6 +15,7 @@ import { verifySessionRemoteExitNodeMiddleware } from "#private/middlewares/veri
 import { Router } from "express";
 import {
     db,
+    logsDb,
     exitNodes,
     Resource,
     ResourcePassword,
@@ -1885,7 +1886,7 @@ hybridRouter.post(
             const batchSize = 100;
             for (let i = 0; i < logEntries.length; i += batchSize) {
                 const batch = logEntries.slice(i, i + batchSize);
-                await db.insert(requestAuditLog).values(batch);
+                await logsDb.insert(requestAuditLog).values(batch);
             }
 
             return response(res, {

server/private/routers/remoteExitNode/handleRemoteExitNodePingMessage.ts

@@ -38,7 +38,7 @@ export const startRemoteExitNodeOfflineChecker = (): void => {
             );
 
             // Find clients that haven't pinged in the last 2 minutes and mark them as offline
-            const newlyOfflineNodes = await db
+            const offlineNodes = await db
                 .update(exitNodes)
                 .set({ online: false })
                 .where(
@@ -53,32 +53,15 @@ export const startRemoteExitNodeOfflineChecker = (): void => {
                 )
                 .returning();
 
-            // Update the sites to offline if they have not pinged either
-            const exitNodeIds = newlyOfflineNodes.map(
-                (node) => node.exitNodeId
-            );
-
-            const sitesOnNode = await db
-                .select()
-                .from(sites)
-                .where(
-                    and(
-                        eq(sites.online, true),
-                        inArray(sites.exitNodeId, exitNodeIds)
-                    )
+            if (offlineNodes.length > 0) {
+                logger.info(
+                    `checkRemoteExitNodeOffline: Marked ${offlineNodes.length} remoteExitNode client(s) offline due to inactivity`
                 );
 
-            // loop through the sites and process their lastBandwidthUpdate as an iso string and if its more than 1 minute old then mark the site offline
-            for (const site of sitesOnNode) {
-                if (!site.lastBandwidthUpdate) {
-                    continue;
-                }
-                const lastBandwidthUpdate = new Date(site.lastBandwidthUpdate);
-                if (Date.now() - lastBandwidthUpdate.getTime() > 60 * 1000) {
-                    await db
-                        .update(sites)
-                        .set({ online: false })
-                        .where(eq(sites.siteId, site.siteId));
+                for (const offlineClient of offlineNodes) {
+                    logger.debug(
+                        `checkRemoteExitNodeOffline: Client ${offlineClient.exitNodeId} marked offline (lastPing: ${offlineClient.lastPing})`
+                    );
                 }
             }
         } catch (error) {

server/routers/client/listClients.ts

@@ -70,7 +70,7 @@ async function getLatestOlmVersion(): Promise<string | null> {
         tags = tags.filter((version) => !version.name.includes("rc"));
         const latestVersion = tags[0].name;
 
-        olmVersionCache.set("latestOlmVersion", latestVersion);
+        olmVersionCache.set("latestOlmVersion", latestVersion, 3600);
 
         return latestVersion;
     } catch (error: any) {

server/routers/client/listUserDevices.ts

@@ -71,7 +71,7 @@ async function getLatestOlmVersion(): Promise<string | null> {
         tags = tags.filter((version) => !version.name.includes("rc"));
         const latestVersion = tags[0].name;
 
-        olmVersionCache.set("latestOlmVersion", latestVersion);
+        olmVersionCache.set("latestOlmVersion", latestVersion, 3600);
 
         return latestVersion;
     } catch (error: any) {

server/routers/newt/handleNewtDisconnectingMessage.ts

@@ -6,7 +6,9 @@ import logger from "@server/logger";
 /**
  * Handles disconnecting messages from sites to show disconnected in the ui
  */
-export const handleNewtDisconnectingMessage: MessageHandler = async (context) => {
+export const handleNewtDisconnectingMessage: MessageHandler = async (
+    context
+) => {
     const { message, client: c, sendToClient } = context;
     const newt = c as Newt;
 
@@ -27,7 +29,7 @@ export const handleNewtDisconnectingMessage: MessageHandler = async (context) =>
             .set({
                 online: false
             })
-            .where(eq(sites.siteId, sites.siteId));
+            .where(eq(sites.siteId, newt.siteId));
     } catch (error) {
         logger.error("Error handling disconnecting message", { error });
     }

server/routers/site/listSites.ts

@@ -55,7 +55,7 @@ async function getLatestNewtVersion(): Promise<string | null> {
         tags = tags.filter((version) => !version.name.includes("rc"));
         const latestVersion = tags[0].name;
 
-        await cache.set("latestNewtVersion", latestVersion);
+        await cache.set("latestNewtVersion", latestVersion, 3600);
 
         return latestVersion;
     } catch (error: any) {
@@ -180,7 +180,7 @@ registry.registerPath({
     method: "get",
     path: "/org/{orgId}/sites",
     description: "List all sites in an organization",
-    tags: [OpenAPITags.Site],
+    tags: [OpenAPITags.Org, OpenAPITags.Site],
     request: {
         params: listSitesParamsSchema,
         query: listSitesSchema

server/routers/user/inviteUser.ts

@@ -201,7 +201,7 @@ export async function inviteUser(
                 );
             }
 
-            await cache.set(email, attempts + 1);
+            await cache.set("regenerateInvite:" + email, attempts + 1, 3600);
 
             const inviteId = existingInvite[0].inviteId; // Retrieve the original inviteId
             const token = generateRandomString(

src/app/[orgId]/settings/(private)/idp/create/page.tsx

@@ -275,6 +275,8 @@ export default function Page() {
         }
     }
 
+    const disabled = !isPaidUser(tierMatrix.orgOidc);
+
     return (
         <>
             <div className="flex justify-between">
@@ -292,6 +294,9 @@ export default function Page() {
                 </Button>
             </div>
 
+            <PaidFeaturesAlert tiers={tierMatrix.orgOidc} />
+
+            <fieldset disabled={disabled} className={disabled ? "opacity-50 pointer-events-none" : ""}>
             <SettingsContainer>
                 <SettingsSection>
                     <SettingsSectionHeader>
@@ -812,9 +817,10 @@ export default function Page() {
                 </Button>
                 <Button
                     type="submit"
-                    disabled={createLoading || !isPaidUser(tierMatrix.orgOidc)}
+                    disabled={createLoading || disabled}
                     loading={createLoading}
                     onClick={() => {
+                        if (disabled) return;
                         // log any issues with the form
                         console.log(form.formState.errors);
                         form.handleSubmit(onSubmit)();
@@ -823,6 +829,7 @@ export default function Page() {
                     {t("idpSubmit")}
                 </Button>
             </div>
+            </fieldset>
         </>
     );
 }

src/components/OrgSelector.tsx

@@ -29,6 +29,7 @@ import { usePathname, useRouter } from "next/navigation";
 import { useMemo, useState } from "react";
 import { useUserContext } from "@app/hooks/useUserContext";
 import { useTranslations } from "next-intl";
+import { build } from "@server/build";
 
 interface OrgSelectorProps {
     orgId?: string;
@@ -50,6 +51,11 @@ export function OrgSelector({
 
     const selectedOrg = orgs?.find((org) => org.orgId === orgId);
 
+    let canCreateOrg = !env.flags.disableUserCreateOrg || user.serverAdmin;
+    if (build === "saas" && user.type !== "internal") {
+        canCreateOrg = false;
+    }
+
     const sortedOrgs = useMemo(() => {
         if (!orgs?.length) return orgs ?? [];
         return [...orgs].sort((a, b) => {
@@ -161,7 +167,7 @@ export function OrgSelector({
                         </CommandGroup>
                     </CommandList>
                 </Command>
-                {(!env.flags.disableUserCreateOrg || user.serverAdmin) && (
+                {canCreateOrg && (
                     <div className="p-2 border-t border-border">
                         <Button
                             variant="ghost"