[code-coverage] Generate code for matching client timeout wrapper

client/clientfactory.go

@@ -97,11 +97,11 @@ func NewRPCClientFactory(
 }
 
 func (cf *rpcClientFactory) NewHistoryClient() (history.Client, error) {
-	return cf.NewHistoryClientWithTimeout(history.DefaultTimeout)
+	return cf.NewHistoryClientWithTimeout(timeoutwrapper.HistoryDefaultTimeout)
 }
 
 func (cf *rpcClientFactory) NewMatchingClient(domainIDToName DomainIDToNameFunc) (matching.Client, error) {
-	return cf.NewMatchingClientWithTimeout(domainIDToName, matching.DefaultTimeout, matching.DefaultLongPollTimeout)
+	return cf.NewMatchingClientWithTimeout(domainIDToName, timeoutwrapper.MatchingDefaultTimeout, timeoutwrapper.MatchingDefaultLongPollTimeout)
 }
 
 func (cf *rpcClientFactory) NewHistoryClientWithTimeout(timeout time.Duration) (history.Client, error) {
@@ -153,20 +153,18 @@ func (cf *rpcClientFactory) NewMatchingClientWithTimeout(
 	peerResolver := matching.NewPeerResolver(cf.resolver, namedPort)
 
 	client := matching.NewClient(
-		timeout,
-		longPollTimeout,
 		rawClient,
 		peerResolver,
 		matching.NewLoadBalancer(domainIDToName, cf.dynConfig),
 	)
+	client = timeoutwrapper.NewMatchingClient(client, longPollTimeout, timeout)
 	if errorRate := cf.dynConfig.GetFloat64Property(dynamicconfig.MatchingErrorInjectionRate)(); errorRate != 0 {
 		client = errorinjectors.NewMatchingClient(client, errorRate, cf.logger)
 	}
 	if cf.metricsClient != nil {
 		client = metered.NewMatchingClient(client, cf.metricsClient)
 	}
 	return client, nil
-
 }
 
 func (cf *rpcClientFactory) NewAdminClientWithTimeoutAndConfig(

client/history/client.go

@@ -39,11 +39,6 @@ import (
 
 var _ Client = (*clientImpl)(nil)
 
-const (
-	// DefaultTimeout is the default timeout used to make calls
-	DefaultTimeout = time.Second * 30
-)
-
 type (
 	clientImpl struct {
 		numberOfShards    int

client/matching/client.go

@@ -22,7 +22,6 @@ package matching
 
 import (
 	"context"
-	"time"
 
 	"go.uber.org/yarpc"
 
@@ -33,35 +32,22 @@ import (
 
 var _ Client = (*clientImpl)(nil)
 
-const (
-	// DefaultTimeout is the default timeout used to make calls
-	DefaultTimeout = time.Minute
-	// DefaultLongPollTimeout is the long poll default timeout used to make calls
-	DefaultLongPollTimeout = time.Minute * 2
-)
-
 type clientImpl struct {
-	timeout         time.Duration
-	longPollTimeout time.Duration
-	client          Client
-	peerResolver    PeerResolver
-	loadBalancer    LoadBalancer
+	client       Client
+	peerResolver PeerResolver
+	loadBalancer LoadBalancer
 }
 
 // NewClient creates a new history service TChannel client
 func NewClient(
-	timeout time.Duration,
-	longPollTimeout time.Duration,
 	client Client,
 	peerResolver PeerResolver,
 	lb LoadBalancer,
 ) Client {
 	return &clientImpl{
-		timeout:         timeout,
-		longPollTimeout: longPollTimeout,
-		client:          client,
-		peerResolver:    peerResolver,
-		loadBalancer:    lb,
+		client:       client,
+		peerResolver: peerResolver,
+		loadBalancer: lb,
 	}
 }
 
@@ -81,8 +67,6 @@ func (c *clientImpl) AddActivityTask(
 	if err != nil {
 		return err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.AddActivityTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -102,8 +86,6 @@ func (c *clientImpl) AddDecisionTask(
 	if err != nil {
 		return err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.AddDecisionTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -123,8 +105,6 @@ func (c *clientImpl) PollForActivityTask(
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := c.createLongPollContext(ctx)
-	defer cancel()
 	return c.client.PollForActivityTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -144,8 +124,6 @@ func (c *clientImpl) PollForDecisionTask(
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := c.createLongPollContext(ctx)
-	defer cancel()
 	return c.client.PollForDecisionTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -165,8 +143,6 @@ func (c *clientImpl) QueryWorkflow(
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.QueryWorkflow(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -179,8 +155,6 @@ func (c *clientImpl) RespondQueryTaskCompleted(
 	if err != nil {
 		return err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.RespondQueryTaskCompleted(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -193,8 +167,6 @@ func (c *clientImpl) CancelOutstandingPoll(
 	if err != nil {
 		return err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.CancelOutstandingPoll(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -207,8 +179,6 @@ func (c *clientImpl) DescribeTaskList(
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.DescribeTaskList(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -221,8 +191,6 @@ func (c *clientImpl) ListTaskListPartitions(
 	if err != nil {
 		return nil, err
 	}
-	ctx, cancel := c.createContext(ctx)
-	defer cancel()
 	return c.client.ListTaskListPartitions(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
 
@@ -271,21 +239,3 @@ func (c *clientImpl) GetTaskListsByDomain(
 		ActivityTaskListMap: activityTaskListMap,
 	}, nil
 }
-
-func (c *clientImpl) createContext(
-	parent context.Context,
-) (context.Context, context.CancelFunc) {
-	if parent == nil {
-		return context.WithTimeout(context.Background(), c.timeout)
-	}
-	return context.WithTimeout(parent, c.timeout)
-}
-
-func (c *clientImpl) createLongPollContext(
-	parent context.Context,
-) (context.Context, context.CancelFunc) {
-	if parent == nil {
-		return context.WithTimeout(context.Background(), c.longPollTimeout)
-	}
-	return context.WithTimeout(parent, c.longPollTimeout)
-}

client/matching/interface.go

@@ -34,6 +34,7 @@ import (
 //go:generate gowrap gen -g -p . -i Client -t ../templates/errorinjectors.tmpl -o ../wrappers/errorinjectors/matching_generated.go -v client=Matching
 //go:generate gowrap gen -g -p . -i Client -t ../templates/grpc.tmpl -o ../wrappers/grpc/matching_generated.go -v client=Matching -v package=matchingv1 -v path=github.com/uber/cadence/.gen/proto/matching/v1 -v prefix=Matching
 //go:generate gowrap gen -g -p . -i Client -t ../templates/thrift.tmpl -o ../wrappers/thrift/matching_generated.go -v client=Matching -v prefix=Matching
+//go:generate gowrap gen -g -p . -i Client -t ../templates/timeout.tmpl -o ../wrappers/timeout/matching_generated.go -v client=Matching
 
 // Client is the interface exposed by types service client
 type Client interface {

client/templates/timeout.tmpl

@@ -4,8 +4,8 @@
 {{ $decorator := (printf "%s%s" (down $clientName) .Interface.Name) }}
 {{ $Decorator := (printf "%s%s" $ClientName .Interface.Name) }}
 {{$largeTimeoutAPIs := list "adminClient.GetCrossClusterTasks" "adminClient.GetReplicationMessages"}}
-{{$longPollTimeoutAPIs := list "frontendClient.ListArchivedWorkflowExecutions" "frontendClient.PollForActivityTask" "frontendClient.PollForDecisionTask"}}
-{{$noTimeoutAPIs := list "historyClient.GetReplicationMessages" "historyClient.GetDLQReplicationMessages" "historyClient.CountDLQMessages" "historyClient.ReadDLQMessages" "historyClient.PurgeDLQMessages" "historyClient.MergeDLQMessages" "historyClient.GetCrossClusterTasks" "historyClient.GetFailoverInfo"}}
+{{$longPollTimeoutAPIs := list "frontendClient.ListArchivedWorkflowExecutions" "frontendClient.PollForActivityTask" "frontendClient.PollForDecisionTask" "matchingClient.PollForActivityTask" "matchingClient.PollForDecisionTask"}}
+{{$noTimeoutAPIs := list "historyClient.GetReplicationMessages" "historyClient.GetDLQReplicationMessages" "historyClient.CountDLQMessages" "historyClient.ReadDLQMessages" "historyClient.PurgeDLQMessages" "historyClient.MergeDLQMessages" "historyClient.GetCrossClusterTasks" "historyClient.GetFailoverInfo" "matchingClient.GetTaskListsByDomain"}}
 {{/*
  $fieldMap defines a map of the decorator struct fields
  with field name as the key and field type as the value
@@ -15,10 +15,8 @@
     {{$fieldMap = merge $fieldMap (dict  "timeout" "time.Duration" "client" .Interface.Type) }}
 {{ else if eq $ClientName "Admin" }}
     {{$fieldMap = merge $fieldMap (dict  "timeout" "time.Duration" "client" .Interface.Type "largeTimeout" "time.Duration") }}
-{{ else if eq $ClientName "Frontend" }}
+{{ else }}
     {{$fieldMap = merge $fieldMap (dict "timeout" "time.Duration" "client" .Interface.Type "longPollTimeout" "time.Duration") }}
-{{ else if eq $ClientName "Matching" }}
-    {{$fieldMap = merge $fieldMap (dict "timeout" "time.Duration" "client" .Interface.Type "longPollTimeout" "time.Duration" "peerResolver" "matching.PeerResolver" "loadBalancer" "matching.LoadBalancer") }}
 {{ end }}
 
 import (
@@ -69,4 +67,4 @@ func (c * {{$decorator}}) {{$method.Declaration}} {
     {{- end }}
     {{$method.Pass ("c.client.") }}
 }
-{{end}}
+{{end}}

client/wrappers/timeout/timeout.go

@@ -36,6 +36,12 @@ const (
 	FrontendDefaultTimeout = 10 * time.Second
 	// FrontendDefaultLongPollTimeout is the frontend service long poll default timeout used to make calls
 	FrontendDefaultLongPollTimeout = time.Minute * 3
+	// MatchingDefaultTimeout is the default timeout used to make calls
+	MatchingDefaultTimeout = time.Minute
+	// MatchingDefaultLongPollTimeout is the long poll default timeout used to make calls
+	MatchingDefaultLongPollTimeout = time.Minute * 2
+	// HistoryDefaultTimeout is the default timeout used to make calls
+	HistoryDefaultTimeout = time.Second * 30
 )
 
 func createContext(parent context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {