Golang模拟实现连接池

package pool

import (
	"errors"
	"io"
	"log"
	"sync"
)

// 本包是pool包, 用于展示如何使用有缓冲的通道实现资源池, 来管理可以在任意数量的goroutine之间共享
// 及独立使用的资源。这种模式在需要共享一组静态资源的情况（如共享数据库连接或者内存缓冲区）
// 下非常有用。如果goroutine需要从池里得到这些资源中的一个。
// 它可以从池里申请，使用完后归还到资源池里

// Pool管理一组可以安全的在多个goroutine间
// 共享的资源。被管理的资源必须实现io.Closer接口
type Pool struct {
	m 				sync.Mutex
	resources 		chan io.Closer
	factory 		func() (io.Closer, error)
	closed 			bool
}

// ErrPoolClosed表示请求(Acquire)了一个已经关闭的池
var ErrPoolClosed = errors.New("Pool has been closed.")

// New创建一个用来管理资源的池
// 这个池需要一个可以分配新资源的函数,
// 并规定池的大小
func New(fn func() (io.Closer, error), size uint) (*Pool, error) {
	if size <= 0 {
		return nil, errors.New("Size value tool small.")
	}
	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:
			log.Println("Acquire:", "Shared Resource")
			if !ok {
				return nil, ErrPoolClosed
			}
			return r, nil
	    // 因为没有空闲资源, 所以提供一个新资源
		default:
			log.Println("Acquire:", "New Resource")
			return p.factory()
	}
}

// Release将一个使用后的资源放回池里
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()

	// 如果pool已经被关闭, 则什么都不做
	if p.closed {
		return
	}

	// 关闭池子
	p.closed = true

	// 清空通道的资源之前，要将通道关闭
	// 如果不这样做，会发生死锁
	close(p.resources)

	// 关闭资源
	for r := range p.resources {
		r.Close()
	}
}

package main

import (
	"GoPratice/pool"
	"io"
	"log"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

// 这个程序展示如何使用pool包
// 来共享一组模拟的数据库连接

const (
	maxGoroutines = 25 // 要使用的goroutine的数量
	pooledResources = 2 // 连接池中的资源的数量
)

// dbConnection 模拟要共享的资源
type dbConnection struct {
	ID int32
}

// dbConnection对象要实现io.Closer接口
func (dbConn *dbConnection) Close() error {
	log.Println("Close: Connection", dbConn.ID)
	return nil
}

// idCounter用来给每一个连接分配一个独一无二的id
var idCounter int32

// createConnection是一个工厂函数
// 当需要一个新连接的时候, 连接池调用这个函数
func createConnection() (io.Closer, error) {
	id := atomic.AddInt32(&idCounter, 1)
	log.Println("Create: New Connection", id)

	return &dbConnection{id}, nil
}

func init() {
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
}

func main() {
	var wg sync.WaitGroup
	wg.Add(maxGoroutines)

	// 创建用来管理连接的连接池
	p, err := pool.New(createConnection, pooledResources)
	if err != nil {
		log.Println(err)
		return
	}

	// 使用池里的连接来完成查询
	for query := 0; query < maxGoroutines; query++ {
		// 每个goroutine需要自己复制一份
		// 要查询值的副本, 不然所有的查询会共享同一个查询变量
		go func(q int) {
			performQueries(q, p)
			wg.Done()
		}(query)
	}

	wg.Wait()

	// 关闭池
	log.Println("shutdown Program.")
	p.Close()
}

func performQueries(query int, p *pool.Pool) {
	// 从池里获取一个连接
	conn, err := p.Acquire()
	if err != nil {
		log.Println(err)
		return
	}
	// 延迟归还连接
	defer p.Release(conn)

	// 用等待模拟查询响应
	time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
	log.Printf("QID[%d] CID[%d]", query, conn.(*dbConnection).ID)
}