Last active
May 20, 2024 14:18
-
-
Save colm-mchugh/ee3b0b7b062a235a871b to your computer and use it in GitHub Desktop.
Using channel and sync.Pool to maintain a pool of bytes.Buffer
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"bytes" | |
"encoding/json" | |
"flag" | |
"io" | |
"log" | |
"net/http" | |
"os" | |
"strings" | |
"sync" | |
) | |
// service vars - log and buffer pool | |
var ( | |
infoLog *log.Logger | |
bufferPool BufferPool | |
) | |
// BufferPool provides an API for managing | |
// bytes.Buffer objects: | |
type BufferPool interface { | |
GetBuffer() *bytes.Buffer | |
PutBuffer(*bytes.Buffer) | |
} | |
// syncPoolBufPool is an implementation of BufferPool | |
// that uses a sync.Pool to maintain buffers: | |
type syncPoolBufPool struct { | |
pool *sync.Pool | |
makeBuffer func() interface{} | |
} | |
func NewSyncPool(buf_size int) BufferPool { | |
var newPool syncPoolBufPool | |
newPool.makeBuffer = func() interface{} { | |
var b bytes.Buffer | |
b.Grow(buf_size) | |
return &b | |
} | |
newPool.pool = &sync.Pool{} | |
newPool.pool.New = newPool.makeBuffer | |
return &newPool | |
} | |
func (bp *syncPoolBufPool) GetBuffer() (b *bytes.Buffer) { | |
pool_object := bp.pool.Get() | |
b, ok := pool_object.(*bytes.Buffer) | |
if !ok { // explicitly make buffer if sync.Pool returns nil: | |
b = bp.makeBuffer().(*bytes.Buffer) | |
} | |
return | |
} | |
func (bp *syncPoolBufPool) PutBuffer(b *bytes.Buffer) { | |
bp.pool.Put(b) | |
} | |
// chanBufferPool is an implementation of BufferPool | |
// that uses a channel to maintain buffers: | |
type chanBufferPool struct { | |
pool chan *bytes.Buffer | |
makeBuffer func() *bytes.Buffer | |
} | |
func NewChanPool(max_bufs, buf_size int) BufferPool { | |
return &chanBufferPool{ | |
pool: make(chan *bytes.Buffer, max_bufs), | |
makeBuffer: func() *bytes.Buffer { | |
var b bytes.Buffer | |
b.Grow(buf_size) | |
return &b | |
}, | |
} | |
} | |
func (bp *chanBufferPool) GetBuffer() (b *bytes.Buffer) { | |
select { | |
case b = <-bp.pool: // found buffer in pool | |
default: | |
b = bp.makeBuffer() | |
} | |
return | |
} | |
func (bp *chanBufferPool) PutBuffer(b *bytes.Buffer) { | |
select { | |
case bp.pool <- b: // put buffer back in pool | |
default: | |
// buffer is left to the mercy of Go GC | |
} | |
} | |
// Service data type: | |
type ResponseData struct { | |
Name string | |
Stuff []string | |
} | |
// Service Implementation: | |
func doService(w http.ResponseWriter, req *http.Request) { | |
// construct response: | |
resp := ResponseData{req.FormValue("input"), []string{"something", "lipsum-orum-sum`"}} | |
// get a buffer for putting the response in: | |
buf := bufferPool.GetBuffer() | |
// convert response to json: | |
js, err := json.Marshal(resp) | |
if err != nil { // give up if json conversion has an error: | |
http.Error(w, err.Error(), http.StatusInternalServerError) | |
return | |
} | |
// put json response in buffer: | |
buf.Write(js) | |
// Write http response header: | |
w.Header().Set("Content-Type", "application/json") | |
// Write http response body: | |
io.Copy(w, buf) | |
// return buffer to pool: | |
bufferPool.PutBuffer(buf) | |
} | |
// flags | |
var BUFFER_SIZE = flag.Int("buffer-size", 64, "Maximum buffer size (under this, responses are buffered, above this, responses are flushed") | |
var POOL_TYPE = flag.String("pool-type", "sync", "Buffer pool implementation. Choices are \"sync\" or \"chan\"") | |
// Service Set-Up: | |
func main() { | |
flag.Parse() | |
infoLog = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile) | |
switch strings.ToLower(*POOL_TYPE) { | |
case "sync": | |
bufferPool = NewSyncPool(*BUFFER_SIZE) | |
case "chan": | |
bufferPool = NewChanPool(32, *BUFFER_SIZE) | |
default: | |
return | |
} | |
http.HandleFunc("/service", doService) | |
http.ListenAndServe(":8080", nil) | |
// to invoke this service: | |
// curl http://localhost:8080/service?input=somestuff | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment