fix(arrow/flight/session): fix flaky race condition test

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() {

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()))
 

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