Go database/sql 基于临时 channel 传递连接

Go 的 database/sql 包中，连接池的管理是其核心功能之一，其中通过临时创建的 channel 来传递连接的方式是一个非常巧妙的设计。这种设计保证了连接的并发安全，同时灵活地管理了连接的获取和归还。

核心逻辑分为以下几个部分：

连接的获取：

当需要一个连接时，尝试从空闲连接池中获取。
如果池中没有可用连接且当前连接数未达到最大限制，则创建新连接。
如果池满了，会通过临时创建的 channel 等待已有连接归还。

连接的归还：

当使用完连接后，将其放回空闲连接池。
如果有等待获取连接的 goroutine，会通过 channel 通知其中一个。

临时的 channel：

如果连接池没有可用连接，会为每个等待连接的请求创建一个临时 channel。
通过这个 channel，将归还的连接传递给等待的请求。
这样可以避免锁的过多竞争，并且实现了高效的等待和通知机制。

以下是 database/sql 源码中与连接获取和归还相关的简化逻辑（部分代码可能略有调整以便更易理解）：

go 复制代码

type DB struct {
 freeConn     []*driverConn            // 空闲连接池，当前所有的空闲连接
 connRequests map[int]chan connRequest // 若空闲连接池中无空闲连接，存储当前有哪些协程在申请新连接
 reqKey       int

 mu sync.Mutex // protects following fields
}

func NewDB(size int) *DB {
 freeConn := make([]*driverConn, 0, size)
 for i := 1; i <= size; i++ {
  freeConn = append(freeConn, &driverConn{})
 }
 return &DB{
  freeConn:     freeConn,
  connRequests: make(map[int]chan connRequest),
  reqKey:       0,
 }
}

type driverConn struct { // 表示真正需要的连接
}

type connRequest struct { // 表示申请一个新连接的请求
 conn *driverConn
}

func (db *DB) Conn() *driverConn {
 db.mu.Lock()
 numFree := len(db.freeConn)
 if numFree > 0 { // 当连接池中有空闲连接的时候
  conn := db.freeConn[0] // 取 slice 切片的第一个元素，然后将 slice 后面的元素往前挪，最后通过截断来"释放"最后一个元素。
  copy(db.freeConn, db.freeConn[1:])
  db.freeConn = db.freeConn[:numFree-1]
  fmt.Println("空闲连接池中获得")
  db.mu.Unlock()
  return conn
 }
 req := make(chan connRequest, 1) // 当连接池中没有空闲连接的时候
 reqKey := db.reqKey + 1
 db.connRequests[reqKey] = req
 db.mu.Unlock()
 select {
 case <-time.After(10 * time.Second):
  return nil
 case ret := <-req:
  fmt.Println("已有连接归还中获得")
  return ret.conn
 }
}

func (db *DB) PutConn(dc *driverConn) {
 db.mu.Lock()
 defer db.mu.Unlock()
 if len(db.connRequests) > 0 {
  var req chan connRequest
  var reqKey int
  for reqKey, req = range db.connRequests {
   break
  }
  delete(db.connRequests, reqKey)
  req <- connRequest{
   conn: dc,
  }
  fmt.Println("归还到新请求连接中")
  return
 }
 db.freeConn = append(db.freeConn, dc)
 fmt.Println("归还到空闲连接池")
}

func main() {
 db := NewDB(3) // 创建对象，设置空闲连接池中有3个空闲连接

 go func() {
  dc := db.Conn() // 获取连接
  time.Sleep(2 * time.Second)
  db.PutConn(dc) // 归还连接
 }()

 go db.Conn() // 获取连接
 go db.Conn() // 获取连接
 go db.Conn() // 获取连接

 time.Sleep(10 * time.Second)
}

// Output:
// 空闲连接池中获得
// 空闲连接池中获得
// 空闲连接池中获得
// 归还到新请求连接中
// 已有连接归还中获得