chore: add cancellablePromise method to Context

ydb-platform · May 8, 2024 · 2d1960e · 2d1960e
1 parent 3cf7a37
commit 2d1960e
Show file tree

Hide file tree

Showing 6 changed files with 97 additions and 10 deletions.
diff --git a/src/__tests__/unit/context.test.ts b/src/__tests__/unit/context.test.ts
@@ -1,5 +1,5 @@
 // @ts-ignore
-import {Context, setContextIdGenerator} from '../../context/Context';
+import {Context, setContextIdGenerator} from '../../context';
 // @ts-ignore
 import {ensureContext} from '../../context/ensureContext';
 // @ts-ignore
@@ -69,12 +69,12 @@ describe('Context', () => {
             expect(done2).toBeUndefined();
 
             const {ctx: ctx3, cancel: cancel3, dispose: dispose3, done: done3} =
-                ctx1.createChild();
+                ctx1.createChild({force: true});
             expect(ctx3).toBeDefined();
             expect(ctx3.err).toBeUndefined();
             expect(ctx3.onCancel).toBeDefined(); // because parent ctx has cancel
             expect(cancel3).toBeUndefined(); // cancel was not requested through options
-            expect(dispose3).toBeDefined(); // dispose cancel chain
+            expect(dispose3).toBeDefined(); // dispose cancel
             expect(done3).toBeUndefined();
 
             const testCancel = new Error('Test cancel');
@@ -295,6 +295,56 @@ describe('Context', () => {
         expect(ctx2[Symbol1]).toBe('aaa');
     });
 
+    it('keep using old context if possible', () => {
+        const {ctx} = Context.createNew();
+        expect(ctx.createChild().ctx).toBe(ctx);
+
+        expect(ctx.createChild({}).ctx).toBe(ctx);
+        expect(ctx.createChild({timeout: -1}).ctx).toBe(ctx);
+        expect(ctx.createChild({timeout: undefined}).ctx).toBe(ctx);
+
+        expect(ctx.createChild({cancel: true}).ctx).not.toBe(ctx);
+        expect(ctx.createChild({cancel: false}).ctx).not.toBe(ctx);
+        expect(ctx.createChild({timeout: 12}).ctx).not.toBe(ctx);
+        expect(ctx.createChild({done: true}).ctx).not.toBe(ctx);
+    });
+
+    for (const scenario of ['ok', 'failed', 'cancel'])
+        it(`'cancellablePromise: scenario: ${scenario}`, async () => {
+            let promiseResolve: (value: unknown) => void, promiseReject: (reason?: any) => void;
+            const promise = new Promise((resolve, reject) => {
+                promiseResolve = resolve;
+                promiseReject = reject;
+            })
+            const {ctx} = Context.createNew();
+            expect(ctx.cancellablePromise(promise)).toBe(promise);
+
+            const {ctx: ctx2, cancel} = ctx.createChild({
+                cancel: true,
+            });
+
+            const promise2 = ctx2.cancellablePromise(promise);
+            expect(promise2).not.toBe(promise);
+
+            switch (scenario) {
+                case 'ok': {
+                    promiseResolve!(12);
+                    expect(await promise2).toBe(12);
+                }
+                    break;
+                case 'failed': {
+                    promiseReject!(new Error('test'));
+                    await expect(promise2).rejects.toThrow('test');
+                }
+                    break;
+                case 'cancel': {
+                    cancel!();
+                    await expect(promise2).rejects.toThrow('Unknown');
+                }
+                    break;
+            }
+        });
+
     it('make 100% coverage', () => {
         {
             const {ctx} = Context.createNew({id: 'test'});

diff --git a/src/__tests__/unit/retryer.test.ts b/src/__tests__/unit/retryer.test.ts
@@ -104,15 +104,20 @@ describe('retryer', () => {
                     false /* error comes thru RetryLambdaResult */
                 ])
                     for (const isIdempotentOp of [false, true]) {
-                        // skip senseless tests
+                        // when backoff is not specified, nonIdempotent && idempotent do not affect the test
+                        if (backoff === null && !(nonIdempotent && idempotent)) continue;
+                        // makes no sense to retry non-idempotent operations, while do not retry idempotent one
                         if (nonIdempotent && !idempotent) continue;
+                        // with simply thrown error, the information that operation is idempotent or not is not available
                         if (simpleError && isIdempotentOp) continue;
+
                         const testName = `retry: ` +
                             `backoff: ${backoff === null ? null : ['No', 'Fast', 'Slow'][backoff]}; ` +
                             `nonIdempotent: ${Number(nonIdempotent)}; idempotent: ${Number(idempotent)}; ` +
                             `simpleError: ${simpleError}; isIdempotentOp: ${Number(isIdempotentOp)}`;
                         // leave the only test, if specified
                         if (ONLY_TEST && testName !== ONLY_TEST) continue;
+
                         it(testName, async () => {
                             const {ctx} = Context.createNew();
                             // @ts-ignore
@@ -237,4 +242,7 @@ describe('retryer', () => {
 // it('stop on context done', async () => {
 //
 // });
+// it('limit by count for legacy', async () => {
+//
+// });
 })
diff --git a/src/context/Context.ts → src/context/context.ts b/src/context/Context.ts → src/context/context.ts
@@ -41,6 +41,10 @@ interface IContextOpts {
      * cancel context after timeout.
      */
     timeout?: number,
+    /**
+     * Force creation of child context, even if there is no sufficient need.
+     */
+    force?: boolean,
 }
 
 export type CtxDone = () => void;
@@ -113,9 +117,17 @@ export class Context {
 
     /**
      * Creates a child context from the this one.
+     *
+     * Note: If there are no sufficient requirements for a new context the parent context
+     * will be keep using.
      */
     public createChild(opts: IContextOpts = {}): IContextCreateResult {
         if (opts.id) throw new Error('This method cannot change the context id');
+        if (!(
+            opts.hasOwnProperty('cancel') ||
+            opts.timeout! > 0 ||
+            opts.done
+        ) && !opts.force) return {ctx: this};
         const ctx = Object.create(this) as Context;
         const originOpts = opts;
         if (this.onCancel)
@@ -137,6 +149,23 @@ export class Context {
         return res;
     }
 
+    /**
+     * Makes a promise cancellable through context, if the context allows cancel or has a timeout.
+     */
+    public cancellablePromise<T>(promise: Promise<T>): Promise<T> {
+        if (!this.onCancel) return promise;
+        let cancelReject: (reason?: any) => void;
+        const cancelPromise = new Promise((_, reject) => {
+            cancelReject = reject;
+        });
+        const unsub = this.onCancel((cause) => {
+            cancelReject(cause);
+        });
+        return (Promise.race([promise, cancelPromise]) as Promise<T>).finally(() => {
+            unsub();
+        });
+    }
+
     /**
      * True if the reason for canceling is timeout.
      */
@@ -221,7 +250,7 @@ function makeContextCancellable(context: Context) {
 function setContextTimeout(timeout: number, cancel: OnCancelListener) {
     let timer: Timeout | undefined = setTimeout(() => {
         // An error is always created rather than using a constant to have an actual callstack
-        const err = new Error('Timeout');
+        const err = new Error(`Timeout: ${timeout} ms`);
         (err as any).cause = timeoutSymbol;
         cancel(err);
     }, timeout);
@@ -238,7 +267,7 @@ function setContextTimeout(timeout: number, cancel: OnCancelListener) {
 
 function createDone(cancel: OnCancelListener) {
     function done() {
-        // An error is always created rather than using a constant to have an actual callstack
+        // The error is always created rather than using a constant to have an actual callstack
         const err = new Error('Done');
         (err as any).cause = doneSymbol;
         cancel(err);
@@ -251,7 +280,7 @@ function initContext(this: Context, opts: IContextOpts) {
     const res: Omit<IContextCreateResult, 'ctx'> = {};
     let cancel: OnCancelListener;
     if (opts.cancel === true) res.cancel = cancel = makeContextCancellable(this);
-    if (opts.timeout && opts.timeout > 0) res.dispose = setContextTimeout(opts.timeout, cancel! || (cancel = makeContextCancellable(this)));
+    if (opts.timeout! > 0) res.dispose = setContextTimeout(opts.timeout!, cancel! || (cancel = makeContextCancellable(this)));
     if (opts.done) res.done = createDone(cancel! || makeContextCancellable(this));
     return res;
 }
diff --git a/src/context/ensure-context.ts b/src/context/ensure-context.ts
@@ -1,4 +1,4 @@
-import {Context} from "./Context";
+import {Context} from "./context";
 
 /**
  * Decorator that ensures:

diff --git a/src/context/has-object-context.ts b/src/context/has-object-context.ts
@@ -1,4 +1,4 @@
-import {Context} from "./Context";
+import {Context} from "./context";
 
 export interface HasObjectContext {
     /**

diff --git a/src/context/index.ts b/src/context/index.ts
@@ -1,4 +1,4 @@
-export {Context, CtxDispose, CtxCancel, CtxDone, CtxUnsubcribe, CtxIdGenerator, setContextIdGenerator} from './Context';
+export {Context, CtxDispose, CtxCancel, CtxDone, CtxUnsubcribe, CtxIdGenerator, setContextIdGenerator} from './context';
 export {ensureContext} from './ensure-context';
 export {HasObjectContext} from './has-object-context';
 export * as contextSymbols from './symbols';