File tree Expand file tree Collapse file tree 2 files changed +63
-0
lines changed Expand file tree Collapse file tree 2 files changed +63
-0
lines changed Original file line number Diff line number Diff line change 1
1
import argparse
2
2
import asyncio
3
+ import trio
3
4
4
5
from pyredis .server import Server
5
6
from pyredis .asyncserver import RedisServerProtocol
@@ -20,6 +21,11 @@ async def amain(args):
20
21
await server .serve_forever ()
21
22
22
23
24
+ async def tmain (args ):
25
+ server = Server (args .port )
26
+ await server .run ()
27
+
28
+
23
29
def main (args ):
24
30
print (f"Starting PyRedis on port: { args .port } " )
25
31
@@ -38,6 +44,7 @@ def main(args):
38
44
default = REDIS_DEFAULT_PORT ,
39
45
)
40
46
parser .add_argument ("--asyncio" , action = argparse .BooleanOptionalAction )
47
+ parser .add_argument ("--trio" , action = argparse .BooleanOptionalAction )
41
48
parser .add_argument (
42
49
"-v" ,
43
50
"--verbose" ,
@@ -52,6 +59,9 @@ def main(args):
52
59
if args .asyncio :
53
60
logging .info ("Using AsyncIO RedisServerProtocol" )
54
61
asyncio .run (amain (args ))
62
+ elif args .trio :
63
+ logging .info ("Using Trio Stream API" )
64
+ trio .run (tmain , args )
55
65
else :
56
66
logging .info ("Using threading module for multi-threading" )
57
67
main (args )
Original file line number Diff line number Diff line change
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
You can’t perform that action at this time.
0 commit comments