package main

import (
	"sync"
	"sync/atomic"
	"time"

	log "github.com/Sirupsen/logrus"
)

type bufferPool struct {
	// limiter has a "true" placeholder for each in-use buffer.
	limiter chan bool
	// allocated is the number of bytes currently allocated to buffers.
	allocated uint64
	// Pool has unused buffers.
	sync.Pool
}

func newBufferPool(count int, bufSize int) *bufferPool {
	p := bufferPool{}
	p.New = func() interface{} {
		atomic.AddUint64(&p.allocated, uint64(bufSize))
		return make([]byte, bufSize)
	}
	p.limiter = make(chan bool, count)
	return &p
}

func (p *bufferPool) Get(size int) []byte {
	select {
	case p.limiter <- true:
	default:
		t0 := time.Now()
		log.Printf("reached max buffers (%d), waiting", cap(p.limiter))
		p.limiter <- true
		log.Printf("waited %v for a buffer", time.Since(t0))
	}
	buf := p.Pool.Get().([]byte)
	if cap(buf) < size {
		log.Fatalf("bufferPool Get(size=%d) but max=%d", size, cap(buf))
	}
	return buf[:size]
}

func (p *bufferPool) Put(buf []byte) {
	p.Pool.Put(buf)
	<-p.limiter
}

// Alloc returns the number of bytes allocated to buffers.
func (p *bufferPool) Alloc() uint64 {
	return atomic.LoadUint64(&p.allocated)
}

// Cap returns the maximum number of buffers allowed.
func (p *bufferPool) Cap() int {
	return cap(p.limiter)
}

// Len returns the number of buffers in use right now.
func (p *bufferPool) Len() int {
	return len(p.limiter)
}