diff --git a/libs/common/src/platform/misc/reference-counting/rc.spec.ts b/libs/common/src/platform/misc/reference-counting/rc.spec.ts index 094abfe3615..5fe44841285 100644 --- a/libs/common/src/platform/misc/reference-counting/rc.spec.ts +++ b/libs/common/src/platform/misc/reference-counting/rc.spec.ts @@ -25,6 +25,150 @@ describe("Rc", () => { rc = new Rc(value); }); + describe("use", () => { + describe("given Rc has not been marked for disposal", () => { + describe("given callback is synchronous", () => { + it("calls the callback", () => { + const spy = jest.fn(); + + rc.use(() => { + spy(); + }); + + expect(spy).toHaveBeenCalled(); + }); + + it("provides value in callback", () => { + rc.use((v) => { + expect(v).toBe(value); + }); + }); + + it("increases refCount while value is in use", () => { + rc.use(() => { + expect(rc["refCount"]).toBe(1); + }); + + expect(rc["refCount"]).toBe(0); + }); + + it("does not free value when refCount reaches 0 when not marked for disposal", () => { + rc.use(() => {}); + + expect(value.isFreed).toBe(false); + }); + + it("frees value directly when marked for disposal if refCount is 0", () => { + rc.use(() => {}); + + rc.markForDisposal(); + + expect(value.isFreed).toBe(true); + }); + + it("frees value after refCount reaches 0 when rc is marked for disposal while in use", () => { + rc.use(() => { + rc.markForDisposal(); + expect(value.isFreed).toBe(false); + }); + + expect(value.isFreed).toBe(true); + }); + + it("throws error when trying to take a disposed reference", () => { + rc.markForDisposal(); + + expect(() => rc.use(() => {})).toThrow(); + }); + }); + + describe("given callback is asynchronous", () => { + it("awaits the callback", async () => { + const spy = jest.fn(); + + await rc.use(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + spy(); + }); + + expect(spy).toHaveBeenCalled(); + }); + + it("provides value in callback", async () => { + await rc.use(async (v) => { + expect(v).toBe(value); + }); + }); + + it("increases refCount while value is in use", async () => { + let resolveCallback: () => void; + const promise = new Promise((resolve) => { + resolveCallback = resolve; + }); + + const usePromise = rc.use(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + await promise; + }); + + // should be 1 because the callback has not resolved yet + expect(rc["refCount"]).toBe(1); + + resolveCallback(); + await usePromise; + + // should be 0 because the callback has resolved + expect(rc["refCount"]).toBe(0); + }); + + it("does not free value when refCount reaches 0 when not marked for disposal", async () => { + await rc.use(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + }); + + expect(value.isFreed).toBe(false); + }); + + it("frees value directly when marked for disposal if refCount is 0", async () => { + await rc.use(async () => {}); + + rc.markForDisposal(); + + expect(value.isFreed).toBe(true); + }); + + it("frees value after refCount reaches 0 when rc is marked for disposal while in use", async () => { + let resolveCallback: () => void; + const promise = new Promise((resolve) => { + resolveCallback = resolve; + }); + + const usePromise = rc.use(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + await promise; + }); + + rc.markForDisposal(); + + // should not be freed yet because the callback has not resolved + expect(value.isFreed).toBe(false); + + resolveCallback(); + await usePromise; + + // should be freed because the callback has resolved + expect(value.isFreed).toBe(true); + }); + + it("throws error when trying to take a disposed reference", async () => { + rc.markForDisposal(); + + await expect(async () => await rc.use(async () => {})).rejects.toThrow(); + }); + }); + }); + }); + it("should increase refCount when taken", () => { rc.take(); diff --git a/libs/common/src/platform/misc/reference-counting/rc.ts b/libs/common/src/platform/misc/reference-counting/rc.ts index 9be102b43d3..3ce77d19ed6 100644 --- a/libs/common/src/platform/misc/reference-counting/rc.ts +++ b/libs/common/src/platform/misc/reference-counting/rc.ts @@ -2,6 +2,8 @@ import { UsingRequired } from "../using-required"; export type Freeable = { free: () => void }; +type UseReturnValue = T extends (value: any) => Promise ? Promise : void; + /** * Reference counted disposable value. * This class is used to manage the lifetime of a value that needs to be @@ -16,6 +18,39 @@ export class Rc { this.value = value; } + /** + * Use the value in a callback. + * The callback will be called immediately if the value is not marked for disposal. + * If the callback returns a promise, the promise will be awaited. + * + * @example + * ```typescript + * function someFunction(rc: Rc) { + * using reference = rc.use((value) => { + * value.doSomething(); + * }); + * } + * ``` + * + * @param fn The callback to call with the value. + */ + use unknown>(fn: Callback): UseReturnValue { + if (this.markedForDisposal) { + throw new Error("Cannot use a value marked for disposal"); + } + + this.refCount++; + const maybePromise = fn(this.value); + + if (maybePromise instanceof Promise) { + return maybePromise.then(() => { + this.release(); + }) as UseReturnValue; + } else { + this.release(); + } + } + /** * Use this function when you want to use the underlying object. * This will guarantee that you have a reference to the object