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

rightsize_modules/
memcached.rs
