Skip to content

Commit 03adb7d

Browse files
MasterPtatoMasterPtato
committed
fix(tunnel): fix ups race condition
1 parent eb42b41 commit 03adb7d

File tree

13 files changed

+93
-66
lines changed

13 files changed

+93
-66
lines changed

‎docker/dev-host/docker-compose.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ services:
150150
- CMD
151151
- curl
152152
- '-f'
153-
- http://127.0.0.1:6421/health
153+
- http://127.0.0.1:6420/health
154154
interval: 2s
155155
timeout: 10s
156156
retries: 10

‎docker/dev-multidc-multinode/docker-compose.yml‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ services:
184184
- CMD
185185
- curl
186186
- '-f'
187-
- http://127.0.0.1:6421/health
187+
- http://127.0.0.1:6420/health
188188
interval: 2s
189189
timeout: 10s
190190
retries: 10
@@ -224,7 +224,7 @@ services:
224224
- CMD
225225
- curl
226226
- '-f'
227-
- http://127.0.0.1:6421/health
227+
- http://127.0.0.1:6420/health
228228
interval: 2s
229229
timeout: 10s
230230
retries: 10
@@ -264,7 +264,7 @@ services:
264264
- CMD
265265
- curl
266266
- '-f'
267-
- http://127.0.0.1:6421/health
267+
- http://127.0.0.1:6420/health
268268
interval: 2s
269269
timeout: 10s
270270
retries: 10
@@ -444,7 +444,7 @@ services:
444444
- CMD
445445
- curl
446446
- '-f'
447-
- http://127.0.0.1:6421/health
447+
- http://127.0.0.1:6420/health
448448
interval: 2s
449449
timeout: 10s
450450
retries: 10
@@ -484,7 +484,7 @@ services:
484484
- CMD
485485
- curl
486486
- '-f'
487-
- http://127.0.0.1:6421/health
487+
- http://127.0.0.1:6420/health
488488
interval: 2s
489489
timeout: 10s
490490
retries: 10
@@ -524,7 +524,7 @@ services:
524524
- CMD
525525
- curl
526526
- '-f'
527-
- http://127.0.0.1:6421/health
527+
- http://127.0.0.1:6420/health
528528
interval: 2s
529529
timeout: 10s
530530
retries: 10
@@ -702,7 +702,7 @@ services:
702702
- CMD
703703
- curl
704704
- '-f'
705-
- http://127.0.0.1:6421/health
705+
- http://127.0.0.1:6420/health
706706
interval: 2s
707707
timeout: 10s
708708
retries: 10
@@ -742,7 +742,7 @@ services:
742742
- CMD
743743
- curl
744744
- '-f'
745-
- http://127.0.0.1:6421/health
745+
- http://127.0.0.1:6420/health
746746
interval: 2s
747747
timeout: 10s
748748
retries: 10
@@ -782,7 +782,7 @@ services:
782782
- CMD
783783
- curl
784784
- '-f'
785-
- http://127.0.0.1:6421/health
785+
- http://127.0.0.1:6420/health
786786
interval: 2s
787787
timeout: 10s
788788
retries: 10

‎docker/dev-multidc/docker-compose.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,7 @@ services:
182182
- CMD
183183
- curl
184184
- '-f'
185-
- http://127.0.0.1:6421/health
185+
- http://127.0.0.1:6420/health
186186
interval: 2s
187187
timeout: 10s
188188
retries: 10
@@ -330,7 +330,7 @@ services:
330330
- CMD
331331
- curl
332332
- '-f'
333-
- http://127.0.0.1:6421/health
333+
- http://127.0.0.1:6420/health
334334
interval: 2s
335335
timeout: 10s
336336
retries: 10
@@ -476,7 +476,7 @@ services:
476476
- CMD
477477
- curl
478478
- '-f'
479-
- http://127.0.0.1:6421/health
479+
- http://127.0.0.1:6420/health
480480
interval: 2s
481481
timeout: 10s
482482
retries: 10

‎docker/dev-multinode/docker-compose.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ services:
173173
- CMD
174174
- curl
175175
- '-f'
176-
- http://127.0.0.1:6421/health
176+
- http://127.0.0.1:6420/health
177177
interval: 2s
178178
timeout: 10s
179179
retries: 10
@@ -210,7 +210,7 @@ services:
210210
- CMD
211211
- curl
212212
- '-f'
213-
- http://127.0.0.1:6421/health
213+
- http://127.0.0.1:6420/health
214214
interval: 2s
215215
timeout: 10s
216216
retries: 10
@@ -247,7 +247,7 @@ services:
247247
- CMD
248248
- curl
249249
- '-f'
250-
- http://127.0.0.1:6421/health
250+
- http://127.0.0.1:6420/health
251251
interval: 2s
252252
timeout: 10s
253253
retries: 10

‎docker/dev/docker-compose.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ services:
173173
- CMD
174174
- curl
175175
- '-f'
176-
- http://127.0.0.1:6421/health
176+
- http://127.0.0.1:6420/health
177177
interval: 2s
178178
timeout: 10s
179179
retries: 10

‎docker/dev/rivet-engine/config.jsonc‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,13 @@
3232
"postgres": {
3333
"url": "postgresql://postgres:postgres@postgres:5432/rivet_engine"
3434
},
35+
"postgres_notify": {
36+
"url": "postgresql://postgres:postgres@postgres:5432/rivet_engine",
37+
"memory_optimization": false
38+
},
39+
// "memory": {
40+
// "channel": "default"
41+
// },
3542
"cache": {
3643
"driver": "in_memory"
3744
},

‎docker/template/src/docker-compose.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,7 @@ export function generateDockerCompose(context: TemplateContext) {
294294
],
295295
ports: isPrimary && i === 0 ? [`6420:6420`] : undefined,
296296
healthcheck: {
297-
test: ["CMD", "curl", "-f", "http://127.0.0.1:6421/health"],
297+
test: ["CMD", "curl", "-f", "http://127.0.0.1:6420/health"],
298298
interval: "2s",
299299
timeout: "10s",
300300
retries: 10,

‎packages/common/config/src/config/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ impl Root {
193193
{
194194
self.pubsub = Some(PubSub::PostgresNotify(pubsub::Postgres {
195195
url: pg.url.clone(),
196+
memory_optimization: true,
196197
}));
197198
}
198199

‎packages/common/universalpubsub/src/driver/postgres/mod.rs‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,7 @@ impl PubSubDriver for PostgresDriver {
113113
tracing::debug!(%subject, ?lock_id, "calculated advisory lock id");
114114

115115
// Create a single connection for both subscription and lock holding
116-
let (client, mut connection) =
117-
tokio_postgres::connect(&self.conn_str, tokio_postgres::NoTls).await?;
116+
let (client, mut connection) = pg_connect(&self.conn_str).await?;
118117

119118
// Set up message forwarding
120119
let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<String>();
@@ -155,7 +154,7 @@ impl PubSubDriver for PostgresDriver {
155154
let listen_subject = subject_owned.clone();
156155

157156
// Spawn task to handle connection, lock acquisition, and LISTEN
158-
tokio::spawn(async move {
157+
let poll_handle = tokio::spawn(async move {
159158
// First acquire the lock while polling the connection
160159
let lock_sql = format!("SELECT pg_try_advisory_lock_shared({})", lock_id);
161160
let lock_future = client_clone.query_one(&lock_sql, &[]);
@@ -265,6 +264,7 @@ impl PubSubDriver for PostgresDriver {
265264
lock_id,
266265
client,
267266
subject: subject.to_string(),
267+
poll_handle,
268268
}))
269269
}
270270

@@ -419,16 +419,15 @@ impl PubSubDriver for PostgresDriver {
419419
// Create a temporary reply subject and a dedicated listener connection
420420
let reply_subject = format!("_INBOX.{}", uuid::Uuid::new_v4());
421421

422-
let (client, mut connection) =
423-
tokio_postgres::connect(&self.conn_str, tokio_postgres::NoTls).await?;
422+
let (client, mut connection) = pg_connect(&self.conn_str).await?;
424423

425424
// Setup connection and LISTEN in a task
426425
let (listen_done_tx, listen_done_rx) = tokio::sync::oneshot::channel();
427426
let reply_subject_clone = reply_subject.clone();
428427

429428
// Spawn task to handle connection and LISTEN
430429
let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel();
431-
tokio::spawn(async move {
430+
let poll_handle = tokio::spawn(async move {
432431
// Convert subject to base64 hash string because Postgres identifiers can only be 63 bytes
433432
let mut hasher = DefaultHasher::new();
434433
reply_subject_clone.hash(&mut hasher);
@@ -513,14 +512,19 @@ impl PubSubDriver for PostgresDriver {
513512
};
514513

515514
// Apply timeout if specified
516-
if let Some(dur) = timeout {
515+
let res = if let Some(dur) = timeout {
517516
match tokio::time::timeout(dur, response_future).await {
518517
std::result::Result::Ok(resp) => resp,
519518
std::result::Result::Err(_) => Err(errors::Ups::RequestTimeout.build().into()),
520519
}
521520
} else {
522521
response_future.await
523-
}
522+
};
523+
524+
// Stop poll loop
525+
poll_handle.abort();
526+
527+
res
524528
}
525529

526530
async fn send_request_reply(&self, reply: &str, payload: &[u8]) -> Result<()> {
@@ -599,6 +603,7 @@ pub struct PostgresSubscriber {
599603
lock_id: i64,
600604
client: Arc<tokio_postgres::Client>,
601605
subject: String,
606+
poll_handle: tokio::task::JoinHandle<()>,
602607
}
603608

604609
#[async_trait]
@@ -716,5 +721,22 @@ impl Drop for PostgresSubscriber {
716721
.execute("SELECT pg_advisory_unlock_shared($1)", &[&lock_id])
717722
.await;
718723
});
724+
725+
// Stop polling task
726+
self.poll_handle.abort();
719727
}
720728
}
729+
730+
async fn pg_connect(
731+
conn_str: &str,
732+
) -> Result<(
733+
tokio_postgres::Client,
734+
tokio_postgres::Connection<
735+
tokio_postgres::Socket,
736+
<tokio_postgres::NoTls as tokio_postgres::tls::TlsConnect<tokio_postgres::Socket>>::Stream,
737+
>,
738+
)> {
739+
let (client, conn) = tokio_postgres::connect(conn_str, tokio_postgres::NoTls).await?;
740+
741+
Ok((client, conn))
742+
}

‎packages/common/universalpubsub/src/pubsub.rs‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,3 @@ impl Message {
7979
pub struct Response {
8080
pub payload: Vec<u8>,
8181
}
82-
83-
impl Drop for Message {
84-
fn drop(&mut self) {
85-
tracing::info!("dropping message");
86-
}
87-
}

0 commit comments

Comments
 (0)