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)
}