golang200行实现连接池

package main

import (
    "errors"
    "fmt"
    "net"
    "os"
    "sync"
    "time"
)

type (
    Config struct {
        //初始化连接数
        InitCap int
        //最大连接数
        MaxCap   int
        InitConn func() (interface{}, error)
        Close    func(interface{}) error
        Validate func(interface{}) error
        Timeout  time.Duration
    }
    Pool struct {
        mu    sync.Mutex
        conns chan *conn
        //防止一个连接被重复加入
        poolMap  map[interface{}]int
        initConn func() (interface{}, error)
        close    func(interface{}) error
        validate func(interface{}) error
        timeout  time.Duration
    }
    conn struct {
        connect    interface{}
        createTime time.Time
    }
)

var (
    ErrClosed       = errors.New("pool is closed")
    ErrInvalidCap   = errors.New("invalid cap settings")
    ErrInvalidInit  = errors.New("invalid init func settings")
    ErrInvalidClose = errors.New("invalid close func settings")
    ErrInvalidConn  = errors.New("invalid conn")
)

const addr string = "127.0.0.1:8230"

func NewPool(config *Config) (*Pool, error) {
    if config.InitCap < 0 || config.MaxCap <= 0 || config.InitCap > config.MaxCap {
        return nil, ErrInvalidCap
    }
    if config.InitConn == nil {
        return nil, ErrInvalidInit
    }
    if config.Close == nil {
        return nil, ErrInvalidClose
    }
    pool := &Pool{
        conns:    make(chan *conn, config.MaxCap),
        poolMap:  make(map[interface{}]int),
        initConn: config.InitConn,
        close:    config.Close,
        timeout:  config.Timeout,
    }
    if config.Validate != nil {
        pool.validate = config.Validate
    }
    for i := 0; i < config.InitCap; i++ {
        connect, err := pool.initConn()
        pool.poolMap[connect] = 1
        if err != nil {
            pool.Release()
            return nil, ErrInvalidInit
        }
        pool.conns <- &conn{connect: connect, createTime: time.Now()}
    }
    return pool, nil
}

func (pool *Pool) Get() (interface{}, error) {
    if pool.conns == nil {
        return nil, ErrClosed
    }
    for {
        select {
        case wrapConn := <-pool.conns:
            if wrapConn == nil {
                return nil, ErrClosed
            }
            //判断是否超时，超时则丢弃
            pool.mu.Lock()
            if timeout := pool.timeout; timeout > 0 {
                if wrapConn.createTime.Add(timeout).Before(time.Now()) {
                    delete(pool.poolMap, wrapConn.connect)
                    pool.close(wrapConn.connect)
                    pool.mu.Unlock()
                    continue
                }
            }
            if pool.validate != nil {
                if err := pool.validate(wrapConn.connect); err != nil {
                    delete(pool.poolMap, wrapConn.connect)
                    pool.close(wrapConn.connect)
                    pool.mu.Unlock()
                    continue
                }
            }
            delete(pool.poolMap, wrapConn.connect)
            pool.mu.Unlock()
            return wrapConn.connect, nil
        default:
            pool.mu.Lock()
            if pool.initConn == nil {
                pool.mu.Unlock()
                return nil, ErrInvalidInit
            }
            connect, err := pool.initConn()
            pool.mu.Unlock()
            if err != nil {
                return nil, err
            }
            return connect, nil
        }
    }
}

func (pool *Pool) Put(connect interface{}) error {
    if connect == nil {
        return ErrInvalidConn
    }
    if pool.validate != nil {
        if err := pool.validate(connect); err != nil {
            pool.close(connect)
            return ErrInvalidConn
        }
    }
    pool.mu.Lock()
    if pool.conns == nil {
        pool.mu.Unlock()
        err := pool.close(connect)
        return err
    }
    if _, ok := pool.poolMap[connect]; ok {
        pool.mu.Unlock()
        return nil
    }
    select {
    case pool.conns <- &conn{connect: connect, createTime: time.Now()}:
        pool.poolMap[connect] = 1
        pool.mu.Unlock()
        return nil
    default:
        pool.mu.Unlock()
        //连接池已满，直接关闭该连接
        return pool.close(connect)
    }
}

// Release 释放连接池中所有连接
func (pool *Pool) Release() {
    pool.mu.Lock()
    for wrapConn := range pool.conns {
        pool.close(wrapConn.connect)
    }
    close(pool.conns)
    pool.conns = nil
    pool.initConn = nil
    pool.validate = nil
    pool.close = nil
    pool.poolMap = nil
    pool.mu.Unlock()
}

func (pool *Pool) Len() int {
    return len(pool.conns)
}

//以下是测试代码
func main() {
    go server()
    //等待tcp server启动
    time.Sleep(2 * time.Second)
    client()
    fmt.Println("服务退出")
    time.Sleep(20 * time.Second)
}

func client() {
    initConn := func() (interface{}, error) { return net.Dial("tcp", addr) }
    close := func(v interface{}) error { return v.(net.Conn).Close() }

    //创建一个连接池： 初始化5，最大连接30
    poolConfig := &Config{
        InitCap:  5,
        MaxCap:   30,
        InitConn: initConn,
        Close:    close,
        Timeout:  15 * time.Second,
    }
    p, err := NewPool(poolConfig)
    if err != nil {
        fmt.Println("err=", err)
    }
    //从连接池中取得一个连接
    v, _ := p.Get()
    _, _ = p.Get()

    //do something
    conn := v.(net.Conn)
    conn.Write([]byte("guoqiang"))
    //将连接放回连接池中
    p.Put(v)
    p.Put(v)
    p.Put(v)
    //查看当前连接中的数量
    current := p.Len()
    fmt.Println("len=", current)
}

func server() {
    l, err := net.Listen("tcp", addr)
    if err != nil {
        fmt.Println("Error listening: ", err)
        os.Exit(1)
    }
    defer l.Close()
    fmt.Println("Listening on ", addr)
    for {
        conn, err := l.Accept()
        if err != nil {
            fmt.Println("Error accepting: ", err)
        }
        buffer := make([]byte, 20480)
        conn.Read(buffer)
        fmt.Printf("Received message %s -> %s message: %s\n", conn.RemoteAddr(), conn.LocalAddr(), buffer)
    }
}