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.

200 lines
3.6 KiB
Go

package starmap
import (
"errors"
"fmt"
"io"
"os"
"runtime"
"sync"
"sync/atomic"
"time"
)
type StarStack struct {
datas []interface{}
pStart uint64
pEnd uint64
cap uint64
isClose atomic.Value
rmu sync.Mutex
wmu sync.Mutex
}
func NewStarStack(cap uint64) *StarStack {
rtnBuffer := new(StarStack)
rtnBuffer.cap = cap
rtnBuffer.datas = make([]interface{}, cap)
rtnBuffer.isClose.Store(false)
return rtnBuffer
}
func (star *StarStack) init() {
star.cap = 1024
star.datas = make([]interface{}, star.cap)
star.isClose.Store(false)
}
func (star *StarStack) Free() uint64 {
return star.cap - star.Len()
}
func (star *StarStack) Cap() uint64 {
return star.cap
}
func (star *StarStack) Len() uint64 {
if star.pEnd >= star.pStart {
return star.pEnd - star.pStart
}
return star.pEnd - star.pStart + star.cap
}
func (star *StarStack) PopNoWait() (interface{}, error) {
if star.isClose.Load() == nil {
star.init()
}
if star.isClose.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 *StarStack) MustPop() interface{} {
if star.isClose.Load() == nil {
star.init()
}
data, err := star.Pop()
if err != nil {
return nil
}
return data
}
func (star *StarStack) Pop() (interface{}, error) {
if star.isClose.Load() == nil {
star.init()
}
for {
if star.isClose.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 {
time.Sleep(time.Microsecond)
runtime.Gosched()
continue
}
return data, nil
}
}
func (star *StarStack) Push(data interface{}) error {
if star.isClose.Load() == nil {
star.init()
}
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 *StarStack) Close() error {
if star.isClose.Load() == nil {
star.init()
}
star.isClose.Store(true)
return nil
}
func (star *StarStack) Read(buf []interface{}) (int, error) {
if star.isClose.Load() == nil {
star.init()
}
if star.isClose.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.PopNoWait()
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 *StarStack) Write(bts []byte) (int, error) {
if star.isClose.Load() == nil {
star.init()
}
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.Push(bts[i])
if err != nil {
fmt.Println("Write bts err:", err)
return sum, err
}
sum++
}
return sum, nil
}