fix(execstream): do not let a resolution answer nobody reads block the end of the session

commit: 79b18d6ec964e9e511a632cda1221980426a26dd
author: brice · 8d ago · 2026-10-14 03:02 UTC
parent: be35c9f
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

2 files · +105 -1

M internal/execstream/frame.go +37 -1


@@ -19,6 +19,7 @@ import (
19	19		"fmt"
20	20		"io"
21	21		"sync"
	22	+	"time"
22	23
23	24		"google.golang.org/protobuf/proto"
24	25

@@ -106,6 +107,38 @@ func Drive(rw io.ReadWriter, run execpb.Run, cancel <-chan execpb.Cancel, onOu
106	107		return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})
107	108		}
108	109
	110	+	// resolveJoinGrace is how long DriveJob waits for the resolution worker to
	111	+	// finish once the session is over before it interrupts a write that is stuck
	112	+	// (a guest that stopped reading): the write deadline is set to now, and the
	113	+	// stream is closed if it cannot be interrupted otherwise. The session is over at
	114	+	// that point, the caller closes the stream anyway.
	115	+	var resolveJoinGrace = 2 * time.Second
	116	+
	117	+	// joinResolver waits for the worker. A worker in a Send that cannot finish
	118	+	// because the peer does not read would block DriveJob for ever, so the write is
	119	+	// interrupted after the grace period.
	120	+	func joinResolver(rw io.ReadWriter, workers *sync.WaitGroup) {
	121	+	finished := make(chan struct{})
	122	+	go func() { workers.Wait(); close(finished) }()
	123	+	select {
	124	+	case <-finished:
	125	+	return
	126	+	case <-time.After(resolveJoinGrace):
	127	+	}
	128	+	if d, ok := rw.(interface{ SetWriteDeadline(time.Time) error }); ok {
	129	+	_ = d.SetWriteDeadline(time.Now())
	130	+	}
	131	+	select {
	132	+	case <-finished:
	133	+	return
	134	+	case <-time.After(resolveJoinGrace):
	135	+	}
	136	+	if cl, ok := rw.(io.Closer); ok {
	137	+	_ = cl.Close()
	138	+	<-finished
	139	+	}
	140	+	}
	141	+
109	142		// resolveQueue is how many ResolveRequest frames may wait for the worker.
110	143		const resolveQueue = 2
111	144

@@ -167,13 +200,16 @@ func DriveJob(rw io.ReadWriter, run execpb.Run, cancel <-chan execpb.Cancel, h
167	200		case req := <-resolveCh:
168	201		resp := h.Resolve(rctx, req)
169	202		resp.Id = req.Id
	203	+	if rctx.Err() != nil {
	204	+	return // the session is over: nobody reads the answer
	205	+	}
170	206		_ = c.Send(resolveFrame(resp))
171	207		case <-rctx.Done():
172	208		return
173	209		}
174	210		}
175	211		}()
176		-	defer func() { rcancel(); workers.Wait() }()
	212	+	defer func() { rcancel(); joinResolver(rw, &workers) }()
177	213		if cancel != nil {
178	214		go func() {
179	215		select {


@@ -19,6 +19,7 @@ import (
19		"fmt"	19		"fmt"
20		"io"	20		"io"
21		"sync"	21		"sync"
			22	+	"time"
22			23
23		"google.golang.org/protobuf/proto"	24		"google.golang.org/protobuf/proto"
24			25

@@ -106,6 +107,38 @@ func Drive(rw io.ReadWriter, run execpb.Run, cancel <-chan execpb.Cancel, onOu
106		return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})	107		return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})
107		}	108		}
108			109
			110	+	// resolveJoinGrace is how long DriveJob waits for the resolution worker to
			111	+	// finish once the session is over before it interrupts a write that is stuck
			112	+	// (a guest that stopped reading): the write deadline is set to now, and the
			113	+	// stream is closed if it cannot be interrupted otherwise. The session is over at
			114	+	// that point, the caller closes the stream anyway.
			115	+	var resolveJoinGrace = 2 * time.Second
			116	+
			117	+	// joinResolver waits for the worker. A worker in a Send that cannot finish
			118	+	// because the peer does not read would block DriveJob for ever, so the write is
			119	+	// interrupted after the grace period.
			120	+	func joinResolver(rw io.ReadWriter, workers *sync.WaitGroup) {
			121	+	finished := make(chan struct{})
			122	+	go func() { workers.Wait(); close(finished) }()
			123	+	select {
			124	+	case <-finished:
			125	+	return
			126	+	case <-time.After(resolveJoinGrace):
			127	+	}
			128	+	if d, ok := rw.(interface{ SetWriteDeadline(time.Time) error }); ok {
			129	+	_ = d.SetWriteDeadline(time.Now())
			130	+	}
			131	+	select {
			132	+	case <-finished:
			133	+	return
			134	+	case <-time.After(resolveJoinGrace):
			135	+	}
			136	+	if cl, ok := rw.(io.Closer); ok {
			137	+	_ = cl.Close()
			138	+	<-finished
			139	+	}
			140	+	}
			141	+
109		// resolveQueue is how many ResolveRequest frames may wait for the worker.	142		// resolveQueue is how many ResolveRequest frames may wait for the worker.
110		const resolveQueue = 2	143		const resolveQueue = 2
111			144

@@ -167,13 +200,16 @@ func DriveJob(rw io.ReadWriter, run execpb.Run, cancel <-chan execpb.Cancel, h
167		case req := <-resolveCh:	200		case req := <-resolveCh:
168		resp := h.Resolve(rctx, req)	201		resp := h.Resolve(rctx, req)
169		resp.Id = req.Id	202		resp.Id = req.Id
			203	+	if rctx.Err() != nil {
			204	+	return // the session is over: nobody reads the answer
			205	+	}
170		_ = c.Send(resolveFrame(resp))	206		_ = c.Send(resolveFrame(resp))
171		case <-rctx.Done():	207		case <-rctx.Done():
172		return	208		return
173		}	209		}
174		}	210		}
175		}()	211		}()
176	-	defer func() { rcancel(); workers.Wait() }()	212	+	defer func() { rcancel(); joinResolver(rw, &workers) }()
177		if cancel != nil {	213		if cancel != nil {
178		go func() {	214		go func() {
179		select {	215		select {

M internal/execstream/frame_test.go +68 -0


@@ -260,3 +260,71 @@ func TestDriveJobResolveBurst(t *testing.T) {
260	260		}
261	261		close(release)
262	262		}
	263	+
	264	+	// GO-009: the resolution worker is in a Send that cannot finish because the
	265	+	// guest stopped reading, and the guest ends the session with Exit. DriveJob
	266	+	// must return, not wait for the write for ever.
	267	+	func TestDriveJobReturnsWhenTheResolverWriteIsStuck(t *testing.T) {
	268	+	old := resolveJoinGrace
	269	+	resolveJoinGrace = 50 * time.Millisecond
	270	+	defer func() { resolveJoinGrace = old }()
	271	+	hostSide, guestSide := net.Pipe() // synchronous: a write needs a reader
	272	+	defer func() { _ = guestSide.Close() }()
	273	+	gc := NewConn(guestSide)
	274	+	go func() {
	275	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}})
	276	+	if _, err := gc.Recv(); err != nil { // the Run frame
	277	+	return
	278	+	}
	279	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_ResolveRequest{ResolveRequest: &execpb.ResolveRequest{Id: 1, Uses: []string{"a/b@v1"}}}})
	280	+	time.Sleep(100 * time.Millisecond) // the answer is being written: nobody reads it
	281	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}})
	282	+	}()
	283	+	done := make(chan error, 1)
	284	+	go func() {
	285	+	_, exit, err := DriveJob(hostSide, &execpb.Run{}, nil, JobHandler{Resolve: func(context.Context, execpb.ResolveRequest) execpb.ResolveResponse {
	286	+	return &execpb.ResolveResponse{}
	287	+	}})
	288	+	if err == nil && exit == nil {
	289	+	err = errors.New("no exit")
	290	+	}
	291	+	done <- err
	292	+	}()
	293	+	select {
	294	+	case err := <-done:
	295	+	if err != nil {
	296	+	t.Fatal(err)
	297	+	}
	298	+	case <-time.After(10 * time.Second):
	299	+	t.Fatal("DriveJob is blocked by a write nobody reads")
	300	+	}
	301	+	}
	302	+
	303	+	// And a worker that has an answer when the session ends does not write it.
	304	+	func TestDriveJobDoesNotSendAfterTheSession(t *testing.T) {
	305	+	hostSide, guestSide := net.Pipe()
	306	+	defer func() { _ = guestSide.Close() }()
	307	+	gc := NewConn(guestSide)
	308	+	release := make(chan struct{})
	309	+	go func() {
	310	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}})
	311	+	_, _ = gc.Recv()
	312	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_ResolveRequest{ResolveRequest: &execpb.ResolveRequest{Id: 1, Uses: []string{"a/b@v1"}}}})
	313	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}})
	314	+	close(release)
	315	+	}()
	316	+	done := make(chan struct{})
	317	+	go func() {
	318	+	defer close(done)
	319	+	_, _, _ = DriveJob(hostSide, &execpb.Run{}, nil, JobHandler{Resolve: func(ctx context.Context, _ execpb.ResolveRequest) execpb.ResolveResponse {
	320	+	<-ctx.Done() // returns only when the session ends
	321	+	return &execpb.ResolveResponse{}
	322	+	}})
	323	+	}()
	324	+	select {
	325	+	case <-done:
	326	+	case <-time.After(10 * time.Second):
	327	+	t.Fatal("DriveJob did not return")
	328	+	}
	329	+	<-release
	330	+	}


@@ -260,3 +260,71 @@ func TestDriveJobResolveBurst(t *testing.T) {
260	}	260		}
261	close(release)	261		close(release)
262	}	262		}
		263	+
		264	+	// GO-009: the resolution worker is in a Send that cannot finish because the
		265	+	// guest stopped reading, and the guest ends the session with Exit. DriveJob
		266	+	// must return, not wait for the write for ever.
		267	+	func TestDriveJobReturnsWhenTheResolverWriteIsStuck(t *testing.T) {
		268	+	old := resolveJoinGrace
		269	+	resolveJoinGrace = 50 * time.Millisecond
		270	+	defer func() { resolveJoinGrace = old }()
		271	+	hostSide, guestSide := net.Pipe() // synchronous: a write needs a reader
		272	+	defer func() { _ = guestSide.Close() }()
		273	+	gc := NewConn(guestSide)
		274	+	go func() {
		275	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}})
		276	+	if _, err := gc.Recv(); err != nil { // the Run frame
		277	+	return
		278	+	}
		279	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_ResolveRequest{ResolveRequest: &execpb.ResolveRequest{Id: 1, Uses: []string{"a/b@v1"}}}})
		280	+	time.Sleep(100 * time.Millisecond) // the answer is being written: nobody reads it
		281	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}})
		282	+	}()
		283	+	done := make(chan error, 1)
		284	+	go func() {
		285	+	_, exit, err := DriveJob(hostSide, &execpb.Run{}, nil, JobHandler{Resolve: func(context.Context, execpb.ResolveRequest) execpb.ResolveResponse {
		286	+	return &execpb.ResolveResponse{}
		287	+	}})
		288	+	if err == nil && exit == nil {
		289	+	err = errors.New("no exit")
		290	+	}
		291	+	done <- err
		292	+	}()
		293	+	select {
		294	+	case err := <-done:
		295	+	if err != nil {
		296	+	t.Fatal(err)
		297	+	}
		298	+	case <-time.After(10 * time.Second):
		299	+	t.Fatal("DriveJob is blocked by a write nobody reads")
		300	+	}
		301	+	}
		302	+
		303	+	// And a worker that has an answer when the session ends does not write it.
		304	+	func TestDriveJobDoesNotSendAfterTheSession(t *testing.T) {
		305	+	hostSide, guestSide := net.Pipe()
		306	+	defer func() { _ = guestSide.Close() }()
		307	+	gc := NewConn(guestSide)
		308	+	release := make(chan struct{})
		309	+	go func() {
		310	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Hello{Hello: &execpb.Hello{}}})
		311	+	_, _ = gc.Recv()
		312	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_ResolveRequest{ResolveRequest: &execpb.ResolveRequest{Id: 1, Uses: []string{"a/b@v1"}}}})
		313	+	_ = gc.Send(&execpb.Frame{Body: &execpb.Frame_Exit{Exit: &execpb.Exit{}}})
		314	+	close(release)
		315	+	}()
		316	+	done := make(chan struct{})
		317	+	go func() {
		318	+	defer close(done)
		319	+	_, _, _ = DriveJob(hostSide, &execpb.Run{}, nil, JobHandler{Resolve: func(ctx context.Context, _ execpb.ResolveRequest) execpb.ResolveResponse {
		320	+	<-ctx.Done() // returns only when the session ends
		321	+	return &execpb.ResolveResponse{}
		322	+	}})
		323	+	}()
		324	+	select {
		325	+	case <-done:
		326	+	case <-time.After(10 * time.Second):
		327	+	t.Fatal("DriveJob did not return")
		328	+	}
		329	+	<-release
		330	+	}