[PureGo] Add GracefulClose to transport Stream by zlata-stefanovic-db · Pull Request #489 · databricks/zerobus-sdk

zlata-stefanovic-db · 2026-07-09T09:29:43Z

What changes are proposed in this pull request?

Adds GracefulClose(ctx) to the pure-Go transport Stream, encapsulating the clean-shutdown sequence so higher layers don't re-implement it and get the order wrong.

rawstream.go — generic gracefulClose(ctx) on rawStream: half-close the send side (CloseSend), drain and discard remaining responses to io.EOF, then release resources (close). The future Arrow/Flight wireStream inherits it for free. Draining to EOF lets the server see an orderly close; a bare Close cancels the context and the server sees an abrupt reset instead.
ephemeral_stream.go — exposes func (s *Stream) GracefulClose(ctx context.Context) error; CloseSend/Close docs updated to cross-reference it. Close() stays as the hard-abort escape hatch.
ctx bounds the drain: on expiry or a non-EOF stream error it hard-aborts (via close) and returns the cause (ctx error preferred), else nil. The ctx is bridged to close so a caller deadline unblocks the in-flight recv, which otherwise waits on the stream's Close-only context.

Why

Before this, the transport layer exposed only Close(), which fires the context cancel — closing without a preceding drain makes the server observe an abrupt RST_STREAM rather than an orderly END_STREAM. The correct sequence was documented on CloseSend but nothing encapsulated it, so every higher layer (ingest, recovery) would have to re-implement it. This is the same trap the Rust gRPC path had to unwind.

Raised in review of #463.

How is this tested?

transport_test.go:

TestStreamGracefulClose — server ends the stream on half-close; drain reaches a clean io.EOF and GracefulClose returns nil.
TestStreamGracefulCloseDrainsPending — an in-flight ack precedes io.EOF; the drain must discard it and keep going.
TestStreamGracefulCloseHonorsDeadline — server never ends the stream (new hangDrain fake-server flag); GracefulClose returns its ctx error promptly and leaves the stream torn down.

gofmt, go vet, and go test -race ./... all clean.

Notes for reviewers

No changelog entry — all code is under purego/internal/, no public API yet.
Timing — @teodordelibasic-db suggested adding this "once ingest lands." Pulling it slightly forward so the ingestion layer has a correct primitive to call from day one rather than open-coding the sequence; happy to defer into the ingestion PR instead if preferred.

Closes #487. Part of #467.

Encapsulate the clean-shutdown sequence (CloseSend, drain to io.EOF, Close) on rawStream so embedding streams don't re-implement it and get the order wrong. A bare Close cancels the context and the server sees an abrupt reset; draining to EOF lets it see an orderly close. ctx bounds the drain: on expiry or a stream error it hard-aborts and returns the cause, else nil. Close stays as the hard-abort escape hatch. Closes #487. Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

Note the send side is half-closed (no Send may follow) and make the 'later Close is a no-op' guarantee explicit: every return path releases resources first. Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

davidtosovic-db · 2026-07-10T09:58:29Z

+// returns the cause, else nil. Must not be called concurrently with Recv, and no
+// Send may follow (the send side is half-closed). A later Close is a no-op, since
+// GracefulClose always releases resources before it returns.
+func (s *Stream) GracefulClose(ctx context.Context) error { return s.gracefulClose(ctx) }


Nothing guarantees that this context has a deadline. Meaning we could hang here if the server is not sending acks on recv() channel.
For Open(....) you have explicit guards against this, where you check for a deadline and apply a default one if it's not present.

openCtx := ctx if _, ok := ctx.Deadline(); !ok { var cancelOpen context.CancelFunc openCtx, cancelOpen = context.WithTimeout(ctx, defaultHandshakeTimeout) defer cancelOpen() }

For safety we should have the same or equivalent mechanism here.

context.Background() passed to GracefulClose would hang forever if the server stalled during the drain, since the only exit was ctx expiry. Apply defaultHandshakeTimeout internally when no deadline is present, mirroring the same guard in Open. Add TestStreamGracefulCloseDefaultsDeadline to cover the no-deadline path (marked -short-skippable as it takes ~30s to fire). Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

…bound Add TestGracefulCloseCloseSendFailure to rawstream_test.go covering the path where CloseSend itself fails â�� verifies the error is returned and close() is called to hard-abort the stream. Tighten the elapsed-time assertion in TestStreamGracefulCloseHonorsDeadline from a fixed 5s bound to 5x the actual deadline (5*200ms = 1s), so the test catches genuine slowness rather than accepting a 4.9s return when the deadline was 200ms. Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

…bound Add TestGracefulCloseCloseSendFailure to rawstream_test.go covering the path where CloseSend itself fails â�� verifies the error is returned and close() is called to hard-abort the stream. Replace the wall-clock elapsed check in TestStreamGracefulCloseHonorsDeadline with errors.Is(err, context.DeadlineExceeded): asserts what the deadline mechanism actually returned rather than inferring it from timing, removing the flakiness entirely. Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

teodordelibasic-db

First pass just regarding nits and styling.

teodordelibasic-db · 2026-07-10T10:33:36Z

+// half-closed). Every return path calls close first, so a later close is a no-op.
+func (s *rawStream[Req, Resp]) gracefulClose(ctx context.Context) error {
+	if err := s.closeSend(); err != nil {
+		s.close() // can't half-close cleanly; hard-abort


Let's tidy up these inline comments.

teodordelibasic-db · 2026-07-10T10:33:45Z

 		req, err := stream.Recv()
 		if err == io.EOF {
+			if f.hangDrain {
+				<-stream.Context().Done() // withhold the clean end


teodordelibasic-db · 2026-07-10T10:34:06Z

-		if err != nil {
-			t.Fatalf("Recv during drain: got %v, want io.EOF", err)
-		}
+	stream.Close() // idempotent no-op after a graceful shutdown


teodordelibasic-db · 2026-07-10T11:02:53Z

 // rawStream; this type adds EphemeralStream wire types and the two handshake
 // hooks. As with the embedded rawStream, a Stream is not safe for concurrent
 // Send and should use a single writer goroutine.
 type Stream struct {


Unrelated to this PR and we can discuss and do a follow up if we decide to do so, but since we will probably have an ArrowStream, I think leaving this as Stream might be confusing. Maybe more natural would be GrpcStream. What do you think?

Maybe we can name it EphemeralStream as the rpc it wraps? Because Arrow is also a grpc stream technically so it might be confusing - so we'd have rawStream as a shared base and then 1. EphemeralStream and 2. ArrowStream

teodordelibasic-db · 2026-07-10T11:03:35Z

Also unrelated and potential follow up, do you think raw_stream.go might be more per Go file naming standards?

Yes, I can rename it

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

zlata-stefanovic-db self-assigned this Jul 9, 2026

zlata-stefanovic-db marked this pull request as ready for review July 9, 2026 10:43

Merge branch 'main' into purego-graceful-close

a55738a

zlata-stefanovic-db requested a review from a team July 9, 2026 17:28

zlata-stefanovic-db mentioned this pull request Jul 10, 2026

[Go] Design a pure-Go SDK (no cgo/FFI) #467

Open

10 tasks

zlata-stefanovic-db added the enhancement New feature or request label Jul 10, 2026

davidtosovic-db reviewed Jul 10, 2026

View reviewed changes

Comment thread purego/internal/transport/rawstream.go Outdated

davidtosovic-db reviewed Jul 10, 2026

View reviewed changes

Comment thread purego/internal/transport/transport_test.go Outdated

davidtosovic-db previously approved these changes Jul 10, 2026

View reviewed changes

zlata-stefanovic-db added 2 commits July 10, 2026 10:41

zlata-stefanovic-db dismissed davidtosovic-db’s stale review via 4e99e19 July 10, 2026 10:45

zlata-stefanovic-db added 2 commits July 10, 2026 10:50

Merge branch 'main' into purego-graceful-close

212faca

teodordelibasic-db reviewed Jul 10, 2026

View reviewed changes

zlata-stefanovic-db added 2 commits July 10, 2026 11:26

[PureGo] Remove inline comments flagged in review

be8c921

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

[PureGo] Rename rawstream.go to raw_stream.go

7dd3703

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

zlata-stefanovic-db requested review from davidtosovic-db and teodordelibasic-db July 10, 2026 12:14

davidtosovic-db previously approved these changes Jul 10, 2026

View reviewed changes

zlata-stefanovic-db mentioned this pull request Jul 10, 2026

[PureGo] Add UC OAuth token provider with per-table caching #492

Draft

[PureGo] Fix 30s CI hang in graceful-close test; doc/test nits

3050ebc

Signed-off-by: Zlata Stefanovic <zlata.stefanovic@databricks.com>

zlata-stefanovic-db dismissed davidtosovic-db’s stale review via 3050ebc July 10, 2026 14:31

davidtosovic-db approved these changes Jul 13, 2026

View reviewed changes

zlata-stefanovic-db added this pull request to the merge queue Jul 13, 2026

Merged via the queue into main with commit 0c9da6f Jul 13, 2026
18 checks passed

zlata-stefanovic-db deleted the purego-graceful-close branch July 13, 2026 08:43

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[PureGo] Add GracefulClose to transport Stream#489

[PureGo] Add GracefulClose to transport Stream#489
zlata-stefanovic-db merged 10 commits into
mainfrom
purego-graceful-close

zlata-stefanovic-db commented Jul 9, 2026

Uh oh!

davidtosovic-db Jul 10, 2026

Uh oh!

Uh oh!

Uh oh!

teodordelibasic-db left a comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Uh oh!

teodordelibasic-db Jul 10, 2026

Uh oh!

teodordelibasic-db Jul 10, 2026

Uh oh!

teodordelibasic-db Jul 10, 2026

Uh oh!

zlata-stefanovic-db Jul 10, 2026 •

edited

Loading

Uh oh!

teodordelibasic-db Jul 10, 2026

Uh oh!

zlata-stefanovic-db Jul 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

zlata-stefanovic-db commented Jul 9, 2026

What changes are proposed in this pull request?

Why

How is this tested?

Notes for reviewers

Uh oh!

davidtosovic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

teodordelibasic-db left a comment

Choose a reason for hiding this comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

zlata-stefanovic-db Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

teodordelibasic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

zlata-stefanovic-db Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zlata-stefanovic-db Jul 10, 2026 •

edited

Loading