Limit size of heartbeat header

packages/app/src/heartbeatService.ts

@@ -21,6 +21,7 @@ import {
   isIndexedDBAvailable,
   validateIndexedDBOpenable
 } from '@firebase/util';
+import { countBytes, splitHeartbeatsCache } from './heartbeatSize';
 import {
   deleteHeartbeatsFromIndexedDB,
   readHeartbeatsFromIndexedDB,
@@ -33,6 +34,8 @@ import {
   HeartbeatStorage
 } from './types';
 
+const HEADER_SIZE_LIMIT_BYTES = 1000;
+
 export class HeartbeatServiceImpl implements HeartbeatService {
   /**
    * The persistence layer for heartbeats
@@ -120,19 +123,42 @@ export class HeartbeatServiceImpl implements HeartbeatService {
     if (this._heartbeatsCache === null) {
       return '';
     }
-    const headerString = base64Encode(
-      JSON.stringify({ version: 2, heartbeats: this._heartbeatsCache })
-    );
-    this._heartbeatsCache = null;
-    // Do not wait for this, to reduce latency.
-    void this._storage.deleteAll();
+    // Count size of _heartbeatsCache after being converted into a base64
+    // header string.
+    const base64Bytes = countBytes(this._heartbeatsCache);
+    // If it exceeds the limit, split out the oldest portion under the
+    // limit to return. Put the rest back into _heartbeatsCache.
+    let headerString = '';
+    if (base64Bytes > HEADER_SIZE_LIMIT_BYTES) {
+      const { heartbeatsToSend, heartbeatsToKeep } = splitHeartbeatsCache(
+        this._heartbeatsCache,
+        HEADER_SIZE_LIMIT_BYTES
+      );
+      headerString = base64Encode(
+        JSON.stringify({ version: 2, heartbeats: heartbeatsToSend })
+      );
+      // Write the portion not sent back to memory, and then to indexedDB.
+      this._heartbeatsCache = heartbeatsToKeep;
+      // This is more likely than deleteAll() to cause some mixed up state
+      // problems if we don't wait for execution to finish.
+      await this._storage.overwrite(this._heartbeatsCache);
+    } else {
+      // If _heartbeatsCache does not exceed the size limit, send all the
+      // data in a header and delete memory and indexedDB caches.
+      headerString = base64Encode(
+        JSON.stringify({ version: 2, heartbeats: this._heartbeatsCache })
+      );
+      this._heartbeatsCache = null;
+      // Do not wait for this, to reduce latency.
+      void this._storage.deleteAll();
+    }
     return headerString;
   }
 }
 
 function getUTCDateString(): string {
   const today = new Date();
-  return today.toISOString().substring(0,10);
+  return today.toISOString().substring(0, 10);
 }
 
 export class HeartbeatStorageImpl implements HeartbeatStorage {

packages/app/src/heartbeatSize.test.ts

@@ -0,0 +1,114 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { base64Encode } from '@firebase/util';
+import { expect } from 'chai';
+import '../test/setup';
+import {
+  countBytes,
+  countHeartbeatBytes,
+  splitHeartbeatsCache
+} from './heartbeatSize';
+
+function generateUserAgentString(pairs: number): string {
+  let uaString = '';
+  for (let i = 0; i < pairs; i++) {
+    uaString += `test-platform/${i % 10}.${i % 10}.${i % 10}`;
+  }
+  return uaString;
+}
+
+function generateDates(count: number): string[] {
+  let currentTimestamp = Date.now();
+  const dates = [];
+  for (let i = 0; i < count; i++) {
+    dates.push(new Date(currentTimestamp).toISOString().slice(0, 10));
+    currentTimestamp += 24 * 60 * 60 * 1000;
+  }
+  return dates;
+}
+
+describe('splitHeartbeatsCache()', () => {
+  it('returns empty heartbeatsToKeep if it cannot get under maxSize', () => {
+    const heartbeats = [
+      { userAgent: generateUserAgentString(1), dates: generateDates(1) }
+    ];
+    const { heartbeatsToKeep, heartbeatsToSend } = splitHeartbeatsCache(
+      heartbeats,
+      5
+    );
+    expect(heartbeatsToSend.length).to.equal(0);
+    expect(heartbeatsToKeep).to.deep.equal(heartbeats);
+  });
+  it('splits heartbeats array', () => {
+    const heartbeats = [
+      { userAgent: generateUserAgentString(20), dates: generateDates(8) },
+      { userAgent: generateUserAgentString(4), dates: generateDates(10) }
+    ];
+    const heartbeat1Size = countHeartbeatBytes(heartbeats[0]);
+    const { heartbeatsToKeep, heartbeatsToSend } = splitHeartbeatsCache(
+      heartbeats,
+      heartbeat1Size + 1
+    );
+    expect(heartbeatsToSend.length).to.equal(1);
+    expect(heartbeatsToKeep.length).to.equal(1);
+  });
+  it('splits the first heartbeat if needed', () => {
+    const heartbeats = [
+      { userAgent: generateUserAgentString(20), dates: generateDates(50) },
+      { userAgent: generateUserAgentString(4), dates: generateDates(10) }
+    ];
+    const heartbeat1Size = countHeartbeatBytes(heartbeats[0]);
+    const { heartbeatsToKeep, heartbeatsToSend } = splitHeartbeatsCache(
+      heartbeats,
+      heartbeat1Size - 50
+    );
+    expect(heartbeatsToSend.length).to.equal(1);
+    expect(heartbeatsToKeep.length).to.equal(2);
+    expect(
+      heartbeatsToSend[0].dates.length + heartbeatsToKeep[0].dates.length
+    ).to.equal(heartbeats[0].dates.length);
+    expect(heartbeatsToSend[0].userAgent).to.equal(
+      heartbeatsToKeep[0].userAgent
+    );
+  });
+});
+
+describe('countBytes()', () => {
+  it('counts how many bytes there will be in a stringified, encoded header', () => {
+    const heartbeats = [
+      { userAgent: generateUserAgentString(1), dates: generateDates(1) },
+      { userAgent: generateUserAgentString(3), dates: generateDates(2) }
+    ];
+    let size: number = 0;
+    const headerString = base64Encode(
+      JSON.stringify({ version: 2, heartbeats })
+    );
+    console.log(JSON.stringify({ version: 2, heartbeats }));
+    // We don't use this measurement method in the app because user
+    // environments are much more unpredictable while we know the
+    // tests will run in either a standard headless browser or Node.
+    if (typeof Blob !== 'undefined') {
+      const blob = new Blob([headerString]);
+      size = blob.size;
+    } else if (typeof Buffer !== 'undefined') {
+      const buffer = Buffer.from(headerString);
+      size = buffer.byteLength;
+    }
+    expect(countBytes(heartbeats)).to.equal(size);
+  });
+});

packages/app/src/heartbeatSize.ts

@@ -0,0 +1,116 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { base64Encode } from '@firebase/util';
+import { HeartbeatsByUserAgent } from './types';
+
+/**
+ * Calculate byte length of a string. From:
+ * https://codereview.stackexchange.com/questions/37512/count-byte-length-of-string
+ */
+function getByteLength(str: string): number {
+  let byteLength = 0;
+  for (let i = 0; i < str.length; i++) {
+    const c = str.charCodeAt(i);
+    byteLength +=
+      (c & 0xf800) === 0xd800
+        ? 2 // Code point is half of a surrogate pair
+        : c < 1 << 7
+        ? 1
+        : c < 1 << 11
+        ? 2
+        : 3;
+  }
+  return byteLength;
+}
+
+/**
+ * Calculate bytes of a single HeartbeatsByUserAgent object after
+ * being stringified and converted to base64.
+ */
+export function countHeartbeatBytes(heartbeat: HeartbeatsByUserAgent): number {
+  return getByteLength(base64Encode(JSON.stringify(heartbeat)));
+}
+
+/**
+ * Calculate bytes of a HeartbeatsByUserAgent array after being wrapped
+ * in a platform logging header JSON object, stringified, and converted
+ * to base 64.
+ */
+export function countBytes(heartbeatsCache: HeartbeatsByUserAgent[]): number {
+  // heartbeatsCache wrapper properties
+  return getByteLength(
+    base64Encode(JSON.stringify({ version: 2, heartbeats: heartbeatsCache }))
+  );
+}
+
+/**
+ * Split a HeartbeatsByUserAgent array into 2 arrays, one that fits
+ * under `maxSize`, to be sent as a header, and the remainder. If
+ * the first heartbeat in the array is too big by itself, it will
+ * split that heartbeat into two by splitting its `dates` array.
+ */
+export function splitHeartbeatsCache(
+  heartbeatsCache: HeartbeatsByUserAgent[],
+  maxSize: number
+): {
+  heartbeatsToSend: HeartbeatsByUserAgent[];
+  heartbeatsToKeep: HeartbeatsByUserAgent[];
+} {
+  const BYTES_PER_DATE = getByteLength(
+    base64Encode(JSON.stringify('2022-12-12'))
+  );
+  let totalBytes = 0;
+  const heartbeatsToSend = [];
+  const heartbeatsToKeep = [...heartbeatsCache];
+  for (const heartbeat of heartbeatsCache) {
+    totalBytes += countHeartbeatBytes(heartbeat);
+    if (totalBytes > maxSize) {
+      if (heartbeatsToSend.length === 0) {
+        // The first heartbeat is too large and needs to be split or we have
+        // nothing to send.
+        const heartbeatBytes = countHeartbeatBytes(heartbeat);
+        const bytesOverLimit = heartbeatBytes - maxSize;
+        const datesToRemove = Math.ceil(bytesOverLimit / BYTES_PER_DATE);
+        if (datesToRemove >= heartbeat.dates.length) {
+          // If no amount of removing dates can get this heartbeat under
+          // the limit (unlikely scenario), nothing can be sent.
+          break;
+        }
+        const heartbeatToSend = {
+          ...heartbeat,
+          dates: heartbeat.dates.slice(0, -datesToRemove)
+        };
+        const heartbeatToKeep = {
+          ...heartbeat,
+          dates: heartbeat.dates.slice(-datesToRemove)
+        };
+        heartbeatsToSend.push(heartbeatToSend);
+        heartbeatsToKeep[0] = heartbeatToKeep;
+      } else {
+        break;
+      }
+    } else {
+      heartbeatsToSend.push(heartbeat);
+      heartbeatsToKeep.shift();
+    }
+  }
+  return {
+    heartbeatsToSend,
+    heartbeatsToKeep
+  };
+}