You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

150 lines
2.8 KiB
Go

2 years ago
package stario
import (
"errors"
"fmt"
"io"
"os"
"runtime"
"sync"
"sync/atomic"
"time"
)
type StarBuffer struct {
io.Reader
io.Writer
io.Closer
datas []byte
pStart uint64
pEnd uint64
cap uint64
isClose atomic.Value
isEnd atomic.Value
rmu sync.Mutex
wmu sync.Mutex
}
func NewStarBuffer(cap uint64) *StarBuffer {
rtnBuffer := new(StarBuffer)
rtnBuffer.cap = cap
rtnBuffer.datas = make([]byte, cap)
rtnBuffer.isClose.Store(false)
rtnBuffer.isEnd.Store(false)
return rtnBuffer
}
func (star *StarBuffer) Free() uint64 {
return star.cap - star.Len()
}
func (star *StarBuffer) Cap() uint64 {
return star.cap
}
func (star *StarBuffer) Len() uint64 {
if star.pEnd >= star.pStart {
return star.pEnd - star.pStart
}
return star.pEnd - star.pStart + star.cap
}
func (star *StarBuffer) getByte() (byte, error) {
if star.isClose.Load().(bool) || (star.Len() == 0 && star.isEnd.Load().(bool)) {
return 0, io.EOF
}
if star.Len() == 0 {
return 0, os.ErrNotExist
}
nowPtr := star.pStart
nextPtr := star.pStart + 1
if nextPtr >= star.cap {
nextPtr = 0
}
data := star.datas[nowPtr]
ok := atomic.CompareAndSwapUint64(&star.pStart, nowPtr, nextPtr)
if !ok {
return 0, os.ErrInvalid
}
return data, nil
}
func (star *StarBuffer) putByte(data byte) error {
if star.isClose.Load().(bool) {
return io.EOF
}
nowPtr := star.pEnd
kariEnd := nowPtr + 1
if kariEnd == star.cap {
kariEnd = 0
}
if kariEnd == atomic.LoadUint64(&star.pStart) {
for {
time.Sleep(time.Microsecond)
runtime.Gosched()
if kariEnd != atomic.LoadUint64(&star.pStart) {
break
}
}
}
star.datas[nowPtr] = data
if ok := atomic.CompareAndSwapUint64(&star.pEnd, nowPtr, kariEnd); !ok {
return os.ErrInvalid
}
return nil
}
func (star *StarBuffer) Close() error {
star.isClose.Store(true)
return nil
}
func (star *StarBuffer) Read(buf []byte) (int, error) {
if star.isClose.Load().(bool) || (star.Len() == 0 && star.isEnd.Load().(bool)) {
return 0, io.EOF
}
if buf == nil {
return 0, errors.New("buffer is nil")
}
star.rmu.Lock()
defer star.rmu.Unlock()
var sum int = 0
for i := 0; i < len(buf); i++ {
data, err := star.getByte()
if err != nil {
if err == io.EOF {
return sum, err
}
if err == os.ErrNotExist {
i--
continue
}
return sum, nil
}
buf[i] = data
sum++
}
return sum, nil
}
func (star *StarBuffer) Write(bts []byte) (int, error) {
if bts == nil && !star.isEnd.Load().(bool) {
star.isEnd.Store(true)
return 0, nil
}
if bts == nil || star.isClose.Load().(bool) {
return 0, io.EOF
}
star.wmu.Lock()
defer star.wmu.Unlock()
var sum = 0
for i := 0; i < len(bts); i++ {
err := star.putByte(bts[i])
if err != nil {
fmt.Println("Write bts err:", err)
return sum, err
}
sum++
}
return sum, nil
}