@@ -3,6 +3,8 @@ use crate::execution_context::WorkloadType;
33use crate :: host:: module_host:: { EventStatus , ModuleEvent } ;
44use crate :: host:: ArgsTuple ;
55use crate :: messages:: websocket as ws;
6+ use bytes:: { BufMut , Bytes , BytesMut } ;
7+ use bytestring:: ByteString ;
68use derive_more:: From ;
79use spacetimedb_client_api_messages:: websocket:: {
810 BsatnFormat , Compression , FormatSwitch , JsonFormat , OneOffTable , RowListLen , WebsocketFormat ,
@@ -27,36 +29,131 @@ pub trait ToProtocol {
2729pub ( super ) type SwitchedServerMessage = FormatSwitch < ws:: ServerMessage < BsatnFormat > , ws:: ServerMessage < JsonFormat > > ;
2830pub ( super ) type SwitchedDbUpdate = FormatSwitch < ws:: DatabaseUpdate < BsatnFormat > , ws:: DatabaseUpdate < JsonFormat > > ;
2931
32+ /// The initial size of a `serialize` buffer.
33+ /// Currently 4k to align with the linux page size
34+ /// and this should be more than enough in the common case.
35+ const SERIALIZE_BUFFER_INIT_CAP : usize = 4096 ;
36+
37+ /// A buffer used by [`serialize`]
38+ pub struct SerializeBuffer {
39+ uncompressed : BytesMut ,
40+ compressed : BytesMut ,
41+ }
42+
43+ impl SerializeBuffer {
44+ pub fn new ( config : ClientConfig ) -> Self {
45+ let uncompressed_capacity = SERIALIZE_BUFFER_INIT_CAP ;
46+ let compressed_capacity = if config. compression == Compression :: None || config. protocol == Protocol :: Text {
47+ 0
48+ } else {
49+ SERIALIZE_BUFFER_INIT_CAP
50+ } ;
51+ Self {
52+ uncompressed : BytesMut :: with_capacity ( uncompressed_capacity) ,
53+ compressed : BytesMut :: with_capacity ( compressed_capacity) ,
54+ }
55+ }
56+
57+ /// Take the uncompressed message as the one to use.
58+ fn uncompressed ( self ) -> ( InUseSerializeBuffer , Bytes ) {
59+ let uncompressed = self . uncompressed . freeze ( ) ;
60+ let in_use = InUseSerializeBuffer :: Uncompressed {
61+ uncompressed : uncompressed. clone ( ) ,
62+ compressed : self . compressed ,
63+ } ;
64+ ( in_use, uncompressed)
65+ }
66+
67+ /// Write uncompressed data with a leading tag.
68+ fn write_with_tag < F > ( & mut self , tag : u8 , write : F ) -> & [ u8 ]
69+ where
70+ F : FnOnce ( bytes:: buf:: Writer < & mut BytesMut > ) ,
71+ {
72+ self . uncompressed . put_u8 ( tag) ;
73+ write ( ( & mut self . uncompressed ) . writer ( ) ) ;
74+ & self . uncompressed [ 1 ..]
75+ }
76+
77+ /// Compress the data from a `write_with_tag` call, and change the tag.
78+ fn compress_with_tag (
79+ self ,
80+ tag : u8 ,
81+ write : impl FnOnce ( & [ u8 ] , & mut bytes:: buf:: Writer < BytesMut > ) ,
82+ ) -> ( InUseSerializeBuffer , Bytes ) {
83+ let mut writer = self . compressed . writer ( ) ;
84+ writer. get_mut ( ) . put_u8 ( tag) ;
85+ write ( & self . uncompressed [ 1 ..] , & mut writer) ;
86+ let compressed = writer. into_inner ( ) . freeze ( ) ;
87+ let in_use = InUseSerializeBuffer :: Compressed {
88+ uncompressed : self . uncompressed ,
89+ compressed : compressed. clone ( ) ,
90+ } ;
91+ ( in_use, compressed)
92+ }
93+ }
94+
95+ type BytesMutWriter < ' a > = bytes:: buf:: Writer < & ' a mut BytesMut > ;
96+
97+ pub enum InUseSerializeBuffer {
98+ Uncompressed { uncompressed : Bytes , compressed : BytesMut } ,
99+ Compressed { uncompressed : BytesMut , compressed : Bytes } ,
100+ }
101+
102+ impl InUseSerializeBuffer {
103+ pub fn try_reclaim ( self ) -> Option < SerializeBuffer > {
104+ let ( mut uncompressed, mut compressed) = match self {
105+ Self :: Uncompressed {
106+ uncompressed,
107+ compressed,
108+ } => ( uncompressed. try_into_mut ( ) . ok ( ) ?, compressed) ,
109+ Self :: Compressed {
110+ uncompressed,
111+ compressed,
112+ } => ( uncompressed, compressed. try_into_mut ( ) . ok ( ) ?) ,
113+ } ;
114+ uncompressed. clear ( ) ;
115+ compressed. clear ( ) ;
116+ Some ( SerializeBuffer {
117+ uncompressed,
118+ compressed,
119+ } )
120+ }
121+ }
122+
30123/// Serialize `msg` into a [`DataMessage`] containing a [`ws::ServerMessage`].
31124///
32125/// If `protocol` is [`Protocol::Binary`],
33126/// the message will be conditionally compressed by this method according to `compression`.
34- pub fn serialize ( msg : impl ToProtocol < Encoded = SwitchedServerMessage > , config : ClientConfig ) -> DataMessage {
35- // TODO(centril, perf): here we are allocating buffers only to throw them away eventually.
36- // Consider pooling these allocations so that we reuse them.
127+ pub fn serialize (
128+ mut buffer : SerializeBuffer ,
129+ msg : impl ToProtocol < Encoded = SwitchedServerMessage > ,
130+ config : ClientConfig ,
131+ ) -> ( InUseSerializeBuffer , DataMessage ) {
37132 match msg. to_protocol ( config. protocol ) {
38- FormatSwitch :: Json ( msg) => serde_json:: to_string ( & SerializeWrapper :: new ( msg) ) . unwrap ( ) . into ( ) ,
133+ FormatSwitch :: Json ( msg) => {
134+ let out: BytesMutWriter < ' _ > = ( & mut buffer. uncompressed ) . writer ( ) ;
135+ serde_json:: to_writer ( out, & SerializeWrapper :: new ( msg) )
136+ . expect ( "should be able to json encode a `ServerMessage`" ) ;
137+
138+ let ( in_use, out) = buffer. uncompressed ( ) ;
139+ // SAFETY: `serde_json::to_writer` states that:
140+ // > "Serialization guarantees it only feeds valid UTF-8 sequences to the writer."
141+ let msg_json = unsafe { ByteString :: from_bytes_unchecked ( out) } ;
142+ ( in_use, msg_json. into ( ) )
143+ }
39144 FormatSwitch :: Bsatn ( msg) => {
40145 // First write the tag so that we avoid shifting the entire message at the end.
41- let mut msg_bytes = vec ! [ SERVER_MSG_COMPRESSION_TAG_NONE ] ;
42- bsatn:: to_writer ( & mut msg_bytes, & msg) . unwrap ( ) ;
146+ let srv_msg = buffer. write_with_tag ( SERVER_MSG_COMPRESSION_TAG_NONE , |w| {
147+ bsatn:: to_writer ( w. into_inner ( ) , & msg) . unwrap ( )
148+ } ) ;
43149
44150 // Conditionally compress the message.
45- let srv_msg = & msg_bytes[ 1 ..] ;
46- let msg_bytes = match ws:: decide_compression ( srv_msg. len ( ) , config. compression ) {
47- Compression :: None => msg_bytes,
48- Compression :: Brotli => {
49- let mut out = vec ! [ SERVER_MSG_COMPRESSION_TAG_BROTLI ] ;
50- ws:: brotli_compress ( srv_msg, & mut out) ;
51- out
52- }
53- Compression :: Gzip => {
54- let mut out = vec ! [ SERVER_MSG_COMPRESSION_TAG_GZIP ] ;
55- ws:: gzip_compress ( srv_msg, & mut out) ;
56- out
57- }
151+ let ( in_use, msg_bytes) = match ws:: decide_compression ( srv_msg. len ( ) , config. compression ) {
152+ Compression :: None => buffer. uncompressed ( ) ,
153+ Compression :: Brotli => buffer. compress_with_tag ( SERVER_MSG_COMPRESSION_TAG_BROTLI , ws:: brotli_compress) ,
154+ Compression :: Gzip => buffer. compress_with_tag ( SERVER_MSG_COMPRESSION_TAG_GZIP , ws:: gzip_compress) ,
58155 } ;
59- msg_bytes. into ( )
156+ ( in_use , msg_bytes. into ( ) )
60157 }
61158 }
62159}
0 commit comments