hrpc: Add QueueBatch to the RegionClient interface

Removes a type assert in SendBatch.
tsuna · Jul 27, 2023 · 3acbbb4 · 3acbbb4
1 parent 4e81719
commit 3acbbb4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/hrpc/call.go b/hrpc/call.go
@@ -43,6 +43,7 @@ type RegionClient interface {
 	Close()
 	Addr() string
 	QueueRPC(Call)
+	QueueBatch(context.Context, []Call)
 	String() string
 }
 

diff --git a/rpc.go b/rpc.go
@@ -250,14 +250,9 @@ func (c *client) SendBatch(ctx context.Context, batch []hrpc.Call) (
 	)
 	wg.Add(len(rpcByClient))
 	for client, rpcs := range rpcByClient {
-		// TODO: Move this to the RegionClient interface so we don't
-		// need to type assert here
-		qb := client.(interface {
-			QueueBatch(ctx context.Context, rpcs []hrpc.Call)
-		})
 		go func(client hrpc.RegionClient, rpcs []hrpc.Call) {
 			defer wg.Done()
-			qb.QueueBatch(ctx, rpcs)
+			client.QueueBatch(ctx, rpcs)
 			ctx, sp := observability.StartSpan(ctx, "waitForResult")
 			defer sp.End()
 			ok := c.waitForCompletion(ctx, client, rpcs, res, rpcToRes)

diff --git a/test/mock/region/client.go b/test/mock/region/client.go