File: byte_buffer_pool.go

package info (click to toggle)
golang-github-viant-toolbox 0.33.2-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, forky, sid, trixie
  • size: 1,280 kB
  • sloc: makefile: 16
file content (60 lines) | stat: -rw-r--r-- 1,305 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
package bridge

import (
	"io"
	"net/http/httputil"
)

type bytesBufferPool struct {
	channel    chan []byte
	bufferSize int
}

func (p *bytesBufferPool) Get() (result []byte) {
	select {
	case result = <-p.channel:
	default:
		result = make([]byte, p.bufferSize)
	}
	return result
}

func (p *bytesBufferPool) Put(b []byte) {
	select {
	case p.channel <- b:
	default: //If the pool is full, discard the buffer.
	}
}

//NewBytesBufferPool returns new httputil.BufferPool pool.
func NewBytesBufferPool(poolSize, bufferSize int) httputil.BufferPool {
	return &bytesBufferPool{
		channel:    make(chan []byte, poolSize),
		bufferSize: bufferSize,
	}
}

//CopyBuffer copies bytes from passed in source to destination with provided pool
func CopyWithBufferPool(source io.Reader, destination io.Writer, pool httputil.BufferPool) (int64, error) {
	buf := pool.Get()
	defer pool.Put(buf)
	var written int64
	for {
		bytesRead, readError := source.Read(buf)
		if bytesRead > 0 {
			bytesWritten, writeError := destination.Write(buf[:bytesRead])
			if bytesWritten > 0 {
				written += int64(bytesWritten)
			}
			if writeError != nil {
				return written, writeError
			}
			if bytesRead != bytesWritten {
				return written, io.ErrShortWrite
			}
		}
		if readError != nil {
			return written, readError
		}
	}
}