client: set TCP_USER_TIMEOUT socket option for linux

internal/syscall/syscall_linux.go

@@ -23,12 +23,18 @@
 package syscall
 
 import (
+	"errors"
+	"fmt"
+	"net"
 	"syscall"
+	"time"
 
 	"golang.org/x/sys/unix"
 	"google.golang.org/grpc/grpclog"
 )
 
+var GetTCPUserTimeoutNoopError = errors.New("placeholder error")
+
 // GetCPUTime returns the how much CPU time has passed since the start of this process.
 func GetCPUTime() int64 {
 	var ts unix.Timespec
@@ -65,3 +71,47 @@ func CPUTimeDiff(first *Rusage, latest *Rusage) (float64, float64) {
 
 	return uTimeElapsed, sTimeElapsed
 }
+
+// SetTCPUserTimeout sets the TCP user timeout on a connection's socket
+func SetTCPUserTimeout(conn net.Conn, timeout time.Duration) error {
+	tcpconn, ok := conn.(*net.TCPConn)
+	if !ok {
+		// not a TCP connection. exit early
+		return nil
+	}
+	rawConn, err := tcpconn.SyscallConn()
+	if err != nil {
+		return fmt.Errorf("error getting raw connection: %v", err)
+	}
+	err = rawConn.Control(func(fd uintptr) {
+		err = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT, int(timeout/time.Millisecond))
+	})
+	if err != nil {
+		return fmt.Errorf("error setting option on socket: %v", err)
+	}
+
+	return nil
+}
+
+// GetTCPUserTimeout gets the TCP user timeout on a connection's socket
+func GetTCPUserTimeout(conn net.Conn) (opt int, err error) {
+	tcpconn, ok := conn.(*net.TCPConn)
+	if !ok {
+		err = fmt.Errorf("conn is not *net.TCPConn. got %T", conn)
+		return
+	}
+	rawConn, err := tcpconn.SyscallConn()
+	if err != nil {
+		err = fmt.Errorf("error getting raw connection: %v", err)
+		return
+	}
+	err = rawConn.Control(func(fd uintptr) {
+		opt, err = syscall.GetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT)
+	})
+	if err != nil {
+		err = fmt.Errorf("error getting option on socket: %v", err)
+		return
+	}
+
+	return
+}

internal/syscall/syscall_nonlinux.go

@@ -20,7 +20,15 @@
 
 package syscall
 
-import "google.golang.org/grpc/grpclog"
+import (
+	"errors"
+	"net"
+	"time"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var GetTCPUserTimeoutNoopError = errors.New("GetTCPUserTimeout is a no-op on non-linux or appengine environments")
 
 func init() {
 	grpclog.Info("CPU time info is unavailable on non-linux or appengine environment.")
@@ -45,3 +53,13 @@ func GetRusage() (rusage *Rusage) {
 func CPUTimeDiff(first *Rusage, latest *Rusage) (float64, float64) {
 	return 0, 0
 }
+
+// SetTCPUserTimeout is a no-op function under non-linux or appengine environments
+func SetTCPUserTimeout(conn net.Conn, timeout time.Duration) error {
+	return nil
+}
+
+// GetTCPUserTimeout is a no-op function under non-linux or appengine environments
+func GetTCPUserTimeout(conn net.Conn) (int, error) {
+	return 0, GetTCPUserTimeoutNoopError
+}

internal/transport/http2_client.go

@@ -35,6 +35,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/syscall"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
@@ -165,6 +166,21 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr TargetInfo, opts Conne
 			conn.Close()
 		}
 	}(conn)
+	kp := opts.KeepaliveParams
+	// Validate keepalive parameters.
+	if kp.Time == 0 {
+		kp.Time = defaultClientKeepaliveTime
+	}
+	if kp.Timeout == 0 {
+		kp.Timeout = defaultClientKeepaliveTimeout
+	}
+	keepaliveEnabled := false
+	if kp.Time != infinity {
+		if err = syscall.SetTCPUserTimeout(conn, kp.Timeout); err != nil {
+			return nil, connectionErrorf(false, err, "transport: failed to set TCP_USER_TIMEOUT: %v", err)
+		}
+		keepaliveEnabled = true
+	}
 	var (
 		isSecure bool
 		authInfo credentials.AuthInfo
@@ -188,14 +204,6 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr TargetInfo, opts Conne
 		}
 		isSecure = true
 	}
-	kp := opts.KeepaliveParams
-	// Validate keepalive parameters.
-	if kp.Time == 0 {
-		kp.Time = defaultClientKeepaliveTime
-	}
-	if kp.Timeout == 0 {
-		kp.Timeout = defaultClientKeepaliveTimeout
-	}
 	dynamicWindow := true
 	icwz := int32(initialWindowSize)
 	if opts.InitialConnWindowSize >= defaultWindowSize {
@@ -239,6 +247,7 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr TargetInfo, opts Conne
 		czData:                new(channelzData),
 		onGoAway:              onGoAway,
 		onClose:               onClose,
+		keepaliveEnabled:      keepaliveEnabled,
 	}
 	t.controlBuf = newControlBuffer(t.ctxDone)
 	if opts.InitialWindowSize >= defaultWindowSize {
@@ -267,8 +276,7 @@ func newHTTP2Client(connectCtx, ctx context.Context, addr TargetInfo, opts Conne
 	if channelz.IsOn() {
 		t.channelzID = channelz.RegisterNormalSocket(t, opts.ChannelzParentID, "")
 	}
-	if t.kp.Time != infinity {
-		t.keepaliveEnabled = true
+	if t.keepaliveEnabled {
 		go t.keepalive()
 	}
 	// Start the reader goroutine for incoming message. Each transport has

internal/transport/transport_test.go

@@ -41,6 +41,7 @@ import (
 	"golang.org/x/net/http2/hpack"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/internal/leakcheck"
+	"google.golang.org/grpc/internal/syscall"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/status"
 )
@@ -2317,3 +2318,53 @@ func TestHeaderTblSize(t *testing.T) {
 		t.Fatalf("expected len(limits) = 2 within 10s, got != 2")
 	}
 }
+
+func TestTCPUserTimeout(t *testing.T) {
+	tests := []struct {
+		time    time.Duration
+		timeout time.Duration
+	}{
+		{
+			10 * time.Second,
+			10 * time.Second,
+		},
+		{
+			0,
+			0,
+		},
+	}
+	for _, tt := range tests {
+		lis, err := net.Listen("tcp", "localhost:0")
+		if err != nil {
+			t.Fatalf("Failed to listen. Err: %v", err)
+		}
+		defer lis.Close()
+		// TODO(deklerk): we can `defer cancel()` here after we drop Go 1.6 support. Until then,
+		// doing a `defer cancel()` could cause the dialer to become broken:
+		// https://github.com/golang/go/issues/15078, https://github.com/golang/go/issues/15035
+		connectCtx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Second))
+		client, err := newHTTP2Client(connectCtx, context.Background(), TargetInfo{Addr: lis.Addr().String()}, ConnectOptions{
+			KeepaliveParams: keepalive.ClientParameters{
+				Time:    tt.time,
+				Timeout: tt.timeout,
+			},
+		}, func() {}, func(GoAwayReason) {}, func() {})
+		if err != nil {
+			cancel() // Do not cancel in success path.
+			t.Fatalf("error creating client: %v", err)
+		}
+		defer client.Close()
+
+		opt, err := syscall.GetTCPUserTimeout(client.conn)
+		if err != nil {
+			if err == syscall.GetTCPUserTimeoutNoopError {
+				t.Skipf("skipping test on unsupported environment: %v", err)
+			}
+			t.Fatalf("GetTCPUserTimeout error: %v", err)
+		}
+		if timeoutMS := int(tt.timeout / time.Millisecond); timeoutMS != opt {
+			t.Fatalf("wrong TCP_USER_TIMEOUT set on conn. expected %d. got %d",
+				timeoutMS, opt)
+		}
+	}
+}