Skip to content

Commit

Permalink
fix: Use WeakMap in DerivedStateProvider to separate user state c…
Browse files Browse the repository at this point in the history
…aches (#12866)

Bug fix for PM-15914 where switching users would incorrectly share cached
derived states. The `DerivedStateProvider` now uses a `WeakMap` to maintain
separate caches for each user's state `Observable`.

- Modifies `DefaultDerivedStateProvider` to use `WeakMap` for caching
- Each user's state `Observable` gets its own definition cache
- Added test to verify correct behavior during user switching
- Allows proper garbage collection of unused state caches

This fixes issues where:
- Users would see other users' cached states after switching accounts
- Derived states weren't properly isolated between users
- Cache keys didn't distinguish between different user states
  • Loading branch information
addisonbeck authored Jan 15, 2025
1 parent 6f018e1 commit 55e4b5e
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,14 @@ import { DerivedStateProvider } from "../derived-state.provider";
import { DefaultDerivedState } from "./default-derived-state";

export class DefaultDerivedStateProvider implements DerivedStateProvider {
private cache: Record<string, DerivedState<unknown>> = {};
/**
* The cache uses a WeakMap to maintain separate derived states per user.
* Each user's state Observable acts as a unique key, without needing to
* pass around `userId`. Also, when a user's state Observable is cleaned up
* (like during an account swap) their cache is automatically garbage
* collected.
*/
private cache = new WeakMap<Observable<unknown>, Record<string, DerivedState<unknown>>>();

constructor() {}

Expand All @@ -17,16 +24,22 @@ export class DefaultDerivedStateProvider implements DerivedStateProvider {
deriveDefinition: DeriveDefinition<TFrom, TTo, TDeps>,
dependencies: TDeps,
): DerivedState<TTo> {
let stateCache = this.cache.get(parentState$);
if (!stateCache) {
stateCache = {};
this.cache.set(parentState$, stateCache);
}

const cacheKey = deriveDefinition.buildCacheKey();
const existingDerivedState = this.cache[cacheKey];
const existingDerivedState = stateCache[cacheKey];
if (existingDerivedState != null) {
// I have to cast out of the unknown generic but this should be safe if rules
// around domain token are made
return existingDerivedState as DefaultDerivedState<TFrom, TTo, TDeps>;
}

const newDerivedState = this.buildDerivedState(parentState$, deriveDefinition, dependencies);
this.cache[cacheKey] = newDerivedState;
stateCache[cacheKey] = newDerivedState;
return newDerivedState;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { DeriveDefinition } from "../derive-definition";
import { StateDefinition } from "../state-definition";

import { DefaultDerivedState } from "./default-derived-state";
import { DefaultDerivedStateProvider } from "./default-derived-state.provider";

let callCount = 0;
const cleanupDelayMs = 10;
Expand Down Expand Up @@ -182,4 +183,29 @@ describe("DefaultDerivedState", () => {
expect(await firstValueFrom(observable)).toEqual(new Date(newDate));
});
});

describe("account switching", () => {
let provider: DefaultDerivedStateProvider;

beforeEach(() => {
provider = new DefaultDerivedStateProvider();
});

it("should provide a dedicated cache for each account", async () => {
const user1State$ = new Subject<string>();
const user1Derived = provider.get(user1State$, deriveDefinition, deps);
const user1Emissions = trackEmissions(user1Derived.state$);

const user2State$ = new Subject<string>();
const user2Derived = provider.get(user2State$, deriveDefinition, deps);
const user2Emissions = trackEmissions(user2Derived.state$);

user1State$.next("2015-12-30");
user2State$.next("2020-12-29");
await awaitAsync();

expect(user1Emissions).toEqual([new Date("2015-12-30")]);
expect(user2Emissions).toEqual([new Date("2020-12-29")]);
});
});
});

0 comments on commit 55e4b5e

Please sign in to comment.