Golang代码搜集-资源池管理

忘是亡心i 2021-09-14 13:32 415阅读 0赞

这段代码来自《Go语言实战》第7.2节，演示如何编写一个资源池，可用于数据库连接等。原代码资源管理的效果并不明显，笔者对测试用代码进行了修改，使效果更明显。下面直接上代码。

//资源池
    package pool
    
    import (
        "errors"
        "io"
        "log"
        "sync"
    )
    
    // Pool 管理一组可以安全地在多个goroutine间共享的资源
    // 被共享的资源必须实现 io.Closer接口
    type Pool struct {
        m         sync.Mutex
        resources chan io.Closer
        factory   func() (io.Closer, error)
        closed    bool
    }
    
    // 资源池已经关闭
    var ErrPoolClosed = errors.New("Pool has been closed.")
    
    // 输入资源池size太小
    var ErrPoolSizeTooSmall = errors.New("Pool size too small")
    
    // New创建一个用来管事的资源池
    func New(fn func() (io.Closer, error), size uint) (*Pool, error) {
        if size <= 0 {
            return nil, ErrPoolSizeTooSmall
        }
        return &Pool{
            factory:   fn,
            resources: make(chan io.Closer, size),
        }, nil
    }
    
    // Acquire 从池中获取一个资源
    func (p *Pool) Acquire() (io.Closer, error) {
        select {
        case r, ok := <-p.resources:
            if !ok {
                return nil, ErrPoolClosed
            }
            log.Println("Acquire:", " Shared Resource")
            return r, nil
        default:
            //没有空闲资源，提供新资源
            log.Println("Acquire:", " New Resource")
            return p.factory()
        }
    }
    func (p *Pool) Release(r io.Closer) {
        //保证本操作和close操作的安全
        p.m.Lock()
        defer p.m.Unlock()
        //如果资源池已经关闭，则关闭当前资源
        if p.closed {
            r.Close()
            return
        }
        select {
        //试图将这资源放入队列
        case p.resources <- r:
            log.Println("Release:", "In Queue")
        default:
            //如果队列已满，则关闭这个资源
            log.Println("Release:", "Closing")
            r.Close()
        }
    }
    
    //Close 资源池停止工作，并关闭所有现有的资源
    func (p *Pool) Close() {
        //保证本操作和Release操作的安全
        p.m.Lock()
        defer p.m.Unlock()
        //如果已经关闭，则什么也不做
        if p.closed {
            return
        }
        //关闭资源池
        p.closed = true
        //关闭通道,必须先关闭
        close(p.resources)
        //关闭资源
        for r := range p.resources {
            r.Close()
        }
    }

//使用示例
    package main
    
    import (
        "fmt"
        "io"
        "log"
        "math/rand"
        "pool"
        "sync"
        "sync/atomic"
        "time"
    )
    
    var (
        maxGoroutines        = 10
        pooledResources uint = 5
    )
    
    // dbConnection 模拟共享资源
    type dbConnection struct {
        Id int32
    }
    
    // 实现io.Closer接口
    func (db *dbConnection) Close() error {
        log.Println("Close:Connection", db.Id)
        return nil
    }
    
    // 用来给第个链接一个独一无二的id
    var idCounter int32
    
    func createConn() (io.Closer, error) {
        id := atomic.AddInt32(&idCounter, 1)
        log.Println("Create:New Conn", id)
        return &dbConnection{id}, nil
    }
    
    func main() {
        rand.Seed(time.Now().UnixNano())
        fmt.Println("This demo presents how to use the pool")
        //创建资源池
        p, err := pool.New(createConn, pooledResources)
        if err != nil {
            log.Fatal("Create pool failed")
        }
        //关闭资源池
        defer p.Close()
    
        var wg sync.WaitGroup
        wg.Add(maxGoroutines)
        for query := 0; query < maxGoroutines; query++ {
            go performQueries(query, p, &wg)
            //随机等一段时间，否则后而gorountine在获取资源时
            //资源池仍没有资源
            time.Sleep(time.Duration(rand.Intn(2)) * time.Second)
        }
        //等待 结束
        wg.Wait()
        fmt.Printf("Query num:%d, Connection num:%d\n", maxGoroutines, idCounter)
    }
    func performQueries(query int, p *pool.Pool, wg *sync.WaitGroup) {
        defer wg.Done()
        //从池里请求一个链接
        conn, err := p.Acquire()
        if err != nil {
            log.Fatal(err)
        }
        defer p.Release(conn)
        //模拟查询
        time.Sleep(time.Duration(rand.Intn(500)) * time.Millisecond)
        log.Printf("QID[%d] CID[%d]\n", query, conn.(*dbConnection).Id)
    }