Skip to content

Commit c51caae

Browse files
sushanbsushanb
committed
refactor(bigtable): address igor nits on Session vRPC
Five nits from the igor-reviewer-agent review of PR googleapis#20213: - Return stillActive bool from markCancelled and use it in awaitInvokeResult so the ctx.Done branch does one slotMu take instead of two (markCancelled + recordCtxDone's activeVRPC). - Move ctx to the first arg of buildInvokeRequest per Go convention. - Drop the dead `_ = rpc` in processResult; remove rpc from the signature (only awaitInvokeResult calls it, no rpc reads today). - Delete the deliver() one-line wrapper — inline the buffered send at both call sites (routeVRPCFrame's non-cancelled branch, cancelActiveRPCs's non-cancelled branch). Two-word "cap-1, drainSlot serialized this write" comment replaces the paragraph docstring. - Replace time.Sleep(20ms) in TestInvoke_ForceCloseWhileSending_ BoundedReturn with a deterministic waitFor on activeVRPC() == nil, so the test doesn't rely on scheduler cooperation. Also converts the two remaining string-literal call sites in noteRetryAttempt and recordCtxDone to the typed SessionEventKind constants (SessionEventRetry, SessionEventCtxDone) landed in the prior PR-1 nit pass.
1 parent d98959d commit c51caae

2 files changed

Lines changed: 28 additions & 34 deletions

File tree

‎bigtable/internal/transport/session_vrpc.go‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -52,16 +52,19 @@ func (s *Session) claimSlot(rpc *vrpcImpl) bool {
5252
// markCancelled records ctx.Done cancellation of rpc without freeing the
5353
// slot — the caller returns, but activeRPC stays until the server response
5454
// arrives to drain it. First-cancel-wins; a racing drain that clears
55-
// activeRPC makes this a no-op.
56-
func (s *Session) markCancelled(rpc *vrpcImpl, res vrpcResult) {
55+
// activeRPC makes this a no-op. Returns whether rpc was still the active
56+
// slot occupant, so callers that also want the "still in flight?" signal
57+
// don't have to re-acquire slotMu with a separate activeVRPC() call.
58+
func (s *Session) markCancelled(rpc *vrpcImpl, res vrpcResult) (stillActive bool) {
5759
s.slotMu.Lock()
5860
defer s.slotMu.Unlock()
5961
if s.activeRPC != rpc {
60-
return
62+
return false
6163
}
6264
if s.currentCancel == nil {
6365
s.currentCancel = &res
6466
}
67+
return true
6568
}
6669

6770
// drainSlot atomically clears the (activeRPC, currentCancel) pair iff
@@ -134,7 +137,7 @@ func (s *Session) Invoke(ctx context.Context, desc VRpcDescriptor, req interface
134137
if attempt > 1 {
135138
s.noteRetryAttempt(ctx, desc.Method(), attempt)
136139
}
137-
sessionReq := buildInvokeRequest(rpcID, reqBytes, attempt, startTime, ctx)
140+
sessionReq := buildInvokeRequest(ctx, rpcID, reqBytes, attempt, startTime)
138141

139142
// Capture SentAt immediately before the frame is handed to Send so
140143
// downstream metrics can compute client-side blocking latency as
@@ -185,7 +188,7 @@ func (s *Session) noteRetryAttempt(ctx context.Context, method string, attempt i
185188
prevCode := status.Code(prev).String()
186189
s.debugf("retry attempt=%d method=%s prev_code=%s prev_err=%v",
187190
attempt, method, prevCode, prev)
188-
s.recordEvent("retry", "attempt=%d method=%s prev_code=%s prev_err=%v",
191+
s.recordEvent(SessionEventRetry, "attempt=%d method=%s prev_code=%s prev_err=%v",
189192
attempt, method, prevCode, prev)
190193
}
191194

@@ -195,7 +198,7 @@ func (s *Session) noteRetryAttempt(ctx context.Context, method string, attempt i
195198
// so the server measures from receive time rather than an absolute wall
196199
// clock. Omitted when ctx has no deadline or the budget is already
197200
// non-positive (the client-side ctx.Done branch will fire immediately).
198-
func buildInvokeRequest(rpcID int64, reqBytes []byte, attempt int64, startTime time.Time, ctx context.Context) *spb.SessionRequest {
201+
func buildInvokeRequest(ctx context.Context, rpcID int64, reqBytes []byte, attempt int64, startTime time.Time) *spb.SessionRequest {
199202
virtRpc := &spb.VirtualRpcRequest{
200203
RpcId: rpcID,
201204
Payload: reqBytes,
@@ -232,15 +235,15 @@ func (s *Session) awaitInvokeResult(ctx context.Context, rpc *vrpcImpl, desc VRp
232235
case <-ctx.Done():
233236
select {
234237
case res := <-rpc.resultChan:
235-
return s.processResult(rpc, desc, result, res)
238+
return s.processResult(desc, result, res)
236239
default:
237240
}
238-
s.recordCtxDone(ctx, rpc, desc.Method(), result.SentAt)
239241
cancelErr := tagErr(StateTransportFailure, ctx.Err())
240-
s.markCancelled(rpc, vrpcResult{err: cancelErr})
242+
stillActive := s.markCancelled(rpc, vrpcResult{err: cancelErr})
243+
s.recordCtxDone(ctx, rpc, desc.Method(), result.SentAt, stillActive)
241244
return cancelErr
242245
case res := <-rpc.resultChan:
243-
return s.processResult(rpc, desc, result, res)
246+
return s.processResult(desc, result, res)
244247
}
245248
}
246249

@@ -252,7 +255,7 @@ func (s *Session) awaitInvokeResult(ctx context.Context, rpc *vrpcImpl, desc VRp
252255
// The res.resp.RpcId == rpc.id check that used to live here is gone:
253256
// under slotMu, handleVRPCResponse gates the id match BEFORE drainSlot,
254257
// so deliver can only ever put a matching-id response into resultChan.
255-
func (s *Session) processResult(rpc *vrpcImpl, desc VRpcDescriptor, result *InvokeResult, res vrpcResult) error {
258+
func (s *Session) processResult(desc VRpcDescriptor, result *InvokeResult, res vrpcResult) error {
256259
result.TransportLatency = time.Since(result.SentAt)
257260
ci := res.ClusterInfo()
258261
result.ClusterInfo = ci
@@ -274,22 +277,21 @@ func (s *Session) processResult(rpc *vrpcImpl, desc VRpcDescriptor, result *Invo
274277
if res.resp.Stats != nil && res.resp.Stats.BackendLatency != nil {
275278
s.recordLatency(res.resp.Stats.BackendLatency.AsDuration())
276279
}
277-
_ = rpc // rpc kept in signature for future per-rpc metrics; no reads today.
278280
return nil
279281
}
280282

281283
// recordCtxDone emits the debug + sessionz event for a ctx cancellation
282-
// or deadline fire while a vRPC was in flight. Captures whether the RPC
283-
// was still holding the slot at cancel time — useful for spotting races
284-
// between our cancel and a late server response.
285-
func (s *Session) recordCtxDone(ctx context.Context, rpc *vrpcImpl, method string, sentAt time.Time) {
286-
stillActive := s.activeVRPC() == rpc
284+
// or deadline fire while a vRPC was in flight. stillActive comes from
285+
// markCancelled's return so no second slotMu take is needed here — it
286+
// reports whether rpc still held the slot at cancel time, useful for
287+
// spotting races between our cancel and a late server response.
288+
func (s *Session) recordCtxDone(ctx context.Context, rpc *vrpcImpl, method string, sentAt time.Time, stillActive bool) {
287289
sessState := State(s.state.Load())
288290
waited := time.Since(sentAt)
289291
peer := s.peerInfoSummary()
290292
s.debugf("vRPC %s rpc_id=%d ctx.Done waited=%v err=%v session_state=%v still_in_flight=%v %s",
291293
method, rpc.id, waited, ctx.Err(), sessState, stillActive, peer)
292-
s.recordEvent("ctx-done", "method=%s rpc_id=%d waited=%v err=%v session_state=%v still_in_flight=%v %s",
294+
s.recordEvent(SessionEventCtxDone, "method=%s rpc_id=%d waited=%v err=%v session_state=%v still_in_flight=%v %s",
293295
method, rpc.id, waited, ctx.Err(), sessState, stillActive, peer)
294296
}
295297

@@ -353,7 +355,9 @@ func (s *Session) routeVRPCFrame(rpcID int64, frameName, nilTag string, counter
353355
// resultChan. Just count the drain for observability.
354356
recordDebugTag(tagSessionVRPCCancelledDrained)
355357
} else {
356-
s.deliver(drained, result)
358+
// resultChan is cap-1 and drainSlot serialized this write; the
359+
// send never blocks.
360+
drained.resultChan <- result
357361
}
358362
// v3: drainSlot success is the sole "session became free" signal.
359363
// Fires on every drain (not just the cancelled branch) so the pool
@@ -386,17 +390,6 @@ func errorResponseToErr(errResp *spb.ErrorResponse) error {
386390
return st.Err()
387391
}
388392

389-
// deliver writes a result onto the RPC's buffered (cap 1) channel.
390-
// Under slotMu, exactly one caller ever holds a drained rpc (the
391-
// winning drainSlot inside handleVRPCResponse / handleVRPCErrorResponse
392-
// / cancelActiveRPCs), so the two-writers race on resultChan is
393-
// impossible in production. The cap-1 buffer is retained as defense
394-
// against the awaitInvokeResult ctx.Done-vs-response tick race — the
395-
// send must not block if the reader stopped listening.
396-
func (s *Session) deliver(rpc *vrpcImpl, res vrpcResult) {
397-
rpc.resultChan <- res
398-
}
399-
400393
// ForceClose transitions the session straight to StateClosed and cancels
401394
// any in-flight vRPC with a TransportFailure-tagged error. Minimal port
402395
// from the sessionz-debug lifecycle body — the full teardown (setCloseReason,
@@ -442,5 +435,5 @@ func (s *Session) cancelActiveRPCs(err error) {
442435
// / benign shutdown while an RPC was in-flight. Server may or may not
443436
// have processed — TransportFailure classification lets idempotent ops
444437
// retry and prevents non-idempotent ones from double-applying.
445-
s.deliver(drained, vrpcResult{err: tagErr(StateTransportFailure, err)})
438+
drained.resultChan <- vrpcResult{err: tagErr(StateTransportFailure, err)}
446439
}

‎bigtable/internal/transport/session_vrpc_test.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -780,9 +780,10 @@ func TestInvoke_ForceCloseWhileSending_BoundedReturn(t *testing.T) {
780780
Description: "mid-send teardown",
781781
})
782782

783-
// Give ForceClose a moment to run before unblocking Send so the
784-
// race resolves in the intended order.
785-
time.Sleep(20 * time.Millisecond)
783+
// Block on cancelActiveRPCs actually clearing the slot before we
784+
// unblock Send — deterministic sync in place of a scheduler-hoping
785+
// time.Sleep so the test can't flake on a loaded CI box.
786+
waitFor(t, time.Second, func() bool { return s.activeVRPC() == nil }, "slot drained by cancelActiveRPCs")
786787
close(sendGate)
787788

788789
select {

0 commit comments

Comments
 (0)