rpc: add intrinsic support for async iterators

koush · Mar 2, 2023 · 4106185 · 4106185
1 parent d91e625
commit 4106185
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 3 deletions.
diff --git a/server/src/rpc.ts b/server/src/rpc.ts
@@ -94,6 +94,17 @@ class RpcProxy implements PrimitiveProxyHandler<any> {
     }
 
     get(target: any, p: PropertyKey, receiver: any): any {
+        if (p === Symbol.asyncIterator) {
+            if (!this.proxyProps?.[Symbol.asyncIterator.toString()])
+                return;
+            return () => {
+                return {
+                    next: async () => {
+                        return this.apply(() => 'next', undefined)
+                    }
+                }
+            };
+        }
         if (p === RpcPeer.PROPERTY_PROXY_ID)
             return this.entry.id;
         if (p === '__proxy_constructor')
@@ -281,6 +292,15 @@ export class RpcPeer {
         }
     }
 
+    static getProxyProperies(value: any) {
+        if (!value[Symbol.asyncIterator])
+            return value?.[RpcPeer.PROPERTY_PROXY_PROPERTIES];
+        return {
+            [Symbol.asyncIterator.toString()]: true,
+            ...value?.[RpcPeer.PROPERTY_PROXY_PROPERTIES],
+        }
+    }
+
     static readonly PROPERTY_PROXY_ID = '__proxy_id';
     static readonly PROPERTY_PROXY_ONEWAY_METHODS = '__proxy_oneway_methods';
     static readonly PROPERTY_JSON_DISABLE_SERIALIZATION = '__json_disable_serialization';
@@ -456,7 +476,7 @@ export class RpcPeer {
                 __remote_proxy_id: proxiedEntry.id,
                 __remote_proxy_finalizer_id,
                 __remote_constructor_name,
-                __remote_proxy_props: value?.[RpcPeer.PROPERTY_PROXY_PROPERTIES],
+                __remote_proxy_props: RpcPeer.getProxyProperies(value),
                 __remote_proxy_oneway_methods: value?.[RpcPeer.PROPERTY_PROXY_ONEWAY_METHODS],
             }
             return ret;
@@ -481,7 +501,7 @@ export class RpcPeer {
                 __remote_proxy_id: undefined,
                 __remote_proxy_finalizer_id: undefined,
                 __remote_constructor_name,
-                __remote_proxy_props: value?.[RpcPeer.PROPERTY_PROXY_PROPERTIES],
+                __remote_proxy_props: RpcPeer.getProxyProperies(value),
                 __remote_proxy_oneway_methods: value?.[RpcPeer.PROPERTY_PROXY_ONEWAY_METHODS],
                 __serialized_value: serialized,
             }
@@ -500,7 +520,7 @@ export class RpcPeer {
             __remote_proxy_id,
             __remote_proxy_finalizer_id: __remote_proxy_id,
             __remote_constructor_name,
-            __remote_proxy_props: value?.[RpcPeer.PROPERTY_PROXY_PROPERTIES],
+            __remote_proxy_props: RpcPeer.getProxyProperies(value),
             __remote_proxy_oneway_methods: value?.[RpcPeer.PROPERTY_PROXY_ONEWAY_METHODS],
         }
 

diff --git a/server/test/rpc-iterator-test.ts b/server/test/rpc-iterator-test.ts
@@ -0,0 +1,27 @@
+import { RpcPeer } from "../src/rpc";
+
+const p1 = new RpcPeer('p1', 'p2', message => {
+    console.log('message p1 p2', message);
+    p2.handleMessage(message);
+});
+
+const p2 = new RpcPeer('p2', 'p1', message => {
+    console.log('message p2 p1', message);
+    p1.handleMessage(message);
+});
+
+async function* generator() {
+    yield 2;
+    yield 3;
+}
+
+p1.params['thing'] = generator();
+
+async function test() {
+    const foo = await p2.getParam('thing');
+    for await (const n of foo) {
+        console.log(n);
+    }
+}
+
+test();