rightsize_modules/
memcached.rs

1//! A single-node Memcached container, ready-checked with a protocol-level `version`
2//! probe instead of the bare listening-port wait.
3
4use std::time::Duration;
5
6use rightsize::wait::{WaitStrategy, WaitTarget};
7use rightsize::{Container, ContainerGuard, Result};
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::net::TcpStream;
10
11/// A single-node Memcached container.
12pub struct MemcachedContainer(Container);
13
14impl MemcachedContainer {
15    /// The guest port memcached listens on.
16    const PORT: u16 = 11211;
17
18    /// Builds a container from the pinned default image (`memcached:1.6-alpine`).
19    pub fn new() -> Self {
20        Self::with_image("memcached:1.6-alpine")
21    }
22
23    /// Builds a container from a caller-chosen image.
24    pub fn with_image(image: &str) -> Self {
25        Self(
26            Container::new(image)
27                .with_exposed_ports(&[Self::PORT])
28                .waiting_for(MemcachedResponds),
29        )
30    }
31
32    /// Boots the container.
33    pub async fn start(self) -> Result<MemcachedGuard> {
34        crate::register_default_backends();
35        Ok(MemcachedGuard(self.0.start().await?))
36    }
37}
38
39impl Default for MemcachedContainer {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45/// The running guard for a [`MemcachedContainer`].
46pub struct MemcachedGuard(ContainerGuard);
47
48impl MemcachedGuard {
49    /// The `host:port` address of the running container.
50    pub fn address(&self) -> String {
51        format!(
52            "{}:{}",
53            self.0.host(),
54            self.0.get_mapped_port(MemcachedContainer::PORT).unwrap()
55        )
56    }
57
58    /// Stops and removes the container, releasing its host port.
59    pub async fn stop(self) -> Result<()> {
60        self.0.stop().await
61    }
62}
63
64impl std::ops::Deref for MemcachedGuard {
65    type Target = ContainerGuard;
66    fn deref(&self) -> &ContainerGuard {
67        &self.0
68    }
69}
70
71/// Memcached logs nothing on startup and the docker userland proxy (or msb's loopback
72/// forwarder) binds the host port before the server inside is accepting, so a bare
73/// TCP-connect wait can pass while the first real client connection still gets a dead
74/// stream. This strategy proves readiness by speaking the protocol: it sends
75/// `version\r\n` and expects a reply starting with `VERSION`.
76struct MemcachedResponds;
77
78impl MemcachedResponds {
79    async fn probe_once(host: &str, port: u16) -> bool {
80        let Ok(connect) = tokio::time::timeout(
81            Duration::from_millis(1000),
82            TcpStream::connect((host, port)),
83        )
84        .await
85        else {
86            return false;
87        };
88        let Ok(mut stream) = connect else {
89            return false;
90        };
91        if tokio::time::timeout(
92            Duration::from_millis(1000),
93            stream.write_all(b"version\r\n"),
94        )
95        .await
96        .is_err()
97        {
98            return false;
99        }
100        let mut buf = [0u8; 64];
101        let Ok(Ok(n)) =
102            tokio::time::timeout(Duration::from_millis(1000), stream.read(&mut buf)).await
103        else {
104            return false;
105        };
106        if n == 0 {
107            return false;
108        }
109        String::from_utf8_lossy(&buf[..n]).starts_with("VERSION")
110    }
111}
112
113#[async_trait::async_trait]
114impl WaitStrategy for MemcachedResponds {
115    async fn wait_until_ready(&self, target: &dyn WaitTarget) -> Result<()> {
116        let guest_port = target
117            .exposed_guest_ports()
118            .first()
119            .copied()
120            .unwrap_or(MemcachedContainer::PORT);
121        let port = target.mapped_port(guest_port);
122        let host = target.host().to_string();
123        rightsize::wait::poll_until_ready(
124            target,
125            Duration::from_secs(60),
126            "a VERSION reply",
127            || {
128                let host = host.clone();
129                async move { Self::probe_once(&host, port).await }
130            },
131        )
132        .await
133    }
134
135    fn with_startup_timeout(self: Box<Self>, _timeout: Duration) -> Box<dyn WaitStrategy> {
136        // This probe doesn't expose a timeout override; keep the fixed poll budget.
137        self
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use std::sync::Arc;
145    use std::sync::atomic::{AtomicBool, Ordering};
146
147    #[test]
148    fn with_image_smoke() {
149        let _ = MemcachedContainer::new();
150        let _ = MemcachedContainer::with_image("memcached:1.6-alpine");
151    }
152
153    /// A fake `VERSION`-replying socket — proves the wait strategy recognizes a real
154    /// memcached protocol reply and (via `probe_once`) rejects a peer that isn't
155    /// speaking memcached at all.
156    #[tokio::test]
157    async fn probe_once_true_on_a_version_reply_false_otherwise() {
158        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
159        let port = listener.local_addr().unwrap().port();
160        let ready = Arc::new(AtomicBool::new(false));
161        let ready_clone = ready.clone();
162        std::thread::spawn(move || {
163            listener.set_nonblocking(true).unwrap();
164            loop {
165                if let Ok((mut stream, _)) = listener.accept() {
166                    use std::io::{Read, Write};
167                    let mut buf = [0u8; 64];
168                    let _ = stream.read(&mut buf);
169                    if ready_clone.load(Ordering::SeqCst) {
170                        let _ = stream.write_all(b"VERSION 1.6.31\r\n");
171                    } else {
172                        let _ = stream.write_all(b"ERROR\r\n");
173                    }
174                    return;
175                }
176                std::thread::sleep(Duration::from_millis(5));
177            }
178        });
179
180        assert!(!MemcachedResponds::probe_once("127.0.0.1", port).await);
181
182        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
183        let port = listener.local_addr().unwrap().port();
184        ready.store(true, Ordering::SeqCst);
185        std::thread::spawn(move || {
186            listener.set_nonblocking(true).unwrap();
187            loop {
188                if let Ok((mut stream, _)) = listener.accept() {
189                    use std::io::{Read, Write};
190                    let mut buf = [0u8; 64];
191                    let _ = stream.read(&mut buf);
192                    let _ = stream.write_all(b"VERSION 1.6.31\r\n");
193                    return;
194                }
195                std::thread::sleep(Duration::from_millis(5));
196            }
197        });
198        assert!(MemcachedResponds::probe_once("127.0.0.1", port).await);
199    }
200}
rightsize_modules/memcached.rs

rightsize_modules/
memcached.rs
