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() { diff --git a/arrow/flight/session/session.go b/arrow/flight/session/session.go index a2969d9ab..202de38b5 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 } @@ -221,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/stateful_session.go b/arrow/flight/session/stateful_session.go index 322f60db4..ac4e700b4 100644 --- a/arrow/flight/session/stateful_session.go +++ b/arrow/flight/session/stateful_session.go @@ -170,7 +170,21 @@ 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 + } + // 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 { return nil, ErrNoSession