starmap/stack.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
}