shardcache_client_rs/
client.rs

1#[cfg(feature = "redis")]
2use std::collections::VecDeque;
3use std::net::ToSocketAddrs;
4use std::time::Duration;
5
6#[cfg(feature = "tls")]
7use crate::ScnpTlsClientConfig;
8use crate::commands::del::{self, Del};
9use crate::commands::exists::{self, Exists};
10use crate::commands::expire::{self, Expire};
11use crate::commands::get::{self, Get};
12use crate::commands::getex::{self, GetEx};
13#[cfg(feature = "redis")]
14use crate::commands::redis::{
15    self, RedisCommand as OptimizedRedisCommand, RedisCommandKind, RedisCommandRouteKeys,
16    RedisRespCommand, RedisResponse,
17};
18use crate::commands::resp::RespCommand;
19use crate::commands::set::{self, Set};
20use crate::commands::setex::{self, SetEx};
21use crate::commands::ttl::{self, Ttl};
22use crate::connection::ScnpConnection;
23use crate::error::{Result, ShardCacheClientError};
24#[cfg(feature = "redis")]
25use crate::routing::ShardCacheRoute;
26use crate::routing::{ShardCacheDirectRouter, ShardCacheRouteMode};
27
28#[cfg(feature = "redis")]
29#[derive(Debug, Clone, Copy)]
30enum RedisPipelineResponse {
31    Native,
32    Resp,
33}
34
35/// Blocking SCNP client for the ordinary server listener.
36#[derive(Debug)]
37pub struct ShardCacheClient {
38    conn: ScnpConnection,
39    #[cfg(feature = "redis")]
40    redis_pipeline_responses: VecDeque<RedisPipelineResponse>,
41}
42
43/// Topology advertised by a shardcache server bootstrap listener.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct ShardCacheTopology {
46    pub node_id: String,
47    pub shard_count: usize,
48    pub route_mode: String,
49    pub direct_shard_base_port: u16,
50    pub capabilities: Vec<String>,
51}
52
53impl ShardCacheClient {
54    /// Connects to a shardcache server listener that accepts generic SCNP.
55    pub fn connect(addr: impl ToSocketAddrs) -> Result<Self> {
56        Ok(Self {
57            conn: ScnpConnection::connect(addr)?,
58            #[cfg(feature = "redis")]
59            redis_pipeline_responses: VecDeque::new(),
60        })
61    }
62
63    /// Connects with explicit TCP connect and per-operation I/O deadlines.
64    pub fn connect_with_timeouts(
65        addr: impl ToSocketAddrs,
66        connect_timeout: Duration,
67        operation_timeout: Duration,
68    ) -> Result<Self> {
69        Self::connect_with_timeouts_and_auth(addr, connect_timeout, operation_timeout, None)
70    }
71
72    /// Connects with deadlines and authenticates before issuing SCNP commands.
73    pub fn connect_with_timeouts_and_auth(
74        addr: impl ToSocketAddrs,
75        connect_timeout: Duration,
76        operation_timeout: Duration,
77        auth_token: Option<&[u8]>,
78    ) -> Result<Self> {
79        let mut conn =
80            ScnpConnection::connect_with_timeouts(addr, connect_timeout, operation_timeout)?;
81        conn.authenticate(auth_token)?;
82        Ok(Self {
83            conn,
84            #[cfg(feature = "redis")]
85            redis_pipeline_responses: VecDeque::new(),
86        })
87    }
88
89    /// Connects with TLS, deadlines, and optional SCNP token authentication.
90    #[cfg(feature = "tls")]
91    pub fn connect_with_timeouts_auth_and_tls(
92        addr: impl ToSocketAddrs,
93        connect_timeout: Duration,
94        operation_timeout: Duration,
95        auth_token: Option<&[u8]>,
96        tls: &ScnpTlsClientConfig,
97    ) -> Result<Self> {
98        let mut conn = ScnpConnection::connect_with_timeouts_and_tls(
99            addr,
100            connect_timeout,
101            operation_timeout,
102            tls,
103        )?;
104        conn.authenticate(auth_token)?;
105        Ok(Self {
106            conn,
107            #[cfg(feature = "redis")]
108            redis_pipeline_responses: VecDeque::new(),
109        })
110    }
111
112    /// Reads `key` into `out`, returning `true` on hit.
113    pub fn get_into(&mut self, key: &[u8], out: &mut Vec<u8>) -> Result<bool> {
114        self.conn.execute(Get::new(key, out))
115    }
116
117    /// Reads `key` while rejecting response bodies larger than `max_body_len`.
118    pub fn get_into_limited(
119        &mut self,
120        key: &[u8],
121        out: &mut Vec<u8>,
122        max_body_len: usize,
123    ) -> Result<bool> {
124        crate::commands::get::write_request(&mut self.conn, None, key)?;
125        self.conn.flush()?;
126        self.conn.read_value_limited(
127            <Get as crate::commands::ScnpCommand>::NAME,
128            out,
129            max_body_len,
130        )
131    }
132
133    /// Sets `key` to `value`.
134    pub fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
135        self.conn.execute(Set::new(key, value))
136    }
137
138    /// Sets `key` to `value` with a millisecond TTL.
139    pub fn set_ex(&mut self, key: &[u8], value: &[u8], ttl_ms: u64) -> Result<()> {
140        self.conn.execute(SetEx::new(key, value, ttl_ms))
141    }
142
143    /// Reads `key` into `out` and sets a millisecond TTL, returning `true` on hit.
144    pub fn get_ex_into(&mut self, key: &[u8], ttl_ms: u64, out: &mut Vec<u8>) -> Result<bool> {
145        self.conn.execute(GetEx::new(key, ttl_ms, out))
146    }
147
148    /// Deletes `key`, returning `true` when an entry was removed.
149    pub fn del(&mut self, key: &[u8]) -> Result<bool> {
150        self.conn.execute(Del::new(key))
151    }
152
153    /// Returns whether `key` exists.
154    pub fn exists(&mut self, key: &[u8]) -> Result<bool> {
155        self.conn.execute(Exists::new(key))
156    }
157
158    /// Returns Redis-compatible TTL seconds for `key`.
159    pub fn ttl(&mut self, key: &[u8]) -> Result<i64> {
160        self.conn.execute(Ttl::new(key))
161    }
162
163    /// Sets a millisecond TTL on `key`, returning `true` when the TTL changed.
164    pub fn expire(&mut self, key: &[u8], ttl_ms: u64) -> Result<bool> {
165        self.conn.execute(Expire::new(key, ttl_ms))
166    }
167
168    /// Returns the first-party Redis command namespace.
169    #[cfg(feature = "redis")]
170    pub fn redis(&mut self) -> crate::Redis<'_, Self> {
171        crate::Redis::new(self)
172    }
173
174    /// Executes a Redis-compatible command through the compact opcode SCNP wrapper.
175    #[cfg(feature = "redis")]
176    pub fn redis_command(
177        &mut self,
178        command: RedisCommandKind,
179        args: &[&[u8]],
180    ) -> Result<RedisResponse> {
181        self.conn.execute(OptimizedRedisCommand::new(command, args))
182    }
183
184    /// Executes a Redis-compatible command by name through native SCNP.
185    ///
186    /// Commands with compact opcodes use the optimized Redis wrapper. Other
187    /// names use the SCNP command-name wrapper and return decoded RESP.
188    #[cfg(feature = "redis")]
189    pub fn redis_command_by_name(
190        &mut self,
191        command: &[u8],
192        args: &[&[u8]],
193    ) -> Result<RedisResponse> {
194        match RedisCommandKind::from_name(command) {
195            Some(command) => self.redis_command(command, args),
196            None => self.redis_resp_command(command, args),
197        }
198    }
199
200    /// Executes a Redis-compatible command through the SCNP command-name wrapper.
201    ///
202    /// This path is still native SCNP, but it carries the Redis command name in
203    /// the body so it can cover commands that do not have a compact opcode.
204    #[cfg(feature = "redis")]
205    pub fn redis_resp_command(&mut self, command: &[u8], args: &[&[u8]]) -> Result<RedisResponse> {
206        validate_redis_command_name(command)?;
207        self.conn.execute(RedisRespCommand::new(command, args))
208    }
209
210    /// Executes a Redis-compatible command through the generic SCNP wrapper.
211    ///
212    /// The server returns RESP bytes as an SCNP value. `out` receives those raw
213    /// bytes so callers can decode exactly the shape they requested.
214    pub fn resp_command_into(&mut self, parts: &[&[u8]], out: &mut Vec<u8>) -> Result<bool> {
215        self.conn.execute(RespCommand::new(parts, out))
216    }
217
218    /// Runs the global SCNP scan wrapper and returns the RESP scan reply bytes.
219    pub fn scan_resp_into(&mut self, cursor: u64, count: usize, out: &mut Vec<u8>) -> Result<bool> {
220        let cursor = cursor.to_string();
221        let count = count.to_string();
222        self.resp_command_into(
223            &[b"SCNP.SCAN", cursor.as_bytes(), b"COUNT", count.as_bytes()],
224            out,
225        )
226    }
227
228    /// Runs a shard-local SCNP scan. Call this concurrently per shard to avoid
229    /// a server-side fanout scan.
230    pub fn scan_shard_resp_into(
231        &mut self,
232        shard_id: usize,
233        cursor: u64,
234        count: usize,
235        out: &mut Vec<u8>,
236    ) -> Result<bool> {
237        let shard_id = shard_id.to_string();
238        let cursor = cursor.to_string();
239        let count = count.to_string();
240        self.resp_command_into(
241            &[
242                b"SCNP.SCANSHARD",
243                shard_id.as_bytes(),
244                cursor.as_bytes(),
245                b"COUNT",
246                count.as_bytes(),
247            ],
248            out,
249        )
250    }
251
252    /// Reads the server's stable topology and transport capabilities.
253    pub fn topology(&mut self) -> Result<ShardCacheTopology> {
254        let mut response = Vec::new();
255        if !self.resp_command_into(&[b"SCNP.TOPOLOGY"], &mut response)? {
256            return Err(ShardCacheClientError::Protocol(
257                "SCNP.TOPOLOGY returned null".into(),
258            ));
259        }
260        parse_topology_response(&response)
261    }
262
263    /// Writes a GET request without flushing or reading its response.
264    pub fn begin_pipeline_get(&mut self, key: &[u8]) -> Result<()> {
265        get::write_request(&mut self.conn, None, key)
266    }
267
268    /// Writes a SET request without flushing or reading its response.
269    pub fn begin_pipeline_set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
270        set::write_request(&mut self.conn, None, key, value)
271    }
272
273    /// Writes a SETEX request without flushing or reading its response.
274    pub fn begin_pipeline_set_ex(&mut self, key: &[u8], value: &[u8], ttl_ms: u64) -> Result<()> {
275        setex::write_request(&mut self.conn, None, key, value, ttl_ms)
276    }
277
278    /// Writes a GETEX request without flushing or reading its response.
279    pub fn begin_pipeline_get_ex(&mut self, key: &[u8], ttl_ms: u64) -> Result<()> {
280        getex::write_request(&mut self.conn, None, key, ttl_ms)
281    }
282
283    /// Writes a DEL request without flushing or reading its response.
284    pub fn begin_pipeline_del(&mut self, key: &[u8]) -> Result<()> {
285        del::write_request(&mut self.conn, None, key)
286    }
287
288    /// Writes an EXISTS request without flushing or reading its response.
289    pub fn begin_pipeline_exists(&mut self, key: &[u8]) -> Result<()> {
290        exists::write_request(&mut self.conn, None, key)
291    }
292
293    /// Writes a TTL request without flushing or reading its response.
294    pub fn begin_pipeline_ttl(&mut self, key: &[u8]) -> Result<()> {
295        ttl::write_request(&mut self.conn, None, key)
296    }
297
298    /// Writes an EXPIRE request without flushing or reading its response.
299    pub fn begin_pipeline_expire(&mut self, key: &[u8], ttl_ms: u64) -> Result<()> {
300        expire::write_request(&mut self.conn, None, key, ttl_ms)
301    }
302
303    /// Writes a compact Redis command request without flushing or reading its response.
304    #[cfg(feature = "redis")]
305    pub fn begin_pipeline_redis_command(
306        &mut self,
307        command: RedisCommandKind,
308        args: &[&[u8]],
309    ) -> Result<()> {
310        redis::write_request(&mut self.conn, command, None, args)?;
311        self.redis_pipeline_responses
312            .push_back(RedisPipelineResponse::Native);
313        Ok(())
314    }
315
316    /// Writes a Redis command request by name without flushing or reading its response.
317    ///
318    /// Compact-opcode commands use the optimized Redis wrapper. Other command
319    /// names use the SCNP command-name wrapper and decode the RESP payload when
320    /// [`finish_pipeline_redis_command`](Self::finish_pipeline_redis_command)
321    /// is called.
322    #[cfg(feature = "redis")]
323    pub fn begin_pipeline_redis_command_by_name(
324        &mut self,
325        command: &[u8],
326        args: &[&[u8]],
327    ) -> Result<()> {
328        match RedisCommandKind::from_name(command) {
329            Some(command) => self.begin_pipeline_redis_command(command, args),
330            None => self.begin_pipeline_redis_resp_command(command, args),
331        }
332    }
333
334    /// Writes a Redis command-name wrapper request without flushing or reading its response.
335    #[cfg(feature = "redis")]
336    pub fn begin_pipeline_redis_resp_command(
337        &mut self,
338        command: &[u8],
339        args: &[&[u8]],
340    ) -> Result<()> {
341        validate_redis_command_name(command)?;
342        redis::write_resp_request(&mut self.conn, command, args)?;
343        self.redis_pipeline_responses
344            .push_back(RedisPipelineResponse::Resp);
345        Ok(())
346    }
347
348    /// Flushes all queued pipelined requests.
349    pub fn flush_pipeline(&mut self) -> Result<()> {
350        self.conn.flush()
351    }
352
353    /// Reads the next pipelined GET response.
354    pub fn finish_pipeline_get_into(&mut self, out: &mut Vec<u8>) -> Result<bool> {
355        self.conn
356            .read_value(<Get as crate::commands::ScnpCommand>::NAME, out)
357    }
358
359    /// Reads the next pipelined SET response.
360    pub fn finish_pipeline_set(&mut self) -> Result<()> {
361        self.conn
362            .expect_ok(<Set as crate::commands::ScnpCommand>::NAME)
363    }
364
365    /// Reads the next pipelined SETEX response.
366    pub fn finish_pipeline_set_ex(&mut self) -> Result<()> {
367        self.conn
368            .expect_ok(<SetEx as crate::commands::ScnpCommand>::NAME)
369    }
370
371    /// Reads the next pipelined GETEX response.
372    pub fn finish_pipeline_get_ex_into(&mut self, out: &mut Vec<u8>) -> Result<bool> {
373        self.conn
374            .read_value(<GetEx as crate::commands::ScnpCommand>::NAME, out)
375    }
376
377    /// Reads the next pipelined DEL response.
378    pub fn finish_pipeline_del(&mut self) -> Result<bool> {
379        self.conn
380            .read_integer(<Del as crate::commands::ScnpCommand>::NAME)
381            .map(|deleted| deleted != 0)
382    }
383
384    /// Reads the next pipelined EXISTS response.
385    pub fn finish_pipeline_exists(&mut self) -> Result<bool> {
386        self.conn
387            .read_integer(<Exists as crate::commands::ScnpCommand>::NAME)
388            .map(|exists| exists != 0)
389    }
390
391    /// Reads the next pipelined TTL response.
392    pub fn finish_pipeline_ttl(&mut self) -> Result<i64> {
393        self.conn
394            .read_integer(<Ttl as crate::commands::ScnpCommand>::NAME)
395    }
396
397    /// Reads the next pipelined EXPIRE response.
398    pub fn finish_pipeline_expire(&mut self) -> Result<bool> {
399        self.conn
400            .read_integer(<Expire as crate::commands::ScnpCommand>::NAME)
401            .map(|changed| changed != 0)
402    }
403
404    /// Reads the next pipelined compact Redis command response.
405    #[cfg(feature = "redis")]
406    pub fn finish_pipeline_redis_command(&mut self) -> Result<RedisResponse> {
407        match self
408            .redis_pipeline_responses
409            .pop_front()
410            .unwrap_or(RedisPipelineResponse::Native)
411        {
412            RedisPipelineResponse::Native => self.conn.read_native_redis_response("REDIS"),
413            RedisPipelineResponse::Resp => self.conn.read_resp_redis_response("RESP"),
414        }
415    }
416}
417
418fn parse_topology_response(response: &[u8]) -> Result<ShardCacheTopology> {
419    if response.first() != Some(&b'$') {
420        return Err(ShardCacheClientError::Protocol(
421            "SCNP.TOPOLOGY did not return a RESP bulk string".into(),
422        ));
423    }
424    let header_end = response
425        .windows(2)
426        .position(|bytes| bytes == b"\r\n")
427        .ok_or_else(|| ShardCacheClientError::Protocol("invalid topology RESP header".into()))?;
428    let length = std::str::from_utf8(&response[1..header_end])
429        .ok()
430        .and_then(|value| value.parse::<usize>().ok())
431        .ok_or_else(|| ShardCacheClientError::Protocol("invalid topology RESP length".into()))?;
432    let body_start = header_end + 2;
433    let body_end = body_start.checked_add(length).ok_or_else(|| {
434        ShardCacheClientError::Protocol("topology response length overflow".into())
435    })?;
436    if response.len() < body_end + 2 || &response[body_end..body_end + 2] != b"\r\n" {
437        return Err(ShardCacheClientError::Protocol(
438            "truncated topology RESP body".into(),
439        ));
440    }
441    let body = std::str::from_utf8(&response[body_start..body_end])
442        .map_err(|_| ShardCacheClientError::Protocol("topology body is not UTF-8".into()))?;
443    let mut fields = body.split('\t');
444    let node_id = fields.next().unwrap_or_default().to_string();
445    let shard_count = fields
446        .next()
447        .and_then(|value| value.parse::<usize>().ok())
448        .ok_or_else(|| ShardCacheClientError::Protocol("invalid topology shard count".into()))?;
449    let route_mode = fields.next().unwrap_or_default().to_string();
450    let direct_shard_base_port = fields
451        .next()
452        .and_then(|value| value.parse::<u16>().ok())
453        .ok_or_else(|| ShardCacheClientError::Protocol("invalid topology direct port".into()))?;
454    let capabilities = fields.map(str::to_string).collect::<Vec<_>>();
455    Ok(ShardCacheTopology {
456        node_id,
457        shard_count,
458        route_mode,
459        direct_shard_base_port,
460        capabilities,
461    })
462}
463
464impl ShardCacheDirectRouter {
465    /// Connects directly to one shard-owned port.
466    pub fn connect_shard(&self, shard_id: usize) -> Result<ShardCacheDirectShardClient> {
467        Ok(ShardCacheDirectShardClient {
468            router: *self,
469            shard_id,
470            conn: ScnpConnection::connect(self.shard_addr(shard_id)?)?,
471        })
472    }
473
474    /// Connects directly to one shard-owned port with explicit deadlines.
475    pub fn connect_shard_with_timeouts(
476        &self,
477        shard_id: usize,
478        connect_timeout: Duration,
479        operation_timeout: Duration,
480    ) -> Result<ShardCacheDirectShardClient> {
481        self.connect_shard_with_timeouts_and_auth(
482            shard_id,
483            connect_timeout,
484            operation_timeout,
485            None,
486        )
487    }
488
489    /// Connects to one shard-owned port with deadlines and SCNP authentication.
490    pub fn connect_shard_with_timeouts_and_auth(
491        &self,
492        shard_id: usize,
493        connect_timeout: Duration,
494        operation_timeout: Duration,
495        auth_token: Option<&[u8]>,
496    ) -> Result<ShardCacheDirectShardClient> {
497        let mut conn = ScnpConnection::connect_with_timeouts(
498            self.shard_addr(shard_id)?,
499            connect_timeout,
500            operation_timeout,
501        )?;
502        conn.authenticate(auth_token)?;
503        Ok(ShardCacheDirectShardClient {
504            router: *self,
505            shard_id,
506            conn,
507        })
508    }
509
510    /// Connects directly to one shard-owned TLS port with deadlines and auth.
511    #[cfg(feature = "tls")]
512    pub fn connect_shard_with_timeouts_auth_and_tls(
513        &self,
514        shard_id: usize,
515        connect_timeout: Duration,
516        operation_timeout: Duration,
517        auth_token: Option<&[u8]>,
518        tls: &ScnpTlsClientConfig,
519    ) -> Result<ShardCacheDirectShardClient> {
520        let mut conn = ScnpConnection::connect_with_timeouts_and_tls(
521            self.shard_addr(shard_id)?,
522            connect_timeout,
523            operation_timeout,
524            tls,
525        )?;
526        conn.authenticate(auth_token)?;
527        Ok(ShardCacheDirectShardClient {
528            router: *self,
529            shard_id,
530            conn,
531        })
532    }
533}
534
535/// Blocking SCNP client that automatically routes each key to its shard port.
536#[derive(Debug)]
537pub struct ShardCacheDirectClient {
538    router: ShardCacheDirectRouter,
539    conns: Vec<ScnpConnection>,
540}
541
542impl ShardCacheDirectClient {
543    /// Connects to every shard-owned port starting at `addr`.
544    ///
545    /// `addr` must be the first direct shard port, not the fanout port.
546    pub fn connect(addr: impl ToSocketAddrs, shard_count: usize) -> Result<Self> {
547        let router = ShardCacheDirectRouter::new(addr, shard_count)?;
548        Self::connect_with_router(router)
549    }
550
551    /// Connects to every shard-owned port using an explicit route mode.
552    pub fn connect_with_route_mode(
553        addr: impl ToSocketAddrs,
554        shard_count: usize,
555        route_mode: ShardCacheRouteMode,
556    ) -> Result<Self> {
557        let router = ShardCacheDirectRouter::new(addr, shard_count)?.with_route_mode(route_mode);
558        Self::connect_with_router(router)
559    }
560
561    fn connect_with_router(router: ShardCacheDirectRouter) -> Result<Self> {
562        let mut conns = Vec::with_capacity(router.shard_count());
563        for shard_id in 0..router.shard_count() {
564            conns.push(ScnpConnection::connect(router.shard_addr(shard_id)?)?);
565        }
566        Ok(Self { router, conns })
567    }
568
569    /// Reads `key` from its owning shard into `out`, returning `true` on hit.
570    pub fn get_into(&mut self, key: &[u8], out: &mut Vec<u8>) -> Result<bool> {
571        let route = self.router.route_key(key);
572        self.conns[route.shard_id].execute(Get::routed(route, key, out))
573    }
574
575    /// Sets `key` on its owning shard.
576    pub fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
577        let route = self.router.route_key(key);
578        self.conns[route.shard_id].execute(Set::routed(route, key, value))
579    }
580
581    /// Sets `key` on its owning shard with a millisecond TTL.
582    pub fn set_ex(&mut self, key: &[u8], value: &[u8], ttl_ms: u64) -> Result<()> {
583        let route = self.router.route_key(key);
584        self.conns[route.shard_id].execute(SetEx::routed(route, key, value, ttl_ms))
585    }
586
587    /// Reads `key` from its owning shard into `out` and sets a millisecond TTL.
588    pub fn get_ex_into(&mut self, key: &[u8], ttl_ms: u64, out: &mut Vec<u8>) -> Result<bool> {
589        let route = self.router.route_key(key);
590        self.conns[route.shard_id].execute(GetEx::routed(route, key, ttl_ms, out))
591    }
592
593    /// Deletes `key` from its owning shard.
594    pub fn del(&mut self, key: &[u8]) -> Result<bool> {
595        let route = self.router.route_key(key);
596        self.conns[route.shard_id].execute(Del::routed(route, key))
597    }
598
599    /// Returns whether `key` exists on its owning shard.
600    pub fn exists(&mut self, key: &[u8]) -> Result<bool> {
601        let route = self.router.route_key(key);
602        self.conns[route.shard_id].execute(Exists::routed(route, key))
603    }
604
605    /// Returns Redis-compatible TTL seconds for `key` on its owning shard.
606    pub fn ttl(&mut self, key: &[u8]) -> Result<i64> {
607        let route = self.router.route_key(key);
608        self.conns[route.shard_id].execute(Ttl::routed(route, key))
609    }
610
611    /// Sets a millisecond TTL on `key` on its owning shard.
612    pub fn expire(&mut self, key: &[u8], ttl_ms: u64) -> Result<bool> {
613        let route = self.router.route_key(key);
614        self.conns[route.shard_id].execute(Expire::routed(route, key, ttl_ms))
615    }
616
617    /// Returns the first-party Redis command namespace for direct shard routing.
618    #[cfg(feature = "redis")]
619    pub fn redis(&mut self) -> crate::Redis<'_, Self> {
620        crate::Redis::new(self)
621    }
622
623    /// Executes a compact Redis command on the owning direct shard.
624    ///
625    /// Commands that require all shards are rejected; use [`ShardCacheClient`] against
626    /// the fanout listener for those.
627    #[cfg(feature = "redis")]
628    pub fn redis_command(
629        &mut self,
630        command: RedisCommandKind,
631        args: &[&[u8]],
632    ) -> Result<RedisResponse> {
633        let route = redis_direct_route(&self.router, command, args)?;
634        let shard_id = route.map_or(0, |route| route.shard_id);
635        self.conns[shard_id].execute(OptimizedRedisCommand::routed(command, route, args))
636    }
637
638    /// Executes a compact Redis command by name on the owning direct shard.
639    #[cfg(feature = "redis")]
640    pub fn redis_command_by_name(
641        &mut self,
642        command: &[u8],
643        args: &[&[u8]],
644    ) -> Result<RedisResponse> {
645        self.redis_command(redis_command_kind_from_name(command)?, args)
646    }
647
648    /// Runs a shard-local SCNP scan on one direct shard connection. Callers can
649    /// invoke this for different shards from different threads for parallel
650    /// scans.
651    pub fn scan_shard_resp_into(
652        &mut self,
653        shard_id: usize,
654        cursor: u64,
655        count: usize,
656        out: &mut Vec<u8>,
657    ) -> Result<bool> {
658        if shard_id >= self.conns.len() {
659            return Err(ShardCacheClientError::Config(format!(
660                "shard {shard_id} is outside configured shard count {}",
661                self.conns.len()
662            )));
663        }
664        let shard_id_text = shard_id.to_string();
665        let cursor = cursor.to_string();
666        let count = count.to_string();
667        self.conns[shard_id].execute(RespCommand::new(
668            &[
669                b"SCNP.SCANSHARD",
670                shard_id_text.as_bytes(),
671                cursor.as_bytes(),
672                b"COUNT",
673                count.as_bytes(),
674            ],
675            out,
676        ))
677    }
678}
679
680/// Blocking SCNP client pinned to one shard-owned port.
681///
682/// This is useful for thread-per-shard clients that pre-partition work.
683#[derive(Debug)]
684pub struct ShardCacheDirectShardClient {
685    router: ShardCacheDirectRouter,
686    shard_id: usize,
687    conn: ScnpConnection,
688}
689
690impl ShardCacheDirectShardClient {
691    /// Returns the shard this client is connected to.
692    pub fn shard_id(&self) -> usize {
693        self.shard_id
694    }
695
696    /// Reads `key` into `out`, returning `true` on hit.
697    pub fn get_into(&mut self, key: &[u8], out: &mut Vec<u8>) -> Result<bool> {
698        let route = self.checked_route(key)?;
699        self.conn.execute(Get::routed(route, key, out))
700    }
701
702    /// Reads `key` while rejecting response bodies larger than `max_body_len`.
703    pub fn get_into_limited(
704        &mut self,
705        key: &[u8],
706        out: &mut Vec<u8>,
707        max_body_len: usize,
708    ) -> Result<bool> {
709        let route = self.checked_route(key)?;
710        crate::commands::get::write_request(&mut self.conn, Some(route), key)?;
711        self.conn.flush()?;
712        self.conn.read_value_limited(
713            <Get as crate::commands::ScnpCommand>::NAME,
714            out,
715            max_body_len,
716        )
717    }
718
719    /// Sets `key` to `value`.
720    pub fn set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
721        let route = self.checked_route(key)?;
722        self.conn.execute(Set::routed(route, key, value))
723    }
724
725    /// Sets `key` to `value` with a millisecond TTL.
726    pub fn set_ex(&mut self, key: &[u8], value: &[u8], ttl_ms: u64) -> Result<()> {
727        let route = self.checked_route(key)?;
728        self.conn.execute(SetEx::routed(route, key, value, ttl_ms))
729    }
730
731    /// Reads `key` into `out` and sets a millisecond TTL, returning `true` on hit.
732    pub fn get_ex_into(&mut self, key: &[u8], ttl_ms: u64, out: &mut Vec<u8>) -> Result<bool> {
733        let route = self.checked_route(key)?;
734        self.conn.execute(GetEx::routed(route, key, ttl_ms, out))
735    }
736
737    /// Deletes `key`, returning `true` when an entry was removed.
738    pub fn del(&mut self, key: &[u8]) -> Result<bool> {
739        let route = self.checked_route(key)?;
740        self.conn.execute(Del::routed(route, key))
741    }
742
743    /// Returns whether `key` exists.
744    pub fn exists(&mut self, key: &[u8]) -> Result<bool> {
745        let route = self.checked_route(key)?;
746        self.conn.execute(Exists::routed(route, key))
747    }
748
749    /// Returns Redis-compatible TTL seconds for `key`.
750    pub fn ttl(&mut self, key: &[u8]) -> Result<i64> {
751        let route = self.checked_route(key)?;
752        self.conn.execute(Ttl::routed(route, key))
753    }
754
755    /// Sets a millisecond TTL on `key`, returning `true` when the TTL changed.
756    pub fn expire(&mut self, key: &[u8], ttl_ms: u64) -> Result<bool> {
757        let route = self.checked_route(key)?;
758        self.conn.execute(Expire::routed(route, key, ttl_ms))
759    }
760
761    /// Returns the first-party Redis command namespace for this shard.
762    #[cfg(feature = "redis")]
763    pub fn redis(&mut self) -> crate::Redis<'_, Self> {
764        crate::Redis::new(self)
765    }
766
767    /// Executes a compact Redis command on this direct shard.
768    ///
769    /// Commands that require all shards are rejected; use [`ShardCacheClient`] against
770    /// the fanout listener for those.
771    #[cfg(feature = "redis")]
772    pub fn redis_command(
773        &mut self,
774        command: RedisCommandKind,
775        args: &[&[u8]],
776    ) -> Result<RedisResponse> {
777        let route = redis_direct_shard_route(&self.router, self.shard_id, command, args)?;
778        self.conn
779            .execute(OptimizedRedisCommand::routed(command, route, args))
780    }
781
782    /// Executes a compact Redis command by name on this direct shard.
783    #[cfg(feature = "redis")]
784    pub fn redis_command_by_name(
785        &mut self,
786        command: &[u8],
787        args: &[&[u8]],
788    ) -> Result<RedisResponse> {
789        self.redis_command(redis_command_kind_from_name(command)?, args)
790    }
791
792    /// Runs a shard-local SCNP scan on this shard-owned connection.
793    pub fn scan_resp_into(&mut self, cursor: u64, count: usize, out: &mut Vec<u8>) -> Result<bool> {
794        let shard_id = self.shard_id.to_string();
795        let cursor = cursor.to_string();
796        let count = count.to_string();
797        self.conn.execute(RespCommand::new(
798            &[
799                b"SCNP.SCANSHARD",
800                shard_id.as_bytes(),
801                cursor.as_bytes(),
802                b"COUNT",
803                count.as_bytes(),
804            ],
805            out,
806        ))
807    }
808
809    /// Writes a routed GET request without flushing or reading its response.
810    pub fn begin_pipeline_get(&mut self, key: &[u8]) -> Result<()> {
811        let route = self.checked_route(key)?;
812        get::write_request(&mut self.conn, Some(route), key)
813    }
814
815    /// Writes a routed SET request without flushing or reading its response.
816    pub fn begin_pipeline_set(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
817        let route = self.checked_route(key)?;
818        set::write_request(&mut self.conn, Some(route), key, value)
819    }
820
821    /// Writes a routed SETEX request without flushing or reading its response.
822    pub fn begin_pipeline_set_ex(&mut self, key: &[u8], value: &[u8], ttl_ms: u64) -> Result<()> {
823        let route = self.checked_route(key)?;
824        setex::write_request(&mut self.conn, Some(route), key, value, ttl_ms)
825    }
826
827    /// Writes a routed GETEX request without flushing or reading its response.
828    pub fn begin_pipeline_get_ex(&mut self, key: &[u8], ttl_ms: u64) -> Result<()> {
829        let route = self.checked_route(key)?;
830        getex::write_request(&mut self.conn, Some(route), key, ttl_ms)
831    }
832
833    /// Writes a routed DEL request without flushing or reading its response.
834    pub fn begin_pipeline_del(&mut self, key: &[u8]) -> Result<()> {
835        let route = self.checked_route(key)?;
836        del::write_request(&mut self.conn, Some(route), key)
837    }
838
839    /// Writes a routed EXISTS request without flushing or reading its response.
840    pub fn begin_pipeline_exists(&mut self, key: &[u8]) -> Result<()> {
841        let route = self.checked_route(key)?;
842        exists::write_request(&mut self.conn, Some(route), key)
843    }
844
845    /// Writes a routed TTL request without flushing or reading its response.
846    pub fn begin_pipeline_ttl(&mut self, key: &[u8]) -> Result<()> {
847        let route = self.checked_route(key)?;
848        ttl::write_request(&mut self.conn, Some(route), key)
849    }
850
851    /// Writes a routed EXPIRE request without flushing or reading its response.
852    pub fn begin_pipeline_expire(&mut self, key: &[u8], ttl_ms: u64) -> Result<()> {
853        let route = self.checked_route(key)?;
854        expire::write_request(&mut self.conn, Some(route), key, ttl_ms)
855    }
856
857    /// Writes a compact Redis command request without flushing or reading its response.
858    #[cfg(feature = "redis")]
859    pub fn begin_pipeline_redis_command(
860        &mut self,
861        command: RedisCommandKind,
862        args: &[&[u8]],
863    ) -> Result<()> {
864        let route = redis_direct_shard_route(&self.router, self.shard_id, command, args)?;
865        redis::write_request(&mut self.conn, command, route, args)
866    }
867
868    /// Writes a compact Redis command request by name without flushing or reading its response.
869    #[cfg(feature = "redis")]
870    pub fn begin_pipeline_redis_command_by_name(
871        &mut self,
872        command: &[u8],
873        args: &[&[u8]],
874    ) -> Result<()> {
875        self.begin_pipeline_redis_command(redis_command_kind_from_name(command)?, args)
876    }
877
878    /// Flushes all queued pipelined requests.
879    pub fn flush_pipeline(&mut self) -> Result<()> {
880        self.conn.flush()
881    }
882
883    /// Reads the next pipelined GET response.
884    pub fn finish_pipeline_get_into(&mut self, out: &mut Vec<u8>) -> Result<bool> {
885        self.conn
886            .read_value(<Get as crate::commands::ScnpCommand>::NAME, out)
887    }
888
889    /// Reads the next pipelined SET response.
890    pub fn finish_pipeline_set(&mut self) -> Result<()> {
891        self.conn
892            .expect_ok(<Set as crate::commands::ScnpCommand>::NAME)
893    }
894
895    /// Reads the next pipelined SETEX response.
896    pub fn finish_pipeline_set_ex(&mut self) -> Result<()> {
897        self.conn
898            .expect_ok(<SetEx as crate::commands::ScnpCommand>::NAME)
899    }
900
901    /// Reads the next pipelined GETEX response.
902    pub fn finish_pipeline_get_ex_into(&mut self, out: &mut Vec<u8>) -> Result<bool> {
903        self.conn
904            .read_value(<GetEx as crate::commands::ScnpCommand>::NAME, out)
905    }
906
907    /// Reads the next pipelined DEL response.
908    pub fn finish_pipeline_del(&mut self) -> Result<bool> {
909        self.conn
910            .read_integer(<Del as crate::commands::ScnpCommand>::NAME)
911            .map(|deleted| deleted != 0)
912    }
913
914    /// Reads the next pipelined EXISTS response.
915    pub fn finish_pipeline_exists(&mut self) -> Result<bool> {
916        self.conn
917            .read_integer(<Exists as crate::commands::ScnpCommand>::NAME)
918            .map(|exists| exists != 0)
919    }
920
921    /// Reads the next pipelined TTL response.
922    pub fn finish_pipeline_ttl(&mut self) -> Result<i64> {
923        self.conn
924            .read_integer(<Ttl as crate::commands::ScnpCommand>::NAME)
925    }
926
927    /// Reads the next pipelined EXPIRE response.
928    pub fn finish_pipeline_expire(&mut self) -> Result<bool> {
929        self.conn
930            .read_integer(<Expire as crate::commands::ScnpCommand>::NAME)
931            .map(|changed| changed != 0)
932    }
933
934    /// Reads the next pipelined compact Redis command response.
935    #[cfg(feature = "redis")]
936    pub fn finish_pipeline_redis_command(&mut self) -> Result<RedisResponse> {
937        self.conn.read_native_redis_response("REDIS")
938    }
939
940    fn checked_route(&self, key: &[u8]) -> Result<crate::routing::ShardCacheRoute> {
941        let route = self.router.route_key(key);
942        if route.shard_id != self.shard_id {
943            return Err(ShardCacheClientError::Config(format!(
944                "key routes to shard {}, but client is connected to shard {}",
945                route.shard_id, self.shard_id
946            )));
947        }
948        Ok(route)
949    }
950}
951
952#[cfg(feature = "redis")]
953fn redis_command_kind_from_name(command: &[u8]) -> Result<RedisCommandKind> {
954    RedisCommandKind::from_name(command).ok_or_else(|| {
955        ShardCacheClientError::Config(format!(
956            "Redis command `{}` is not available on direct SCNP shard clients; use ShardCacheClient on the fanout listener for command-name fallback",
957            String::from_utf8_lossy(command)
958        ))
959    })
960}
961
962#[cfg(feature = "redis")]
963fn validate_redis_command_name(command: &[u8]) -> Result<()> {
964    if command.is_empty() {
965        return Err(ShardCacheClientError::Config(
966            "Redis command name cannot be empty".into(),
967        ));
968    }
969    if command.iter().any(|byte| byte.is_ascii_whitespace()) {
970        return Err(ShardCacheClientError::Config(format!(
971            "Redis command name cannot contain whitespace: `{}`",
972            String::from_utf8_lossy(command)
973        )));
974    }
975    Ok(())
976}
977
978#[cfg(feature = "redis")]
979fn redis_direct_route(
980    router: &ShardCacheDirectRouter,
981    command: RedisCommandKind,
982    args: &[&[u8]],
983) -> Result<Option<ShardCacheRoute>> {
984    let keys = match command.route_keys(args) {
985        RedisCommandRouteKeys::None => return Ok(None),
986        RedisCommandRouteKeys::AllShards => {
987            return Err(ShardCacheClientError::Config(format!(
988                "{} requires all shards; use ShardCacheClient on the fanout listener",
989                command.name()
990            )));
991        }
992        RedisCommandRouteKeys::Keys(keys) if keys.is_empty() => return Ok(None),
993        RedisCommandRouteKeys::Keys(keys) => keys,
994    };
995
996    let first_route = router.route_key(keys[0]);
997    for key in keys.iter().skip(1) {
998        let route = router.route_key(key);
999        if route.shard_id != first_route.shard_id {
1000            return Err(ShardCacheClientError::Config(format!(
1001                "{} keys span multiple direct shards",
1002                command.name()
1003            )));
1004        }
1005    }
1006    Ok(Some(first_route))
1007}
1008
1009#[cfg(feature = "redis")]
1010fn redis_direct_shard_route(
1011    router: &ShardCacheDirectRouter,
1012    shard_id: usize,
1013    command: RedisCommandKind,
1014    args: &[&[u8]],
1015) -> Result<Option<ShardCacheRoute>> {
1016    let route = redis_direct_route(router, command, args)?;
1017    if let Some(route) = route
1018        && route.shard_id != shard_id
1019    {
1020        return Err(ShardCacheClientError::Config(format!(
1021            "{} routes to shard {}, but client is connected to shard {}",
1022            command.name(),
1023            route.shard_id,
1024            shard_id
1025        )));
1026    }
1027    Ok(route)
1028}
shardcache_client_rs/client.rs

shardcache_client_rs/
client.rs
