Golang你一定要懂的连接池实现

type Transport struct {
  //操作空闲连接需要获取锁
  idleMu    sync.Mutex
  //空闲连接池，key为协议目标地址等组合
  idleConn   map[connectMethodKey][]*persistConn // most recently used at end
  //等待空闲连接的队列，基于切片实现，队列大小无限制
  idleConnWait map[connectMethodKey]wantConnQueue // waiting getConns
  
  //排队等待建立连接需要获取锁
  connsPerHostMu  sync.Mutex
  //每个host建立的连接数
  connsPerHost   map[connectMethodKey]int
  //等待建立连接的队列，同样基于切片实现，队列大小无限制
  connsPerHostWait map[connectMethodKey]wantConnQueue // waiting getConns
  
  //最大空闲连接数
  MaxIdleConns int
  //每个目标host最大空闲连接数；默认为2（注意默认值）
  MaxIdleConnsPerHost int
  //每个host可建立的最大连接数
  MaxConnsPerHost int
  //连接多少时间没有使用则被关闭
  IdleConnTimeout time.Duration
  
  //禁用长连接，使用短连接
  DisableKeepAlives bool
}

func (t *Transport) getConn(treq *transportRequest, cm connectMethod) (pc *persistConn, err error) {}

func (t *Transport) tryPutIdleConn(pconn *persistConn) error {}

//getConn方法内部实现

if delivered := t.queueForIdleConn(w); delivered {
  return pc, nil
}
  
t.queueForDial(w)

select {
  case <-w.ready:
    return w.pc, w.err
    
  //超时被取消
  case <-req.Cancel:
    return nil, errRequestCanceledConn
  ……
}

var errRequestCanceledConn = errors.New("net/http: request canceled while waiting for connection") // TODO: unify"htmlcode">

func (t *Transport) queueForIdleConn(w *wantConn) (delivered bool) {
  //如果配置了空闲超时时间，获取到连接需要检测，超时则关闭连接
  if t.IdleConnTimeout > 0 {
    oldTime = time.Now().Add(-t.IdleConnTimeout)
  }
  
  if list, ok := t.idleConn[w.key]; ok {
    for len(list) > 0 && !stop {
      pconn := list[len(list)-1]
      tooOld := !oldTime.IsZero() && pconn.idleAt.Round(0).Before(oldTime)
      //超时了，关闭连接
      if tooOld {
        go pconn.closeConnIfStillIdle()
      }
      
      //分发连接到wantConn
      delivered = w.tryDeliver(pconn, nil)
    }
  }
  
  //排队等待空闲连接
  q := t.idleConnWait[w.key]
  q.pushBack(w)
  t.idleConnWait[w.key] = q
}


排队等待新建连接的逻辑如下：


func (t *Transport) queueForDial(w *wantConn) {
  //如果没有限制最大连接数，直接建立连接
  if t.MaxConnsPerHost <= 0 {
    go t.dialConnFor(w)
    return
  }
  
  //如果没超过连接数限制，直接建立连接
  if n := t.connsPerHost[w.key]; n < t.MaxConnsPerHost {
    go t.dialConnFor(w)
    return
  }
  
  //排队等待连接建立
  q := t.connsPerHostWait[w.key]
  q.pushBack(w)
  t.connsPerHostWait[w.key] = q
}



连接建立完成后，同样会调用tryDeliver分发连接到wantConn，同时关闭通道w.ready，这样主协程纠接触阻塞了。


func (w *wantConn) tryDeliver(pc *persistConn, err error) bool {
  w.pc = pc
  close(w.ready)
}



请求处理完成后，通过tryPutIdleConn将连接放回连接池；这时候如果存在等待空闲连接的协程，则需要分发复用该连接。另外，在回收连接时，还需要校验空闲连接数目是否超过限制：


func (t *Transport) tryPutIdleConn(pconn *persistConn) error {
  //禁用长连接；或者最大空闲连接数不合法
  if t.DisableKeepAlives || t.MaxIdleConnsPerHost < 0 {
    return errKeepAlivesDisabled
  }
  
  if q, ok := t.idleConnWait[key]; ok {
    //如果等待队列不为空，分发连接
    for q.len() > 0 {
      w := q.popFront()
      if w.tryDeliver(pconn, nil) {
        done = true
        break
      }
    }
  }
  
  //空闲连接数目超过限制，默认为DefaultMaxIdleConnsPerHost=2
  idles := t.idleConn[key]
  if len(idles) >= t.maxIdleConnsPerHost() {
    return errTooManyIdleHost
  }

}


空闲连接超时关闭


Golang HTTP连接池如何实现空闲连接的超时关闭逻辑呢？从上述queueForIdleConn逻辑可以看到，每次在获取到空闲连接时，都会检测是否已经超时，超时则关闭连接。
那如果没有业务请求到达，一直不需要获取连接，空闲连接就不会超时关闭吗？其实在将空闲连接添加到连接池时，Golang同时还设置了定时器，定时器到期后，自然会关闭该连接。


pconn.idleTimer = time.AfterFunc(t.IdleConnTimeout, pconn.closeConnIfStillIdle)


排队队列怎么实现


怎么实现队列模型呢？很简单，可以基于切片：


queue  []*wantConn

//入队
queue = append(queue, w)

//出队
v := queue[0]
queue[0] = nil
queue = queue[1:]




这样有什么问题吗？随着频繁的入队与出队操作，切片queue的底层数组，会有大量空间无法复用而造成浪费。除非该切片执行了扩容操作。
Golang在实现队列时，使用了两个切片head和tail；head切片用于出队操作，tail切片用于入队操作；出队时，如果head切片为空，则交换head与tail。通过这种方式，Golang实现了底层数组空间的复用。


func (q *wantConnQueue) pushBack(w *wantConn) {
  q.tail = append(q.tail, w)
}

func (q *wantConnQueue) popFront() *wantConn {
  if q.headPos >= len(q.head) {
    if len(q.tail) == 0 {
      return nil
    }
    // Pick up tail as new head, clear tail.
    q.head, q.headPos, q.tail = q.tail, 0, q.head[:0]
  }
  w := q.head[q.headPos]
  q.head[q.headPos] = nil
  q.headPos++
  return w
}