Aws Lambda Extensions: Build Custom Monitoring Without Touching Application Code

# Lambda execution lifecycle with extensions:
# 1. INIT: Runtime starts → Extensions register → Your init code runs
# 2. INVOKE: Extension blocks on /next → Function runs → Extension does post-invoke work
# 3. SHUTDOWN: SIGTERM sent → Extension has 2s to flush → SIGKILL if not done

# Key: Extension MUST respond to /next within the function timeout
# Slow extension = slow function invocation

// extensions/secrets-cache/index.js — runs as separate process
const http = require('http');
const { SSMClient, GetParameterCommand } = require('@aws-sdk/client-ssm');

const ssm = new SSMClient({ region: process.env.AWS_REGION });
const cache = new Map();
const TTL = 5 * 60 * 1000; // 5 minutes

async function getSecret(name) {
  const hit = cache.get(name);
  if (hit && Date.now() - hit.at < TTL) return hit.value;
  const resp = await ssm.send(new GetParameterCommand({ Name: name, WithDecryption: true }));
  cache.set(name, { value: resp.Parameter.Value, at: Date.now() });
  return resp.Parameter.Value;
}

// HTTP server — function calls localhost:2773 instead of SSM on every invoke
http.createServer(async (req, res) => {
  const name = new URL(req.url, 'http://localhost').searchParams.get('name');
  try {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ value: await getSecret(name) }));
  } catch (e) {
    res.writeHead(500); res.end(JSON.stringify({ error: e.message }));
  }
}).listen(2773);

// Register with Extensions API
const API = 'http://' + process.env.AWS_LAMBDA_RUNTIME_API;
const reg = await fetch(API + '/2020-01-01/extension/register', {
  method: 'POST',
  headers: { 'Lambda-Extension-Name': 'secrets-cache' },
  body: JSON.stringify({ events: ['INVOKE', 'SHUTDOWN'] }),
});
const id = reg.headers.get('Lambda-Extension-Identifier');

// Event loop — must keep calling /next
while (true) {
  const evt = await (await fetch(API + '/2020-01-01/extension/event/next',
    { headers: { 'Lambda-Extension-Identifier': id } })).json();
  if (evt.eventType === 'SHUTDOWN') process.exit(0);
}

// In your Lambda function — zero SSM calls per invocation:
const { value } = await (await fetch('http://localhost:2773/?name=/prod/db/pass')).json();

// Receive function logs in real time via Telemetry API
// Use to forward logs to Datadog, Splunk, or custom destination

const http = require('http');

// 1. Start HTTP receiver
http.createServer((req, res) => {
  let body = '';
  req.on('data', c => body += c);
  req.on('end', () => {
    JSON.parse(body).forEach(log => {
      if (log.type === 'function') sendToSplunk(log.record);
    });
    res.writeHead(200); res.end();
  });
}).listen(4243);

// 2. Subscribe after registration
await fetch(API + '/2022-07-01/telemetry', {
  method: 'PUT',
  headers: { 'Lambda-Extension-Identifier': extensionId },
  body: JSON.stringify({
    schemaVersion: '2022-12-13',
    destination: { protocol: 'HTTP', URI: 'http://sandbox.localdomain:4243' },
    types: ['function', 'platform'],
    buffering: { timeoutMs: 1000, maxBytes: 262144, maxItems: 1000 },
  }),
});

# Must live in /extensions/ directory inside the layer ZIP
mkdir -p layer/extensions
cp my-extension layer/extensions/
chmod +x layer/extensions/my-extension
cd layer && zip -r ../extension.zip extensions/

aws lambda publish-layer-version   --layer-name my-extension   --zip-file fileb://extension.zip   --compatible-runtimes nodejs20.x python3.12