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        crate::register_default_backends();
44        Ok(MongoDbGuard(self.0.start().await?))
45    }
46}
47
48impl Default for MongoDbContainer {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54/// Retries `rs.initiate()` (via `rs.status()` first, so a retry after a partial
55/// initiate doesn't re-initiate) through the proxy-accepts-before-mongod-listens race:
56/// the listening-port wait can return before `mongod` is far enough along to
57/// accept a real client command, so the first few `mongosh` invocations may fail —
58/// that's expected and retried, not a fatal error.
59async fn initiate_replica_set(guard: &ContainerGuard) -> Result<()> {
60    poll_until(guard, "rs.initiate to succeed", |guard| {
61        Box::pin(async move {
62            let result = guard
63                .exec(&[
64                    "mongosh",
65                    "--quiet",
66                    "--eval",
67                    "try { rs.status() } catch (e) { rs.initiate() }",
68                ])
69                .await;
70            matches!(result, Ok(r) if r.exit_code == 0)
71        })
72    })
73    .await
74}
75
76/// Polls until `db.hello().isWritablePrimary` reports `true`.
77async fn await_primary_elected(guard: &ContainerGuard) -> Result<()> {
78    poll_until(guard, "a PRIMARY to be elected", |guard| {
79        Box::pin(async move {
80            let result = guard
81                .exec(&[
82                    "mongosh",
83                    "--quiet",
84                    "--eval",
85                    "db.hello().isWritablePrimary",
86                ])
87                .await;
88            matches!(result, Ok(r) if r.stdout.trim().ends_with("true"))
89        })
90    })
91    .await
92}
93
94/// A tiny deadline/poll-interval loop, local to this module. `cond` returning
95/// `false` (including on an `exec` error, swallowed here) means "not yet"; the loop
96/// keeps retrying until `cond` is `true` or the deadline passes.
97async fn poll_until<F>(guard: &ContainerGuard, what: &str, mut cond: F) -> Result<()>
98where
99    F: FnMut(&ContainerGuard) -> BoxFuture<'_, bool>,
100{
101    let deadline = tokio::time::Instant::now() + REPLICA_SET_TIMEOUT;
102    loop {
103        if cond(guard).await {
104            return Ok(());
105        }
106        if tokio::time::Instant::now() >= deadline {
107            break;
108        }
109        tokio::time::sleep(POLL_INTERVAL).await;
110    }
111    Err(RightsizeError::ContainerLaunch(format!(
112        "Mongo replica set on {}:{} did not reach '{what}' within {}s",
113        guard.host(),
114        guard.get_mapped_port(MongoDbContainer::PORT).unwrap_or(0),
115        REPLICA_SET_TIMEOUT.as_secs(),
116    )))
117}
118
119/// The running guard for a [`MongoDbContainer`].
120pub struct MongoDbGuard(ContainerGuard);
121
122impl MongoDbGuard {
123    /// A `mongodb://` connection string for the running container's `test` database.
124    pub fn connection_string(&self) -> String {
125        format!(
126            "mongodb://{}:{}/test?directConnection=true",
127            self.0.host(),
128            self.0.get_mapped_port(MongoDbContainer::PORT).unwrap()
129        )
130    }
131
132    /// Alias for [`Self::connection_string`]; the container is always a (single-node)
133    /// replica set.
134    pub fn replica_set_url(&self) -> String {
135        self.connection_string()
136    }
137
138    /// Stops and removes the container, releasing its host port.
139    pub async fn stop(self) -> Result<()> {
140        self.0.stop().await
141    }
142}
143
144impl std::ops::Deref for MongoDbGuard {
145    type Target = ContainerGuard;
146    fn deref(&self) -> &ContainerGuard {
147        &self.0
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn with_image_smoke() {
157        let _ = MongoDbContainer::new();
158        let _ = MongoDbContainer::with_image("mongo:8.0");
159    }
160}
rightsize_modules/mongodb.rs

rightsize_modules/
mongodb.rs
