Skip to content

Commit 3176d09

Browse files
ParidelPooyaPooya Paridel
andcommitted
refactor: replace Telemetry with typed contexts and modernize logger
- Remove Telemetry interface and StructuredLogger in favor of OperationContext pattern - Add StepContext, WaitForConditionContext, and WaitForCallbackContext as type aliases - Replace structured-logger with context-logger for better step-aware logging - Extract createDefaultLogger from durable-context to utils/logger/default-logger - Unify attempt numbering in wait-for-condition handler (remove duplicate variables) - Use step's built-in logger in wait-for-callback handler instead of manual creation - Update all function signatures from Telemetry to specific context types - Rename configureLogger to setCustomLogger for consistency - Add tests for new logger utilities - use execution_arn instead of execution_id By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license. ### Checklist - [Y] I have filled out every section of the PR template - [Y] I have thoroughly tested this change ### Testing #### Unit Tests Yes Co-authored-by: Pooya Paridel <parpooya@amazon.com>
1 parent a64ecbc commit 3176d09

File tree

18 files changed

+521
-496
lines changed

18 files changed

+521
-496
lines changed

packages/lambda-durable-functions-sdk-js/bundle-size-history.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,4 +194,4 @@
194194
"size": 214426,
195195
"gitCommit": "9801dc79f2a4a980f971e2de1b85aa23b4c7a579"
196196
}
197-
]
197+
]

packages/lambda-durable-functions-sdk-js/src/context/durable-context/durable-context.test.ts

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,8 @@ describe("Durable Context", () => {
136136
mockExecutionContext,
137137
mockCheckpointHandler,
138138
mockParentContext,
139-
expect.any(Function),
139+
expect.any(Function), // createStepId
140+
expect.any(Function), // createContextLogger
140141
);
141142
expect(mockStepHandler).toHaveBeenCalledWith("test-step", stepFn, options);
142143
});
@@ -428,7 +429,7 @@ describe("Durable Context", () => {
428429
mockExecutionContext,
429430
mockParentContext,
430431
);
431-
expect(typeof durableContext.configureLogger).toBe("function");
432+
expect(typeof durableContext.setCustomLogger).toBe("function");
432433
});
433434

434435
it("should configure custom logger through DurableContext", () => {
@@ -446,13 +447,49 @@ describe("Durable Context", () => {
446447
};
447448

448449
// Configure custom logger
449-
durableContext.configureLogger(mockCustomLogger);
450+
durableContext.setCustomLogger(mockCustomLogger);
450451

451452
// Verify that the custom logger was set by checking the setCustomLogger was called
452453
// Since the step handler is mocked, we can't test the full integration here
453454
// but we can verify the method exists and doesn't throw
454455
expect(() =>
455-
durableContext.configureLogger(mockCustomLogger),
456+
durableContext.setCustomLogger(mockCustomLogger),
456457
).not.toThrow();
457458
});
459+
460+
it("should use default logger when no custom logger is set", () => {
461+
const consoleSpy = jest.spyOn(console, "log").mockImplementation();
462+
463+
// Temporarily restore the real step handler to test actual logger usage
464+
const originalCreateStepHandler = jest.requireActual(
465+
"../../handlers/step-handler/step-handler",
466+
).createStepHandler;
467+
(createStepHandler as jest.Mock).mockImplementationOnce(
468+
originalCreateStepHandler,
469+
);
470+
471+
const durableContext = createDurableContext(
472+
mockExecutionContext,
473+
mockParentContext,
474+
);
475+
476+
// This will use the real step handler which will call the context logger
477+
const stepFn = jest.fn().mockImplementation(async (ctx) => {
478+
// Use all logger methods to ensure full coverage of createDefaultLogger
479+
ctx.logger.log("custom", "test message");
480+
ctx.logger.info("info test");
481+
ctx.logger.error("error test", new Error("test"));
482+
ctx.logger.warn("warn test");
483+
ctx.logger.debug("debug test");
484+
return "result";
485+
});
486+
487+
// This should trigger the default logger
488+
durableContext.step("test-step", stepFn);
489+
490+
// Verify console.log was called
491+
expect(consoleSpy).toHaveBeenCalled();
492+
493+
consoleSpy.mockRestore();
494+
});
458495
});

packages/lambda-durable-functions-sdk-js/src/context/durable-context/durable-context.ts

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,29 @@ import { createMapHandler } from "../../handlers/map-handler/map-handler";
2929
import { createParallelHandler } from "../../handlers/parallel-handler/parallel-handler";
3030
import { createPromiseHandler } from "../../handlers/promise-handler/promise-handler";
3131
import { createConcurrentExecutionHandler } from "../../handlers/concurrent-execution-handler/concurrent-execution-handler";
32-
import { setCustomLogger } from "../../utils/logger/structured-logger";
32+
import { createContextLoggerFactory } from "../../utils/logger/context-logger";
33+
import { createDefaultLogger } from "../../utils/logger/default-logger";
3334

3435
export const createDurableContext = (
3536
executionContext: ExecutionContext,
3637
parentContext: Context,
3738
stepPrefix?: string,
3839
checkpointToken?: string,
3940
): DurableContext => {
41+
// Local logger state for this context instance
42+
let contextLogger: Logger | null = null;
43+
44+
// Local getter function for this context
45+
const getLogger = (): Logger => {
46+
return contextLogger || createDefaultLogger();
47+
};
48+
49+
// Create context logger factory
50+
const createContextLogger = createContextLoggerFactory(
51+
executionContext,
52+
getLogger,
53+
);
54+
4055
let stepCounter = 0;
4156
const checkpoint = createCheckpoint(executionContext, checkpointToken || "");
4257

@@ -55,6 +70,7 @@ export const createDurableContext = (
5570
checkpoint,
5671
parentContext,
5772
createStepId,
73+
createContextLogger,
5874
);
5975
return stepHandler(nameOrFn, fnOrOptions, maybeOptions);
6076
};
@@ -153,8 +169,8 @@ export const createDurableContext = (
153169

154170
const promise = createPromiseHandler(step);
155171

156-
const configureLogger = (logger: Logger): void => {
157-
setCustomLogger(logger);
172+
const setCustomLogger = (logger: Logger): void => {
173+
contextLogger = logger;
158174
};
159175

160176
return {
@@ -168,13 +184,14 @@ export const createDurableContext = (
168184
executionContext,
169185
checkpoint,
170186
createStepId,
187+
createContextLogger,
171188
),
172189
createCallback,
173190
waitForCallback,
174191
map,
175192
parallel,
176193
executeConcurrently,
177194
promise,
178-
configureLogger,
195+
setCustomLogger,
179196
};
180197
};

packages/lambda-durable-functions-sdk-js/src/handlers/step-handler/step-handler.test.ts

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,25 +67,36 @@ describe("Step Handler", () => {
6767
mockCheckpoint = createMockCheckpoint();
6868
mockParentContext = { awsRequestId: "mock-request-id" };
6969
createStepId = jest.fn().mockReturnValue("test-step-id");
70+
71+
const mockLogger = {
72+
log: jest.fn(),
73+
info: jest.fn(),
74+
error: jest.fn(),
75+
warn: jest.fn(),
76+
debug: jest.fn(),
77+
};
78+
const createMockEnrichedLogger = () => mockLogger;
79+
7080
stepHandler = createStepHandler(
7181
mockExecutionContext,
7282
mockCheckpoint,
7383
mockParentContext,
7484
createStepId,
85+
createMockEnrichedLogger,
7586
);
7687

7788
// Reset the mock for retryPresets.default
7889
(retryPresets.default as jest.Mock).mockReset();
7990
});
8091

81-
test("should execute step function with Telemetry", async () => {
92+
test("should execute step function with StepContext", async () => {
8293
const stepFn = jest.fn().mockResolvedValue("step-result");
8394

8495
const result = await stepHandler("test-step", stepFn);
8596

8697
expect(result).toBe("step-result");
8798
expect(stepFn).toHaveBeenCalledTimes(1);
88-
// Verify that Telemetry was passed to the step function
99+
// Verify that StepContext was passed to the step function
89100
expect(stepFn.mock.calls[0].length).toBe(1);
90101
expect(stepFn.mock.calls[0][0]).toHaveProperty("logger");
91102
});

packages/lambda-durable-functions-sdk-js/src/handlers/step-handler/step-handler.ts

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import {
55
RetryDecision,
66
StepSemantics,
77
OperationSubType,
8-
Telemetry,
8+
StepContext,
9+
Logger,
910
} from "../../types";
1011
import { Context } from "aws-lambda";
1112
import {
@@ -26,7 +27,6 @@ import {
2627
import { isUnrecoverableError } from "../../errors/unrecoverable-error/unrecoverable-error";
2728
import { OperationInterceptor } from "../../mocks/operation-interceptor";
2829
import { createErrorObjectFromError } from "../../utils/error-object/error-object";
29-
import { createStructuredLogger } from "../../utils/logger/structured-logger";
3030

3131
const waitForTimer = <T>(
3232
context: ExecutionContext,
@@ -39,14 +39,15 @@ const waitForTimer = <T>(
3939
reason: TerminationReason.RETRY_SCHEDULED,
4040
message: `Retry scheduled for ${name || stepId}`,
4141
});
42-
return new Promise<T>(() => { });
42+
return new Promise<T>(() => {});
4343
};
4444

4545
export const createStepHandler = (
4646
context: ExecutionContext,
4747
checkpoint: ReturnType<typeof createCheckpoint>,
4848
parentContext: Context,
4949
createStepId: () => string,
50+
createContextLogger: (stepId: string, attempt?: number) => Logger,
5051
) => {
5152
return async <T>(
5253
nameOrFn: string | undefined | StepFunc<T>,
@@ -158,7 +159,15 @@ export const createStepHandler = (
158159
// READY: Timer completed, execute step function
159160
// STARTED: Retry after error (AtLeastOncePerRetry semantics), execute step function
160161
// undefined: First execution, execute step function
161-
return executeStep(context, checkpoint, stepId, name, fn, options);
162+
return executeStep(
163+
context,
164+
checkpoint,
165+
stepId,
166+
name,
167+
fn,
168+
createContextLogger,
169+
options,
170+
);
162171
};
163172
};
164173

@@ -195,6 +204,7 @@ export const executeStep = async <T>(
195204
stepId: string,
196205
name: string | undefined,
197206
fn: StepFunc<T>,
207+
createContextLogger: (stepId: string, attempt?: number) => Logger,
198208
options?: StepConfig<T>,
199209
): Promise<T> => {
200210
// Determine step semantics (default to AT_LEAST_ONCE_PER_RETRY if not specified)
@@ -228,23 +238,19 @@ export const executeStep = async <T>(
228238
}
229239

230240
try {
231-
// Get current attempt number for logger
241+
// Get current attempt number for logger enrichment
232242
const stepData = context.getStepData(stepId);
233243
const currentAttempt = stepData?.StepDetails?.Attempt || 0;
234244

235-
// Create telemetry with logger
236-
const telemetry: Telemetry = {
237-
logger: createStructuredLogger({
238-
stepId,
239-
executionId: context.durableExecutionArn,
240-
attempt: currentAttempt,
241-
}),
245+
// Create step context with enriched logger
246+
const stepContext: StepContext = {
247+
logger: createContextLogger(stepId, currentAttempt),
242248
};
243249

244-
// Execute the step function with stepUtil
250+
// Execute the step function with stepContext
245251
const result = await OperationInterceptor.forExecution(
246252
context.durableExecutionArn,
247-
).execute(name, () => fn(telemetry));
253+
).execute(name, () => fn(stepContext));
248254

249255
// Serialize the result for consistency
250256
const serializedResult = await safeSerialize(
@@ -309,7 +315,7 @@ export const executeStep = async <T>(
309315
});
310316

311317
// Return a never-resolving promise to ensure the execution doesn't continue
312-
return new Promise<T>(() => { });
318+
return new Promise<T>(() => {});
313319
}
314320

315321
const stepData = context.getStepData(stepId);

packages/lambda-durable-functions-sdk-js/src/handlers/wait-for-callback-handler/wait-for-callback-handler.test.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,15 @@ describe("waitForCallback handler", () => {
1010
let stepIdCounter: number;
1111
let mockRunInChildContext: jest.Mock;
1212

13+
const createMockLogger = () => ({
14+
log: jest.fn(),
15+
info: jest.fn(),
16+
error: jest.fn(),
17+
warn: jest.fn(),
18+
debug: jest.fn(),
19+
});
20+
const createMockEnrichedLogger = () => createMockLogger();
21+
1322
beforeEach(() => {
1423
stepIdCounter = 0;
1524

packages/lambda-durable-functions-sdk-js/src/handlers/wait-for-callback-handler/wait-for-callback-handler.ts

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@ import {
55
CreateCallbackConfig,
66
DurableContext,
77
OperationSubType,
8-
Telemetry,
8+
WaitForCallbackContext,
9+
StepContext,
10+
Logger,
911
} from "../../types";
1012
import { log } from "../../utils/logger/logger";
1113

@@ -69,12 +71,17 @@ export const createWaitForCallbackHandler = (
6971
});
7072

7173
// Execute the submitter step (submitter is now mandatory)
72-
await childCtx.step(async (telemetry: Telemetry) => {
74+
await childCtx.step(async (stepContext: StepContext) => {
75+
// Use the step's built-in logger instead of creating a new one
76+
const callbackContext: WaitForCallbackContext = {
77+
logger: stepContext.logger,
78+
};
79+
7380
log(context.isVerbose, "📤", "Executing submitter:", {
7481
callbackId,
7582
name,
7683
});
77-
await submitter(callbackId, telemetry);
84+
await submitter(callbackId, callbackContext);
7885
log(context.isVerbose, "✅", "Submitter completed:", {
7986
callbackId,
8087
name,

packages/lambda-durable-functions-sdk-js/src/handlers/wait-for-condition-handler/wait-for-condition-handler.test.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,20 @@ describe("WaitForCondition Handler", () => {
6767
mockExecutionRunner,
6868
);
6969

70+
const mockLogger = {
71+
log: jest.fn(),
72+
info: jest.fn(),
73+
error: jest.fn(),
74+
warn: jest.fn(),
75+
debug: jest.fn(),
76+
};
77+
const createMockEnrichedLogger = () => mockLogger;
78+
7079
waitForConditionHandler = createWaitForConditionHandler(
7180
mockExecutionContext,
7281
mockCheckpoint,
7382
createStepId,
83+
createMockEnrichedLogger,
7484
);
7585
});
7686

0 commit comments

Comments
 (0)