services/keepstore/bufferpool.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package main
   6
   7 import (
   8         "sync"
   9         "sync/atomic"
  10         "time"
  11
  12         log "github.com/Sirupsen/logrus"
  13 )
  14
  15 type bufferPool struct {
  16         // limiter has a "true" placeholder for each in-use buffer.
  17         limiter chan bool
  18         // allocated is the number of bytes currently allocated to buffers.
  19         allocated uint64
  20         // Pool has unused buffers.
  21         sync.Pool
  22 }
  23
  24 func newBufferPool(count int, bufSize int) *bufferPool {
  25         p := bufferPool{}
  26         p.New = func() interface{} {
  27                 atomic.AddUint64(&p.allocated, uint64(bufSize))
  28                 return make([]byte, bufSize)
  29         }
  30         p.limiter = make(chan bool, count)
  31         return &p
  32 }
  33
  34 func (p *bufferPool) Get(size int) []byte {
  35         select {
  36         case p.limiter <- true:
  37         default:
  38                 t0 := time.Now()
  39                 log.Printf("reached max buffers (%d), waiting", cap(p.limiter))
  40                 p.limiter <- true
  41                 log.Printf("waited %v for a buffer", time.Since(t0))
  42         }
  43         buf := p.Pool.Get().([]byte)
  44         if cap(buf) < size {
  45                 log.Fatalf("bufferPool Get(size=%d) but max=%d", size, cap(buf))
  46         }
  47         return buf[:size]
  48 }
  49
  50 func (p *bufferPool) Put(buf []byte) {
  51         p.Pool.Put(buf)
  52         <-p.limiter
  53 }
  54
  55 // Alloc returns the number of bytes allocated to buffers.
  56 func (p *bufferPool) Alloc() uint64 {
  57         return atomic.LoadUint64(&p.allocated)
  58 }
  59
  60 // Cap returns the maximum number of buffers allowed.
  61 func (p *bufferPool) Cap() int {
  62         return cap(p.limiter)
  63 }
  64
  65 // Len returns the number of buffers in use right now.
  66 func (p *bufferPool) Len() int {
  67         return len(p.limiter)
  68 }