From 7b06acc46911d721282e0f2b837f012e0c772b62 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 11:41:56 -0500 Subject: [PATCH 1/7] fix(arrow/flight/session): fix flaky race condition test --- arrow/flight/session/stateful_session.go | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/arrow/flight/session/stateful_session.go b/arrow/flight/session/stateful_session.go index 322f60db4..dff7ee92f 100644 --- a/arrow/flight/session/stateful_session.go +++ b/arrow/flight/session/stateful_session.go @@ -170,7 +170,15 @@ func (manager *statefulServerSessionManager) GetSession(ctx context.Context) (Se sessionID, err := getSessionIDFromIncomingCookie(ctx) if err == nil { - return manager.store.Get(sessionID) + session, err := manager.store.Get(sessionID) + if err != nil { + // If the session isn't in the store (e.g., it was closed/removed), + // treat it as if no session exists so a new one can be created. + // This handles the race condition where a client sends a stale + // cookie before processing the session deletion trailer. + return nil, ErrNoSession + } + return session, nil } if err == http.ErrNoCookie { return nil, ErrNoSession From ac22d3efee5c1340a2a72188aae5027b781c677a Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 12:01:27 -0500 Subject: [PATCH 2/7] another race condition test failure --- arrow/flight/session/stateful_session.go | 6 ++++++ arrow/flight/session/stateless_session.go | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/arrow/flight/session/stateful_session.go b/arrow/flight/session/stateful_session.go index dff7ee92f..ac4e700b4 100644 --- a/arrow/flight/session/stateful_session.go +++ b/arrow/flight/session/stateful_session.go @@ -178,6 +178,12 @@ func (manager *statefulServerSessionManager) GetSession(ctx context.Context) (Se // cookie before processing the session deletion trailer. return nil, ErrNoSession } + // Also check if the session has been marked as closed but not yet removed. + // This handles the race between CallCompleted removing the session and + // StartCall looking it up from the cookie. + if session.Closed() { + return nil, ErrNoSession + } return session, nil } if err == http.ErrNoCookie { diff --git a/arrow/flight/session/stateless_session.go b/arrow/flight/session/stateless_session.go index 9e4afb60f..105b91453 100644 --- a/arrow/flight/session/stateless_session.go +++ b/arrow/flight/session/stateless_session.go @@ -51,6 +51,12 @@ func (manager *statelessServerSessionManager) GetSession(ctx context.Context) (S session, err = getSessionFromIncomingCookie(ctx) if err == nil { + // Check if the session has been marked as closed. + // This handles the race between CallCompleted marking the session closed + // and StartCall looking it up from the cookie. + if session.Closed() { + return nil, ErrNoSession + } return session, err } if err == http.ErrNoCookie { From f7f247d071fffc9b4d8a2fe048dab8201a0bb1cf Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 12:26:25 -0500 Subject: [PATCH 3/7] more flaky fixes --- arrow/flight/client.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/arrow/flight/client.go b/arrow/flight/client.go index 9e5da87c9..96eb0e6b3 100644 --- a/arrow/flight/client.go +++ b/arrow/flight/client.go @@ -152,8 +152,13 @@ func CreateClientMiddleware(middleware CustomClientMiddleware) ClientMiddleware post.CallCompleted(csCtx, err) } if isHdrs { - hdrmd, _ := cs.Header() - hdrs.HeadersReceived(csCtx, metadata.Join(hdrmd, cs.Trailer())) + // Only retrieve headers and trailers if the stream completed successfully + // or with io.EOF. If the stream was cancelled or had an error before + // completion, headers/trailers may not be available. + if err == nil || errors.Is(err, io.EOF) { + hdrmd, _ := cs.Header() + hdrs.HeadersReceived(csCtx, metadata.Join(hdrmd, cs.Trailer())) + } } } go func() { From 4a97a7f53232b8a061020d8532854dd5a6c58018 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 13:09:32 -0500 Subject: [PATCH 4/7] remove unecessary check, fix race --- arrow/flight/session/stateless_session.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/arrow/flight/session/stateless_session.go b/arrow/flight/session/stateless_session.go index 105b91453..9e4afb60f 100644 --- a/arrow/flight/session/stateless_session.go +++ b/arrow/flight/session/stateless_session.go @@ -51,12 +51,6 @@ func (manager *statelessServerSessionManager) GetSession(ctx context.Context) (S session, err = getSessionFromIncomingCookie(ctx) if err == nil { - // Check if the session has been marked as closed. - // This handles the race between CallCompleted marking the session closed - // and StartCall looking it up from the cookie. - if session.Closed() { - return nil, ErrNoSession - } return session, err } if err == http.ErrNoCookie { From 0f8c830d1356016cef6ae9a43fba2c112f6c9d21 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 13:48:53 -0500 Subject: [PATCH 5/7] properly synchronize between Close and Closed --- arrow/flight/session/session.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/arrow/flight/session/session.go b/arrow/flight/session/session.go index a2969d9ab..65ae63b6b 100644 --- a/arrow/flight/session/session.go +++ b/arrow/flight/session/session.go @@ -141,12 +141,16 @@ func (session *serverSession) EraseSessionOption(name string) { } func (session *serverSession) Close() error { + session.mu.Lock() + defer session.mu.Unlock() session.options = nil session.closed = true return nil } func (session *serverSession) Closed() bool { + session.mu.RLock() + defer session.mu.RUnlock() return session.closed } From b84984c9a558147f4bdae10df71fb5d863877e4b Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 14:09:01 -0500 Subject: [PATCH 6/7] next attempt --- arrow/flight/session/stateless_session.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/arrow/flight/session/stateless_session.go b/arrow/flight/session/stateless_session.go index 9e4afb60f..4e3446333 100644 --- a/arrow/flight/session/stateless_session.go +++ b/arrow/flight/session/stateless_session.go @@ -51,6 +51,14 @@ func (manager *statelessServerSessionManager) GetSession(ctx context.Context) (S session, err = getSessionFromIncomingCookie(ctx) if err == nil { + // Check if the session has been marked as closed. + // While freshly decoded sessions always have closed=false, this check + // provides defense against edge cases where a session might be cached + // or reused. It also handles the race condition where a client sends + // a request with a stale cookie before processing the deletion trailer. + if session.Closed() { + return nil, ErrNoSession + } return session, err } if err == http.ErrNoCookie { From 7ae18ca675d86c26ca1bebc9b8a4dba4bd2f46b6 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 4 Mar 2026 14:27:33 -0500 Subject: [PATCH 7/7] another attempt... --- arrow/flight/session/session.go | 6 +++++- arrow/flight/session/stateless_session.go | 8 -------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/arrow/flight/session/session.go b/arrow/flight/session/session.go index 65ae63b6b..202de38b5 100644 --- a/arrow/flight/session/session.go +++ b/arrow/flight/session/session.go @@ -225,7 +225,11 @@ func (middleware *serverSessionMiddleware) CallCompleted(ctx context.Context, _ } if session.Closed() { - // Invalidate the client's cookie + // Invalidate the client's cookie by sending back the incoming cookie + // with MaxAge=-1. This follows the HTTP cookie protocol for deletion. + // Note: For stateless sessions, there's an inherent race condition where + // the client might send the old cookie again before processing this trailer. + // This is acceptable as the session data is client-controlled anyway. clientCookie.MaxAge = -1 grpc.SetTrailer(ctx, metadata.Pairs("Set-Cookie", clientCookie.String())) diff --git a/arrow/flight/session/stateless_session.go b/arrow/flight/session/stateless_session.go index 4e3446333..9e4afb60f 100644 --- a/arrow/flight/session/stateless_session.go +++ b/arrow/flight/session/stateless_session.go @@ -51,14 +51,6 @@ func (manager *statelessServerSessionManager) GetSession(ctx context.Context) (S session, err = getSessionFromIncomingCookie(ctx) if err == nil { - // Check if the session has been marked as closed. - // While freshly decoded sessions always have closed=false, this check - // provides defense against edge cases where a session might be cached - // or reused. It also handles the race condition where a client sends - // a request with a stale cookie before processing the deletion trailer. - if session.Closed() { - return nil, ErrNoSession - } return session, err } if err == http.ErrNoCookie {