From 09697b7679e0d934d51ef1ee55192e236c2c02f1 Mon Sep 17 00:00:00 2001
From: Dreamacro <8615343+Dreamacro@users.noreply.github.com>
Date: Fri, 23 Jul 2021 00:30:23 +0800
Subject: [PATCH] Chore: adjust batch

---
 adapter/provider/healthcheck.go |  2 +-
 common/batch/batch.go           | 12 +++---------
 common/batch/batch_test.go      |  7 ++++---
 3 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/adapter/provider/healthcheck.go b/adapter/provider/healthcheck.go
index 4d2711a5..f41a4788 100644
--- a/adapter/provider/healthcheck.go
+++ b/adapter/provider/healthcheck.go
@@ -59,7 +59,7 @@ func (hc *HealthCheck) touch() {
 }
 
 func (hc *HealthCheck) check() {
-	b := batch.New(batch.WithConcurrencyNum(10))
+	b, _ := batch.New(context.Background(), batch.WithConcurrencyNum(10))
 	for _, proxy := range hc.proxies {
 		p := proxy
 		b.Go(p.Name(), func() (interface{}, error) {
diff --git a/common/batch/batch.go b/common/batch/batch.go
index 7fbca421..ce20ed61 100644
--- a/common/batch/batch.go
+++ b/common/batch/batch.go
@@ -89,7 +89,9 @@ func (b *Batch) Result() map[string]Result {
 	return copy
 }
 
-func New(opts ...Option) *Batch {
+func New(ctx context.Context, opts ...Option) (*Batch, context.Context) {
+	ctx, cancel := context.WithCancel(ctx)
+
 	b := &Batch{
 		result: map[string]Result{},
 	}
@@ -98,14 +100,6 @@ func New(opts ...Option) *Batch {
 		o(b)
 	}
 
-	return b
-}
-
-func WithContext(ctx context.Context, opts ...Option) (*Batch, context.Context) {
-	ctx, cancel := context.WithCancel(ctx)
-
-	b := New(opts...)
 	b.cancel = cancel
-
 	return b, ctx
 }
diff --git a/common/batch/batch_test.go b/common/batch/batch_test.go
index 4fcdbe81..9465dbad 100644
--- a/common/batch/batch_test.go
+++ b/common/batch/batch_test.go
@@ -11,7 +11,7 @@ import (
 )
 
 func TestBatch(t *testing.T) {
-	b := New()
+	b, _ := New(context.Background())
 
 	now := time.Now()
 	b.Go("foo", func() (interface{}, error) {
@@ -37,7 +37,8 @@ func TestBatch(t *testing.T) {
 }
 
 func TestBatchWithConcurrencyNum(t *testing.T) {
-	b := New(
+	b, _ := New(
+		context.Background(),
 		WithConcurrencyNum(3),
 	)
 
@@ -61,7 +62,7 @@ func TestBatchWithConcurrencyNum(t *testing.T) {
 }
 
 func TestBatchContext(t *testing.T) {
-	b, ctx := WithContext(context.Background())
+	b, ctx := New(context.Background())
 
 	b.Go("error", func() (interface{}, error) {
 		time.Sleep(time.Millisecond * 100)