fix(rpc): cancel cannot work

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

@@ -1,19 +1,8 @@
 /* eslint-disable no-console */
 import { BinaryReader } from '@furyjs/fury/dist/lib/reader';
-import { BinaryWriter } from '@furyjs/fury/dist/lib/writer';
 
 import { LengthFieldBasedFrameDecoder, indicator } from '../../src/common/connection/drivers/frame-decoder';
 
-const writer = BinaryWriter({});
-
-export function prependLengthField(content: Uint8Array) {
-  writer.reset();
-  writer.buffer(indicator);
-  writer.uint32(content.byteLength);
-  writer.buffer(content);
-  return writer.dump();
-}
-
 function round(x: number, count: number) {
   return Math.round(x * 10 ** count) / 10 ** count;
 }
@@ -46,7 +35,7 @@ console.timeEnd('createPayload');
 // 1m
 const pressure = 1024 * 1024;
 
-const purePackets = [p1k, p64k, p128k, p5m, p10m].map((v) => [prependLengthField(v), v] as const);
+const purePackets = [p1k, p64k, p128k, p5m, p10m].map((v) => [LengthFieldBasedFrameDecoder.construct(v), v] as const);
 
 const size = purePackets.reduce((acc, v) => acc + v[0].byteLength, 0);
 
@@ -59,7 +48,7 @@ purePackets.forEach((v) => {
 });
 
 const mixedPackets = [p1m, p5m].map((v) => {
-  const sumiPacket = prependLengthField(v);
+  const sumiPacket = LengthFieldBasedFrameDecoder.construct(v);
   const newPacket = createPayload(1024 + sumiPacket.byteLength);
   newPacket.set(sumiPacket, 1024);
   return [newPacket, v] as const;
@@ -70,7 +59,7 @@ const packets = [...purePackets, ...mixedPackets];
 describe('frame decoder', () => {
   it('can create frame', () => {
     const content = new Uint8Array([1, 2, 3]);
-    const packet = prependLengthField(content);
+    const packet = LengthFieldBasedFrameDecoder.construct(content);
     const reader = BinaryReader({});
 
     reader.reset(packet);
@@ -127,7 +116,7 @@ describe('frame decoder', () => {
 
   it('can decode a stream it has no valid length info', (done) => {
     const v = createPayload(1024);
-    const sumiPacket = prependLengthField(v);
+    const sumiPacket = LengthFieldBasedFrameDecoder.construct(v);
 
     const decoder = new LengthFieldBasedFrameDecoder();
     decoder.onData((data) => {

packages/connection/__test__/common/rpc/sumi-rpc.test.ts

@@ -1,7 +1,7 @@
 import { Type } from '@furyjs/fury';
 
 import { METHOD_NOT_REGISTERED } from '@opensumi/ide-connection/lib/common/constants';
-import { Deferred } from '@opensumi/ide-core-common';
+import { CancellationToken, CancellationTokenSource, Deferred, sleep } from '@opensumi/ide-core-common';
 import { IReadableStream, listenReadable } from '@opensumi/ide-utils/lib/stream';
 
 import { test } from './common-tester';
@@ -96,4 +96,24 @@ describe('sumi rpc only', () => {
     await deferred.promise;
     expect(msg).toBe(message);
   });
+
+  it('can cancel a call', (done) => {
+    const method = 'cancelTest';
+    pair.connection2.onRequest(method, async (t: CancellationToken) => {
+      while (!t.isCancellationRequested) {
+        // regular cancellation requires async for it to work
+        await sleep(0);
+      }
+
+      done();
+    });
+
+    pair.connection1.listen();
+    pair.connection2.listen();
+
+    const source = new CancellationTokenSource();
+
+    pair.connection1.sendRequest(method, source.token);
+    source.cancel();
+  });
 });

packages/connection/src/common/connection/drivers/frame-decoder.ts

@@ -1,3 +1,5 @@
+import { BinaryWriter } from '@furyjs/fury/dist/lib/writer';
+
 import { Emitter, readUInt32LE } from '@opensumi/ide-core-common';
 
 import { Buffers } from '../../buffers/buffers';
@@ -155,4 +157,14 @@ export class LengthFieldBasedFrameDecoder {
     this.dataEmitter.dispose();
     this.buffers.dispose();
   }
+
+  static writer = BinaryWriter({});
+
+  static construct(content: Uint8Array) {
+    LengthFieldBasedFrameDecoder.writer.reset();
+    LengthFieldBasedFrameDecoder.writer.buffer(indicator);
+    LengthFieldBasedFrameDecoder.writer.uint32(content.byteLength);
+    LengthFieldBasedFrameDecoder.writer.buffer(content);
+    return LengthFieldBasedFrameDecoder.writer.dump();
+  }
 }

packages/connection/src/common/connection/drivers/stream.ts

@@ -1,9 +1,7 @@
-import { IDisposable, writeUInt32LE } from '@opensumi/ide-core-common';
-
-import { pool } from '../../buffers';
+import { IDisposable } from '@opensumi/ide-core-common';
 
 import { BaseConnection } from './base';
-import { LengthFieldBasedFrameDecoder, indicator } from './frame-decoder';
+import { LengthFieldBasedFrameDecoder } from './frame-decoder';
 
 import type { Readable, Writable } from 'stream';
 
@@ -23,13 +21,10 @@ export class StreamConnection extends BaseConnection<Uint8Array> {
   }
 
   send(data: Uint8Array): void {
-    this.writable.write(indicator);
-    const buf = pool.alloc(4);
-    writeUInt32LE(buf, data.byteLength, 0);
-    this.writable.write(buf, () => {
-      pool.free(buf);
+    const result = LengthFieldBasedFrameDecoder.construct(data);
+    this.writable.write(result, () => {
+      // TODO: logger error
     });
-    this.writable.write(data);
   }
 
   onMessage(cb: (data: Uint8Array) => void): IDisposable {

packages/connection/src/common/fury-extends/any.ts

@@ -23,6 +23,8 @@ export class AnySerializer {
   write(data: any) {
     const { writer } = this;
     const type = typeof data;
+    writer.reserve(1);
+
     switch (type) {
       case 'undefined':
         writer.uint8(ProtocolType.Undefined);
@@ -32,10 +34,12 @@ export class AnySerializer {
         writer.stringOfVarUInt32(data);
         break;
       case 'boolean':
+        writer.reserve(1);
         writer.uint8(ProtocolType.Boolean);
         writer.uint8(data ? 1 : 0);
         break;
       case 'number':
+        writer.reserve(8);
         if ((data | 0) === data) {
           writer.uint8(ProtocolType.Int32);
           writer.int32(data);
@@ -45,19 +49,22 @@ export class AnySerializer {
         }
         break;
       case 'bigint':
+        writer.reserve(8);
         writer.uint8(ProtocolType.BigInt);
         writer.int64(data);
         break;
       case 'object':
         if (data === null) {
           writer.uint8(ProtocolType.Null);
         } else if (Array.isArray(data)) {
+          writer.reserve(4);
           writer.uint8(ProtocolType.Array);
           writer.varUInt32(data.length);
           for (const element of data) {
             this.write(element);
           }
         } else if (isUint8Array(data)) {
+          writer.reserve(4);
           writer.uint8(ProtocolType.Buffer);
           writer.varUInt32(data.byteLength);
           writer.buffer(data);

packages/connection/src/common/rpc/connection.ts

@@ -175,9 +175,8 @@ export class SumiConnection implements IDisposable {
       // skip version, currently only have version 1
       reader.skip(1);
 
-      const opType = reader.uint8();
+      const opType = reader.uint8() as OperationType;
       const requestId = reader.uint32();
-      const method = reader.stringOfVarUInt32();
 
       if (this._timeoutHandles.has(requestId)) {
         // Ignore some jest test scenarios where clearTimeout is not defined.
@@ -190,6 +189,7 @@ export class SumiConnection implements IDisposable {
 
       switch (opType) {
         case OperationType.Response: {
+          const method = reader.stringOfVarUInt32();
           const status = reader.uint16();
 
           const runCallback = (headers: IResponseHeaders, error?: any, result?: any) => {
@@ -246,6 +246,7 @@ export class SumiConnection implements IDisposable {
         case OperationType.Notification:
         // fall through
         case OperationType.Request: {
+          const method = reader.stringOfVarUInt32();
           const headers = this.io.requestHeadersSerializer.read() as IRequestHeaders;
           const args = this.io.getProcessor(method).readRequest();