rightsize_modules/
mongodb.rs

1//! A single-node MongoDB container running as a one-member replica set (required for
2//! transactions/change streams). The [`Container::with_post_start`] hook initiates the
3//! replica set and waits for a primary to be elected before `start()` returns, so
4//! [`MongoDbGuard::connection_string`] is always usable immediately after `start()`.
5
6use std::time::Duration;
7
8use rightsize::{BoxFuture, Container, ContainerGuard, Result, RightsizeError, Wait};
9
10const REPLICA_SET_TIMEOUT: Duration = Duration::from_secs(60);
11const POLL_INTERVAL: Duration = Duration::from_millis(500);
12
13/// A single-node MongoDB container, started as a one-member replica set named
14/// `docker-rs`.
15pub struct MongoDbContainer(Container);
16
17impl MongoDbContainer {
18    /// The guest port `mongod` listens on.
19    const PORT: u16 = 27017;
20
21    /// Builds a container from the pinned default image (`mongo:8.0`).
22    pub fn new() -> Self {
23        Self::with_image("mongo:8.0")
24    }
25
26    /// Builds a container from a caller-chosen image.
27    pub fn with_image(image: &str) -> Self {
28        let container = Container::new(image)
29            .with_exposed_ports(&[Self::PORT])
30            .with_command(&["mongod", "--replSet", "docker-rs", "--bind_ip_all"])
31            .waiting_for(Wait::for_listening_port())
32            .with_post_start(|guard: &ContainerGuard| -> BoxFuture<'_, Result<()>> {
33                Box::pin(async move {
34                    initiate_replica_set(guard).await?;
35                    await_primary_elected(guard).await
36                })
37            });
38        Self(container)
39    }
40
41    /// Boots the container. Does not return until the replica set has a primary.
42    pub async fn start(self) -> Result<MongoDbGuard> {
43        Ok(MongoDbGuard(self.0.start().await?))
44    }
45}
46
47impl Default for MongoDbContainer {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53/// Retries `rs.initiate()` (via `rs.status()` first, so a retry after a partial
54/// initiate doesn't re-initiate) through the proxy-accepts-before-mongod-listens race:
55/// the listening-port wait can return before `mongod` is far enough along to
56/// accept a real client command, so the first few `mongosh` invocations may fail —
57/// that's expected and retried, not a fatal error.
58async fn initiate_replica_set(guard: &ContainerGuard) -> Result<()> {
59    poll_until(guard, "rs.initiate to succeed", |guard| {
60        Box::pin(async move {
61            let result = guard
62                .exec(&[
63                    "mongosh",
64                    "--quiet",
65                    "--eval",
66                    "try { rs.status() } catch (e) { rs.initiate() }",
67                ])
68                .await;
69            matches!(result, Ok(r) if r.exit_code == 0)
70        })
71    })
72    .await
73}
74
75/// Polls until `db.hello().isWritablePrimary` reports `true`.
76async fn await_primary_elected(guard: &ContainerGuard) -> Result<()> {
77    poll_until(guard, "a PRIMARY to be elected", |guard| {
78        Box::pin(async move {
79            let result = guard
80                .exec(&[
81                    "mongosh",
82                    "--quiet",
83                    "--eval",
84                    "db.hello().isWritablePrimary",
85                ])
86                .await;
87            matches!(result, Ok(r) if r.stdout.trim().ends_with("true"))
88        })
89    })
90    .await
91}
92
93/// A tiny deadline/poll-interval loop, local to this module. `cond` returning
94/// `false` (including on an `exec` error, swallowed here) means "not yet"; the loop
95/// keeps retrying until `cond` is `true` or the deadline passes.
96async fn poll_until<F>(guard: &ContainerGuard, what: &str, mut cond: F) -> Result<()>
97where
98    F: FnMut(&ContainerGuard) -> BoxFuture<'_, bool>,
99{
100    let deadline = tokio::time::Instant::now() + REPLICA_SET_TIMEOUT;
101    loop {
102        if cond(guard).await {
103            return Ok(());
104        }
105        if tokio::time::Instant::now() >= deadline {
106            break;
107        }
108        tokio::time::sleep(POLL_INTERVAL).await;
109    }
110    Err(RightsizeError::ContainerLaunch(format!(
111        "Mongo replica set on {}:{} did not reach '{what}' within {}s",
112        guard.host(),
113        guard.get_mapped_port(MongoDbContainer::PORT).unwrap_or(0),
114        REPLICA_SET_TIMEOUT.as_secs(),
115    )))
116}
117
118/// The running guard for a [`MongoDbContainer`].
119pub struct MongoDbGuard(ContainerGuard);
120
121impl MongoDbGuard {
122    /// A `mongodb://` connection string for the running container's `test` database.
123    pub fn connection_string(&self) -> String {
124        format!(
125            "mongodb://{}:{}/test?directConnection=true",
126            self.0.host(),
127            self.0.get_mapped_port(MongoDbContainer::PORT).unwrap()
128        )
129    }
130
131    /// Alias for [`Self::connection_string`]; the container is always a (single-node)
132    /// replica set.
133    pub fn replica_set_url(&self) -> String {
134        self.connection_string()
135    }
136
137    /// Stops and removes the container, releasing its host port.
138    pub async fn stop(self) -> Result<()> {
139        self.0.stop().await
140    }
141}
142
143impl std::ops::Deref for MongoDbGuard {
144    type Target = ContainerGuard;
145    fn deref(&self) -> &ContainerGuard {
146        &self.0
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    #[test]
155    fn with_image_smoke() {
156        let _ = MongoDbContainer::new();
157        let _ = MongoDbContainer::with_image("mongo:8.0");
158    }
159}
rightsize_modules/mongodb.rs

rightsize_modules/
mongodb.rs
