Properly handle grpc dial errors in the throttler metric aggregation (vitessio#18073) (#155)

Signed-off-by: Arthur Schreiber <[email protected]>
Signed-off-by: Mohamed Hamza <[email protected]>
Co-authored-by: Arthur Schreiber <[email protected]>

Author: mhamza15

go/vt/vttablet/tabletserver/throttle/base/throttle_metric.go

@@ -43,7 +43,10 @@ package base
 
 import (
 	"errors"
-	"strings"
+	"net"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // MetricResult is what we expect our probes to return. This can be a numeric result, or
@@ -66,11 +69,20 @@ var ErrNoSuchMetric = errors.New("No such metric")
 var ErrInvalidCheckType = errors.New("Unknown throttler check type")
 
 // IsDialTCPError sees if the given error indicates a TCP issue
-func IsDialTCPError(e error) bool {
-	if e == nil {
+func IsDialTCPError(err error) bool {
+	if err == nil {
 		return false
 	}
-	return strings.HasPrefix(e.Error(), "dial tcp")
+
+	if s, ok := status.FromError(err); ok {
+		return s.Code() == codes.Unavailable || s.Code() == codes.DeadlineExceeded
+	}
+
+	switch err := err.(type) {
+	case *net.OpError:
+		return err.Op == "dial" && err.Net == "tcp"
+	}
+	return false
 }
 
 type noHostsMetricResult struct{}

go/vt/vttablet/tabletserver/throttle/throttler_test.go

@@ -28,10 +28,15 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"vitess.io/vitess/go/vt/grpcclient"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/vtenv"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
+	"vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/base"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/config"
 	"vitess.io/vitess/go/vt/vttablet/tabletserver/throttle/mysql"
 	"vitess.io/vitess/go/vt/vttablet/tmclient"
@@ -157,6 +162,21 @@ func newTestThrottler() *Throttler {
 	return throttler
 }
 
+func TestIsDialTCPError(t *testing.T) {
+	// Verify that IsDialTCPError actually recognizes grpc dial errors
+	cc, err := grpcclient.DialContext(context.Background(), ":0", true, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	defer cc.Close()
+
+	err = cc.Invoke(context.Background(), "/Fail", nil, nil)
+
+	require.True(t, base.IsDialTCPError(err))
+	require.True(t, base.IsDialTCPError(fmt.Errorf("wrapped: %w", err)))
+
+	nonDialErr := fmt.Errorf("rpc error: code = NotFound desc = method not found")
+	require.False(t, base.IsDialTCPError(nonDialErr))
+}
+
 func TestIsAppThrottled(t *testing.T) {
 	throttler := Throttler{
 		throttledApps:   cache.New(cache.NoExpiration, 0),