WPB-22967: migrate budget store to PostGreSQL

changelog.d/5-internal/WPB-22967

@@ -0,0 +1 @@
+Migration of budget store from cassandra to postgres

charts/wire-server/values.yaml

@@ -96,6 +96,7 @@ galley:
       blockList: cassandra
       passwordReset: cassandra
       customBackend: cassandra
+      budget: cassandra
     settings:
       httpPoolSize: 128
       maxTeamSize: 10000

docs/src/developer/reference/config-options.md

@@ -2332,6 +2332,7 @@ galley:
       blockList: postgresql
       passwordReset: postgresql
       customBackend: postgresql
+      budget: postgresql
 background-worker:
   config:
     migrateConversations: false
@@ -2375,6 +2376,7 @@ The current settings and their background-worker flags are:
 - `blockList` -> `migrateBlockList`
 - `passwordReset` -> `migratePasswordReset`
 - `customBackend` -> `migrateCustomBackend`
+- `budget` -> none (ephemeral TTL'd rate-limit counters; no backfill worker, no `wire_budget_migration_finished` metric — Cassandra rows age out on their own while new writes go to PostgreSQL under `migration-to-postgresql`). Because there is no backfill, keep `budget` on `migration-to-postgresql` for at least the longest configured `limitFailedLogins.timeout` (with PostgreSQL mirror writes healthy) before cutting over to `postgresql`; otherwise counters that only exist in Cassandra are dropped and affected users get one extra rate-limit window.
 
 **Migration pattern per migration setting**
 
@@ -2398,6 +2400,7 @@ The current settings and their background-worker flags are:
          blockList: cassandra
          passwordReset: cassandra
          customBackend: cassandra
+         budget: cassandra
    background-worker:
      config:
        migrateConversations: false
@@ -2471,6 +2474,7 @@ The current settings and their background-worker flags are:
          blockList: cassandra
          passwordReset: postgresql
          customBackend: postgresql
+         budget: postgresql
    background-worker:
      config:
        migrateConversations: false

hack/helm_vars/common.yaml.gotmpl

@@ -22,6 +22,7 @@ passwordReset: {{ $preferredStore }}
 customBackendStore: {{ $preferredStore }}
 userStore: {{ $preferredStore }}
 blockListStore: {{ $preferredStore }}
+budgetStore: {{ $preferredStore }}
 
 {{- if (eq (env "UPLOAD_XML_S3_BASE_URL") "") }}
 uploadXml: {}

hack/helm_vars/wire-server/values.yaml.gotmpl

@@ -313,6 +313,7 @@ galley:
       customBackend: {{ .Values.customBackendStore }}
       user: {{ .Values.userStore }}
       blockList: {{ .Values.blockListStore }}
+      budget: {{ .Values.budgetStore }}
     settings:
       maxConvAndTeamSize: 16
       maxTeamSize: 32

integration/integration.cabal

@@ -178,6 +178,7 @@ library
     Test.Meetings
     Test.MessageTimer
     Test.Migration.BlockList
+    Test.Migration.Budget
     Test.Migration.Conversation
     Test.Migration.ConversationCodes
     Test.Migration.CustomBackend

integration/test/Test/Migration/Budget.hs

@@ -0,0 +1,114 @@
+-- This file is part of the Wire Server implementation.
+--
+-- Copyright (C) 2026 Wire Swiss GmbH <opensource@wire.com>
+--
+-- This program is free software: you can redistribute it and/or modify it under
+-- the terms of the GNU Affero General Public License as published by the Free
+-- Software Foundation, either version 3 of the License, or (at your option) any
+-- later version.
+--
+-- This program is distributed in the hope that it will be useful, but WITHOUT
+-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+-- details.
+--
+-- You should have received a copy of the GNU Affero General Public License along
+-- with this program. If not, see <https://www.gnu.org/licenses/>.
+
+module Test.Migration.Budget (testBudgetMigration) where
+
+import API.Brig (login)
+import API.Common (defPassword)
+import Control.Monad.Codensity
+import Control.Monad.Reader
+import qualified Data.ByteString.Char8 as BSChar8
+import SetupHelpers
+import Testlib.Prelude
+import Testlib.ResourcePool
+import Text.Read (readMaybe)
+
+-- | Drives the login rate-limit (budget) store through the cutover lifecycle
+-- (cassandra -> migration-to-postgresql -> postgresql). Budget rows are
+-- ephemeral TTL'd counters, so unlike the other store migrations there is no
+-- backfill worker: Cassandra rows age out on their own while new writes flow
+-- to Postgres. The 60s login-block timeout is what keeps rows alive across
+-- backend restarts within this test.
+--
+-- P2 proves reads stay Cassandra-sourced during dual-write (u1's Cassandra row
+-- still blocks). P3 proves a row dual-written under migration-to-postgresql is
+-- visible to the pure Postgres interpreter (u2 still blocked), and that a
+-- Cassandra-only row does NOT survive cutover (u1 unblocked) — the documented
+-- no-backfill behavior.
+testBudgetMigration :: (HasCallStack) => App ()
+testBudgetMigration = do
+  resourcePool <- asks (.resourcePool)
+  runCodensity (acquireResources 1 resourcePool) $ \[backend] -> do
+    let domain = backend.berDomain
+
+    -- P1 cassandra: budget rows are written only to Cassandra
+    u1 <-
+      runCodensity (startDynamicBackend backend (conf "cassandra")) $ \_ -> do
+        u1 <- randomUser domain def
+        exhaustLogins domain u1
+        pure u1
+
+    -- P2 migration-to-postgresql: reads still come from Cassandra (source of
+    -- truth), writes are mirrored to Postgres
+    u2 <-
+      runCodensity (startDynamicBackend backend (conf "migration-to-postgresql")) $ \_ -> do
+        -- u1 is still blocked: its budget row lives in Cassandra and reads
+        -- are served from there during migration.
+        assertBlocked domain u1
+        u2 <- randomUser domain def
+        exhaustLogins domain u2
+        pure u2
+
+    -- P3 postgresql: reads and writes are served exclusively from Postgres
+    runCodensity (startDynamicBackend backend (conf "postgresql")) $ \_ -> do
+      -- u2 is still blocked: the dual-written row is visible to the
+      -- Postgres-only interpreter.
+      assertBlocked domain u2
+      -- u1 is NOT blocked: its budget row only ever existed in Cassandra and
+      -- was never backfilled (no backfill worker for ephemeral data).
+      assertUnblocked domain u1
+      -- Write and read via the pure Postgres interpreter.
+      u3 <- randomUser domain def
+      exhaustLogins domain u3
+  where
+    conf db =
+      def
+        { brigCfg =
+            setField @_ @Int "optSettings.setLimitFailedLogins.timeout" 60
+              >=> setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
+              >=> setField "postgresMigration.budget" db
+        }
+
+-- | Exhaust the failed-login budget of a user: 5 wrong passwords get plain
+-- login-failure responses, and the 6th attempt (even with the correct
+-- password) is blocked with a @Retry-After@ header.
+exhaustLogins :: (HasCallStack) => String -> Value -> App ()
+exhaustLogins domain u = do
+  email <- asString (u %. "email")
+  for_ [1 :: Int .. 5] $ \_ ->
+    login domain email "wrong-password" `bindResponse` \resp -> do
+      resp.status `shouldMatchInt` 403
+      resp.json %. "label" `shouldMatch` "invalid-credentials"
+  assertBlockedEmail domain email
+
+assertBlocked :: (HasCallStack) => String -> Value -> App ()
+assertBlocked domain u = asString (u %. "email") >>= assertBlockedEmail domain
+
+assertBlockedEmail :: (HasCallStack) => String -> String -> App ()
+assertBlockedEmail domain email =
+  login domain email defPassword `bindResponse` \resp -> do
+    resp.status `shouldMatchInt` 403
+    resp.json %. "label" `shouldMatch` "client-error"
+    -- The Retry-After header distinguishes LoginBlocked from a plain
+    -- login-failure.
+    let Just (_retryAfter :: Int) = readMaybe . BSChar8.unpack =<< lookup (fromString "Retry-After") resp.headers
+    pure ()
+
+assertUnblocked :: (HasCallStack) => String -> Value -> App ()
+assertUnblocked domain u = do
+  email <- asString (u %. "email")
+  login domain email defPassword >>= assertSuccess

libs/wire-subsystems/postgres-migrations/20261020201733-budget.sql

@@ -0,0 +1,10 @@
+CREATE TABLE budget (
+  key text NOT NULL,
+  budget int4 NOT NULL,
+  expires_at timestamptz NOT NULL,
+  PRIMARY KEY (key)
+);
+
+-- Serves the (future) cleanup DELETE ... WHERE expires_at <= now(); lookups
+-- are PK-only, and expired rows are filtered by expires_at > now().
+CREATE INDEX budget_expires_at_idx ON budget (expires_at);

libs/wire-subsystems/src/Wire/BudgetStore/DualWrite.hs

@@ -0,0 +1,43 @@
+-- This file is part of the Wire Server implementation.
+--
+-- Copyright (C) 2026 Wire Swiss GmbH <opensource@wire.com>
+--
+-- This program is free software: you can redistribute it and/or modify it under
+-- the terms of the GNU Affero General Public License as published by the Free
+-- Software Foundation, either version 3 of the License, or (at your option) any
+-- later version.
+--
+-- This program is distributed in the hope that it will be useful, but WITHOUT
+-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+-- details.
+--
+-- You should have received a copy of the GNU Affero General Public License along
+-- with this program. If not, see <https://www.gnu.org/licenses/>.
+
+module Wire.BudgetStore.DualWrite
+  ( interpretBudgetStoreToCassandraAndPostgres,
+  )
+where
+
+import Cassandra (MonadClient)
+import Imports
+import Polysemy
+import Wire.BudgetStore
+import Wire.BudgetStore.Cassandra (budgetStoreToCassandra)
+import Wire.BudgetStore.Postgres (interpretBudgetStoreToPostgres)
+import Wire.Postgres (PGConstraints)
+
+-- | Cassandra is the source of truth during migration (reads stay Cassandra to
+-- preserve in-flight TTL semantics); writes are mirrored to Postgres. Budget
+-- rows are ephemeral TTL'd rate-limit counters, so there is no backfill loop.
+interpretBudgetStoreToCassandraAndPostgres ::
+  forall m r.
+  (MonadClient m, Member (Embed m) r, PGConstraints r) =>
+  InterpreterFor BudgetStore r
+interpretBudgetStoreToCassandraAndPostgres = interpret $ \case
+  LookupBudget k ->
+    budgetStoreToCassandra @m (lookupBudget k)
+  InsertBudget k b -> do
+    budgetStoreToCassandra @m (insertBudget k b)
+    interpretBudgetStoreToPostgres (insertBudget k b)

libs/wire-subsystems/src/Wire/BudgetStore/Postgres.hs

@@ -0,0 +1,66 @@
+-- This file is part of the Wire Server implementation.
+--
+-- Copyright (C) 2026 Wire Swiss GmbH <opensource@wire.com>
+--
+-- This program is free software: you can redistribute it and/or modify it under
+-- the terms of the GNU Affero General Public License as published by the Free
+-- Software Foundation, either version 3 of the License, or (at your option) any
+-- later version.
+--
+-- This program is distributed in the hope that it will be useful, but WITHOUT
+-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+-- details.
+--
+-- You should have received a copy of the GNU Affero General Public License along
+-- with this program. If not, see <https://www.gnu.org/licenses/>.
+
+module Wire.BudgetStore.Postgres
+  ( interpretBudgetStoreToPostgres,
+  )
+where
+
+import Hasql.Statement (Statement)
+import Hasql.TH (maybeStatement, resultlessStatement)
+import Imports
+import Polysemy
+import Wire.BudgetStore
+import Wire.Postgres
+
+interpretBudgetStoreToPostgres ::
+  (PGConstraints r) =>
+  InterpreterFor BudgetStore r
+interpretBudgetStoreToPostgres =
+  interpret $ \case
+    LookupBudget k -> lookupBudgetP k
+    InsertBudget k b -> insertBudgetP k b
+
+lookupBudgetP :: (PGConstraints r) => BudgetKey -> Sem r (Maybe Budget)
+lookupBudgetP (BudgetKey k') =
+  fmap mk <$> runStatement k' lookupBudgetStatement
+  where
+    mk (val, ttl) = Budget (fromIntegral ttl) val
+
+lookupBudgetStatement :: Statement Text (Maybe (Int32, Int32))
+lookupBudgetStatement =
+  -- Mirrors Cassandra's ttl(budget): remaining whole seconds, 0 once the row
+  -- is expired. GREATEST(0, ...) only guards the theoretical now() == row
+  -- boundary; the expires_at > now() filter already excludes expired rows.
+  [maybeStatement|
+      SELECT budget :: int4, GREATEST(0, FLOOR(EXTRACT(EPOCH FROM (expires_at - now())))) :: int4
+      FROM budget WHERE key = ($1 :: text) AND expires_at > now()
+    |]
+
+insertBudgetP :: (PGConstraints r) => BudgetKey -> Budget -> Sem r ()
+insertBudgetP (BudgetKey k') (Budget ttl val) =
+  runStatement (k', val, round ttl) insertBudgetStatement
+
+insertBudgetStatement :: Statement (Text, Int32, Int32) ()
+insertBudgetStatement =
+  [resultlessStatement|
+      INSERT INTO budget (key, budget, expires_at)
+      VALUES (($1 :: text), ($2 :: int4), now() + make_interval(secs => $3 :: int))
+      ON CONFLICT (key) DO UPDATE
+      SET budget = ($2 :: int4),
+          expires_at = now() + make_interval(secs => $3 :: int)
+    |]

libs/wire-subsystems/src/Wire/PostgresMigrationOpts.hs

@@ -59,7 +59,8 @@ data PostgresMigrationOpts = PostgresMigrationOpts
     user :: StorageLocation,
     blockList :: StorageLocation,
     passwordReset :: StorageLocation,
-    customBackend :: StorageLocation
+    customBackend :: StorageLocation,
+    budget :: StorageLocation
   }
   deriving (Show)
 
@@ -74,3 +75,4 @@ instance FromJSON PostgresMigrationOpts where
       <*> o .: "blockList"
       <*> o .: "passwordReset"
       <*> o .: "customBackend"
+      <*> o .: "budget"

libs/wire-subsystems/wire-subsystems.cabal

@@ -252,6 +252,8 @@ library
     Wire.BrigAPIAccess.Rpc
     Wire.BudgetStore
     Wire.BudgetStore.Cassandra
+    Wire.BudgetStore.DualWrite
+    Wire.BudgetStore.Postgres
     Wire.ClientStore
     Wire.ClientStore.Cassandra
     Wire.ClientStore.DynamoDB

postgres-schema.sql

@@ -1370,6 +1370,19 @@ CREATE TABLE public.bot_conv (
 
 ALTER TABLE public.bot_conv OWNER TO "wire-server";
 
+--
+-- Name: budget; Type: TABLE; Schema: public; Owner: wire-server
+--
+
+CREATE TABLE public.budget (
+    key text NOT NULL,
+    budget integer NOT NULL,
+    expires_at timestamp with time zone NOT NULL
+);
+
+
+ALTER TABLE public.budget OWNER TO "wire-server";
+
 --
 -- Name: collaborators; Type: TABLE; Schema: public; Owner: wire-server
 --
@@ -1989,6 +2002,14 @@ ALTER TABLE ONLY public.bot_conv
     ADD CONSTRAINT bot_conv_pkey PRIMARY KEY (id);
 
 
+--
+-- Name: budget budget_pkey; Type: CONSTRAINT; Schema: public; Owner: wire-server
+--
+
+ALTER TABLE ONLY public.budget
+    ADD CONSTRAINT budget_pkey PRIMARY KEY (key);
+
+
 --
 -- Name: collaborators collaborators_pkey; Type: CONSTRAINT; Schema: public; Owner: wire-server
 --
@@ -2533,6 +2554,13 @@ CREATE INDEX bot_conv_conv_idx ON public.bot_conv USING btree (conv);
 CREATE INDEX bot_conv_team_idx ON public.bot_conv USING btree (conv_team);
 
 
+--
+-- Name: budget_expires_at_idx; Type: INDEX; Schema: public; Owner: wire-server
+--
+
+CREATE INDEX budget_expires_at_idx ON public.budget USING btree (expires_at);
+
+
 --
 -- Name: collaborators_team_id_idx; Type: INDEX; Schema: public; Owner: wire-server
 --

services/background-worker/background-worker.integration.yaml

@@ -99,3 +99,4 @@ postgresMigration:
   user: postgresql
   blockList: postgresql
   customBackend: postgresql
+  budget: postgresql