From 08a56ef4802a2b953571b68e518b482336a272cb Mon Sep 17 00:00:00 2001
From: Sam Rose <hello@samwho.dev>
Date: Wed, 14 Aug 2024 15:41:26 +0100
Subject: [PATCH] Cache feature flags per-request, set default values flag to
 false by default.

---
 packages/backend-core/src/context/types.ts    |  3 ++
 packages/backend-core/src/features/index.ts   | 35 +++++++++++++++----
 .../src/features/tests/features.spec.ts       |  4 ++-
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/packages/backend-core/src/context/types.ts b/packages/backend-core/src/context/types.ts
index 6694320978..fe6072e85c 100644
--- a/packages/backend-core/src/context/types.ts
+++ b/packages/backend-core/src/context/types.ts
@@ -18,4 +18,7 @@ export type ContextMap = {
     oauthClient: OAuth2Client
     clients: Record<string, GoogleSpreadsheet>
   }
+  featureFlagCache?: {
+    [key: string]: Record<string, any>
+  }
 }
diff --git a/packages/backend-core/src/features/index.ts b/packages/backend-core/src/features/index.ts
index a95ef5c714..b03641ec90 100644
--- a/packages/backend-core/src/features/index.ts
+++ b/packages/backend-core/src/features/index.ts
@@ -18,6 +18,10 @@ export function init(opts?: PostHogOptions) {
   }
 }
 
+export function shutdown() {
+  posthog?.shutdown()
+}
+
 export abstract class Flag<T> {
   static boolean(defaultValue: boolean): Flag<boolean> {
     return new BooleanFlag(defaultValue)
@@ -87,7 +91,14 @@ class NumberFlag extends Flag<number> {
 }
 
 export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
-  constructor(private readonly flagSchema: T) {}
+  // This is used to safely cache flags sets in the current request context.
+  // Because multiple sets could theoretically exist, we don't want the cache of
+  // one to leak into another.
+  private readonly setId: string
+
+  constructor(private readonly flagSchema: T) {
+    this.setId = crypto.randomUUID()
+  }
 
   defaults(): FlagValues<T> {
     return Object.keys(this.flagSchema).reduce((acc, key) => {
@@ -119,6 +130,15 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
 
   async fetch(ctx?: UserCtx): Promise<FlagValues<T>> {
     return await tracer.trace("features.fetch", async span => {
+      const requestContext = context.getCurrentContext()
+      const cachedFlags = requestContext?.featureFlagCache?.[this.setId] as
+        | FlagValues<T>
+        | undefined
+      if (cachedFlags) {
+        span?.addTags({ fromCache: true })
+        return cachedFlags
+      }
+
       const tags: Record<string, any> = {}
       const flagValues = this.defaults()
       const currentTenantId = context.getTenantId()
@@ -187,10 +207,7 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
       tags[`identity.tenantId`] = identity?.tenantId
       tags[`identity._id`] = identity?._id
 
-      // Until we're confident this performs well, we're only enabling it in QA
-      // and test environments.
-      const usePosthog = env.isTest() || env.isQA()
-      if (usePosthog && posthog && identity?.type === IdentityType.USER) {
+      if (posthog && identity?.type === IdentityType.USER) {
         tags[`readFromPostHog`] = true
 
         const personProperties: Record<string, string> = {}
@@ -204,7 +221,6 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
             personProperties,
           }
         )
-        console.log("posthog flags", JSON.stringify(posthogFlags))
 
         for (const [name, value] of Object.entries(posthogFlags.featureFlags)) {
           if (!this.isFlagName(name)) {
@@ -236,6 +252,11 @@ export class FlagSet<V extends Flag<any>, T extends { [key: string]: V }> {
         }
       }
 
+      if (requestContext) {
+        requestContext.featureFlagCache ??= {}
+        requestContext.featureFlagCache[this.setId] = flagValues
+      }
+
       for (const [key, value] of Object.entries(flagValues)) {
         tags[`flags.${key}.value`] = value
       }
@@ -255,5 +276,5 @@ export const flags = new FlagSet({
   GOOGLE_SHEETS: Flag.boolean(false),
   USER_GROUPS: Flag.boolean(false),
   ONBOARDING_TOUR: Flag.boolean(false),
-  DEFAULT_VALUES: Flag.boolean(true),
+  DEFAULT_VALUES: Flag.boolean(false),
 })
diff --git a/packages/backend-core/src/features/tests/features.spec.ts b/packages/backend-core/src/features/tests/features.spec.ts
index 1b6a6e041b..1e8e25654a 100644
--- a/packages/backend-core/src/features/tests/features.spec.ts
+++ b/packages/backend-core/src/features/tests/features.spec.ts
@@ -1,5 +1,5 @@
 import { IdentityContext, IdentityType, UserCtx } from "@budibase/types"
-import { Flag, FlagSet, FlagValues, init } from "../"
+import { Flag, FlagSet, FlagValues, init, shutdown } from "../"
 import { context } from "../.."
 import environment, { withEnv } from "../../environment"
 import nodeFetch from "node-fetch"
@@ -197,6 +197,8 @@ describe("feature flags", () => {
             throw new Error("No expected value")
           }
         })
+
+        shutdown()
       })
     }
   )