Merge #49362

49362: cli/quit: proceed with hard shutdown even with short --drain-wait r=tbg a=knz

Found while investigating #49359.
:facepalm: on my side for not testing this.
Also :zap: on grpc for cooking their own error protocol that's incompatible with Go's.

Release note (bug fix): When the value passed to `--drain-wait` is
very small, but non-zero, `cockroach quit` in certain cases would
not proceed to perform a hard shutdown. This has been corrected.
This bug existed since v19.1.9, v19.2.7 and v20.1.1.

Co-authored-by: Raphael 'kena' Poss <[email protected]>

pkg/cli/interactive_tests/test_quit.tcl

@@ -0,0 +1,21 @@
+#! /usr/bin/env expect -f
+#
+source [file join [file dirname $argv0] common.tcl]
+
+start_server $argv
+
+spawn /bin/bash
+send "PS1=':''/# '\r"
+eexpect ":/# "
+
+start_test "Test that quit with a very short timeout still proceeds with hard shutdown"
+
+send "$argv quit --insecure --drain-wait=1ns\r"
+eexpect "drain did not complete successfully"
+eexpect "hard shutdown"
+eexpect "ok"
+eexpect ":/# "
+
+end_test
+
+stop_server $argv

pkg/cli/quit.go

@@ -128,7 +128,7 @@ func doDrain(
 		hardError, remainingWork, err = doDrainNoTimeout(ctx, c)
 		return err
 	})
-	if errors.HasType(err, (*contextutil.TimeoutError)(nil)) {
+	if errors.HasType(err, (*contextutil.TimeoutError)(nil)) || grpcutil.IsTimeout(err) {
 		log.Infof(ctx, "drain timed out: %v", err)
 		err = errors.New("drain timeout")
 	}
@@ -160,7 +160,7 @@ func doDrainNoTimeout(
 		})
 		if err != nil {
 			fmt.Fprintf(stderr, "\n") // finish the line started above.
-			return true, remainingWork, errors.Wrap(err, "error sending drain request")
+			return !grpcutil.IsTimeout(err), remainingWork, errors.Wrap(err, "error sending drain request")
 		}
 		for {
 			resp, err := stream.Recv()

pkg/util/grpcutil/grpc_util.go

@@ -41,6 +41,19 @@ func IsLocalRequestContext(ctx context.Context) bool {
 	return ctx.Value(localRequestKey{}) != nil
 }
 
+// IsTimeout returns true if err's Cause is a gRPC timeout, or the request
+// was canceled by a context timeout.
+func IsTimeout(err error) bool {
+	if errors.Is(err, context.DeadlineExceeded) {
+		return true
+	}
+	err = errors.Cause(err)
+	if s, ok := status.FromError(err); ok {
+		return s.Code() == codes.DeadlineExceeded
+	}
+	return false
+}
+
 // IsClosedConnection returns true if err's Cause is an error produced by gRPC
 // on closed connections.
 func IsClosedConnection(err error) bool {