fix(pb): backport serverless fix (#4611)

# Description

Please include a summary of the changes and the related issue. Please also include relevant motivation and context.

## Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes.

## Checklist:

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes

Author: MasterPtato

engine/packages/pegboard/src/workflows/serverless/conn.rs

@@ -156,16 +156,13 @@ enum OutboundReqOutput {
 #[timeout = u64::MAX]
 async fn outbound_req(ctx: &ActivityCtx, input: &OutboundReqInput) -> Result<OutboundReqOutput> {
 	let mut term_signal = TermSignal::get();
-	let mut drain_sub = ctx
-		.subscribe::<Drain>(("workflow_id", ctx.workflow_id()))
-		.await?;
 
 	loop {
 		metrics::SERVERLESS_OUTBOUND_REQ_ACTIVE
 			.with_label_values(&[&input.namespace_id.to_string(), &input.runner_name])
 			.inc();
 
-		let res = outbound_req_inner(ctx, input, &mut term_signal, &mut drain_sub).await;
+		let res = outbound_req_inner(ctx, input, &mut term_signal).await;
 
 		metrics::SERVERLESS_OUTBOUND_REQ_ACTIVE
 			.with_label_values(&[&input.namespace_id.to_string(), &input.runner_name])
@@ -203,7 +200,6 @@ async fn outbound_req_inner(
 	ctx: &ActivityCtx,
 	input: &OutboundReqInput,
 	term_signal: &mut TermSignal,
-	drain_sub: &mut message::SubscriptionHandle<Drain>,
 ) -> Result<OutboundReqOutput> {
 	if is_runner_draining(ctx, input.receiver_wf_id).await? {
 		return Ok(OutboundReqOutput::Draining { drain_sent: false });
@@ -438,7 +434,6 @@ async fn outbound_req_inner(
 			}
 		},
 		_ = tokio::time::sleep(sleep_until_drain) => {}
-		_ = drain_sub.next() => {}
 		_ = term_signal.recv() => {}
 	};
 

engine/packages/pegboard/src/workflows/serverless/receiver.rs

@@ -48,11 +48,7 @@ pub async fn pegboard_serverless_receiver(ctx: &mut WorkflowCtx, input: &Input)
 		.send()
 		.await?;
 
-	// if the connection is currently running an outbound req, this will be received
-	ctx.msg(conn::Drain {})
-		.topic(("workflow_id", conn_wf_id))
-		.send()
-		.await?;
+	ctx.removed::<Message<conn::Drain>>().await?;
 
 	// Wait for connection wf to complete so this wf's state remains readable
 	ctx.workflow::<conn::Input>(conn_wf_id).output().await?;