Block large payloads inside DocumentDeltaConnection, size due to the 1MB Kafka limit

api-report/telemetry-utils.api.md

@@ -247,8 +247,8 @@ export class TelemetryUTLogger implements ITelemetryLogger {
 // @public
 export class ThresholdCounter {
     constructor(threshold: number, logger: ITelemetryLogger, thresholdMultiple?: number);
-    send(eventName: string, value: number): void;
-    sendIfMultiple(eventName: string, value: number): void;
+    send(eventName: string, value: number, metadata?: any): boolean;
+    sendIfMultiple(eventName: string, value: number, metadata?: any): boolean;
     }
 
 // @public

packages/drivers/driver-base/.eslintrc.js

@@ -7,6 +7,9 @@ module.exports = {
     "extends": [
         "@fluidframework/eslint-config-fluid/eslint7"
     ],
+    "parserOptions": {
+        "project": ["./tsconfig.json", "./src/test/tsconfig.json"]
+    },
     "rules": {
         "@typescript-eslint/strict-boolean-expressions": "off"
     }

packages/drivers/driver-base/package.json

@@ -12,18 +12,24 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "build": "npm run build:genver && concurrently npm:build:compile npm:lint && npm run build:docs",
-    "build:compile": "concurrently npm:tsc npm:build:esnext",
+    "build:commonjs": "npm run tsc && npm run build:test",
+    "build:compile": "concurrently npm:build:commonjs npm:build:esnext",
     "build:docs": "api-extractor run --local --typescript-compiler-folder ../../../node_modules/typescript && copyfiles -u 1 ./_api-extractor-temp/doc-models/* ../../../_api-extractor-temp/",
     "build:esnext": "tsc --project ./tsconfig.esnext.json",
     "build:full": "npm run build",
     "build:full:compile": "npm run build:compile",
     "build:genver": "gen-version",
+    "build:test": "tsc --project ./src/test/tsconfig.json",
     "ci:build:docs": "api-extractor run --typescript-compiler-folder ../../../node_modules/typescript && copyfiles -u 1 ./_api-extractor-temp/* ../../../_api-extractor-temp/",
     "clean": "rimraf dist lib *.tsbuildinfo *.build.log",
     "eslint": "eslint --format stylish src",
     "eslint:fix": "eslint --format stylish src --fix",
     "lint": "npm run eslint",
     "lint:fix": "npm run eslint:fix",
+    "test": "npm run test:mocha",
+    "test:coverage": "nyc npm test -- --reporter xunit --reporter-option output=nyc/junit-report.xml",
+    "test:mocha": "mocha --recursive dist/test -r node_modules/@fluidframework/mocha-test-setup --unhandled-rejections=strict",
+    "test:mocha:verbose": "cross-env FLUID_TEST_VERBOSE=1 npm run test:mocha",
     "tsc": "tsc",
     "tsfmt": "tsfmt --verify",
     "tsfmt:fix": "tsfmt --replace"
@@ -39,7 +45,9 @@
   "devDependencies": {
     "@fluidframework/build-common": "^0.23.0",
     "@fluidframework/eslint-config-fluid": "^0.24.0",
+    "@fluidframework/mocha-test-setup": "^0.54.0",
     "@microsoft/api-extractor": "^7.16.1",
+    "@types/mocha": "^8.2.2",
     "@types/node": "^12.19.0",
     "@types/socket.io-client": "^1.4.32",
     "@typescript-eslint/eslint-plugin": "~4.14.0",
@@ -53,6 +61,8 @@
     "eslint-plugin-prefer-arrow": "~1.2.2",
     "eslint-plugin-react": "~7.22.0",
     "eslint-plugin-unicorn": "~26.0.1",
+    "mocha": "^8.4.0",
+    "nyc": "^15.0.0",
     "rimraf": "^2.6.2",
     "typescript": "~4.1.3",
     "typescript-formatter": "7.1.0"

packages/drivers/driver-base/src/documentDeltaConnection.ts

@@ -9,7 +9,7 @@ import {
     IDocumentDeltaConnectionEvents,
     DriverError,
 } from "@fluidframework/driver-definitions";
-import { createGenericNetworkError } from "@fluidframework/driver-utils";
+import { createGenericNetworkError, createWriteError } from "@fluidframework/driver-utils";
 import {
     ConnectionMode,
     IClientConfiguration,
@@ -24,9 +24,12 @@ import {
 } from "@fluidframework/protocol-definitions";
 import { IDisposable, ITelemetryLogger } from "@fluidframework/common-definitions";
 import { ChildLogger } from "@fluidframework/telemetry-utils";
+import { MessageSizeValidator } from "./messageSizeValidator";
 
 // Local storage key to disable the BatchManager
 const batchManagerDisabledKey = "FluidDisableBatchManager";
+// Local storage key to enable message size tracking
+const limitPayloadSizeKey = "FluidEnablePayloadSizeLimit";
 
 // See #8129.
 // Need to move to common-utils (tracked as #8165)
@@ -105,6 +108,10 @@ export class DocumentDeltaConnection
     protected _disposed: boolean = false;
     protected readonly logger: ITelemetryLogger;
     protected readonly isBatchManagerDisabled: boolean = false;
+    private messageSizeValidator: MessageSizeValidator | undefined;
+    // Close to 1MB with a small buffer to allow for redundant string escaping during transfer
+    private readonly maxPayloadSizeInBytes = 900000;
+    private readonly shouldLimitPayloadSize: boolean = false;
 
     public get details(): IConnected {
         if (!this._details) {
@@ -157,7 +164,8 @@ export class DocumentDeltaConnection
             }
         });
 
-        this.isBatchManagerDisabled = DocumentDeltaConnection.disabledBatchManagerFeatureGate;
+        this.isBatchManagerDisabled = DocumentDeltaConnection.booleanFeature(batchManagerDisabledKey);
+        this.shouldLimitPayloadSize = DocumentDeltaConnection.booleanFeature(limitPayloadSizeKey);
     }
 
     /**
@@ -280,6 +288,11 @@ export class DocumentDeltaConnection
     }
 
     protected emitMessages(type: string, messages: IDocumentMessage[][]) {
+        if (this.messageSizeValidator && !this.messageSizeValidator.isPayloadValid(messages)) {
+            this.emit("error", createWriteError("Payload too large"));
+            return;
+        }
+
         // Although the implementation here disconnects the socket and does not reuse it, other subclasses
         // (e.g. OdspDocumentDeltaConnection) may reuse the socket.  In these cases, we need to avoid emitting
         // on the still-live socket.
@@ -288,11 +301,11 @@ export class DocumentDeltaConnection
         }
     }
 
-    private static get disabledBatchManagerFeatureGate() {
+    private static booleanFeature(key: string): boolean {
         try {
             return localStorage !== undefined
                 && typeof localStorage === "object"
-                && localStorage.getItem(batchManagerDisabledKey) === "1";
+                && localStorage.getItem(key) === "1";
         } catch (e) { }
         return false;
     }
@@ -470,6 +483,12 @@ export class DocumentDeltaConnection
             }, timeout + 2000);
         });
 
+        if (this.shouldLimitPayloadSize) {
+            this.messageSizeValidator = new MessageSizeValidator(
+                this.maxMessageSize + 1,
+                this.maxPayloadSizeInBytes,
+                ChildLogger.create(this.logger, "MessageSizeValidator"));
+        }
         assert(!this.disposed, 0x246 /* "checking consistency of socket & _disposed flags" */);
     }
 

packages/drivers/driver-base/src/messageSizeValidator.ts

@@ -0,0 +1,75 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { ITelemetryLogger } from "@fluidframework/common-definitions";
+import { IDocumentMessage } from "@fluidframework/protocol-definitions";
+import { ThresholdCounter } from "@fluidframework/telemetry-utils";
+
+export class MessageSizeValidator {
+    private readonly payloadSizeCountersWithEvents = [
+        // The order here matters, in order to save telemetry quota.
+        // The first counter to exceed its limit will short-circuit
+        // event publishing.
+        {
+            counter: new ThresholdCounter(this.maxPayloadSizeInBytes, this.logger),
+            eventName: "OpsPayloadSizeLimitExceeded",
+        },
+        {
+            counter: new ThresholdCounter(this.maxPayloadSizeInBytes / 2, this.logger),
+            eventName: "OpsPayloadSize50PcOfMax",
+        },
+        {
+            counter: new ThresholdCounter(this.maxPayloadSizeInBytes / 4, this.logger),
+            eventName: "OpsPayloadSize25PcOfMax",
+        },
+    ];
+
+    private readonly messageSizeCounter = new ThresholdCounter(this.maxMessageSizeInBytes, this.logger);
+    private readonly messageSizeEvent = "OpSizeLimitExceeded";
+
+    constructor(
+        private readonly maxMessageSizeInBytes: number,
+        private readonly maxPayloadSizeInBytes: number,
+        private readonly logger: ITelemetryLogger,
+    ) {
+    }
+
+    private trackPayload(payloadSizeInBytes: number) {
+        for (const x of this.payloadSizeCountersWithEvents) {
+            if (x.counter.send(x.eventName, payloadSizeInBytes, { max: this.maxPayloadSizeInBytes })) {
+                break;
+            }
+        }
+    }
+
+    private trackMessage(messageSizeInBytes: number) {
+        this.messageSizeCounter.send(this.messageSizeEvent, messageSizeInBytes, { max: this.maxMessageSizeInBytes });
+    }
+
+    public isPayloadValid(messages: IDocumentMessage[][]): boolean {
+        let payloadSizeInBytes = 0;
+        let allMessagesUnderLimit = true;
+
+        for (const inner of messages) {
+            for (const message of inner) {
+                const messageSize = MessageSizeValidator.sizeInBytes(message);
+                allMessagesUnderLimit &&= messageSize < this.maxMessageSizeInBytes;
+                payloadSizeInBytes = payloadSizeInBytes + messageSize;
+                this.trackMessage(messageSize);
+            }
+        }
+
+        this.trackPayload(payloadSizeInBytes);
+        return allMessagesUnderLimit && payloadSizeInBytes < this.maxPayloadSizeInBytes;
+    }
+
+    public static sizeInBytes(message: IDocumentMessage): number {
+        const { contents, ...restOfObject } = message;
+        // `contents` is already stringified. Re-stringifying the whole message will
+        // lead to additional escape characters which will increase the size artificially.
+        return new TextEncoder().encode(message.contents).length
+            + new TextEncoder().encode(JSON.stringify(restOfObject)).length;
+    }
+}

packages/drivers/driver-base/src/test/messageSizeValidator.spec.ts

@@ -0,0 +1,102 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { strict as assert } from "assert";
+import { IDocumentMessage } from "@fluidframework/protocol-definitions";
+import { MockLogger } from "@fluidframework/telemetry-utils";
+import { MessageSizeValidator } from "../messageSizeValidator";
+
+const generateStringOfSize = (sizeInBytes: number): string => new Array(sizeInBytes + 1).join("0");
+const generateMessageOfSize = (sizeInBytes: number): IDocumentMessage => {
+    const envelope = {
+        clientSequenceNumber: 1,
+        metadata: {
+            meta: "data",
+            other: "data",
+        },
+        referenceSequenceNumber: 0,
+        type: "test",
+    };
+
+    const message = (envelope as IDocumentMessage);
+    message.contents = generateStringOfSize(sizeInBytes - new TextEncoder().encode(JSON.stringify(message)).length);
+    return message;
+};
+
+describe("Message size validation", () => {
+    let logger: MockLogger;
+    let validator: MessageSizeValidator;
+    const maxMessageSizeInBytes = 10 * 1000;
+    const maxPayloadSizeInBytes = 100 * maxMessageSizeInBytes;
+
+    beforeEach(() => {
+        logger = new MockLogger();
+        validator = new MessageSizeValidator(maxMessageSizeInBytes, maxPayloadSizeInBytes, logger);
+    });
+
+    it("Should fail when message size is max", () => {
+        assert(!validator.isPayloadValid([[generateMessageOfSize(maxMessageSizeInBytes)]]));
+        assert(logger.matchEvents([{
+            eventName: "OpSizeLimitExceeded",
+            category: "performance",
+            value: maxMessageSizeInBytes,
+            max: maxMessageSizeInBytes,
+        }]));
+    });
+
+    it("Should succeed when payload is lower than 25% of max", () => {
+        assert(validator.isPayloadValid([[generateMessageOfSize(maxMessageSizeInBytes - 1)]]));
+        assert(logger.matchEvents([]));
+    });
+
+    it("Should succeed when payload is lower than 50% of max", async () => {
+        const size = 26 * (maxMessageSizeInBytes - 1);
+        assert(validator.isPayloadValid([Array(26).fill(generateMessageOfSize(maxMessageSizeInBytes - 1))]));
+        assert(logger.matchEvents([{
+            eventName: "OpsPayloadSize25PcOfMax",
+            category: "performance",
+            value: size,
+            max: maxPayloadSizeInBytes,
+        }]));
+    });
+
+    it("Should succeed when payload is between 50% and 100% of max", () => {
+        const size = 51 * (maxMessageSizeInBytes - 1);
+        assert(validator.isPayloadValid([Array(51).fill(generateMessageOfSize(maxMessageSizeInBytes - 1))]));
+        assert(logger.matchEvents([{
+            eventName: "OpsPayloadSize50PcOfMax",
+            category: "performance",
+            value: size,
+            max: maxPayloadSizeInBytes,
+        }]));
+    });
+
+    it("Should fail when payload size is higher than max", () => {
+        const size = 101 * (maxMessageSizeInBytes - 1);
+        assert(!validator.isPayloadValid([Array(101).fill(generateMessageOfSize(maxMessageSizeInBytes - 1))]));
+        assert(logger.matchEvents([{
+            eventName: "OpsPayloadSizeLimitExceeded",
+            category: "performance",
+            value: size,
+            max: maxPayloadSizeInBytes,
+        }]));
+    });
+
+    it("Should fail when payload size is higher than max and message is larger than max", () => {
+        const size = 100 * (maxMessageSizeInBytes);
+        assert(!validator.isPayloadValid([Array(100).fill(generateMessageOfSize(maxMessageSizeInBytes))]));
+        assert(logger.matchEvents([{
+            eventName: "OpSizeLimitExceeded",
+            category: "performance",
+            value: maxMessageSizeInBytes,
+            max: maxMessageSizeInBytes,
+        }, {
+            eventName: "OpsPayloadSizeLimitExceeded",
+            category: "performance",
+            value: size,
+            max: maxPayloadSizeInBytes,
+        }]));
+    });
+});

packages/drivers/driver-base/src/test/tsconfig.json

@@ -0,0 +1,21 @@
+{
+    "extends": "@fluidframework/build-common/ts-common-config.json",
+    "compilerOptions": {
+        "rootDir": "./",
+        "outDir": "../../dist/test",
+        "types": [
+            "node",
+            "mocha"
+        ],
+        "declaration": false,
+        "declarationMap": false
+    },
+    "include": [
+        "./**/*"
+    ],
+    "references": [
+        {
+            "path": "../.."
+        }
+    ]
+}

packages/drivers/driver-base/tsconfig.json

@@ -2,16 +2,18 @@
     "extends": "@fluidframework/build-common/ts-common-config.json",
     "exclude": [
         "dist",
-        "node_modules"
+        "node_modules",
+        "src/test/**/*"
     ],
     "compilerOptions": {
         "rootDir": "./src",
         "outDir": "./dist",
         "types": [
             "@types/socket.io-client"
-        ]
+        ],
+        "composite": true
     },
     "include": [
         "src/**/*"
     ]
-}
+}

packages/drivers/local-driver/src/test/tsconfig.json

@@ -18,4 +18,4 @@
             "path": "../.."
         }
     ]
-}
+}

packages/loader/container-loader/src/connectionManager.ts

@@ -78,7 +78,11 @@ function getNackReconnectInfo(nackContent: INackContent) {
 const createReconnectError = (fluidErrorCode: string, err: any) =>
     wrapError(
         err,
-        (errorMessage: string) => new GenericNetworkError(fluidErrorCode, errorMessage, true /* canRetry */),
+        (errorMessage: string) => new GenericNetworkError(
+            fluidErrorCode,
+            errorMessage,
+            err?.canRetry === true || err?.canRetry === undefined, // unless explicitly specified, this will retry
+        ),
     );
 
 /**