143 lines
2.4 KiB
Go
143 lines
2.4 KiB
Go
package runtimex
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"sync/atomic"
|
|
)
|
|
|
|
var (
|
|
ErrStackClosed = errors.New("stack closed")
|
|
ErrStackFull = errors.New("stack full")
|
|
)
|
|
|
|
type ChanStack struct {
|
|
data chan interface{}
|
|
cap uint64
|
|
current uint64
|
|
isClose atomic.Value
|
|
}
|
|
|
|
func NewChanStack(cap uint64) *ChanStack {
|
|
rtnBuffer := new(ChanStack)
|
|
rtnBuffer.cap = cap
|
|
rtnBuffer.isClose.Store(false)
|
|
rtnBuffer.data = make(chan interface{}, cap)
|
|
return rtnBuffer
|
|
}
|
|
|
|
func (s *ChanStack) init() {
|
|
s.cap = 1024
|
|
s.data = make(chan interface{}, s.cap)
|
|
s.isClose.Store(false)
|
|
}
|
|
|
|
func (s *ChanStack) Free() uint64 {
|
|
return s.cap - atomic.LoadUint64(&s.current)
|
|
}
|
|
|
|
func (s *ChanStack) Cap() uint64 {
|
|
return s.cap
|
|
}
|
|
|
|
func (s *ChanStack) Len() uint64 {
|
|
return atomic.LoadUint64(&s.current)
|
|
}
|
|
|
|
func (s *ChanStack) Pop() (interface{}, error) {
|
|
if s.isClose.Load() == nil {
|
|
s.init()
|
|
}
|
|
if s.isClose.Load().(bool) {
|
|
return nil, io.EOF
|
|
}
|
|
data, ok := <-s.data
|
|
if !ok {
|
|
s.isClose.Store(true)
|
|
return nil, io.EOF
|
|
}
|
|
for {
|
|
current := atomic.LoadUint64(&s.current)
|
|
if current == 0 {
|
|
break
|
|
}
|
|
if atomic.CompareAndSwapUint64(&s.current, current, current-1) {
|
|
break
|
|
}
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
func (s *ChanStack) Push(data interface{}) error {
|
|
if s.isClose.Load() == nil {
|
|
s.init()
|
|
}
|
|
if s.isClose.Load().(bool) {
|
|
return io.EOF
|
|
}
|
|
if err := func() (err error) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
err = io.EOF
|
|
}
|
|
}()
|
|
s.data <- data
|
|
return nil
|
|
}(); err != nil {
|
|
return err
|
|
}
|
|
for {
|
|
current := atomic.LoadUint64(&s.current)
|
|
if atomic.CompareAndSwapUint64(&s.current, current, current+1) {
|
|
break
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *ChanStack) TryPush(data interface{}) error {
|
|
if s.isClose.Load() == nil {
|
|
s.init()
|
|
}
|
|
if s.isClose.Load().(bool) {
|
|
return io.EOF
|
|
}
|
|
if err := func() (err error) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
err = io.EOF
|
|
}
|
|
}()
|
|
select {
|
|
case s.data <- data:
|
|
return nil
|
|
default:
|
|
return ErrStackFull
|
|
}
|
|
}(); err != nil {
|
|
return err
|
|
}
|
|
for {
|
|
current := atomic.LoadUint64(&s.current)
|
|
if atomic.CompareAndSwapUint64(&s.current, current, current+1) {
|
|
break
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *ChanStack) Close() error {
|
|
if s.isClose.Load() == nil {
|
|
s.init()
|
|
}
|
|
if s.isClose.Load().(bool) {
|
|
return ErrStackClosed
|
|
}
|
|
s.isClose.Store(true)
|
|
defer func() {
|
|
recover()
|
|
}()
|
|
close(s.data)
|
|
return nil
|
|
}
|