Skip to content

Commit a02ca33

Browse files
committed
Add trio for step 4
1 parent 2293ebf commit a02ca33

File tree

2 files changed

+63
-0
lines changed

2 files changed

+63
-0
lines changed

pyredis/__main__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import argparse
22
import asyncio
3+
import trio
34

45
from pyredis.server import Server
56
from pyredis.asyncserver import RedisServerProtocol
@@ -20,6 +21,11 @@ async def amain(args):
2021
await server.serve_forever()
2122

2223

24+
async def tmain(args):
25+
server = Server(args.port)
26+
await server.run()
27+
28+
2329
def main(args):
2430
print(f"Starting PyRedis on port: {args.port}")
2531

@@ -38,6 +44,7 @@ def main(args):
3844
default=REDIS_DEFAULT_PORT,
3945
)
4046
parser.add_argument("--asyncio", action=argparse.BooleanOptionalAction)
47+
parser.add_argument("--trio", action=argparse.BooleanOptionalAction)
4148
parser.add_argument(
4249
"-v",
4350
"--verbose",
@@ -52,6 +59,9 @@ def main(args):
5259
if args.asyncio:
5360
logging.info("Using AsyncIO RedisServerProtocol")
5461
asyncio.run(amain(args))
62+
elif args.trio:
63+
logging.info("Using Trio Stream API")
64+
trio.run(tmain, args)
5565
else:
5666
logging.info("Using threading module for multi-threading")
5767
main(args)

pyredis/trioserver.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
from trio import SocketListener, serve_tcp, SocketStream
2+
import logging
3+
import trio
4+
5+
from pyredis.protocol import extract_frame_from_buffer, encode_message
6+
from pyredis.commands import handle_command
7+
from pyredis.datastore import DataStore
8+
9+
RECV_SIZE = 2048
10+
log = logging.getLogger("pyredis")
11+
12+
13+
class Server:
14+
def __init__(self, port) -> None:
15+
self.port = port
16+
self._running = False
17+
self._datastore = DataStore()
18+
19+
async def run(self):
20+
self._running = True
21+
22+
async with trio.open_nursery() as nursery:
23+
await nursery.start(
24+
serve_tcp,
25+
self.handle_client_connection,
26+
port=self.port,
27+
host="127.0.0.1",
28+
)
29+
30+
async def handle_client_connection(self, client_stream: SocketStream):
31+
buffer = bytearray()
32+
try:
33+
while True:
34+
data = await client_stream.receive_some(RECV_SIZE)
35+
log.info("Received data from client")
36+
if not data:
37+
log.info("Readched EOF")
38+
break
39+
buffer.extend(data)
40+
frame, frame_size = extract_frame_from_buffer(buffer)
41+
log.info("Extracted one frame from received data")
42+
if frame:
43+
buffer = buffer[frame_size:]
44+
log.info("Processing one frame")
45+
result = handle_command(frame, self._datastore)
46+
await client_stream.send_all(encode_message(result))
47+
48+
finally:
49+
log.info("Attempt to close stream")
50+
await client_stream.aclose()
51+
52+
def stop(self):
53+
self._running = False

0 commit comments

Comments
 (0)