From 0ee46fab7ad27f9a53b654ade5fd00058a975f1b Mon Sep 17 00:00:00 2001
From: Zellyn Hunter <zellyn@squareup.com>
Date: Fri, 9 Jul 2021 16:46:18 -0400
Subject: [PATCH] add cancel-after: cancel without timing out

---
 cmd/grpcurl/grpcurl.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/cmd/grpcurl/grpcurl.go b/cmd/grpcurl/grpcurl.go
index b0e69a9..8b91e42 100644
--- a/cmd/grpcurl/grpcurl.go
+++ b/cmd/grpcurl/grpcurl.go
@@ -126,6 +126,9 @@ var (
 		The maximum total time the operation can take, in seconds. This is
 		useful for preventing batch jobs that use grpcurl from hanging due to
 		slow or bad network links or due to incorrect stream method usage.`))
+	cancelAfter = flags.Float64("cancel-after", 0, prettify(`
+		Time after which the context (and hence request) will be canceled.
+		This is useful for testing request cancellation.`))
 	maxMsgSz = flags.Int("max-msg-sz", 0, prettify(`
 		The maximum encoded size of a response message, in bytes, that grpcurl
 		will accept. If not specified, defaults to 4,194,304 (4 megabytes).`))
@@ -278,6 +281,9 @@ func main() {
 	if *maxTime < 0 {
 		fail(nil, "The -max-time argument must not be negative.")
 	}
+	if *cancelAfter < 0 {
+		fail(nil, "The -cancel-after argument must not be negative.")
+	}
 	if *maxMsgSz < 0 {
 		fail(nil, "The -max-msg-sz argument must not be negative.")
 	}
@@ -388,6 +394,17 @@ func main() {
 		defer cancel()
 	}
 
+	if *cancelAfter > 0 {
+		timeout := time.Duration(*cancelAfter * float64(time.Second))
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
+		defer cancel()
+		go func() {
+			time.Sleep(timeout)
+			cancel()
+		}()
+	}
+
 	dial := func() *grpc.ClientConn {
 		dialTime := 10 * time.Second
 		if *connectTimeout > 0 {