aws-samples / aws-samples/sample-collaborative-ai-dlc

[Feature]: Harden Gremlin connection handling in lambdas

Open
#43 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
JavaScript
Stars
75
Forks
23
Avg merge
3d 17h
Merged PRs (30d)
24

Description

### Description

Our Gremlin lambda handlers use a simplified connection pattern: open a `DriverRemoteConnection` inside the handler, run queries in a `try`, close in `finally`. No retries, no reconnect on websocket close, no websocket-close handler.

AWS's [official Lambda + Neptune JavaScript example](https://docs.aws.amazon.com/neptune/latest/userguide/lambda-functions-examples.html#lambda-functions-examples-javascript) documents a more defensive pattern that handles failure modes our current code doesn't.

### Use case

Failure modes the current pattern is vulnerable to:

- **Silent websocket disconnects.** The gremlin-javascript driver doesn't raise an exception when the connection drops mid-query — the promise resolves with `null`. Queries appear to succeed but return no data. AWS's recommended fix is a `'ws close'` event handler that throws, forcing the invocation to fail so the caller retries.
- **No retries on transient errors.** AWS enumerates specific retryable errors (`ConcurrentModificationException`, `ReadOnlyViolationException` during primary failover, `WebSocket is not open`, `Connection reset by peer`, several more). We currently fail on first occurrence.
- **No reconnect on retry.** When one of those errors fires, AWS recommends closing and recreating the connection before retrying, not reusing the broken one.

### Area

Backend (Lambda)

### Additional context

**Before — current pattern (`lambda/timeline-events/index.js`):**

```js
const getConnection = async () => {
const host = process.env.NEPTUNE_ENDPOINT;
const credentials = await fromNodeProviderChain()();
credentials.region = process.env.AWS_REGION || 'us-east-1';
const connInfo = getUrlAndHeaders(host, '8182', credentials, '/gremlin', 'wss');
return new DriverRemoteConnection(connInfo.url, { headers: connInfo.headers });
};

exports.handler = async (event) => {
let conn;
try {
conn = await getConnection();
const g = traversal().withRemote(conn);
const events = await g.V().hasLabel('TimelineEvent').valueMap().toList();
return { statusCode: 200, body: JSON.stringify(events) };
} finally {
if (conn) try { await conn.close(); } catch (e) {}
}
};
```

Same shape in `lambda/tasks/`, `lambda/purge-neptune/`, `lambda/artifacts/`, `lambda/agents/`, `lambda/answer-question/`, etc. — ~15 files total.

**After — option 1: shared in-repo wrapper.**

Call sites become:

```js
import { createGremlinClient } from '../shared/gremlin.js';

const query = createGremlinClient({
host: process.env.NEPTUNE_ENDPOINT,
port: 8182,
useIam: true,
});

export const handler = async (event) => {
const events = await query(g => g.V().hasLabel('TimelineEvent').valueMap().toList());
return { statusCode: 200, body: JSON.stringify(events) };
};
```

But we own the implementation in `lambda/shared/gremlin.js`:

```js
import gremlin from 'gremlin';
import { getUrlAndHeaders } from 'gremlin-aws-sigv4/lib/utils.js';
import { fromNodeProviderChain } from '@aws-sdk/credential-providers';

const { driver: { DriverRemoteConnection }, process: { AnonymousTraversalSource } } = gremlin;
const { traversal } = AnonymousTraversalSource;

const RETRYABLE = [
/ConcurrentModificationException/,
/ReadOnlyViolationException/,
/WebSocket is not open/,
/Connection reset by peer/,
/Timed out while waiting for an available host/,
/Connection to server is no longer active/,
/SSLEngine closed already/,
/Broken pipe/,
];

const isRetryable = (err) => RETRYABLE.some(re => re.test(err?.message ?? ''));
const sleep = (ms) => new Promise(r => setTimeout(r, ms));
const backoff = (attempt) => Math.min(1000 * 2 ** attempt, 8000) + Math.random() * 250;

export function createGremlinClient({ host, port, useIam, maxRetries = 5 }) {
let conn = null;

const openConnection = async () => {
let url, headers;
if (useIam) {
const creds = await fromNodeProviderChain()();
creds.region = process.env.AWS_REGION;
({ url, headers } = getUrlAndHeaders(host, port, creds, '/gremlin', 'wss'));
} else {
url = `wss://${host}:${port}/gremlin`;
}
const c = new DriverRemoteConnection(url, { headers });
// AWS-recommended: promote silent ws close to an exception so in-flight
// queries fail instead of resolving with null.
c._client._connection._ws.on('close', () => {
c._rejectInFlight?.(new Error('WebSocket is not open'));
});
return c;
};

const closeConnection = async () => {
if (!conn) return;
try { await conn.close(); } catch {}
conn = null;
};

return async function query(fn) {
for (let attempt = 0; attempt <= maxRetries; attempt++) {
try {
if (!conn) conn = await openConnection();
const g = traversal().withRemote(conn);
return await fn(g);
} catch (err) {
await closeConnection();
if (attempt === maxRetries || !isRetryable(err)) throw err;
await sleep(backoff(attempt));
}
}
};
}
```

Plus tests for: retry classification, backoff behavior, connection recreation on retry, ws-close promotion, IAM vs non-IAM paths, max-retries exhaustion.

**After — option 2: use [`neptune-lambda-client`](https://github.com/svozza/neptune-lambda-client).**

A small library I wrote that implements this pattern. It's been in production in the [Workload Discovery](https://github.com/awslabs/workload-discovery-on-aws) codebase for ~5 years, so battle-tested against real Neptune failover and concurrent-write conditions.

```js
import { create } from 'neptune-lambda-client';

const query = create({
host: process.env.NEPTUNE_ENDPOINT,
port: 8182,
useIam: true,
});

export const handler = async (event) => {
const events = await query(g => g.V().hasLabel('TimelineEvent').valueMap().toList());
return { statusCode: 200, body: JSON.stringify(events) };
};
```

Also drops the `gremlin-aws-sigv4` dependency from ~19 lambdas — SigV4 signing is handled internally when `useIam: true`.

Contributor guide

Open the contributing guide

Research direction

Start with the AWS Lambda and Neptune example linked in the issue, then inspect the repeated connection patterns in lambda/timeline-events/index.js, lambda/tasks/, lambda/purge-neptune/, lambda/artifacts/, lambda/agents/, and lambda/answer-question/. Compare the two proposed approaches and establish the shared wrapper or dependency choice. Done means the Lambda call sites use hardened connection handling and tests cover retry classification, backoff, reconnection, websocket-close promotion, IAM paths, and retry exhaustion.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, javascript
Domain
backend, cloud, databases
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.