From dda0b22870a512e86d2bbfc0d82293c8e2d881d6 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Sun, 30 May 2021 17:19:20 -0400
Subject: [PATCH] Try out `msgspec` in our msgpack stream channel

Can only really use an encoder currently since there is no streaming api
in `msgspec` as of currently. See jcrist/msgspec#27.

Not sure if any encoding speedups are currently noticeable especially
without any validation going on yet XD.

First experiments toward #196
---
 tractor/_ipc.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tractor/_ipc.py b/tractor/_ipc.py
index 08057e9..6051a15 100644
--- a/tractor/_ipc.py
+++ b/tractor/_ipc.py
@@ -4,10 +4,11 @@ Inter-process comms abstractions
 """
 import platform
 import typing
-from typing import Any, Tuple, Optional
+from typing import Any, Tuple, Optional, Callable
 from functools import partial
 
 import msgpack
+import msgspec
 import trio
 from async_generator import asynccontextmanager
 
@@ -27,6 +28,9 @@ except ImportError:
     Unpacker = partial(msgpack.Unpacker, strict_map_key=False)
 
 
+ms_decode = msgspec.Encoder().encode
+
+
 class MsgpackTCPStream:
     '''A ``trio.SocketStream`` delivering ``msgpack`` formatted data
     using ``msgpack-python``.
@@ -35,15 +39,16 @@ class MsgpackTCPStream:
     def __init__(
         self,
         stream: trio.SocketStream,
-
     ) -> None:
 
         self.stream = stream
         assert self.stream.socket
+
         # should both be IP sockets
         lsockname = stream.socket.getsockname()
         assert isinstance(lsockname, tuple)
         self._laddr = lsockname[:2]
+
         rsockname = stream.socket.getpeername()
         assert isinstance(rsockname, tuple)
         self._raddr = rsockname[:2]
@@ -61,6 +66,7 @@ class MsgpackTCPStream:
             raw=False,
             use_list=False,
         )
+        # decoder = msgspec.Decoder()  #dict[str, Any])
         while True:
             try:
                 data = await self.stream.receive_some(2**10)
@@ -95,6 +101,7 @@ class MsgpackTCPStream:
                     f'transport {self} was already closed prior ro read'
                 )
 
+            # yield decoder.decode(data)
             unpacker.feed(data)
             for packet in unpacker:
                 yield packet
@@ -111,7 +118,9 @@ class MsgpackTCPStream:
     async def send(self, data: Any) -> None:
         async with self._send_lock:
             return await self.stream.send_all(
-                msgpack.dumps(data, use_bin_type=True))
+                # msgpack.dumps(data, use_bin_type=True))
+                ms_decode(data)
+            )
 
     async def recv(self) -> Any:
         return await self._agen.asend(None)