Adjust Kafka message estimator on MESSAGE_TOO_LARGE (#6374)

graphql-hive · Jan 22, 2025 · 393ece7 · 393ece7
1 parent 0e4be14
commit 393ece7
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 2 deletions.
diff --git a/.changeset/ka-f-ka.md b/.changeset/ka-f-ka.md
@@ -0,0 +1,5 @@
+---
+'hive': patch
+---
+
+Adjust the Kafka message size estimation only when Kafka gives back `MESSAGE_TOO_LARGE` error
diff --git a/packages/services/usage/__tests__/buffer.spec.ts b/packages/services/usage/__tests__/buffer.spec.ts
@@ -34,6 +34,9 @@ test('increase the defaultBytesPerOperation estimation by 5% when over 100 calls
     limitInBytes: eventHubLimitInBytes,
     useEstimator: true,
     onRetry,
+    isTooLargePayloadError() {
+      return true;
+    },
     calculateReportSize(report) {
       return report.size;
     },
@@ -150,6 +153,9 @@ test('buffer should split the report into multiple reports when the estimated si
     interval,
     limitInBytes: eventHubLimitInBytes,
     useEstimator: true,
+    isTooLargePayloadError() {
+      return true;
+    },
     calculateReportSize(report) {
       return report.size;
     },
@@ -262,6 +268,9 @@ test('buffer create two chunks out of one buffer when actual buffer size is too
     interval,
     limitInBytes: eventHubLimitInBytes,
     useEstimator: true,
+    isTooLargePayloadError() {
+      return true;
+    },
     calculateReportSize(report) {
       return report.size;
     },

diff --git a/packages/services/usage/src/buffer.ts b/packages/services/usage/src/buffer.ts
@@ -131,6 +131,7 @@ export function createKVBuffer<T>(config: {
   calculateReportSize(report: T): number;
   split(report: T, numOfChunks: number): readonly T[];
   onRetry(reports: readonly T[]): void;
+  isTooLargePayloadError(error: unknown): boolean;
   sender(
     reports: readonly T[],
     estimatedSizeInBytes: number,
@@ -254,8 +255,10 @@ export function createKVBuffer<T>(config: {
         await flushBuffer(reports, size, batchId);
       } catch (error) {
         logger.error(error);
-        // the payload size was most likely too big
-        estimator.overflowed(batchId);
+        if (config.isTooLargePayloadError(error)) {
+          // the payload size was most likely too big
+          estimator.overflowed(batchId);
+        }
       }
     }
 

diff --git a/packages/services/usage/src/usage.ts b/packages/services/usage/src/usage.ts
@@ -156,6 +156,9 @@ export function createUsage(config: {
     interval: config.kafka.buffer.interval,
     limitInBytes: 990_000, // 1MB is the limit of a single request to EventHub, let's keep it below that
     useEstimator: config.kafka.buffer.dynamic,
+    isTooLargePayloadError(error) {
+      return error instanceof Error && 'type' in error && error.type === 'MESSAGE_TOO_LARGE';
+    },
     calculateReportSize(report) {
       return Object.keys(report.map).length;
     },