From f5c114d546bd98c03dcd14fa4cbc1355f528f737 Mon Sep 17 00:00:00 2001
From: Anish Adhikari <aanish@uber.com>
Date: Mon, 15 Jun 2020 01:09:39 -0700
Subject: [PATCH] Fix ListTaskListPartition command

---
 service/matching/matchingEngine.go | 33 +++++++++++++++---------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/service/matching/matchingEngine.go b/service/matching/matchingEngine.go
index 526a0c52d1d..d8f00046942 100644
--- a/service/matching/matchingEngine.go
+++ b/service/matching/matchingEngine.go
@@ -593,11 +593,12 @@ func (e *matchingEngineImpl) ListTaskListPartitions(
 	if err != nil {
 		return nil, err
 	}
-	resp := workflow.ListTaskListPartitionsResponse{
+	resp := &workflow.ListTaskListPartitionsResponse{
 		ActivityTaskListPartitions: activityTaskListInfo,
 		DecisionTaskListPartitions: decisionTaskListInfo,
 	}
-	return &resp, nil
+
+	return resp, nil
 }
 
 func (e *matchingEngineImpl) listTaskListPartitions(
@@ -605,8 +606,7 @@ func (e *matchingEngineImpl) listTaskListPartitions(
 	taskListType int,
 ) ([]*workflow.TaskListPartitionMetadata, error) {
 	partitions, err := e.getAllPartitions(
-		request.GetDomain(),
-		*request.TaskList,
+		request,
 		taskListType,
 	)
 	if err != nil {
@@ -615,13 +615,12 @@ func (e *matchingEngineImpl) listTaskListPartitions(
 
 	var partitionHostInfo []*workflow.TaskListPartitionMetadata
 	for _, partition := range partitions {
-		if host, err := e.getHostInfo(partition); err != nil {
-			partitionHostInfo = append(partitionHostInfo,
-				&workflow.TaskListPartitionMetadata{
-					Key:           common.StringPtr(partition),
-					OwnerHostName: common.StringPtr(host),
-				})
-		}
+		host, _ := e.getHostInfo(partition)
+		partitionHostInfo = append(partitionHostInfo,
+			&workflow.TaskListPartitionMetadata{
+				Key:           common.StringPtr(partition),
+				OwnerHostName: common.StringPtr(host),
+			})
 	}
 	return partitionHostInfo, nil
 }
@@ -635,22 +634,22 @@ func (e *matchingEngineImpl) getHostInfo(partitionKey string) (string, error) {
 }
 
 func (e *matchingEngineImpl) getAllPartitions(
-	domain string,
-	taskList workflow.TaskList,
+	request *m.ListTaskListPartitionsRequest,
 	taskListType int,
 ) ([]string, error) {
 	var partitionKeys []string
-	domainID, err := e.domainCache.GetDomainID(domain)
+	domainID, err := e.domainCache.GetDomainID(request.GetDomain())
 	if err != nil {
 		return partitionKeys, err
 	}
-	taskListID, err := newTaskListID(domainID, taskList.GetName(), persistence.TaskListTypeDecision)
+	taskList := request.GetTaskList()
+	taskListID, err := newTaskListID(domainID, taskList.GetName(), taskListType)
 	rootPartition := taskListID.GetRoot()
 
 	partitionKeys = append(partitionKeys, rootPartition)
 
-	nWritePartitions := e.config.GetTasksBatchSize
-	n := nWritePartitions(domain, rootPartition, taskListType)
+	nWritePartitions := e.config.NumTasklistWritePartitions
+	n := nWritePartitions(request.GetDomain(), rootPartition, taskListType)
 	if n <= 0 {
 		return partitionKeys, nil
 	}