Add TaskListPartitionConfig message to proto

Shaddoll · Oct 14, 2024 · f13e95d · f13e95d
1 parent a32159c
commit f13e95d
Show file tree

Hide file tree

Showing 38 changed files with 1,126 additions and 481 deletions.
diff --git a/.gen/proto/matching/v1/service.pb.go b/.gen/proto/matching/v1/service.pb.go
diff --git a/.gen/proto/matching/v1/service.pb.yarpc.go b/.gen/proto/matching/v1/service.pb.yarpc.go
diff --git a/client/matching/client.go b/client/matching/client.go
@@ -55,7 +55,7 @@ func (c *clientImpl) AddActivityTask(
 	ctx context.Context,
 	request *types.AddActivityTaskRequest,
 	opts ...yarpc.CallOption,
-) error {
+) (*types.AddActivityTaskResponse, error) {
 	partition := c.loadBalancer.PickWritePartition(
 		request.GetDomainUUID(),
 		*request.GetTaskList(),
@@ -65,7 +65,7 @@ func (c *clientImpl) AddActivityTask(
 	request.TaskList.Name = partition
 	peer, err := c.peerResolver.FromTaskList(request.TaskList.GetName())
 	if err != nil {
-		return err
+		return nil, err
 	}
 	return c.client.AddActivityTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }
@@ -74,7 +74,7 @@ func (c *clientImpl) AddDecisionTask(
 	ctx context.Context,
 	request *types.AddDecisionTaskRequest,
 	opts ...yarpc.CallOption,
-) error {
+) (*types.AddDecisionTaskResponse, error) {
 	partition := c.loadBalancer.PickWritePartition(
 		request.GetDomainUUID(),
 		*request.GetTaskList(),
@@ -84,7 +84,7 @@ func (c *clientImpl) AddDecisionTask(
 	request.TaskList.Name = partition
 	peer, err := c.peerResolver.FromTaskList(request.TaskList.GetName())
 	if err != nil {
-		return err
+		return nil, err
 	}
 	return c.client.AddDecisionTask(ctx, request, append(opts, yarpc.WithShardKey(peer))...)
 }

diff --git a/client/matching/client_test.go b/client/matching/client_test.go
@@ -58,74 +58,6 @@ func TestClient_withoutResponse(t *testing.T) {
 		mock      func(*MockPeerResolver, *MockLoadBalancer, *MockClient)
 		wantError bool
 	}{
-		{
-			name: "AddActivityTask",
-			op: func(c Client) error {
-				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
-				c.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(nil)
-			},
-		},
-		{
-			name: "AddActivityTask - Error in resolving peer",
-			op: func(c Client) error {
-				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", assert.AnError)
-			},
-			wantError: true,
-		},
-		{
-			name: "AddActivityTask - Error while adding activity task",
-			op: func(c Client) error {
-				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
-				c.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(assert.AnError)
-			},
-			wantError: true,
-		},
-		{
-			name: "AddDecisionTask",
-			op: func(c Client) error {
-				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
-				c.EXPECT().AddDecisionTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(nil)
-			},
-		},
-		{
-			name: "AddDecisionTask - Error in resolving peer",
-			op: func(c Client) error {
-				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", assert.AnError)
-			},
-			wantError: true,
-		},
-		{
-			name: "AddDecisionTask - Error while adding decision task",
-			op: func(c Client) error {
-				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
-			},
-			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
-				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
-				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
-				c.EXPECT().AddDecisionTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(assert.AnError)
-			},
-			wantError: true,
-		},
 		{
 			name: "RespondQueryTaskCompleted",
 			op: func(c Client) error {
@@ -220,6 +152,76 @@ func TestClient_withResponse(t *testing.T) {
 		want      any
 		wantError bool
 	}{
+		{
+			name: "AddActivityTask",
+			op: func(c Client) (any, error) {
+				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
+				c.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(&types.AddActivityTaskResponse{}, nil)
+			},
+			want: &types.AddActivityTaskResponse{},
+		},
+		{
+			name: "AddActivityTask - Error in resolving peer",
+			op: func(c Client) (any, error) {
+				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", assert.AnError)
+			},
+			wantError: true,
+		},
+		{
+			name: "AddActivityTask - Error while adding activity task",
+			op: func(c Client) (any, error) {
+				return c.AddActivityTask(context.Background(), testAddActivityTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeActivity, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
+				c.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(nil, assert.AnError)
+			},
+			wantError: true,
+		},
+		{
+			name: "AddDecisionTask",
+			op: func(c Client) (any, error) {
+				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
+				c.EXPECT().AddDecisionTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(&types.AddDecisionTaskResponse{}, nil)
+			},
+			want: &types.AddDecisionTaskResponse{},
+		},
+		{
+			name: "AddDecisionTask - Error in resolving peer",
+			op: func(c Client) (any, error) {
+				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", assert.AnError)
+			},
+			wantError: true,
+		},
+		{
+			name: "AddDecisionTask - Error while adding decision task",
+			op: func(c Client) (any, error) {
+				return c.AddDecisionTask(context.Background(), testAddDecisionTaskRequest())
+			},
+			mock: func(p *MockPeerResolver, balancer *MockLoadBalancer, c *MockClient) {
+				balancer.EXPECT().PickWritePartition(_testDomainUUID, types.TaskList{Name: _testTaskList}, persistence.TaskListTypeDecision, "").Return(_testPartition)
+				p.EXPECT().FromTaskList(_testPartition).Return("peer0", nil)
+				c.EXPECT().AddDecisionTask(gomock.Any(), gomock.Any(), []yarpc.CallOption{yarpc.WithShardKey("peer0")}).Return(nil, assert.AnError)
+			},
+			wantError: true,
+		},
 		{
 			name: "PollForActivityTask",
 			op: func(c Client) (any, error) {

diff --git a/client/matching/interface.go b/client/matching/interface.go
@@ -38,8 +38,8 @@ import (
 
 // Client is the interface exposed by types service client
 type Client interface {
-	AddActivityTask(context.Context, *types.AddActivityTaskRequest, ...yarpc.CallOption) error
-	AddDecisionTask(context.Context, *types.AddDecisionTaskRequest, ...yarpc.CallOption) error
+	AddActivityTask(context.Context, *types.AddActivityTaskRequest, ...yarpc.CallOption) (*types.AddActivityTaskResponse, error)
+	AddDecisionTask(context.Context, *types.AddDecisionTaskRequest, ...yarpc.CallOption) (*types.AddDecisionTaskResponse, error)
 	CancelOutstandingPoll(context.Context, *types.CancelOutstandingPollRequest, ...yarpc.CallOption) error
 	DescribeTaskList(context.Context, *types.MatchingDescribeTaskListRequest, ...yarpc.CallOption) (*types.DescribeTaskListResponse, error)
 	ListTaskListPartitions(context.Context, *types.MatchingListTaskListPartitionsRequest, ...yarpc.CallOption) (*types.ListTaskListPartitionsResponse, error)

diff --git a/client/matching/interface_mock.go b/client/matching/interface_mock.go
diff --git a/client/templates/thrift.tmpl b/client/templates/thrift.tmpl
@@ -18,6 +18,12 @@ import (
 func (g {{$decorator}}) {{$method.Declaration}} {
 	{{- if eq $method.Name "CountDLQMessages"}}
 		return nil, thrift.ToError(&types.BadRequestError{Message: "Feature not supported on TChannel"})
+	{{- else if or (eq $method.Name "AddDecisionTask") (eq $method.Name "AddActivityTask")}}
+		{{(index $method.Results 1).Name}} = g.c.{{$method.Name}}({{(index $method.Params 0).Name}}, thrift.From{{$prefix}}{{$Request}}({{(index $method.Params 1).Name}}), {{(index $method.Params 2).Pass}})
+		if {{(index $method.Results 1).Name}} != nil {
+			return nil, {{(index $method.Results 1).Name}}
+		}
+		return &types.{{$method.Name}}Response{}, nil
 	{{- else}}
 	{{- if eq (len $method.Params) 2}}
 	{{- if eq (len $method.Results) 1}}

diff --git a/client/wrappers/errorinjectors/matching_generated.go b/client/wrappers/errorinjectors/matching_generated.go
diff --git a/client/wrappers/grpc/matching_generated.go b/client/wrappers/grpc/matching_generated.go
diff --git a/client/wrappers/metered/matching_generated.go b/client/wrappers/metered/matching_generated.go
diff --git a/client/wrappers/metered/metered_test.go b/client/wrappers/metered/metered_test.go
@@ -88,13 +88,13 @@ func TestMatching(t *testing.T) {
 		metricsClient := metrics.NewClient(testScope, metrics.ServiceIdx(0))
 
 		clientMock.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), gomock.Any()).
-			Return(nil).Times(1)
+			Return(&types.AddActivityTaskResponse{}, nil).Times(1)
 
 		retryableClient := NewMatchingClient(
 			clientMock,
 			metricsClient)
 
-		err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
+		_, err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
 			ForwardedFrom: "test",
 			TaskList:      &types.TaskList{Name: "test"},
 		})
@@ -132,13 +132,13 @@ func TestMatching(t *testing.T) {
 		metricsClient := metrics.NewClient(testScope, metrics.ServiceIdx(0))
 
 		clientMock.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), gomock.Any()).
-			Return(nil).Times(1)
+			Return(&types.AddActivityTaskResponse{}, nil).Times(1)
 
 		retryableClient := NewMatchingClient(
 			clientMock,
 			metricsClient)
 
-		err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
+		_, err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
 			ForwardedFrom: "",
 			TaskList:      &types.TaskList{Name: "test"},
 		})
@@ -153,13 +153,13 @@ func TestMatching(t *testing.T) {
 		metricsClient := metrics.NewClient(testScope, metrics.ServiceIdx(0))
 
 		clientMock.EXPECT().AddActivityTask(gomock.Any(), gomock.Any(), gomock.Any()).
-			Return(nil).Times(1)
+			Return(&types.AddActivityTaskResponse{}, nil).Times(1)
 
 		retryableClient := NewMatchingClient(
 			clientMock,
 			metricsClient)
 
-		err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
+		_, err := retryableClient.AddActivityTask(context.Background(), &types.AddActivityTaskRequest{
 			ForwardedFrom: "",
 			TaskList:      &types.TaskList{Name: common.ReservedTaskListPrefix + "test"},
 		})