fix(aws-serverless): Keep the Lambda extension polling past 300s invocations

packages/aws-serverless/src/lambda-extension/aws-lambda-extension.ts

@@ -10,6 +10,95 @@ import {
 } from '@sentry/core';
 import { DEBUG_BUILD } from './debug-build';
 
+const POLL_RETRY_BASE_MS = 100;
+const POLL_RETRY_MAX_MS = 5_000;
+/** Bounded so a permanently unreachable API exits instead of logging every 5s forever. */
+const POLL_MAX_CONSECUTIVE_FAILURES = 20;
+
+/** The body lands in an error message that a failing poll writes to the console. */
+const ERROR_BODY_MAX_LENGTH = 200;
+
+/**
+ * Detects a peer that went away without a FIN/RST, which a request deadline cannot do here: the
+ * poll is open across the environment's frozen idle time, which is unbounded, and a socket
+ * deadline runs on real time and would fire on thaw after a long idle — destroying a poll that
+ * was about to be answered. Keep-alive probes only travel while the environment is running.
+ */
+const POLL_KEEPALIVE_MS = 30_000;
+
+/** 408 and 429 are the retryable ones; the rest of 4xx means the poll itself is refused. */
+const RETRYABLE_CLIENT_ERRORS = [408, 429];
+
+interface ExtensionEvent {
+  eventType?: string;
+}
+
+interface ExtensionsApiResponse {
+  statusCode: number;
+  body: string;
+}
+
+export class ExtensionsApiError extends Error {
+  public constructor(
+    message: string,
+    public readonly statusCode: number,
+  ) {
+    super(message);
+    this.name = 'ExtensionsApiError';
+  }
+}
+
+/**
+ * Structural rather than `instanceof`: the check has to hold for an error that crossed a
+ * module boundary, and a transport failure carries `code`, never `statusCode`.
+ */
+function isClientError(err: unknown): boolean {
+  const statusCode = (err as { statusCode?: unknown } | null)?.statusCode;
+  return (
+    typeof statusCode === 'number' &&
+    statusCode >= 400 &&
+    statusCode < 500 &&
+    !RETRYABLE_CLIENT_ERRORS.includes(statusCode)
+  );
+}
+
+/**
+ * Exported only for testing purposes.
+ *
+ * `fetch` cannot be used for the long poll: Node's implementation applies undici's 300s
+ * `headersTimeout`, and lifting it would mean passing a dispatcher and depending on `undici`
+ * directly. `http.request` has no default timeout, and the Extensions API is plain HTTP on
+ * localhost.
+ */
+export function request(url: string, headers: Record<string, string>): Promise<ExtensionsApiResponse> {
+  return new Promise((resolve, reject) => {
+    const req = http.request(url, { headers }, res => {
+      const chunks: Buffer[] = [];
+      res.on('data', (chunk: Buffer) => chunks.push(chunk));
+      res.on('end', () => resolve({ statusCode: res.statusCode ?? 0, body: Buffer.concat(chunks).toString() }));
+      res.on('error', err => {
+        req.destroy();
+        reject(err);
+      });
+    });
+
+    req.on('socket', socket => socket.setKeepAlive(true, POLL_KEEPALIVE_MS));
+
+    req.on('error', reject);
+    req.end();
+  });
+}
+
+function truncate(body: string): string {
+  return body.length > ERROR_BODY_MAX_LENGTH ? `${body.slice(0, ERROR_BODY_MAX_LENGTH)}...` : body;
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise(resolve => {
+    setTimeout(resolve, ms);
+  });
+}
+
 /**
  * The Extension API Client.
  */
@@ -42,25 +131,85 @@ export class AwsLambdaExtension {
     }
 
     this._extensionId = res.headers.get('lambda-extension-identifier');
+
+    if (!this._extensionId) {
+      throw new Error('Extensions API accepted the registration without returning an extension identifier');
+    }
   }
 
   /**
-   * Advances the extension to the next event.
+   * Advances the extension to the next event and returns it.
    */
-  public async next(): Promise<void> {
+  public async next(): Promise<ExtensionEvent> {
     if (!this._extensionId) {
       throw new Error('Extension ID is not set');
     }
 
-    const res = await fetch(`${this._baseUrl}/event/next`, {
-      headers: {
-        'Lambda-Extension-Identifier': this._extensionId,
-        'Content-Type': 'application/json',
-      },
+    // This request blocks until the next event arrives, so it stays open for the whole
+    // duration of the current invocation. Under `fetch` that is capped at 300s, so any
+    // invocation that runs longer than that loses the extension partway through.
+    const res = await request(`${this._baseUrl}/event/next`, {
+      'Lambda-Extension-Identifier': this._extensionId,
+      'Content-Type': 'application/json',
     });
 
-    if (!res.ok) {
-      throw new Error(`Failed to advance to next event: ${await res.text()}`);
+    if (res.statusCode < 200 || res.statusCode > 299) {
+      throw new ExtensionsApiError(`Failed to advance to next event: ${truncate(res.body)}`, res.statusCode);
+    }
+
+    try {
+      return JSON.parse(res.body) as ExtensionEvent;
+    } catch {
+      // Not an empty event: `run` reads `eventType` to decide when to stop, so a body it cannot
+      // read has to be a failed poll. Returning `{}` would look like an INVOKE — resetting the
+      // backoff and re-polling with no delay, which spins the loop on any endpoint answering
+      // 200 with something that is not JSON, and skips the SHUTDOWN exit.
+      throw new Error(`Failed to parse the event from the Extensions API: ${truncate(res.body)}`);
+    }
+  }
+
+  /**
+   * Polls the Extensions API until the environment shuts down.
+   *
+   * A failed poll is retried rather than ending the loop. Lambda only completes an invocation
+   * once the runtime and every registered extension have asked for the next event, so an
+   * extension that stops polling does not fail loudly — it leaves every later invocation on
+   * that execution environment running until the function timeout kills it.
+   */
+  public async run(): Promise<void> {
+    let consecutiveFailures = 0;
+
+    for (;;) {
+      try {
+        const event = await this.next();
+        consecutiveFailures = 0;
+
+        // The runtime API is torn down right after this, so polling again would only produce
+        // errors on the way out.
+        if (event.eventType === 'SHUTDOWN') {
+          return;
+        }
+      } catch (err) {
+        // A poll the API refuses outright is not going to start working; retrying only buries
+        // the reason under a console error every few seconds for the life of the environment.
+        if (isClientError(err)) {
+          throw err;
+        }
+
+        consecutiveFailures++;
+
+        // Same reasoning once a recoverable-looking failure stops recovering.
+        if (consecutiveFailures >= POLL_MAX_CONSECUTIVE_FAILURES) {
+          throw err;
+        }
+
+        consoleSandbox(() => {
+          // eslint-disable-next-line no-console
+          console.error('Sentry Lambda extension: polling the Extensions API failed, retrying.', err);
+        });
+
+        await sleep(Math.min(POLL_RETRY_BASE_MS * 2 ** (consecutiveFailures - 1), POLL_RETRY_MAX_MS));
+      }
     }
   }
 

packages/aws-serverless/src/lambda-extension/index.ts

@@ -1,21 +1,37 @@
 #!/usr/bin/env node
-import { debug } from '@sentry/core';
+import { consoleSandbox } from '@sentry/core';
 import { AwsLambdaExtension } from './aws-lambda-extension';
-import { DEBUG_BUILD } from './debug-build';
 
-async function main(): Promise<void> {
-  const extension = new AwsLambdaExtension();
+const extension = new AwsLambdaExtension();
 
+async function main(): Promise<void> {
   await extension.register();
 
   extension.startSentryTunnel();
 
-  // eslint-disable-next-line no-constant-condition
-  while (true) {
-    await extension.next();
-  }
+  // Returns on SHUTDOWN. The process is left to idle rather than exiting, so envelopes the
+  // tunnel is still forwarding get their chance to land before Lambda reaps the environment.
+  await extension.run();
 }
 
-main().catch(err => {
-  DEBUG_BUILD && debug.error('Error in Lambda Extension', err);
+main().catch(async err => {
+  // The debug logger is only enabled from `Sentry.init`, and this process never calls it, so
+  // nothing reported through the logger from here would ever be visible.
+  consoleSandbox(() => {
+    // eslint-disable-next-line no-console
+    console.error('Sentry Lambda extension: stopped, events will no longer be tunnelled.', err);
+  });
+
+  // Reporting lets Lambda recycle the environment; `error` rethrows, and a registration that
+  // never completed has no id to report with, so neither path should mask the exit.
+  await extension.error('exit', err as Error).catch(() => undefined);
+
+  // Exiting here is not optional: the tunnel server holds a referenced handle, so the process
+  // would otherwise stay alive and registered while never asking for another event — and Lambda
+  // holds every later invocation on this execution environment open until the function timeout.
+  //
+  // Deferred by one turn of the loop because `process.exit` does not wait for stderr, which is a
+  // pipe under Lambda — exiting straight from the microtask above truncates the message written
+  // there to a single pipe buffer.
+  setImmediate(() => process.exit(1));
 });