Fix batched streams [#3439]

AccumulateNetwork · Oct 4, 2023 · 93e74f4 · 93e74f4
1 parent c0516ce
commit 93e74f4
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/pkg/api/v3/batch.go b/pkg/api/v3/batch.go
@@ -9,11 +9,13 @@ package api
 import "context"
 
 type BatchData struct {
-	values map[any]any
+	context context.Context
+	values  map[any]any
 }
 
-func (d *BatchData) Get(k any) any { return d.values[k] }
-func (d *BatchData) Put(k, v any)  { d.values[k] = v }
+func (d *BatchData) Context() context.Context { return d.context }
+func (d *BatchData) Get(k any) any            { return d.values[k] }
+func (d *BatchData) Put(k, v any)             { d.values[k] = v }
 
 type contextKeyBatch struct{}
 
@@ -29,9 +31,10 @@ func ContextWithBatchData(ctx context.Context) (context.Context, context.CancelF
 		return ctx, func() {}, v
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
 	bd := new(BatchData)
+	bd.context = ctx
 	bd.values = map[any]any{}
-	ctx, cancel := context.WithCancel(ctx)
 	ctx = context.WithValue(ctx, contextKeyBatch{}, bd)
 	return ctx, cancel, bd
 }

diff --git a/pkg/api/v3/p2p/dial_self.go b/pkg/api/v3/p2p/dial_self.go
@@ -46,6 +46,9 @@ func (d *selfDialer) Dial(ctx context.Context, addr multiaddr.Multiaddr) (messag
 }
 
 func handleLocally(ctx context.Context, service *serviceHandler) message.Stream {
+	if batch := api.GetBatchData(ctx); batch != nil {
+		ctx = batch.Context()
+	}
 	p, q := message.DuplexPipe(ctx)
 	go func() {
 		// Panic protection

diff --git a/pkg/api/v3/p2p/stream.go b/pkg/api/v3/p2p/stream.go
@@ -32,7 +32,14 @@ func openStreamFor(ctx context.Context, host dialerHost, peer peer.ID, sa *api.S
 	}
 
 	// Close the stream when the context is canceled
-	go func() { <-ctx.Done(); _ = conn.Close() }()
+	go func() {
+		ctx := ctx
+		if batch := api.GetBatchData(ctx); batch != nil {
+			ctx = batch.Context()
+		}
+		<-ctx.Done()
+		_ = conn.Close()
+	}()
 
 	s := new(stream)
 	s.peer = peer