Golang代码搜集-资源池管理

这段代码来自《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)
}