增加pool,简易客户端连接池管理

pool.go

@@ -0,0 +1,72 @@
+package tdx
+
+import (
+	"errors"
+	"github.com/injoyai/base/safe"
+)
+
+func NewPool(dial func() (*Client, error), number int) (*Pool, error) {
+	ch := make(chan *Client, number)
+	p := &Pool{
+		ch: ch,
+		Closer: safe.NewCloser().SetCloseFunc(func(err error) error {
+			close(ch)
+			return nil
+		}),
+	}
+	for i := 0; i < number; i++ {
+		c, err := dial()
+		if err != nil {
+			return nil, err
+		}
+		p.ch <- c
+	}
+	return p, nil
+}
+
+type Pool struct {
+	ch chan *Client
+	*safe.Closer
+}
+
+func (this *Pool) Get() (*Client, error) {
+	select {
+	case <-this.Done():
+		return nil, this.Err()
+	case c, ok := <-this.ch:
+		if !ok {
+			return nil, errors.New("已关闭")
+		}
+		return c, nil
+	}
+}
+
+func (this *Pool) Put(c *Client) {
+	select {
+	case <-this.Done():
+		c.Close()
+		return
+	case this.ch <- c:
+	}
+}
+
+func (this *Pool) Do(fn func(c *Client) error) error {
+	c, err := this.Get()
+	if err != nil {
+		return err
+	}
+	defer this.Put(c)
+	return fn(c)
+}
+
+func (this *Pool) Go(fn func(c *Client)) error {
+	c, err := this.Get()
+	if err != nil {
+		return err
+	}
+	go func(c *Client) {
+		defer this.Put(c)
+		fn(c)
+	}(c)
+	return nil
+}