kvserver: handle context cancellation when receiving snapshots

Following a3fd4fb, context errors are now propagated back up the
stack when receiving snapshots. However, this triggered a
`maybeFatalOnRaftReadyErr` assertion which crashed the node.

`handleRaftReadyRaftMuLocked` (which is called directly when applying
snapshots) does not appear prepared to safely handle arbitrary context
cancellation, as it is typically called with the Raft scheduler's
background context. Furthermore, by the time we call it we have already
received the entire snapshot, so it does not seem useful to abort
snapshot application just because the caller goes away.

This patch therefore uses a new background context for applying
snapshots, disconnected from the client's context, once the entire
snapshot has been received.

Release note: None

Author: erikgrinaker

pkg/kv/kvserver/store_raft.go

@@ -338,14 +338,14 @@ func (s *Store) processRaftSnapshotRequest(
 		var expl string
 		var err error
 		stats, expl, err = r.handleRaftReadyRaftMuLocked(ctx, inSnap)
+		maybeFatalOnRaftReadyErr(ctx, expl, err)
 		if !stats.snap.applied {
 			// This line would be hit if a snapshot was sent when it isn't necessary
 			// (i.e. follower was able to catch up via the log in the interim) or when
 			// multiple snapshots raced (as is possible when raft leadership changes
 			// and both the old and new leaders send snapshots).
 			log.Infof(ctx, "ignored stale snapshot at index %d", snapHeader.RaftMessageRequest.Message.Snapshot.Metadata.Index)
 		}
-		maybeFatalOnRaftReadyErr(ctx, expl, err)
 		return nil
 	})
 }

pkg/kv/kvserver/store_snapshot.go

@@ -660,7 +660,13 @@ func (s *Store) receiveSnapshot(
 		return err
 	}
 	inSnap.placeholder = placeholder
-	if err := s.processRaftSnapshotRequest(ctx, header, inSnap); err != nil {
+
+	// Use a background context for applying the snapshot, as handleRaftReady is
+	// not prepared to deal with for arbitrary context cancellation. We've already
+	// received the entire snapshot at this point too, so there's no point in
+	// abandoning application half-way through if the caller goes away.
+	applyCtx := s.AnnotateCtx(context.Background())
+	if err := s.processRaftSnapshotRequest(applyCtx, header, inSnap); err != nil {
 		return sendSnapshotError(stream, errors.Wrap(err.GoError(), "failed to apply snapshot"))
 	}
 	return stream.Send(&SnapshotResponse{Status: SnapshotResponse_APPLIED})