feat: optimize stream packet parse performance (#3282)

* feat: optimize stream packet parse performance

* test: fix testcase

* test: fix testcase

* test: fix testcase

* chore: update code

Author: bytemain

packages/connection/__test__/common/buffers.test.ts

@@ -1,7 +1,7 @@
-import { Buffers } from '../../src/common/connection/buffers';
+import { Buffers, copy } from '../../src/common/connection/buffers';
 
 describe('Buffers', () => {
-  it('can append and slice', () => {
+  it('can push and slice', () => {
     const list = new Buffers();
     list.push(new Uint8Array([1, 2, 3]));
     list.push(new Uint8Array([4, 5, 6]));
@@ -18,6 +18,25 @@ describe('Buffers', () => {
     expect(list.slice(2, 6)).toEqual(new Uint8Array([3, 4, 5, 6]));
   });
 
+  it('slice', () => {
+    const xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    const splits = [[4, 2, 3, 1], [2, 2, 2, 2, 2], [1, 6, 3, 1], [9, 2], [10], [5, 5]];
+
+    splits.forEach(function (split) {
+      const bufs = create(xs, split);
+      expect(new Uint8Array(xs)).toEqual(bufs.slice());
+
+      for (let i = 0; i < xs.length; i++) {
+        for (let j = i; j < xs.length; j++) {
+          const a = bufs.slice(i, j);
+          const b = new Uint8Array(xs.slice(i, j));
+
+          expect(a).toEqual(b);
+        }
+      }
+    });
+  });
+
   it('can splice', () => {
     const list = new Buffers();
     list.push(new Uint8Array([1, 2, 3]));
@@ -44,10 +63,54 @@ describe('Buffers', () => {
       new Uint8Array([4, 5, 6]),
     ]);
   });
+
+  it('splice', () => {
+    const xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    const splits = [[4, 2, 3, 1], [2, 2, 2, 2, 2], [1, 6, 3, 1], [9, 2], [10], [5, 5]];
+
+    splits.forEach(function (split) {
+      for (let i = 0; i < xs.length; i++) {
+        for (let j = i; j < xs.length; j++) {
+          const bufs = create(xs, split);
+          const xs_ = xs.slice();
+
+          const a_ = bufs.splice(i, j);
+          const a = [].slice.call(a_.slice());
+          const b = xs_.splice(i, j);
+          expect(a).toEqual(b);
+          expect(bufs.slice()).toEqual(new Uint8Array(xs_));
+        }
+      }
+    });
+  });
+
+  it('spliceRep', () => {
+    const xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    const splits = [[4, 2, 3, 1], [2, 2, 2, 2, 2], [1, 6, 3, 1], [9, 2], [10], [5, 5]];
+    const reps = [[], [1], [5, 6], [3, 1, 3, 3, 7], [9, 8, 7, 6, 5, 4, 3, 2, 1, 2, 3, 4, 5]];
+
+    splits.forEach(function (split) {
+      reps.forEach(function (rep) {
+        for (let i = 0; i < xs.length; i++) {
+          for (let j = i; j < xs.length; j++) {
+            const bufs = create(xs, split);
+            const xs_ = xs.slice();
+
+            const a_ = bufs.splice(i, j, new Uint8Array(rep));
+            const a = [].slice.call(a_.slice());
+            const b = xs_.splice(i, j, ...rep);
+
+            expect(a).toEqual(b);
+            expect(bufs.slice()).toEqual(new Uint8Array(xs_));
+          }
+        }
+      });
+    });
+  });
+
   it('can copy', () => {
     const list = new Buffers();
     list.push(new Uint8Array([1, 2, 3]));
-
     list.push(new Uint8Array([4, 5, 6]));
 
     const target = new Uint8Array(7);
@@ -72,5 +135,132 @@ describe('Buffers', () => {
     list.set(1, 10);
     expect(list.get(0)).toEqual(9);
     expect(list.get(1)).toEqual(10);
+
+    const xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    const splits = [[4, 2, 3, 1], [2, 2, 2, 2, 2], [1, 6, 3, 1], [9, 2], [10], [5, 5]];
+
+    splits.forEach(function (split) {
+      const bufs = create(xs, split);
+      const buf = new Uint8Array(xs);
+
+      for (let i = 0; i < xs.length; i++) {
+        for (let j = i; j < xs.length; j++) {
+          const t0 = new Uint8Array(j - i);
+          const t1 = new Uint8Array(j - i);
+
+          expect(bufs.copy(t0, 0, i, j)).toEqual(copy(buf, t1, 0, i, j));
+          expect([].slice.call(t0)).toEqual([].slice.call(t1));
+        }
+      }
+    });
+  });
+
+  it('unshift', () => {
+    const bufs = new Buffers();
+    bufs.unshift(new Uint8Array([6, 7, 8, 9]));
+    bufs.unshift(new Uint8Array([4, 5]));
+    bufs.unshift(new Uint8Array([1, 2, 3]));
+    bufs.unshift(new Uint8Array([0]));
+
+    expect([].slice.call(bufs.slice())).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    expect(bufs.buffers.length).toEqual(4);
+    expect(bufs.byteLength).toEqual(10);
+  });
+
+  it('cursor should work', () => {
+    const list = new Buffers();
+    list.push(new Uint8Array([1, 2, 3, 4]));
+    list.push(new Uint8Array([4, 5, 6, 7, 8, 9]));
+    list.push(new Uint8Array([7, 8, 9]));
+
+    const cursor = list.cursor(0);
+    let count = 0;
+    for (const a of cursor.iterator()) {
+      expect(a).toEqual(list.get(count));
+      if (count === 4) {
+        break;
+      }
+
+      count++;
+    }
+    expect(count).toEqual(4);
+    expect(cursor.value).toEqual(list.get(5));
+    expect(cursor.offset).toEqual(5);
+    expect(cursor.line).toEqual(1);
+    expect(cursor.lineOffset).toEqual(1);
+    expect(cursor.lineWidth).toEqual(6);
+    expect(list.pos(cursor.offset)).toEqual({ buf: cursor.line, offset: cursor.lineOffset });
+
+    count++;
+
+    // reentrant
+    for (const a of cursor.iterator()) {
+      expect(a).toEqual(list.get(count));
+      count++;
+    }
+
+    // because here we always let count plus 1
+    expect(count).toEqual(list.byteLength);
+    expect(cursor.offset).toEqual(list.byteLength);
+    expect(cursor.line).toEqual(3);
+    expect(cursor.lineOffset).toEqual(0);
+  });
+
+  it('cursor can break', () => {
+    const list = new Buffers();
+    list.push(new Uint8Array([1, 2, 3, 4]));
+    list.push(new Uint8Array([4, 5, 6, 7, 8, 9]));
+    list.push(new Uint8Array([7, 8, 9]));
+
+    const cursor = list.cursor(0);
+    let count = 0;
+    const iter = cursor.iterator();
+    for (const a of iter) {
+      expect(a).toEqual(list.get(count));
+      if (count === 4) {
+        iter.return();
+      } else {
+        count++;
+      }
+    }
+
+    expect(count).toEqual(4);
+
+    expect(cursor.offset).toEqual(5);
+    expect(cursor.value).toEqual(list.get(cursor.offset));
+    expect(cursor.line).toEqual(1);
+    expect(cursor.lineOffset).toEqual(1);
+    expect(list.pos(cursor.offset)).toEqual({ buf: cursor.line, offset: cursor.lineOffset });
+
+    cursor.move(1);
+    expect(cursor.offset).toEqual(6);
+    expect(cursor.value).toEqual(list.get(cursor.offset));
+    expect(cursor.line).toEqual(1);
+    expect(cursor.lineOffset).toEqual(2);
+    expect(list.pos(cursor.offset)).toEqual({ buf: cursor.line, offset: cursor.lineOffset });
+
+    cursor.move(2);
+    expect(cursor.offset).toEqual(8);
+    expect(cursor.value).toEqual(list.get(cursor.offset));
+    expect(cursor.line).toEqual(1);
+    expect(cursor.lineOffset).toEqual(4);
+    expect(list.pos(cursor.offset)).toEqual({ buf: cursor.line, offset: cursor.lineOffset });
+
+    cursor.move(3);
+    expect(cursor.offset).toEqual(11);
+    expect(cursor.value).toEqual(list.get(cursor.offset));
+    expect(cursor.line).toEqual(2);
+    expect(cursor.lineOffset).toEqual(1);
+    expect(list.pos(cursor.offset)).toEqual({ buf: cursor.line, offset: cursor.lineOffset });
   });
 });
+
+function create(xs: number[], split: number[]) {
+  const bufs = new Buffers();
+  let offset = 0;
+  split.forEach(function (i) {
+    bufs.push(new Uint8Array(xs.slice(offset, offset + i)));
+    offset += i;
+  });
+  return bufs;
+}

packages/connection/__test__/common/stream-decoder.test.ts

@@ -51,6 +51,16 @@ const mixedPackets = [p1m, p5m].map((v) => {
   return [newPacket, v] as const;
 });
 
+const size = purePackets.reduce((acc, v) => acc + v[0].byteLength, 0);
+
+let offset = 0;
+const bigPayload = createPayload(size);
+purePackets.forEach((v) => {
+  const sumiPacket = v[0];
+  bigPayload.set(sumiPacket, offset);
+  offset += sumiPacket.byteLength;
+});
+
 const packets = [...purePackets, ...mixedPackets];
 
 describe('stream-packet', () => {
@@ -69,14 +79,7 @@ describe('stream-packet', () => {
       const decoder = new StreamPacketDecoder();
 
       decoder.onData((data) => {
-        expect(data.byteLength).toEqual(expected.byteLength);
-        for (let i = 0; i < 10; i++) {
-          // 随机选一些数据(<= 100字节)，对比是否正确，对比整个数组的话，超大 buffer 会很耗时
-          const start = Math.floor(Math.random() * data.byteLength);
-          const end = Math.floor(Math.random() * 1024);
-
-          expect(data.subarray(start, end)).toEqual(expected.subarray(start, end));
-        }
+        fastExpectBufferEqual(data, expected);
         decoder.dispose();
         done();
       });
@@ -91,6 +94,50 @@ describe('stream-packet', () => {
       logMemoryUsage();
     });
   });
+
+  it('can decode a stream payload contains multiple packets', (done) => {
+    const decoder = new StreamPacketDecoder();
+    const expectCount = purePackets.length;
+    let count = 0;
+    decoder.onData((data) => {
+      const expected = purePackets[count][1];
+      fastExpectBufferEqual(data, expected);
+
+      count++;
+      if (count === expectCount) {
+        decoder.dispose();
+        done();
+      }
+    });
+
+    console.log('write chunk', bigPayload.byteLength);
+    // write chunk by ${pressure} bytes
+    for (let i = 0; i < bigPayload.byteLength; i += pressure) {
+      decoder.push(bigPayload.subarray(i, i + pressure));
+      logMemoryUsage();
+    }
+
+    logMemoryUsage();
+  });
+
+  it('can decode a stream it header and payload are separated', (done) => {
+    const v = createPayload(1024);
+    const sumiPacket = createStreamPacket(v);
+
+    const decoder = new StreamPacketDecoder();
+    decoder.onData((data) => {
+      fastExpectBufferEqual(data, v);
+      done();
+    });
+
+    console.log('write chunk', sumiPacket.byteLength);
+
+    // use pressure = 2 to simulate the header and payload are separated
+    const pressure = 2;
+    for (let i = 0; i < sumiPacket.byteLength; i += pressure) {
+      decoder.push(sumiPacket.subarray(i, i + pressure));
+    }
+  });
 });
 
 function logMemoryUsage() {
@@ -103,3 +150,16 @@ function logMemoryUsage() {
 
   console.log(text);
 }
+
+function fastExpectBufferEqual(data: Uint8Array, expected: Uint8Array, confidenceLevel = 10) {
+  expect(data.byteLength).toEqual(expected.byteLength);
+
+  for (let i = 0; i < confidenceLevel; i++) {
+    // 如果对比整个 Uint8Array 的话，超大 Uint8Array 会很耗时
+    // 随机选一些数据(<= 1024 字节)，对比是否正确，
+    const start = Math.floor(Math.random() * data.byteLength);
+    const end = Math.floor(Math.random() * 1024);
+
+    expect(data.subarray(start, end)).toEqual(expected.subarray(start, end));
+  }
+}

packages/connection/__test__/node/ws-channel.test.ts

@@ -2,6 +2,7 @@ import net from 'net';
 
 import { WSChannel } from '@opensumi/ide-connection';
 import { normalizedIpcHandlerPathAsync } from '@opensumi/ide-core-common/lib/utils/ipc';
+
 const total = 1000;
 
 describe('ws channel node', () => {