Compare commits
15 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
348ff7418b
|
|||
|
c8facb5a03
|
|||
|
3add9183b3
|
|||
| 3c83d1d39f | |||
| d528a583f4 | |||
| a986fafa09 | |||
| 8e09b7eec0 | |||
| 609bdf7079 | |||
|
95503ec65a
|
|||
|
538445a201
|
|||
|
beaae47d8c
|
|||
|
51c0d54943
|
|||
|
14203b8fa3
|
|||
| 95e2899e8a | |||
| b568f51676 |
@@ -0,0 +1,7 @@
|
|||||||
|
.sentrux/
|
||||||
|
agent_readme.md
|
||||||
|
target.md
|
||||||
|
.gocache
|
||||||
|
.gocache/
|
||||||
|
.tmp_*/
|
||||||
|
.idea
|
||||||
@@ -2,127 +2,204 @@ package stario
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
|
||||||
"io"
|
"io"
|
||||||
"sync"
|
"sync"
|
||||||
"time"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var ErrStarBufferInvalidCapacity = errors.New("star buffer capacity must be greater than zero")
|
||||||
|
var ErrStarBufferClosed = errors.New("star buffer closed")
|
||||||
|
var ErrStarBufferWriteClosed = errors.New("star buffer write closed")
|
||||||
|
|
||||||
|
// StarBuffer is a blocking ring buffer that implements stream-style reads and writes.
|
||||||
|
//
|
||||||
|
// Close marks the write side finished after all payload bytes are sent.
|
||||||
|
// Abort aborts both sides immediately but still allows buffered bytes to be drained.
|
||||||
type StarBuffer struct {
|
type StarBuffer struct {
|
||||||
io.Reader
|
|
||||||
io.Writer
|
|
||||||
io.Closer
|
|
||||||
datas []byte
|
datas []byte
|
||||||
pStart int
|
pStart uint64
|
||||||
pEnd int
|
pEnd uint64
|
||||||
cap int
|
size uint64
|
||||||
|
cap uint64
|
||||||
isClose bool
|
isClose bool
|
||||||
isEnd bool
|
isWriteEnd bool
|
||||||
rmu sync.Mutex
|
mu sync.Mutex
|
||||||
wmu sync.Mutex
|
notEmpty *sync.Cond
|
||||||
|
notFull *sync.Cond
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewStarBuffer(cap int) *StarBuffer {
|
func NewStarBuffer(cap uint64) (*StarBuffer, error) {
|
||||||
rtnBuffer := new(StarBuffer)
|
if cap == 0 {
|
||||||
rtnBuffer.cap = cap
|
return nil, ErrStarBufferInvalidCapacity
|
||||||
rtnBuffer.datas = make([]byte, cap)
|
}
|
||||||
return rtnBuffer
|
rtnBuffer := &StarBuffer{
|
||||||
|
cap: cap,
|
||||||
|
datas: make([]byte, cap),
|
||||||
|
}
|
||||||
|
rtnBuffer.notEmpty = sync.NewCond(&rtnBuffer.mu)
|
||||||
|
rtnBuffer.notFull = sync.NewCond(&rtnBuffer.mu)
|
||||||
|
return rtnBuffer, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) Free() int {
|
// Free returns the remaining writable capacity.
|
||||||
return star.cap - star.Len()
|
func (star *StarBuffer) Free() uint64 {
|
||||||
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
|
return star.cap - star.size
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) Cap() int {
|
// Cap returns the fixed buffer capacity.
|
||||||
|
func (star *StarBuffer) Cap() uint64 {
|
||||||
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
return star.cap
|
return star.cap
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) Len() int {
|
// Len returns the currently buffered byte count.
|
||||||
length := star.pEnd - star.pStart
|
func (star *StarBuffer) Len() uint64 {
|
||||||
if length < 0 {
|
star.mu.Lock()
|
||||||
return star.cap + length - 1
|
defer star.mu.Unlock()
|
||||||
}
|
return star.size
|
||||||
return length
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) getByte() (byte, error) {
|
func (star *StarBuffer) getByte() (byte, error) {
|
||||||
if star.isClose || (star.isEnd && star.Len() == 0) {
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
|
for star.size == 0 && !star.isWriteEnd && !star.isClose {
|
||||||
|
star.notEmpty.Wait()
|
||||||
|
}
|
||||||
|
if star.size == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
}
|
}
|
||||||
if star.Len() == 0 {
|
|
||||||
return 0, errors.New("no byte available now")
|
|
||||||
}
|
|
||||||
data := star.datas[star.pStart]
|
data := star.datas[star.pStart]
|
||||||
star.pStart++
|
star.pStart = (star.pStart + 1) % star.cap
|
||||||
if star.pStart == star.cap {
|
star.size--
|
||||||
star.pStart = 0
|
star.notFull.Broadcast()
|
||||||
}
|
|
||||||
return data, nil
|
return data, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) putByte(data byte) error {
|
func (star *StarBuffer) putByte(data byte) error {
|
||||||
if star.isClose || star.isEnd {
|
star.mu.Lock()
|
||||||
return io.EOF
|
defer star.mu.Unlock()
|
||||||
}
|
for star.size == star.cap && !star.isClose && !star.isWriteEnd {
|
||||||
kariEnd := star.pEnd + 1
|
star.notFull.Wait()
|
||||||
if kariEnd == star.cap {
|
|
||||||
kariEnd = 0
|
|
||||||
}
|
|
||||||
if kariEnd == star.pStart {
|
|
||||||
for {
|
|
||||||
time.Sleep(time.Microsecond)
|
|
||||||
if kariEnd != star.pStart {
|
|
||||||
break
|
|
||||||
}
|
}
|
||||||
|
if star.isClose {
|
||||||
|
return ErrStarBufferClosed
|
||||||
}
|
}
|
||||||
|
if star.isWriteEnd {
|
||||||
|
return ErrStarBufferWriteClosed
|
||||||
}
|
}
|
||||||
star.datas[star.pEnd] = data
|
star.datas[star.pEnd] = data
|
||||||
star.pEnd = kariEnd
|
star.pEnd = (star.pEnd + 1) % star.cap
|
||||||
|
star.size++
|
||||||
|
star.notEmpty.Broadcast()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Close closes only the write side and satisfies the usual io.Closer-style
|
||||||
|
// "producer finished" semantics.
|
||||||
|
//
|
||||||
|
// Buffered bytes remain readable until drained; afterwards reads return io.EOF.
|
||||||
func (star *StarBuffer) Close() error {
|
func (star *StarBuffer) Close() error {
|
||||||
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
|
return star.closeWriteLocked()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Abort aborts the buffer and wakes blocked readers/writers immediately.
|
||||||
|
//
|
||||||
|
// Buffered bytes remain readable until drained; subsequent writes fail with
|
||||||
|
// ErrStarBufferClosed.
|
||||||
|
func (star *StarBuffer) Abort() error {
|
||||||
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
star.isClose = true
|
star.isClose = true
|
||||||
|
star.isWriteEnd = true
|
||||||
|
star.notEmpty.Broadcast()
|
||||||
|
star.notFull.Broadcast()
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
func (star *StarBuffer) Read(buf []byte) (int, error) {
|
|
||||||
|
func (star *StarBuffer) closeWriteLocked() error {
|
||||||
if star.isClose {
|
if star.isClose {
|
||||||
|
return ErrStarBufferClosed
|
||||||
|
}
|
||||||
|
star.isWriteEnd = true
|
||||||
|
star.notEmpty.Broadcast()
|
||||||
|
star.notFull.Broadcast()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (star *StarBuffer) Read(buf []byte) (int, error) {
|
||||||
|
if len(buf) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
star.mu.Lock()
|
||||||
|
defer star.mu.Unlock()
|
||||||
|
for star.size == 0 && !star.isWriteEnd && !star.isClose {
|
||||||
|
star.notEmpty.Wait()
|
||||||
|
}
|
||||||
|
if star.size == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
}
|
}
|
||||||
if buf == nil {
|
sum := minInt(len(buf), int(star.size))
|
||||||
return 0, errors.New("buffer is nil")
|
first := minInt(sum, int(star.cap-star.pStart))
|
||||||
}
|
copy(buf, star.datas[star.pStart:star.pStart+uint64(first)])
|
||||||
star.rmu.Lock()
|
second := sum - first
|
||||||
defer star.rmu.Unlock()
|
if second > 0 {
|
||||||
var sum int = 0
|
copy(buf[first:], star.datas[:second])
|
||||||
for i := 0; i < len(buf); i++ {
|
|
||||||
data, err := star.getByte()
|
|
||||||
if err != nil {
|
|
||||||
if err == io.EOF {
|
|
||||||
return sum, err
|
|
||||||
}
|
|
||||||
return sum, nil
|
|
||||||
}
|
|
||||||
buf[i] = data
|
|
||||||
sum++
|
|
||||||
}
|
}
|
||||||
|
star.pStart = (star.pStart + uint64(sum)) % star.cap
|
||||||
|
star.size -= uint64(sum)
|
||||||
|
star.notFull.Broadcast()
|
||||||
return sum, nil
|
return sum, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (star *StarBuffer) Write(bts []byte) (int, error) {
|
func (star *StarBuffer) Write(bts []byte) (int, error) {
|
||||||
if bts == nil || star.isClose {
|
if len(bts) == 0 {
|
||||||
star.isEnd = true
|
return 0, nil
|
||||||
return 0, io.EOF
|
|
||||||
}
|
}
|
||||||
star.wmu.Lock()
|
star.mu.Lock()
|
||||||
defer star.wmu.Unlock()
|
defer star.mu.Unlock()
|
||||||
var sum = 0
|
sum := 0
|
||||||
for i := 0; i < len(bts); i++ {
|
for sum < len(bts) {
|
||||||
err := star.putByte(bts[i])
|
for star.size == star.cap && !star.isClose && !star.isWriteEnd {
|
||||||
if err != nil {
|
star.notFull.Wait()
|
||||||
fmt.Println("Write bts err:", err)
|
|
||||||
return sum, err
|
|
||||||
}
|
}
|
||||||
sum++
|
if star.isClose {
|
||||||
|
if sum == 0 {
|
||||||
|
return 0, ErrStarBufferClosed
|
||||||
|
}
|
||||||
|
return sum, ErrStarBufferClosed
|
||||||
|
}
|
||||||
|
if star.isWriteEnd {
|
||||||
|
if sum == 0 {
|
||||||
|
return 0, ErrStarBufferWriteClosed
|
||||||
|
}
|
||||||
|
return sum, ErrStarBufferWriteClosed
|
||||||
|
}
|
||||||
|
space := int(star.cap - star.size)
|
||||||
|
if space == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
n := minInt(len(bts)-sum, space)
|
||||||
|
first := minInt(n, int(star.cap-star.pEnd))
|
||||||
|
copy(star.datas[star.pEnd:star.pEnd+uint64(first)], bts[sum:sum+first])
|
||||||
|
second := n - first
|
||||||
|
if second > 0 {
|
||||||
|
copy(star.datas[:second], bts[sum+first:sum+n])
|
||||||
|
}
|
||||||
|
star.pEnd = (star.pEnd + uint64(n)) % star.cap
|
||||||
|
star.size += uint64(n)
|
||||||
|
sum += n
|
||||||
|
star.notEmpty.Broadcast()
|
||||||
}
|
}
|
||||||
return sum, nil
|
return sum, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func minInt(a int, b int) int {
|
||||||
|
if a < b {
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,44 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func BenchmarkStarBufferByteRoundTrip(b *testing.B) {
|
||||||
|
buf, err := NewStarBuffer(4096)
|
||||||
|
if err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(1)
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if err := buf.putByte('a'); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := buf.getByte(); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkStarBufferChunkRoundTrip(b *testing.B) {
|
||||||
|
buf, err := NewStarBuffer(8192)
|
||||||
|
if err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
payload := []byte("hello world b612 hello world b612 b612 b612 b612 b612 b612")
|
||||||
|
scratch := make([]byte, len(payload))
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(payload)))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if _, err := buf.Write(payload); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := io.ReadFull(buf, scratch); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
+102
@@ -0,0 +1,102 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewStarBufferRejectsZeroCapacity(t *testing.T) {
|
||||||
|
buf, err := NewStarBuffer(0)
|
||||||
|
if err != ErrStarBufferInvalidCapacity {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if buf != nil {
|
||||||
|
t.Fatal("expected nil buffer when capacity is invalid")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarBufferCloseDrainsThenEOF(t *testing.T) {
|
||||||
|
buf, err := NewStarBuffer(4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := buf.Write([]byte("abcd")); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := buf.Close(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
got := make([]byte, 4)
|
||||||
|
n, err := buf.Read(got)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n != 4 || !bytes.Equal(got[:n], []byte("abcd")) {
|
||||||
|
t.Fatalf("unexpected payload: n=%d data=%q", n, got[:n])
|
||||||
|
}
|
||||||
|
|
||||||
|
n, err = buf.Read(got)
|
||||||
|
if n != 0 || err != io.EOF {
|
||||||
|
t.Fatalf("expected EOF after draining buffer, got n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := buf.Write([]byte("x")); err != ErrStarBufferWriteClosed {
|
||||||
|
t.Fatalf("unexpected write error after Close: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarBufferAbortAllowsDrain(t *testing.T) {
|
||||||
|
buf, err := NewStarBuffer(4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := buf.Write([]byte("ab")); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := buf.Abort(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
got := make([]byte, 2)
|
||||||
|
n, err := buf.Read(got)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n != 2 || !bytes.Equal(got[:n], []byte("ab")) {
|
||||||
|
t.Fatalf("unexpected payload after abort: n=%d data=%q", n, got[:n])
|
||||||
|
}
|
||||||
|
|
||||||
|
n, err = buf.Read(got)
|
||||||
|
if n != 0 || err != io.EOF {
|
||||||
|
t.Fatalf("expected EOF after draining aborted buffer, got n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := buf.Write([]byte("x")); err != ErrStarBufferClosed {
|
||||||
|
t.Fatalf("unexpected write error after Abort: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarBufferNilReadIsNoOp(t *testing.T) {
|
||||||
|
buf, err := NewStarBuffer(4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n, err := buf.Read(nil); n != 0 || err != nil {
|
||||||
|
t.Fatalf("expected nil read to be a no-op, got n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarBufferNilWriteIsNoOp(t *testing.T) {
|
||||||
|
buf, err := NewStarBuffer(4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n, err := buf.Write(nil); n != 0 || err != nil {
|
||||||
|
t.Fatalf("expected nil write to be a no-op, got n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
if _, err := buf.Write([]byte("ab")); err != nil {
|
||||||
|
t.Fatalf("nil write must not end the write side, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,111 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ERR_TIMEOUT is the legacy timeout sentinel used by WaitUntilTimeout*.
|
||||||
|
var ERR_TIMEOUT = errors.New("TIME OUT")
|
||||||
|
|
||||||
|
// WaitUntilContext runs fn and returns either its result or the context error,
|
||||||
|
// whichever happens first.
|
||||||
|
func WaitUntilContext(ctx context.Context, fn func(context.Context) error) error {
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
finished := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
finished <- fn(ctx)
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case err := <-finished:
|
||||||
|
return err
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilContextFinished is the asynchronous form of WaitUntilContext.
|
||||||
|
func WaitUntilContextFinished(ctx context.Context, fn func(context.Context) error) <-chan error {
|
||||||
|
result := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
result <- WaitUntilContext(ctx, fn)
|
||||||
|
close(result)
|
||||||
|
}()
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilContextDone adapts a done-channel worker to a context-based wait.
|
||||||
|
func WaitUntilContextDone(ctx context.Context, fn func(<-chan struct{}) error) error {
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
return WaitUntilContext(ctx, func(context.Context) error {
|
||||||
|
return fn(ctx.Done())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilContextDoneFinished is the asynchronous form of WaitUntilContextDone.
|
||||||
|
func WaitUntilContextDoneFinished(ctx context.Context, fn func(<-chan struct{}) error) <-chan error {
|
||||||
|
return WaitUntilContextFinished(ctx, func(ctx context.Context) error {
|
||||||
|
return fn(ctx.Done())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilTimeout is a legacy timeout helper kept for compatibility.
|
||||||
|
//
|
||||||
|
// The provided stop channel must be treated as receive-only by callers. New
|
||||||
|
// code should prefer WaitUntilContext or WaitUntilContextDone.
|
||||||
|
func WaitUntilTimeout(tm time.Duration, fn func(chan struct{}) error) error {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), tm)
|
||||||
|
defer cancel()
|
||||||
|
err := WaitUntilContextDone(ctx, func(done <-chan struct{}) error {
|
||||||
|
return fn(bridgeDoneChan(done))
|
||||||
|
})
|
||||||
|
if errors.Is(err, context.DeadlineExceeded) {
|
||||||
|
return ERR_TIMEOUT
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilFinished runs fn asynchronously and returns its eventual result.
|
||||||
|
func WaitUntilFinished(fn func() error) <-chan error {
|
||||||
|
return WaitUntilContextFinished(context.Background(), func(ctx context.Context) error {
|
||||||
|
return fn()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// WaitUntilTimeoutFinished is the asynchronous form of WaitUntilTimeout.
|
||||||
|
//
|
||||||
|
// The provided stop channel must be treated as receive-only by callers. New
|
||||||
|
// code should prefer WaitUntilContextFinished or WaitUntilContextDoneFinished.
|
||||||
|
func WaitUntilTimeoutFinished(tm time.Duration, fn func(chan struct{}) error) <-chan error {
|
||||||
|
result := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), tm)
|
||||||
|
defer cancel()
|
||||||
|
err := WaitUntilContextDone(ctx, func(done <-chan struct{}) error {
|
||||||
|
return fn(bridgeDoneChan(done))
|
||||||
|
})
|
||||||
|
if errors.Is(err, context.DeadlineExceeded) {
|
||||||
|
err = ERR_TIMEOUT
|
||||||
|
}
|
||||||
|
result <- err
|
||||||
|
close(result)
|
||||||
|
}()
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func bridgeDoneChan(done <-chan struct{}) chan struct{} {
|
||||||
|
stop := make(chan struct{})
|
||||||
|
if done == nil {
|
||||||
|
return stop
|
||||||
|
}
|
||||||
|
go func() {
|
||||||
|
<-done
|
||||||
|
close(stop)
|
||||||
|
}()
|
||||||
|
return stop
|
||||||
|
}
|
||||||
+174
@@ -0,0 +1,174 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestWaitUntilContextReturnsWorkerError(t *testing.T) {
|
||||||
|
want := errors.New("worker failed")
|
||||||
|
err := WaitUntilContext(context.Background(), func(ctx context.Context) error {
|
||||||
|
return want
|
||||||
|
})
|
||||||
|
if !errors.Is(err, want) {
|
||||||
|
t.Fatalf("unexpected error: got %v want %v", err, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilContextReturnsDeadlineExceeded(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
workerReturned := make(chan struct{})
|
||||||
|
err := WaitUntilContext(ctx, func(ctx context.Context) error {
|
||||||
|
<-ctx.Done()
|
||||||
|
close(workerReturned)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if !errors.Is(err, context.DeadlineExceeded) {
|
||||||
|
t.Fatalf("unexpected context error: %v", err)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-workerReturned:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("worker did not return after context deadline")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilContextFinishedReturnsCanceled(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
result := WaitUntilContextFinished(ctx, func(ctx context.Context) error {
|
||||||
|
<-ctx.Done()
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
err, ok := <-result
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("result channel closed without value")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("unexpected context error: %v", err)
|
||||||
|
}
|
||||||
|
if _, ok := <-result; ok {
|
||||||
|
t.Fatal("result channel should be closed after delivering the outcome")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilContextDoneBridgesDoneChannel(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
doneClosed := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
time.Sleep(20 * time.Millisecond)
|
||||||
|
cancel()
|
||||||
|
}()
|
||||||
|
err := WaitUntilContextDone(ctx, func(done <-chan struct{}) error {
|
||||||
|
<-done
|
||||||
|
close(doneClosed)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("unexpected context error: %v", err)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-doneClosed:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("done channel was not closed when context canceled")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilContextDoneFinishedReturnsCanceled(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
doneClosed := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
time.Sleep(20 * time.Millisecond)
|
||||||
|
cancel()
|
||||||
|
}()
|
||||||
|
result := WaitUntilContextDoneFinished(ctx, func(done <-chan struct{}) error {
|
||||||
|
<-done
|
||||||
|
close(doneClosed)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
err, ok := <-result
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("result channel closed without value")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("unexpected context error: %v", err)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-doneClosed:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("done channel was not closed when context canceled")
|
||||||
|
}
|
||||||
|
if _, ok := <-result; ok {
|
||||||
|
t.Fatal("result channel should be closed after delivering the outcome")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilTimeoutReturnsWorkerError(t *testing.T) {
|
||||||
|
want := errors.New("worker failed")
|
||||||
|
err := WaitUntilTimeout(time.Second, func(stop chan struct{}) error {
|
||||||
|
return want
|
||||||
|
})
|
||||||
|
if !errors.Is(err, want) {
|
||||||
|
t.Fatalf("unexpected error: got %v want %v", err, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilTimeoutTimesOutWithoutBlockingWorkerReturn(t *testing.T) {
|
||||||
|
workerReturned := make(chan struct{})
|
||||||
|
err := WaitUntilTimeout(20*time.Millisecond, func(stop chan struct{}) error {
|
||||||
|
<-stop
|
||||||
|
close(workerReturned)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if !errors.Is(err, ERR_TIMEOUT) {
|
||||||
|
t.Fatalf("unexpected timeout error: %v", err)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-workerReturned:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("worker did not return after timeout signal")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilFinishedReturnsWorkerError(t *testing.T) {
|
||||||
|
want := errors.New("worker failed")
|
||||||
|
err := <-WaitUntilFinished(func() error {
|
||||||
|
return want
|
||||||
|
})
|
||||||
|
if !errors.Is(err, want) {
|
||||||
|
t.Fatalf("unexpected error: got %v want %v", err, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilTimeoutFinishedReturnsTimeout(t *testing.T) {
|
||||||
|
result := WaitUntilTimeoutFinished(20*time.Millisecond, func(stop chan struct{}) error {
|
||||||
|
<-stop
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
err, ok := <-result
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("result channel closed without value")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, ERR_TIMEOUT) {
|
||||||
|
t.Fatalf("unexpected timeout error: %v", err)
|
||||||
|
}
|
||||||
|
if _, ok := <-result; ok {
|
||||||
|
t.Fatal("result channel should be closed after delivering the outcome")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitUntilTimeoutFinishedReturnsWorkerError(t *testing.T) {
|
||||||
|
want := errors.New("worker failed")
|
||||||
|
err, ok := <-WaitUntilTimeoutFinished(time.Second, func(stop chan struct{}) error {
|
||||||
|
return want
|
||||||
|
})
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("result channel closed without value")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, want) {
|
||||||
|
t.Fatalf("unexpected error: got %v want %v", err, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
+305
@@ -0,0 +1,305 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DefaultFrameReaderBufferSize is the default transport read chunk size used by
|
||||||
|
// FrameReader.
|
||||||
|
const DefaultFrameReaderBufferSize = 32 * 1024
|
||||||
|
|
||||||
|
var framePayloadPool sync.Pool
|
||||||
|
|
||||||
|
type frameReaderConnKey struct{}
|
||||||
|
|
||||||
|
// FrameWriter adapts StarQueue framing helpers to an io.Writer.
|
||||||
|
type FrameWriter struct {
|
||||||
|
writer io.Writer
|
||||||
|
queue *StarQueue
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewFrameWriter creates a framing writer backed by queue. When queue is nil, a
|
||||||
|
// default StarQueue is created.
|
||||||
|
func NewFrameWriter(writer io.Writer, queue *StarQueue) *FrameWriter {
|
||||||
|
if queue == nil {
|
||||||
|
queue = NewQueue()
|
||||||
|
}
|
||||||
|
return &FrameWriter{
|
||||||
|
writer: writer,
|
||||||
|
queue: queue,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFrame writes one framed payload.
|
||||||
|
func (writer *FrameWriter) WriteFrame(payload []byte) error {
|
||||||
|
if writer == nil || writer.writer == nil || writer.queue == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.queue.WriteFrame(writer.writer, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFrameBuffers writes one framed payload using net.Buffers when possible.
|
||||||
|
func (writer *FrameWriter) WriteFrameBuffers(payload []byte) error {
|
||||||
|
if writer == nil || writer.writer == nil || writer.queue == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.queue.WriteFrameBuffers(writer.writer, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFramesBuffers writes multiple framed payloads in one batch when
|
||||||
|
// possible.
|
||||||
|
func (writer *FrameWriter) WriteFramesBuffers(payloads ...[]byte) error {
|
||||||
|
if writer == nil || writer.writer == nil || writer.queue == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.queue.WriteFramesBuffers(writer.writer, payloads...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// FrameReader adapts StarQueue parsing helpers to an io.Reader.
|
||||||
|
type FrameReader struct {
|
||||||
|
reader io.Reader
|
||||||
|
queue *StarQueue
|
||||||
|
connKey interface{}
|
||||||
|
readSize int
|
||||||
|
|
||||||
|
header [queHeaderSize]byte
|
||||||
|
headerRead int
|
||||||
|
payload []byte
|
||||||
|
payloadRead int
|
||||||
|
release func()
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewFrameReader creates a framing reader backed by queue. When queue is nil, a
|
||||||
|
// default StarQueue is created.
|
||||||
|
func NewFrameReader(reader io.Reader, queue *StarQueue) *FrameReader {
|
||||||
|
if queue == nil {
|
||||||
|
queue = NewQueue()
|
||||||
|
}
|
||||||
|
return &FrameReader{
|
||||||
|
reader: reader,
|
||||||
|
queue: queue,
|
||||||
|
connKey: &frameReaderConnKey{},
|
||||||
|
readSize: DefaultFrameReaderBufferSize,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetReadBufferSize updates the underlying transport read chunk size.
|
||||||
|
func (reader *FrameReader) SetReadBufferSize(size int) {
|
||||||
|
if reader == nil || size <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
reader.readSize = size
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetConnKey overrides the internal queue connection key.
|
||||||
|
func (reader *FrameReader) SetConnKey(conn interface{}) error {
|
||||||
|
if reader == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if conn == nil {
|
||||||
|
reader.connKey = &frameReaderConnKey{}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := validateConnKey(conn); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
reader.connKey = conn
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Next returns the next framed payload.
|
||||||
|
func (reader *FrameReader) Next() ([]byte, error) {
|
||||||
|
if reader == nil || reader.reader == nil || reader.queue == nil {
|
||||||
|
return nil, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
payload, release, err := reader.NextPooled()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if release == nil {
|
||||||
|
return payload, nil
|
||||||
|
}
|
||||||
|
owned := cloneBytes(payload)
|
||||||
|
release()
|
||||||
|
return owned, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// NextPooled returns the next frame payload. The caller must call release when
|
||||||
|
// it is non-nil.
|
||||||
|
func (reader *FrameReader) NextPooled() ([]byte, func(), error) {
|
||||||
|
if reader == nil || reader.reader == nil || reader.queue == nil {
|
||||||
|
return nil, nil, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if err := reader.readHeader(); err != nil {
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
if err := reader.ensurePayloadBuffer(true); err != nil {
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
if len(reader.payload) > 0 {
|
||||||
|
if err := reader.readInto(reader.payload, &reader.payloadRead); err != nil {
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return reader.finishFrame()
|
||||||
|
}
|
||||||
|
|
||||||
|
// NextView reads the next frame and exposes its payload only for the duration
|
||||||
|
// of fn.
|
||||||
|
func (reader *FrameReader) NextView(fn func(FrameView) error) error {
|
||||||
|
if fn == nil {
|
||||||
|
return ErrQueueFrameHandlerNil
|
||||||
|
}
|
||||||
|
payload, release, err := reader.NextPooled()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if release != nil {
|
||||||
|
defer release()
|
||||||
|
}
|
||||||
|
return fn(FrameView{
|
||||||
|
Payload: payload,
|
||||||
|
Conn: reader.connKey,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinFrameReaderError(parseErr error, readErr error) error {
|
||||||
|
switch {
|
||||||
|
case parseErr == nil:
|
||||||
|
return readErr
|
||||||
|
case readErr == nil:
|
||||||
|
return parseErr
|
||||||
|
default:
|
||||||
|
return errors.Join(parseErr, readErr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) normalizeReadErr(parseErr error, readErr error) error {
|
||||||
|
if errors.Is(readErr, io.EOF) && reader.hasBufferedState() {
|
||||||
|
reader.clearBufferedState()
|
||||||
|
readErr = io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
return joinFrameReaderError(parseErr, readErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) hasBufferedState() bool {
|
||||||
|
if reader == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return reader.headerRead > 0 || reader.payloadRead > 0 || len(reader.payload) > 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) clearBufferedState() {
|
||||||
|
if reader == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
reader.headerRead = 0
|
||||||
|
reader.payloadRead = 0
|
||||||
|
if reader.release != nil {
|
||||||
|
reader.release()
|
||||||
|
}
|
||||||
|
reader.release = nil
|
||||||
|
reader.payload = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) readHeader() error {
|
||||||
|
if reader.headerRead == queHeaderSize {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return reader.readInto(reader.header[:], &reader.headerRead)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) ensurePayloadBuffer(pooled bool) error {
|
||||||
|
if reader.payload != nil || reader.headerRead < queHeaderSize {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
header, err := parseHeaderBytes(reader.header[:], reader.queue.maxLength)
|
||||||
|
if err != nil {
|
||||||
|
reader.clearBufferedState()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if header.Length == 0 {
|
||||||
|
reader.payload = []byte{}
|
||||||
|
reader.release = nil
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
payloadLen := int(header.Length)
|
||||||
|
switch {
|
||||||
|
case reader.queue.Encode && reader.queue.DecodeFunc != nil:
|
||||||
|
reader.payload = make([]byte, payloadLen)
|
||||||
|
case pooled:
|
||||||
|
buf := getFramePayloadBuffer(payloadLen)
|
||||||
|
reader.payload = buf
|
||||||
|
reader.release = func() {
|
||||||
|
putFramePayloadBuffer(buf)
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
reader.payload = make([]byte, payloadLen)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) finishFrame() ([]byte, func(), error) {
|
||||||
|
payload := reader.payload
|
||||||
|
release := reader.release
|
||||||
|
reader.payload = nil
|
||||||
|
reader.payloadRead = 0
|
||||||
|
reader.release = nil
|
||||||
|
reader.headerRead = 0
|
||||||
|
if reader.queue.Encode && reader.queue.DecodeFunc != nil {
|
||||||
|
decoded := reader.queue.DecodeFunc(payload)
|
||||||
|
if release != nil {
|
||||||
|
release()
|
||||||
|
release = nil
|
||||||
|
}
|
||||||
|
payload = decoded
|
||||||
|
}
|
||||||
|
return payload, release, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *FrameReader) readInto(dst []byte, offset *int) error {
|
||||||
|
for *offset < len(dst) {
|
||||||
|
maxRead := len(dst) - *offset
|
||||||
|
if reader.readSize > 0 && maxRead > reader.readSize {
|
||||||
|
maxRead = reader.readSize
|
||||||
|
}
|
||||||
|
n, err := reader.reader.Read(dst[*offset : *offset+maxRead])
|
||||||
|
if n > 0 {
|
||||||
|
*offset += n
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, io.EOF) {
|
||||||
|
if *offset == 0 {
|
||||||
|
return io.EOF
|
||||||
|
}
|
||||||
|
if *offset < len(dst) {
|
||||||
|
return io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if n == 0 {
|
||||||
|
return io.ErrNoProgress
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func getFramePayloadBuffer(size int) []byte {
|
||||||
|
if size <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if pooled, ok := framePayloadPool.Get().([]byte); ok && cap(pooled) >= size {
|
||||||
|
return pooled[:size]
|
||||||
|
}
|
||||||
|
return make([]byte, size)
|
||||||
|
}
|
||||||
|
|
||||||
|
func putFramePayloadBuffer(buf []byte) {
|
||||||
|
if cap(buf) == 0 || cap(buf) > 32*1024*1024 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
framePayloadPool.Put(buf[:0])
|
||||||
|
}
|
||||||
+225
@@ -0,0 +1,225 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
type chunkedReader struct {
|
||||||
|
data []byte
|
||||||
|
max int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *chunkedReader) Read(p []byte) (int, error) {
|
||||||
|
if len(reader.data) == 0 {
|
||||||
|
return 0, io.EOF
|
||||||
|
}
|
||||||
|
if reader.max > 0 && len(p) > reader.max {
|
||||||
|
p = p[:reader.max]
|
||||||
|
}
|
||||||
|
n := copy(p, reader.data)
|
||||||
|
reader.data = reader.data[n:]
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type trackingReader struct {
|
||||||
|
data []byte
|
||||||
|
maxRequested int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *trackingReader) Read(p []byte) (int, error) {
|
||||||
|
if len(p) > reader.maxRequested {
|
||||||
|
reader.maxRequested = len(p)
|
||||||
|
}
|
||||||
|
if len(reader.data) == 0 {
|
||||||
|
return 0, io.EOF
|
||||||
|
}
|
||||||
|
n := copy(p, reader.data)
|
||||||
|
reader.data = reader.data[n:]
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type stagedReaderStep struct {
|
||||||
|
data []byte
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
type stagedReader struct {
|
||||||
|
steps []stagedReaderStep
|
||||||
|
index int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *stagedReader) Read(p []byte) (int, error) {
|
||||||
|
if reader.index >= len(reader.steps) {
|
||||||
|
return 0, io.EOF
|
||||||
|
}
|
||||||
|
step := &reader.steps[reader.index]
|
||||||
|
n := copy(p, step.data)
|
||||||
|
step.data = step.data[n:]
|
||||||
|
if len(step.data) == 0 {
|
||||||
|
err := step.err
|
||||||
|
reader.index++
|
||||||
|
return n, err
|
||||||
|
}
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameWriterReaderRoundTrip(t *testing.T) {
|
||||||
|
var wire bytes.Buffer
|
||||||
|
writer := NewFrameWriter(&wire, nil)
|
||||||
|
if err := writer.WriteFrame([]byte("one")); err != nil {
|
||||||
|
t.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
if err := writer.WriteFramesBuffers([][]byte{[]byte("two"), []byte("three")}...); err != nil {
|
||||||
|
t.Fatalf("WriteFramesBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
reader := NewFrameReader(bytes.NewReader(wire.Bytes()), nil)
|
||||||
|
cases := []string{"one", "two", "three"}
|
||||||
|
for _, want := range cases {
|
||||||
|
got, err := reader.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next returned error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != want {
|
||||||
|
t.Fatalf("unexpected frame: got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, err := reader.Next(); !errors.Is(err, io.EOF) {
|
||||||
|
t.Fatalf("expected EOF after draining frames, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderHandlesPartialTransportReads(t *testing.T) {
|
||||||
|
var wire bytes.Buffer
|
||||||
|
writer := NewFrameWriter(&wire, nil)
|
||||||
|
if err := writer.WriteFrame([]byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
reader := NewFrameReader(&chunkedReader{data: wire.Bytes(), max: 3}, nil)
|
||||||
|
reader.SetReadBufferSize(3)
|
||||||
|
got, err := reader.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next returned error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != "hello" {
|
||||||
|
t.Fatalf("unexpected frame: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderSetReadBufferSizeLimitsUnderlyingReadSize(t *testing.T) {
|
||||||
|
var wire bytes.Buffer
|
||||||
|
writer := NewFrameWriter(&wire, nil)
|
||||||
|
if err := writer.WriteFrame([]byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
source := &trackingReader{data: wire.Bytes()}
|
||||||
|
reader := NewFrameReader(source, nil)
|
||||||
|
reader.SetReadBufferSize(3)
|
||||||
|
|
||||||
|
got, err := reader.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Next returned error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != "hello" {
|
||||||
|
t.Fatalf("unexpected frame: %q", got)
|
||||||
|
}
|
||||||
|
if got, want := source.maxRequested, 3; got != want {
|
||||||
|
t.Fatalf("max requested read size = %d, want %d", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderNextPooledAndNextView(t *testing.T) {
|
||||||
|
queue := NewQueue()
|
||||||
|
wire := append(queue.BuildMessage([]byte("alpha")), queue.BuildMessage([]byte("beta"))...)
|
||||||
|
reader := NewFrameReader(bytes.NewReader(wire), queue)
|
||||||
|
|
||||||
|
first, release, err := reader.NextPooled()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NextPooled returned error: %v", err)
|
||||||
|
}
|
||||||
|
if string(first) != "alpha" {
|
||||||
|
t.Fatalf("unexpected pooled frame: %q", first)
|
||||||
|
}
|
||||||
|
if release == nil {
|
||||||
|
t.Fatal("NextPooled should return a release func")
|
||||||
|
}
|
||||||
|
release()
|
||||||
|
|
||||||
|
var second string
|
||||||
|
if err := reader.NextView(func(view FrameView) error {
|
||||||
|
second = string(view.Payload)
|
||||||
|
return nil
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("NextView returned error: %v", err)
|
||||||
|
}
|
||||||
|
if second != "beta" {
|
||||||
|
t.Fatalf("unexpected view frame: %q", second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderPreservesPartialFrameAcrossDeadline(t *testing.T) {
|
||||||
|
queue := NewQueue()
|
||||||
|
frame := queue.BuildMessage([]byte("hello"))
|
||||||
|
reader := NewFrameReader(&stagedReader{
|
||||||
|
steps: []stagedReaderStep{
|
||||||
|
{data: append([]byte(nil), frame[:5]...), err: os.ErrDeadlineExceeded},
|
||||||
|
{data: append([]byte(nil), frame[5:]...), err: nil},
|
||||||
|
},
|
||||||
|
}, queue)
|
||||||
|
|
||||||
|
if _, _, err := reader.NextPooled(); !errors.Is(err, os.ErrDeadlineExceeded) {
|
||||||
|
t.Fatalf("expected deadline exceeded on partial frame, got %v", err)
|
||||||
|
}
|
||||||
|
got, release, err := reader.NextPooled()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NextPooled after deadline returned error: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != "hello" {
|
||||||
|
t.Fatalf("unexpected frame after deadline: %q", got)
|
||||||
|
}
|
||||||
|
if release != nil {
|
||||||
|
release()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameWriterNilWriterFails(t *testing.T) {
|
||||||
|
writer := NewFrameWriter(nil, nil)
|
||||||
|
if err := writer.WriteFrame([]byte("hello")); !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("expected io.ErrClosedPipe, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderTruncatedFrameReturnsUnexpectedEOF(t *testing.T) {
|
||||||
|
queue := NewQueue()
|
||||||
|
frame := queue.BuildMessage([]byte("hello"))
|
||||||
|
reader := NewFrameReader(bytes.NewReader(frame[:len(frame)-1]), nil)
|
||||||
|
|
||||||
|
if _, err := reader.Next(); !errors.Is(err, io.ErrUnexpectedEOF) {
|
||||||
|
t.Fatalf("expected io.ErrUnexpectedEOF, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFrameReaderReturnsUnexpectedEOFAfterPendingFrames(t *testing.T) {
|
||||||
|
queue := NewQueue()
|
||||||
|
first := queue.BuildMessage([]byte("one"))
|
||||||
|
second := queue.BuildMessage([]byte("two"))
|
||||||
|
wire := append(append([]byte{}, first...), second[:len(second)-1]...)
|
||||||
|
reader := NewFrameReader(bytes.NewReader(wire), nil)
|
||||||
|
|
||||||
|
got, err := reader.Next()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error on first frame: %v", err)
|
||||||
|
}
|
||||||
|
if string(got) != "one" {
|
||||||
|
t.Fatalf("unexpected first frame: %q", got)
|
||||||
|
}
|
||||||
|
if _, err := reader.Next(); !errors.Is(err, io.ErrUnexpectedEOF) {
|
||||||
|
t.Fatalf("expected io.ErrUnexpectedEOF after draining pending frame, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
module b612.me/stario
|
||||||
|
|
||||||
|
go 1.18
|
||||||
|
|
||||||
|
require golang.org/x/term v0.23.0
|
||||||
|
|
||||||
|
require golang.org/x/sys v0.23.0 // indirect
|
||||||
@@ -0,0 +1,4 @@
|
|||||||
|
golang.org/x/sys v0.23.0 h1:YfKFowiIMvtgl1UERQoTPPToxltDeZfbj4H7dVUCwmM=
|
||||||
|
golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||||
|
golang.org/x/term v0.23.0 h1:F6D4vR+EHoL9/sWAWgAR1H2DcHr4PareCbAaCo1RpuU=
|
||||||
|
golang.org/x/term v0.23.0/go.mod h1:DgV24QBUrK6jhZXl+20l6UWznPlwAHm1Q1mGHtydmSk=
|
||||||
@@ -3,53 +3,334 @@ package stario
|
|||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"golang.org/x/term"
|
||||||
|
"io"
|
||||||
"os"
|
"os"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"golang.org/x/crypto/ssh/terminal"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type InputMsg struct {
|
type InputMsg struct {
|
||||||
msg string
|
msg string
|
||||||
err error
|
err error
|
||||||
|
skipSliceSigErr bool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type rawInputSignalMode uint8
|
||||||
|
|
||||||
|
const (
|
||||||
|
rawInputSignalIgnore rawInputSignalMode = iota
|
||||||
|
rawInputSignalExit
|
||||||
|
rawInputSignalReturnError
|
||||||
|
)
|
||||||
|
|
||||||
|
type rawTerminalSession struct {
|
||||||
|
fd int
|
||||||
|
state *term.State
|
||||||
|
reader *bufio.Reader
|
||||||
|
input io.Closer
|
||||||
|
redrawHint string
|
||||||
|
printNewline bool
|
||||||
|
mu sync.Mutex
|
||||||
|
}
|
||||||
|
|
||||||
|
var rawTerminalSessionFactory = newRawTerminalSession
|
||||||
|
var inputSignalHandler = signal
|
||||||
|
|
||||||
|
// Passwd reads one password-style line in raw mode.
|
||||||
|
//
|
||||||
|
// When the user presses an input signal such as Ctrl+C, this compatibility
|
||||||
|
// entry exits the current flow and returns an empty message with a nil error.
|
||||||
func Passwd(hint string, defaultVal string) InputMsg {
|
func Passwd(hint string, defaultVal string) InputMsg {
|
||||||
return passwd(hint, defaultVal, "")
|
return passwd(hint, defaultVal, "", rawInputSignalExit)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// PasswdWithMask is like Passwd but echoes the provided mask string.
|
||||||
func PasswdWithMask(hint string, defaultVal string, mask string) InputMsg {
|
func PasswdWithMask(hint string, defaultVal string, mask string) InputMsg {
|
||||||
return passwd(hint, defaultVal, mask)
|
return passwd(hint, defaultVal, mask, rawInputSignalExit)
|
||||||
}
|
}
|
||||||
|
|
||||||
func passwd(hint string, defaultVal string, mask string) InputMsg {
|
// PasswdResponseSignal is like Passwd but preserves input-signal errors for
|
||||||
var ioBuf []byte
|
// callers that need to distinguish Ctrl+C / Ctrl+Z style exits.
|
||||||
|
func PasswdResponseSignal(hint string, defaultVal string) InputMsg {
|
||||||
|
return passwd(hint, defaultVal, "", rawInputSignalReturnError)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PasswdResponseSignalWithMask is like PasswdResponseSignal but echoes the
|
||||||
|
// provided mask string.
|
||||||
|
func PasswdResponseSignalWithMask(hint string, defaultVal string, mask string) InputMsg {
|
||||||
|
return passwd(hint, defaultVal, mask, rawInputSignalReturnError)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessageBoxRaw reads one line in raw mode without treating control keys as
|
||||||
|
// exit signals.
|
||||||
|
func MessageBoxRaw(hint string, defaultVal string) InputMsg {
|
||||||
|
return messageBox(hint, defaultVal)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRawTerminalSession(hint string, printNewline bool) (*rawTerminalSession, error) {
|
||||||
if hint != "" {
|
if hint != "" {
|
||||||
fmt.Print(hint)
|
fmt.Print(hint)
|
||||||
}
|
}
|
||||||
state, err := terminal.MakeRaw(int(os.Stdin.Fd()))
|
input, err := openRawTerminalInput()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return InputMsg{"", err}
|
return nil, err
|
||||||
}
|
}
|
||||||
defer terminal.Restore(0, state)
|
fd := int(input.Fd())
|
||||||
inputReader := bufio.NewReader(os.Stdin)
|
state, err := term.MakeRaw(fd)
|
||||||
for {
|
|
||||||
b, err := inputReader.ReadByte()
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return InputMsg{"", err}
|
_ = input.Close()
|
||||||
|
return nil, err
|
||||||
}
|
}
|
||||||
if b == 0x0d {
|
return &rawTerminalSession{
|
||||||
|
fd: fd,
|
||||||
|
state: state,
|
||||||
|
reader: bufio.NewReader(input),
|
||||||
|
input: input,
|
||||||
|
redrawHint: promptRedrawHint(hint),
|
||||||
|
printNewline: printNewline,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (session *rawTerminalSession) Close() {
|
||||||
|
if session == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
session.mu.Lock()
|
||||||
|
defer session.mu.Unlock()
|
||||||
|
if session.state != nil {
|
||||||
|
_ = term.Restore(session.fd, session.state)
|
||||||
|
session.state = nil
|
||||||
|
}
|
||||||
|
if session.printNewline {
|
||||||
fmt.Println()
|
fmt.Println()
|
||||||
return InputMsg{strings.TrimSpace(string(ioBuf)), nil}
|
session.printNewline = false
|
||||||
}
|
}
|
||||||
if mask != "" {
|
if session.input != nil {
|
||||||
fmt.Print(mask)
|
_ = session.input.Close()
|
||||||
}
|
session.input = nil
|
||||||
ioBuf = append(ioBuf, b)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (session *rawTerminalSession) Restore() error {
|
||||||
|
if session == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
session.mu.Lock()
|
||||||
|
defer session.mu.Unlock()
|
||||||
|
if session.state == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err := term.Restore(session.fd, session.state); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
session.state = nil
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (session *rawTerminalSession) Abort() {
|
||||||
|
session.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func promptRedrawHint(hint string) string {
|
||||||
|
if strings.Index(hint, "\n") >= 0 {
|
||||||
|
hint = hint[strings.LastIndex(hint, "\n")+1:]
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(hint)
|
||||||
|
}
|
||||||
|
|
||||||
|
func finalizeInputValue(raw string, defaultVal string) string {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if len(raw) == 0 {
|
||||||
|
return defaultVal
|
||||||
|
}
|
||||||
|
return raw
|
||||||
|
}
|
||||||
|
|
||||||
|
func renderRawEcho(ioBuf []rune, mask string) string {
|
||||||
|
if mask == "" {
|
||||||
|
return string(ioBuf)
|
||||||
|
}
|
||||||
|
return strings.Repeat(mask, len(ioBuf))
|
||||||
|
}
|
||||||
|
|
||||||
|
func rawEchoRenderUnit(r rune, mask string, maskWidth int) (string, int, bool) {
|
||||||
|
if mask != "" {
|
||||||
|
if maskWidth <= 0 {
|
||||||
|
return "", 0, false
|
||||||
|
}
|
||||||
|
return mask, maskWidth, true
|
||||||
|
}
|
||||||
|
width := runeDisplayWidth(r)
|
||||||
|
if width <= 0 {
|
||||||
|
return "", 0, false
|
||||||
|
}
|
||||||
|
return string(r), width, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func redrawPromptLine(hint string, echo string, lastWidth int) int {
|
||||||
|
nowWidth := stringDisplayWidth(hint) + stringDisplayWidth(echo)
|
||||||
|
clearWidth := lastWidth
|
||||||
|
if nowWidth > clearWidth {
|
||||||
|
clearWidth = nowWidth
|
||||||
|
}
|
||||||
|
fmt.Print("\r")
|
||||||
|
if clearWidth > 0 {
|
||||||
|
fmt.Print(strings.Repeat(" ", clearWidth))
|
||||||
|
fmt.Print("\r")
|
||||||
|
}
|
||||||
|
if hint != "" {
|
||||||
|
fmt.Print(hint)
|
||||||
|
}
|
||||||
|
if echo != "" {
|
||||||
|
fmt.Print(echo)
|
||||||
|
}
|
||||||
|
return nowWidth
|
||||||
|
}
|
||||||
|
|
||||||
|
func erasePromptTail(width int) {
|
||||||
|
if width <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
backtrack := strings.Repeat("\b", width)
|
||||||
|
fmt.Print(backtrack)
|
||||||
|
fmt.Print(strings.Repeat(" ", width))
|
||||||
|
fmt.Print(backtrack)
|
||||||
|
}
|
||||||
|
|
||||||
|
func redrawPromptEcho(hint string, ioBuf []rune, mask string, lastWidth int) (int, int) {
|
||||||
|
echo := renderRawEcho(ioBuf, mask)
|
||||||
|
echoWidth := stringDisplayWidth(echo)
|
||||||
|
return echoWidth, redrawPromptLine(hint, echo, lastWidth)
|
||||||
|
}
|
||||||
|
|
||||||
|
func stringDisplayWidth(text string) int {
|
||||||
|
width := 0
|
||||||
|
for _, r := range text {
|
||||||
|
width += runeDisplayWidth(r)
|
||||||
|
}
|
||||||
|
return width
|
||||||
|
}
|
||||||
|
|
||||||
|
func runeDisplayWidth(r rune) int {
|
||||||
|
switch {
|
||||||
|
case r == '\t':
|
||||||
|
return 4
|
||||||
|
case r < 0x20 || (r >= 0x7f && r < 0xa0):
|
||||||
|
return 0
|
||||||
|
case isWideRune(r):
|
||||||
|
return 2
|
||||||
|
default:
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func isWideRune(r rune) bool {
|
||||||
|
return r >= 0x1100 && (r <= 0x115f ||
|
||||||
|
r == 0x2329 || r == 0x232a ||
|
||||||
|
(r >= 0x2e80 && r <= 0xa4cf && r != 0x303f) ||
|
||||||
|
(r >= 0xac00 && r <= 0xd7a3) ||
|
||||||
|
(r >= 0xf900 && r <= 0xfaff) ||
|
||||||
|
(r >= 0xfe10 && r <= 0xfe19) ||
|
||||||
|
(r >= 0xfe30 && r <= 0xfe6f) ||
|
||||||
|
(r >= 0xff00 && r <= 0xff60) ||
|
||||||
|
(r >= 0xffe0 && r <= 0xffe6) ||
|
||||||
|
(r >= 0x1f300 && r <= 0x1f64f) ||
|
||||||
|
(r >= 0x1f900 && r <= 0x1f9ff) ||
|
||||||
|
(r >= 0x20000 && r <= 0x3fffd))
|
||||||
|
}
|
||||||
|
|
||||||
|
func signalInputResult(mode rawInputSignalMode, err error) InputMsg {
|
||||||
|
switch mode {
|
||||||
|
case rawInputSignalExit:
|
||||||
|
return InputMsg{msg: "", err: nil}
|
||||||
|
case rawInputSignalReturnError:
|
||||||
|
return InputMsg{msg: "", err: err}
|
||||||
|
default:
|
||||||
|
return InputMsg{msg: "", err: nil}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func rawLineInput(hint string, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
|
||||||
|
session, err := rawTerminalSessionFactory(hint, true)
|
||||||
|
if err != nil {
|
||||||
|
return InputMsg{msg: "", err: err}
|
||||||
|
}
|
||||||
|
defer session.Close()
|
||||||
|
return rawLineInputSession(session, defaultVal, mask, signalMode)
|
||||||
|
}
|
||||||
|
|
||||||
|
func rawLineInputSession(session *rawTerminalSession, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
|
||||||
|
if session == nil || session.reader == nil {
|
||||||
|
return InputMsg{msg: "", err: io.ErrClosedPipe}
|
||||||
|
}
|
||||||
|
ioBuf := make([]rune, 0, 16)
|
||||||
|
promptWidth := stringDisplayWidth(session.redrawHint)
|
||||||
|
maskWidth := stringDisplayWidth(mask)
|
||||||
|
echoWidth := 0
|
||||||
|
lastWidth := promptWidth
|
||||||
|
for {
|
||||||
|
b, _, err := session.reader.ReadRune()
|
||||||
|
if err != nil {
|
||||||
|
return InputMsg{msg: "", err: err}
|
||||||
|
}
|
||||||
|
if signalMode != rawInputSignalIgnore && isSignal(b) {
|
||||||
|
session.Close()
|
||||||
|
if signalMode == rawInputSignalExit {
|
||||||
|
return signalInputResult(signalMode, nil)
|
||||||
|
}
|
||||||
|
return signalInputResult(signalMode, inputSignalHandler(b))
|
||||||
|
}
|
||||||
|
switch b {
|
||||||
|
case 0x0d, 0x0a:
|
||||||
|
return InputMsg{msg: finalizeInputValue(string(ioBuf), defaultVal), err: nil}
|
||||||
|
case 0x08, 0x7F:
|
||||||
|
if len(ioBuf) > 0 {
|
||||||
|
removed := ioBuf[len(ioBuf)-1]
|
||||||
|
ioBuf = ioBuf[:len(ioBuf)-1]
|
||||||
|
if _, removedWidth, ok := rawEchoRenderUnit(removed, mask, maskWidth); ok {
|
||||||
|
erasePromptTail(removedWidth)
|
||||||
|
echoWidth -= removedWidth
|
||||||
|
if echoWidth < 0 {
|
||||||
|
echoWidth = 0
|
||||||
|
}
|
||||||
|
lastWidth = promptWidth + echoWidth
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
ioBuf = append(ioBuf, b)
|
||||||
|
if appendText, appendWidth, ok := rawEchoRenderUnit(b, mask, maskWidth); ok {
|
||||||
|
fmt.Print(appendText)
|
||||||
|
echoWidth += appendWidth
|
||||||
|
lastWidth = promptWidth + echoWidth
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
echoWidth, lastWidth = redrawPromptEcho(session.redrawHint, ioBuf, mask, lastWidth)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func messageBox(hint string, defaultVal string) InputMsg {
|
||||||
|
return rawLineInput(hint, defaultVal, "", rawInputSignalIgnore)
|
||||||
|
}
|
||||||
|
|
||||||
|
func isSignal(s rune) bool {
|
||||||
|
switch s {
|
||||||
|
case 0x03, 0x1a, 0x1c:
|
||||||
|
return true
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func passwd(hint string, defaultVal string, mask string, signalMode rawInputSignalMode) InputMsg {
|
||||||
|
return rawLineInput(hint, defaultVal, mask, signalMode)
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessageBox reads one line in cooked mode and falls back to defaultVal when
|
||||||
|
// the trimmed input is empty.
|
||||||
func MessageBox(hint string, defaultVal string) InputMsg {
|
func MessageBox(hint string, defaultVal string) InputMsg {
|
||||||
if hint != "" {
|
if hint != "" {
|
||||||
fmt.Print(hint)
|
fmt.Print(hint)
|
||||||
@@ -57,9 +338,15 @@ func MessageBox(hint string, defaultVal string) InputMsg {
|
|||||||
inputReader := bufio.NewReader(os.Stdin)
|
inputReader := bufio.NewReader(os.Stdin)
|
||||||
str, err := inputReader.ReadString('\n')
|
str, err := inputReader.ReadString('\n')
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return InputMsg{"", err}
|
return InputMsg{msg: str, err: err}
|
||||||
}
|
}
|
||||||
return InputMsg{strings.TrimSpace(str), nil}
|
str = finalizeInputValue(str, defaultVal)
|
||||||
|
return InputMsg{msg: str, err: err}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im *InputMsg) IgnoreSliceParseError(i bool) *InputMsg {
|
||||||
|
im.skipSliceSigErr = i
|
||||||
|
return im
|
||||||
}
|
}
|
||||||
|
|
||||||
func (im InputMsg) String() (string, error) {
|
func (im InputMsg) String() (string, error) {
|
||||||
@@ -74,6 +361,38 @@ func (im InputMsg) MustString() string {
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceString(sep string) ([]string, error) {
|
||||||
|
if im.err != nil {
|
||||||
|
return nil, im.err
|
||||||
|
}
|
||||||
|
if len(strings.TrimSpace(im.msg)) == 0 {
|
||||||
|
return []string{}, nil
|
||||||
|
}
|
||||||
|
return strings.Split(im.msg, sep), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceString(sep string) []string {
|
||||||
|
res, _ := im.SliceString(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) sliceFn(sep string, fn func(string) (interface{}, error)) ([]interface{}, error) {
|
||||||
|
var res []interface{}
|
||||||
|
data, err := im.SliceString(sep)
|
||||||
|
if err != nil {
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
for _, v := range data {
|
||||||
|
code, err := fn(strings.TrimSpace(v))
|
||||||
|
if err != nil && !im.skipSliceSigErr {
|
||||||
|
return nil, err
|
||||||
|
} else if err == nil {
|
||||||
|
res = append(res, code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return res, nil
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) Int() (int, error) {
|
func (im InputMsg) Int() (int, error) {
|
||||||
if im.err != nil {
|
if im.err != nil {
|
||||||
return 0, im.err
|
return 0, im.err
|
||||||
@@ -81,6 +400,22 @@ func (im InputMsg) Int() (int, error) {
|
|||||||
return strconv.Atoi(im.msg)
|
return strconv.Atoi(im.msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceInt(sep string) ([]int, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
return strconv.Atoi(v)
|
||||||
|
})
|
||||||
|
var res []int
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(int))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceInt(sep string) []int {
|
||||||
|
res, _ := im.SliceInt(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) MustInt() int {
|
func (im InputMsg) MustInt() int {
|
||||||
res, _ := im.Int()
|
res, _ := im.Int()
|
||||||
return res
|
return res
|
||||||
@@ -98,6 +433,22 @@ func (im InputMsg) MustInt64() int64 {
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceInt64(sep string) ([]int64, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
return strconv.ParseInt(v, 10, 64)
|
||||||
|
})
|
||||||
|
var res []int64
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(int64))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceInt64(sep string) []int64 {
|
||||||
|
res, _ := im.SliceInt64(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) Uint64() (uint64, error) {
|
func (im InputMsg) Uint64() (uint64, error) {
|
||||||
if im.err != nil {
|
if im.err != nil {
|
||||||
return 0, im.err
|
return 0, im.err
|
||||||
@@ -109,6 +460,21 @@ func (im InputMsg) MustUint64() uint64 {
|
|||||||
res, _ := im.Uint64()
|
res, _ := im.Uint64()
|
||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
func (im InputMsg) SliceUint64(sep string) ([]uint64, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
return strconv.ParseUint(v, 10, 64)
|
||||||
|
})
|
||||||
|
var res []uint64
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(uint64))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceUint64(sep string) []uint64 {
|
||||||
|
res, _ := im.SliceUint64(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) Bool() (bool, error) {
|
func (im InputMsg) Bool() (bool, error) {
|
||||||
if im.err != nil {
|
if im.err != nil {
|
||||||
@@ -122,6 +488,22 @@ func (im InputMsg) MustBool() bool {
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceBool(sep string) ([]bool, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
return strconv.ParseBool(v)
|
||||||
|
})
|
||||||
|
var res []bool
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(bool))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceBool(sep string) []bool {
|
||||||
|
res, _ := im.SliceBool(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) Float64() (float64, error) {
|
func (im InputMsg) Float64() (float64, error) {
|
||||||
if im.err != nil {
|
if im.err != nil {
|
||||||
return 0, im.err
|
return 0, im.err
|
||||||
@@ -134,6 +516,22 @@ func (im InputMsg) MustFloat64() float64 {
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceFloat64(sep string) ([]float64, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
return strconv.ParseFloat(v, 64)
|
||||||
|
})
|
||||||
|
var res []float64
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(float64))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceFloat64(sep string) []float64 {
|
||||||
|
res, _ := im.SliceFloat64(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func (im InputMsg) Float32() (float32, error) {
|
func (im InputMsg) Float32() (float32, error) {
|
||||||
if im.err != nil {
|
if im.err != nil {
|
||||||
return 0, im.err
|
return 0, im.err
|
||||||
@@ -147,57 +545,123 @@ func (im InputMsg) MustFloat32() float32 {
|
|||||||
return res
|
return res
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) SliceFloat32(sep string) ([]float32, error) {
|
||||||
|
data, err := im.sliceFn(sep, func(v string) (interface{}, error) {
|
||||||
|
f, err := strconv.ParseFloat(v, 32)
|
||||||
|
return float32(f), err
|
||||||
|
})
|
||||||
|
var res []float32
|
||||||
|
for _, v := range data {
|
||||||
|
res = append(res, v.(float32))
|
||||||
|
}
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (im InputMsg) MustSliceFloat32(sep string) []float32 {
|
||||||
|
res, _ := im.SliceFloat32(sep)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
func YesNo(hint string, defaults bool) bool {
|
func YesNo(hint string, defaults bool) bool {
|
||||||
res := strings.ToUpper(MessageBox(hint, "").MustString())
|
res, err := YesNoE(hint, defaults)
|
||||||
if res == "" {
|
if err != nil {
|
||||||
return defaults
|
return defaults
|
||||||
}
|
}
|
||||||
res = res[0:1]
|
return res
|
||||||
if res == "Y" {
|
}
|
||||||
return true
|
|
||||||
} else if res == "N" {
|
func parseYesNoValue(raw string, defaults bool) (bool, bool) {
|
||||||
return false
|
raw = strings.TrimSpace(strings.ToUpper(raw))
|
||||||
} else {
|
if raw == "" {
|
||||||
return defaults
|
return defaults, true
|
||||||
|
}
|
||||||
|
switch []rune(raw)[0] {
|
||||||
|
case 'Y':
|
||||||
|
return true, true
|
||||||
|
case 'N':
|
||||||
|
return false, true
|
||||||
|
default:
|
||||||
|
return false, false
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func StopUntil(hint string, trigger string, repeat bool) error {
|
func YesNoE(hint string, defaults bool) (bool, error) {
|
||||||
pressLen := len(trigger)
|
for {
|
||||||
if trigger == "" {
|
res, err := MessageBox(hint, "").String()
|
||||||
pressLen = 1
|
if err != nil {
|
||||||
} else {
|
return false, err
|
||||||
pressLen = len(trigger)
|
|
||||||
}
|
}
|
||||||
state, err := terminal.MakeRaw(int(os.Stdin.Fd()))
|
if answer, ok := parseYesNoValue(res, defaults); ok {
|
||||||
|
return answer, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildTriggerPrefixTable(triggerRunes []rune) []int {
|
||||||
|
if len(triggerRunes) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
prefix := make([]int, len(triggerRunes))
|
||||||
|
for i := 1; i < len(triggerRunes); i++ {
|
||||||
|
j := prefix[i-1]
|
||||||
|
for j > 0 && triggerRunes[i] != triggerRunes[j] {
|
||||||
|
j = prefix[j-1]
|
||||||
|
}
|
||||||
|
if triggerRunes[i] == triggerRunes[j] {
|
||||||
|
j++
|
||||||
|
}
|
||||||
|
prefix[i] = j
|
||||||
|
}
|
||||||
|
return prefix
|
||||||
|
}
|
||||||
|
|
||||||
|
func advanceTriggerIndex(triggerRunes []rune, prefix []int, current int, input rune) (int, bool) {
|
||||||
|
if len(triggerRunes) == 0 {
|
||||||
|
return 0, true
|
||||||
|
}
|
||||||
|
for current > 0 && input != triggerRunes[current] {
|
||||||
|
current = prefix[current-1]
|
||||||
|
}
|
||||||
|
if input == triggerRunes[current] {
|
||||||
|
current++
|
||||||
|
if current == len(triggerRunes) {
|
||||||
|
return current, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return current, false
|
||||||
|
}
|
||||||
|
|
||||||
|
// StopUntil keeps reading raw input until trigger is matched.
|
||||||
|
// When trigger == "", it returns after the first key press, which is used for
|
||||||
|
// "press any key to continue" style prompts.
|
||||||
|
func StopUntil(hint string, trigger string, repeat bool) error {
|
||||||
|
triggerRunes := []rune(trigger)
|
||||||
|
prefix := buildTriggerPrefixTable(triggerRunes)
|
||||||
|
session, err := rawTerminalSessionFactory(hint, false)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
defer terminal.Restore(0, state)
|
defer session.Close()
|
||||||
inputReader := bufio.NewReader(os.Stdin)
|
|
||||||
//ioBuf := make([]byte, pressLen)
|
|
||||||
if hint != "" && !repeat {
|
|
||||||
fmt.Print(hint)
|
|
||||||
}
|
|
||||||
i := 0
|
i := 0
|
||||||
for {
|
for {
|
||||||
b, err := inputReader.ReadByte()
|
b, _, err := session.reader.ReadRune()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if trigger == "" {
|
if trigger == "" {
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
if b == trigger[i] {
|
next, complete := advanceTriggerIndex(triggerRunes, prefix, i, b)
|
||||||
i++
|
if complete {
|
||||||
if i == pressLen {
|
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
if next > 0 {
|
||||||
|
i = next
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
i = 0
|
i = 0
|
||||||
if hint != "" && repeat {
|
if hint != "" && repeat {
|
||||||
fmt.Print("\n")
|
fmt.Print("\r\n")
|
||||||
fmt.Print(hint)
|
fmt.Print(hint)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,81 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func benchmarkRawRedrawAppend(b *testing.B, current []rune, next rune, mask string) {
|
||||||
|
b.Run("legacy", func(b *testing.B) {
|
||||||
|
ioBuf := make([]rune, len(current)+1)
|
||||||
|
copy(ioBuf, current)
|
||||||
|
ioBuf[len(current)] = next
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
echo := renderRawEcho(ioBuf, mask)
|
||||||
|
_ = stringDisplayWidth(echo)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("fast", func(b *testing.B) {
|
||||||
|
maskWidth := stringDisplayWidth(mask)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_, _, ok := rawEchoRenderUnit(next, mask, maskWidth)
|
||||||
|
if !ok {
|
||||||
|
b.Fatal("fast path rejected append rune")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func benchmarkRawRedrawBackspace(b *testing.B, current []rune, mask string) {
|
||||||
|
if len(current) == 0 {
|
||||||
|
b.Fatal("backspace benchmark requires non-empty input")
|
||||||
|
}
|
||||||
|
last := current[len(current)-1]
|
||||||
|
trimmed := current[:len(current)-1]
|
||||||
|
|
||||||
|
b.Run("legacy", func(b *testing.B) {
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
echo := renderRawEcho(trimmed, mask)
|
||||||
|
_ = stringDisplayWidth(echo)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
b.Run("fast", func(b *testing.B) {
|
||||||
|
maskWidth := stringDisplayWidth(mask)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_, _, ok := rawEchoRenderUnit(last, mask, maskWidth)
|
||||||
|
if !ok {
|
||||||
|
b.Fatal("fast path rejected backspace rune")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkRawRedrawAppendPlainLong(b *testing.B) {
|
||||||
|
current := []rune(strings.Repeat("a", 4096))
|
||||||
|
benchmarkRawRedrawAppend(b, current, 'b', "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkRawRedrawAppendMaskLong(b *testing.B) {
|
||||||
|
current := []rune(strings.Repeat("a", 4096))
|
||||||
|
benchmarkRawRedrawAppend(b, current, 'b', "[]")
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkRawRedrawBackspacePlainLong(b *testing.B) {
|
||||||
|
current := []rune(strings.Repeat("a", 4096))
|
||||||
|
benchmarkRawRedrawBackspace(b, current, "")
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkRawRedrawBackspaceMaskLong(b *testing.B) {
|
||||||
|
current := []rune(strings.Repeat("a", 4096))
|
||||||
|
benchmarkRawRedrawBackspace(b, current, "[]")
|
||||||
|
}
|
||||||
+177
@@ -0,0 +1,177 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
)
|
||||||
|
|
||||||
|
const defaultCopyContextBufferSize = 32 * 1024
|
||||||
|
|
||||||
|
type readContextResult struct {
|
||||||
|
data []byte
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
type writeContextResult struct {
|
||||||
|
n int
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReadFullContext reads exactly len(buf) bytes unless the context is canceled
|
||||||
|
// or the underlying reader returns an error.
|
||||||
|
//
|
||||||
|
// If ctx is canceled while the underlying Read call is blocked, ReadFullContext
|
||||||
|
// returns ctx.Err() without waiting for that call to finish. The underlying
|
||||||
|
// reader may still complete asynchronously afterwards.
|
||||||
|
func ReadFullContext(ctx context.Context, reader io.Reader, buf []byte) (int, error) {
|
||||||
|
if reader == nil {
|
||||||
|
return 0, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
total := 0
|
||||||
|
for total < len(buf) {
|
||||||
|
chunk, err := readContext(ctx, reader, len(buf)-total)
|
||||||
|
if len(chunk) > 0 {
|
||||||
|
total += copy(buf[total:], chunk)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, io.EOF) {
|
||||||
|
if total > 0 {
|
||||||
|
return total, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
return total, io.EOF
|
||||||
|
}
|
||||||
|
return total, err
|
||||||
|
}
|
||||||
|
if len(chunk) == 0 {
|
||||||
|
return total, io.ErrNoProgress
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return total, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFullContext writes the full payload unless the context is canceled or
|
||||||
|
// the underlying writer returns an error.
|
||||||
|
//
|
||||||
|
// If ctx is canceled while the underlying Write call is blocked,
|
||||||
|
// WriteFullContext returns ctx.Err() without waiting for that call to finish.
|
||||||
|
// The underlying writer may still complete asynchronously afterwards.
|
||||||
|
func WriteFullContext(ctx context.Context, writer io.Writer, data []byte) (int, error) {
|
||||||
|
if writer == nil {
|
||||||
|
return 0, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
total := 0
|
||||||
|
for total < len(data) {
|
||||||
|
written, err := writeContext(ctx, writer, data[total:])
|
||||||
|
if written > 0 {
|
||||||
|
total += written
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return total, err
|
||||||
|
}
|
||||||
|
if written == 0 {
|
||||||
|
return total, io.ErrNoProgress
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return total, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CopyContext copies from src to dst until EOF, context cancellation, or a
|
||||||
|
// non-EOF error occurs.
|
||||||
|
//
|
||||||
|
// If ctx is canceled while the current read or write is blocked, CopyContext
|
||||||
|
// returns ctx.Err() without waiting for that operation to finish.
|
||||||
|
func CopyContext(ctx context.Context, dst io.Writer, src io.Reader) (int64, error) {
|
||||||
|
if dst == nil || src == nil {
|
||||||
|
return 0, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
var copied int64
|
||||||
|
for {
|
||||||
|
chunk, readErr := readContext(ctx, src, defaultCopyContextBufferSize)
|
||||||
|
if len(chunk) > 0 {
|
||||||
|
written, writeErr := WriteFullContext(ctx, dst, chunk)
|
||||||
|
copied += int64(written)
|
||||||
|
if writeErr != nil {
|
||||||
|
return copied, writeErr
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if readErr != nil {
|
||||||
|
if errors.Is(readErr, io.EOF) {
|
||||||
|
return copied, nil
|
||||||
|
}
|
||||||
|
return copied, readErr
|
||||||
|
}
|
||||||
|
if len(chunk) == 0 {
|
||||||
|
return copied, io.ErrNoProgress
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func readContext(ctx context.Context, reader io.Reader, size int) ([]byte, error) {
|
||||||
|
if size <= 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
resultCh := make(chan readContextResult, 1)
|
||||||
|
go func() {
|
||||||
|
tmp := make([]byte, size)
|
||||||
|
n, err := reader.Read(tmp)
|
||||||
|
if n > 0 {
|
||||||
|
tmp = tmp[:n]
|
||||||
|
} else {
|
||||||
|
tmp = nil
|
||||||
|
}
|
||||||
|
resultCh <- readContextResult{data: tmp, err: err}
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result.data, result.err
|
||||||
|
case <-ctx.Done():
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result.data, result.err
|
||||||
|
default:
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeContext(ctx context.Context, writer io.Writer, data []byte) (int, error) {
|
||||||
|
if len(data) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
payload := append([]byte(nil), data...)
|
||||||
|
resultCh := make(chan writeContextResult, 1)
|
||||||
|
go func() {
|
||||||
|
n, err := writer.Write(payload)
|
||||||
|
resultCh <- writeContextResult{n: n, err: err}
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result.n, result.err
|
||||||
|
case <-ctx.Done():
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result.n, result.err
|
||||||
|
default:
|
||||||
|
return 0, ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,153 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type chunkedWriter struct {
|
||||||
|
buf bytes.Buffer
|
||||||
|
max int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (writer *chunkedWriter) Write(p []byte) (int, error) {
|
||||||
|
if writer.max <= 0 || len(p) <= writer.max {
|
||||||
|
return writer.buf.Write(p)
|
||||||
|
}
|
||||||
|
return writer.buf.Write(p[:writer.max])
|
||||||
|
}
|
||||||
|
|
||||||
|
type blockingReader struct {
|
||||||
|
started chan struct{}
|
||||||
|
release chan struct{}
|
||||||
|
data []byte
|
||||||
|
once sync.Once
|
||||||
|
}
|
||||||
|
|
||||||
|
func (reader *blockingReader) Read(p []byte) (int, error) {
|
||||||
|
reader.once.Do(func() { close(reader.started) })
|
||||||
|
<-reader.release
|
||||||
|
n := copy(p, reader.data)
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type blockingWriter struct {
|
||||||
|
started chan struct{}
|
||||||
|
release chan struct{}
|
||||||
|
once sync.Once
|
||||||
|
}
|
||||||
|
|
||||||
|
func (writer *blockingWriter) Write(p []byte) (int, error) {
|
||||||
|
writer.once.Do(func() { close(writer.started) })
|
||||||
|
<-writer.release
|
||||||
|
return len(p), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadFullContext(t *testing.T) {
|
||||||
|
buf := make([]byte, 5)
|
||||||
|
n, err := ReadFullContext(context.Background(), bytes.NewBufferString("hello"), buf)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadFullContext returned error: %v", err)
|
||||||
|
}
|
||||||
|
if n != 5 || string(buf) != "hello" {
|
||||||
|
t.Fatalf("unexpected payload: n=%d data=%q", n, buf)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadFullContextReturnsUnexpectedEOF(t *testing.T) {
|
||||||
|
buf := make([]byte, 5)
|
||||||
|
n, err := ReadFullContext(context.Background(), bytes.NewBufferString("hey"), buf)
|
||||||
|
if !errors.Is(err, io.ErrUnexpectedEOF) {
|
||||||
|
t.Fatalf("expected unexpected EOF, got %v", err)
|
||||||
|
}
|
||||||
|
if n != 3 || string(buf[:n]) != "hey" {
|
||||||
|
t.Fatalf("unexpected payload: n=%d data=%q", n, buf[:n])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadFullContextCanceledWhileBlocked(t *testing.T) {
|
||||||
|
reader := &blockingReader{
|
||||||
|
started: make(chan struct{}),
|
||||||
|
release: make(chan struct{}),
|
||||||
|
data: []byte("hello"),
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
buf := make([]byte, 5)
|
||||||
|
_, err := ReadFullContext(ctx, reader, buf)
|
||||||
|
done <- err
|
||||||
|
}()
|
||||||
|
|
||||||
|
<-reader.started
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("expected context canceled, got %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("ReadFullContext did not return after cancel")
|
||||||
|
}
|
||||||
|
|
||||||
|
close(reader.release)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWriteFullContext(t *testing.T) {
|
||||||
|
writer := &chunkedWriter{max: 2}
|
||||||
|
n, err := WriteFullContext(context.Background(), writer, []byte("hello"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("WriteFullContext returned error: %v", err)
|
||||||
|
}
|
||||||
|
if n != 5 || writer.buf.String() != "hello" {
|
||||||
|
t.Fatalf("unexpected write result: n=%d data=%q", n, writer.buf.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWriteFullContextCanceledWhileBlocked(t *testing.T) {
|
||||||
|
writer := &blockingWriter{
|
||||||
|
started: make(chan struct{}),
|
||||||
|
release: make(chan struct{}),
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
_, err := WriteFullContext(ctx, writer, []byte("hello"))
|
||||||
|
done <- err
|
||||||
|
}()
|
||||||
|
|
||||||
|
<-writer.started
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("expected context canceled, got %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("WriteFullContext did not return after cancel")
|
||||||
|
}
|
||||||
|
|
||||||
|
close(writer.release)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCopyContext(t *testing.T) {
|
||||||
|
var dst bytes.Buffer
|
||||||
|
written, err := CopyContext(context.Background(), &dst, bytes.NewBufferString("hello world"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CopyContext returned error: %v", err)
|
||||||
|
}
|
||||||
|
if written != int64(len("hello world")) || dst.String() != "hello world" {
|
||||||
|
t.Fatalf("unexpected copy result: written=%d data=%q", written, dst.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
+234
@@ -0,0 +1,234 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func installRawInputStub(t *testing.T, input string, signalErr error) {
|
||||||
|
t.Helper()
|
||||||
|
prevFactory := rawTerminalSessionFactory
|
||||||
|
prevSignalHandler := inputSignalHandler
|
||||||
|
rawTerminalSessionFactory = func(hint string, printNewline bool) (*rawTerminalSession, error) {
|
||||||
|
return &rawTerminalSession{
|
||||||
|
reader: bufio.NewReader(strings.NewReader(input)),
|
||||||
|
redrawHint: promptRedrawHint(hint),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
inputSignalHandler = func(sigtype rune) error {
|
||||||
|
return signalErr
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
rawTerminalSessionFactory = prevFactory
|
||||||
|
inputSignalHandler = prevSignalHandler
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPromptRedrawHint(t *testing.T) {
|
||||||
|
got := promptRedrawHint("头部提示\n 中文确认: ")
|
||||||
|
if got != "中文确认:" {
|
||||||
|
t.Fatalf("unexpected redraw hint: got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStringDisplayWidth(t *testing.T) {
|
||||||
|
got := stringDisplayWidth("中a[]")
|
||||||
|
if got != 5 {
|
||||||
|
t.Fatalf("unexpected display width: got %d want 5", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRawEchoRenderUnitPlainRune(t *testing.T) {
|
||||||
|
text, width, ok := rawEchoRenderUnit('中', "", 0)
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected plain wide rune to use fast path")
|
||||||
|
}
|
||||||
|
if text != "中" || width != 2 {
|
||||||
|
t.Fatalf("unexpected render unit: got (%q, %d) want (%q, %d)", text, width, "中", 2)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRawEchoRenderUnitMaskedRune(t *testing.T) {
|
||||||
|
text, width, ok := rawEchoRenderUnit('a', "[]", stringDisplayWidth("[]"))
|
||||||
|
if !ok {
|
||||||
|
t.Fatal("expected masked rune to use fast path")
|
||||||
|
}
|
||||||
|
if text != "[]" || width != 2 {
|
||||||
|
t.Fatalf("unexpected render unit: got (%q, %d) want (%q, %d)", text, width, "[]", 2)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRawEchoRenderUnitFallsBackForControlRune(t *testing.T) {
|
||||||
|
if _, _, ok := rawEchoRenderUnit('\x00', "", 0); ok {
|
||||||
|
t.Fatal("expected control rune to fall back to full redraw")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSignalInputResultExitSuppressesError(t *testing.T) {
|
||||||
|
got := signalInputResult(rawInputSignalExit, ErrSignalInterrupt)
|
||||||
|
if got.err != nil {
|
||||||
|
t.Fatalf("expected nil error, got %v", got.err)
|
||||||
|
}
|
||||||
|
if got.msg != "" {
|
||||||
|
t.Fatalf("expected empty message, got %q", got.msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSignalInputResultReturnErrorPreservesSignal(t *testing.T) {
|
||||||
|
got := signalInputResult(rawInputSignalReturnError, ErrSignalInterrupt)
|
||||||
|
if !errors.Is(got.err, ErrSignalInterrupt) {
|
||||||
|
t.Fatalf("expected signal error, got %v", got.err)
|
||||||
|
}
|
||||||
|
if got.msg != "" {
|
||||||
|
t.Fatalf("expected empty message, got %q", got.msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPasswdSuppressesSignalError(t *testing.T) {
|
||||||
|
installRawInputStub(t, string([]rune{0x03}), ErrSignalInterrupt)
|
||||||
|
got := Passwd("", "fallback")
|
||||||
|
if got.err != nil {
|
||||||
|
t.Fatalf("expected nil error, got %v", got.err)
|
||||||
|
}
|
||||||
|
if got.msg != "" {
|
||||||
|
t.Fatalf("expected empty message after signal exit, got %q", got.msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPasswdResponseSignalPreservesSignalError(t *testing.T) {
|
||||||
|
installRawInputStub(t, string([]rune{0x03}), ErrSignalInterrupt)
|
||||||
|
got := PasswdResponseSignal("", "fallback")
|
||||||
|
if !errors.Is(got.err, ErrSignalInterrupt) {
|
||||||
|
t.Fatalf("expected interrupt error, got %v", got.err)
|
||||||
|
}
|
||||||
|
if got.msg != "" {
|
||||||
|
t.Fatalf("expected empty message after signal exit, got %q", got.msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStopUntilEmptyTriggerReturnsAfterFirstKey(t *testing.T) {
|
||||||
|
installRawInputStub(t, "abc", nil)
|
||||||
|
if err := StopUntil("", "", false); err != nil {
|
||||||
|
t.Fatalf("StopUntil returned error: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseYesNoValue(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
input string
|
||||||
|
defaults bool
|
||||||
|
want bool
|
||||||
|
ok bool
|
||||||
|
}{
|
||||||
|
{name: "default", input: " ", defaults: true, want: true, ok: true},
|
||||||
|
{name: "yes", input: "yes", defaults: false, want: true, ok: true},
|
||||||
|
{name: "no", input: "No", defaults: true, want: false, ok: true},
|
||||||
|
{name: "invalid", input: "maybe", defaults: false, want: false, ok: false},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
got, ok := parseYesNoValue(tc.input, tc.defaults)
|
||||||
|
if got != tc.want || ok != tc.ok {
|
||||||
|
t.Fatalf("%s: got (%v, %v) want (%v, %v)", tc.name, got, ok, tc.want, tc.ok)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSliceFloat32(t *testing.T) {
|
||||||
|
data := InputMsg{msg: "1.5,2.25", err: nil}
|
||||||
|
got, err := data.SliceFloat32(",")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("SliceFloat32 returned error: %v", err)
|
||||||
|
}
|
||||||
|
if len(got) != 2 || got[0] != float32(1.5) || got[1] != float32(2.25) {
|
||||||
|
t.Fatalf("unexpected float32 slice: %#v", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTypedSliceParsingTrimsTokenWhitespace(t *testing.T) {
|
||||||
|
ints, err := (InputMsg{msg: "1, 2, 3"}).SliceInt(",")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("SliceInt returned error: %v", err)
|
||||||
|
}
|
||||||
|
if len(ints) != 3 || ints[0] != 1 || ints[1] != 2 || ints[2] != 3 {
|
||||||
|
t.Fatalf("unexpected int slice: %#v", ints)
|
||||||
|
}
|
||||||
|
|
||||||
|
bools, err := (InputMsg{msg: "true, false, true"}).SliceBool(",")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("SliceBool returned error: %v", err)
|
||||||
|
}
|
||||||
|
if len(bools) != 3 || !bools[0] || bools[1] || !bools[2] {
|
||||||
|
t.Fatalf("unexpected bool slice: %#v", bools)
|
||||||
|
}
|
||||||
|
|
||||||
|
float64s, err := (InputMsg{msg: "1.25, 2.5, 3.75"}).SliceFloat64(",")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("SliceFloat64 returned error: %v", err)
|
||||||
|
}
|
||||||
|
if len(float64s) != 3 || float64s[0] != 1.25 || float64s[1] != 2.5 || float64s[2] != 3.75 {
|
||||||
|
t.Fatalf("unexpected float64 slice: %#v", float64s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAdvanceTriggerIndexHandlesOverlap(t *testing.T) {
|
||||||
|
trigger := []rune("aba")
|
||||||
|
prefix := buildTriggerPrefixTable(trigger)
|
||||||
|
index := 0
|
||||||
|
complete := false
|
||||||
|
for _, r := range []rune("aaba") {
|
||||||
|
index, complete = advanceTriggerIndex(trigger, prefix, index, r)
|
||||||
|
if complete {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !complete {
|
||||||
|
t.Fatal("expected overlapped trigger to complete")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func Test_Slice(t *testing.T) {
|
||||||
|
var data = InputMsg{
|
||||||
|
msg: "true,false,true,true,false,0,1,hello",
|
||||||
|
err: nil,
|
||||||
|
skipSliceSigErr: false,
|
||||||
|
}
|
||||||
|
res, err := data.IgnoreSliceParseError(true).SliceBool(",")
|
||||||
|
if err != nil {
|
||||||
|
fmt.Println(res)
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(res) == 0 {
|
||||||
|
t.Fatal(res)
|
||||||
|
}
|
||||||
|
fmt.Println(res)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSliceMsg(t *testing.T) {
|
||||||
|
var data = InputMsg{
|
||||||
|
msg: "",
|
||||||
|
err: nil,
|
||||||
|
skipSliceSigErr: false,
|
||||||
|
}
|
||||||
|
res, err := data.SliceString(",")
|
||||||
|
if err != nil {
|
||||||
|
fmt.Println(res)
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(res) != 0 {
|
||||||
|
t.Fatal(res)
|
||||||
|
}
|
||||||
|
fmt.Println(len(res))
|
||||||
|
res2, err := data.SliceInt64(",")
|
||||||
|
if err != nil {
|
||||||
|
fmt.Println(res2)
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(res2) != 0 {
|
||||||
|
t.Fatal(res2)
|
||||||
|
}
|
||||||
|
fmt.Println(len(res2))
|
||||||
|
}
|
||||||
@@ -0,0 +1,71 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import "io"
|
||||||
|
|
||||||
|
// StarPipeReader is the read side returned by NewStarPipe.
|
||||||
|
type StarPipeReader struct {
|
||||||
|
buf *StarBuffer
|
||||||
|
}
|
||||||
|
|
||||||
|
// StarPipeWriter is the write side returned by NewStarPipe.
|
||||||
|
type StarPipeWriter struct {
|
||||||
|
buf *StarBuffer
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewStarPipe creates a buffered in-memory pipe backed by StarBuffer.
|
||||||
|
//
|
||||||
|
// The writer side uses Close to signal a graceful end-of-stream and Abort to
|
||||||
|
// fail both sides immediately.
|
||||||
|
func NewStarPipe(capacity uint64) (*StarPipeReader, *StarPipeWriter, error) {
|
||||||
|
buf, err := NewStarBuffer(capacity)
|
||||||
|
if err != nil {
|
||||||
|
return nil, nil, err
|
||||||
|
}
|
||||||
|
return &StarPipeReader{buf: buf}, &StarPipeWriter{buf: buf}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Read reads buffered bytes from the pipe.
|
||||||
|
func (reader *StarPipeReader) Read(p []byte) (int, error) {
|
||||||
|
if reader == nil || reader.buf == nil {
|
||||||
|
return 0, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return reader.buf.Read(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close aborts the pipe from the read side and wakes blocked writers.
|
||||||
|
func (reader *StarPipeReader) Close() error {
|
||||||
|
if reader == nil || reader.buf == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return reader.buf.Abort()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Abort aborts the pipe from the read side and wakes blocked writers.
|
||||||
|
func (reader *StarPipeReader) Abort() error {
|
||||||
|
return reader.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Write writes bytes into the pipe buffer.
|
||||||
|
func (writer *StarPipeWriter) Write(p []byte) (int, error) {
|
||||||
|
if writer == nil || writer.buf == nil {
|
||||||
|
return 0, io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.buf.Write(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close gracefully closes the write side. Buffered bytes remain readable until
|
||||||
|
// drained.
|
||||||
|
func (writer *StarPipeWriter) Close() error {
|
||||||
|
if writer == nil || writer.buf == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.buf.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Abort aborts the pipe immediately.
|
||||||
|
func (writer *StarPipeWriter) Abort() error {
|
||||||
|
if writer == nil || writer.buf == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
return writer.buf.Abort()
|
||||||
|
}
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestStarPipeRoundTrip(t *testing.T) {
|
||||||
|
reader, writer, err := NewStarPipe(16)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
want := []byte("hello world")
|
||||||
|
go func() {
|
||||||
|
_, _ = writer.Write(want)
|
||||||
|
_ = writer.Close()
|
||||||
|
}()
|
||||||
|
|
||||||
|
got, err := io.ReadAll(reader)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadAll failed: %v", err)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("unexpected payload: got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarPipeReaderCloseAbortsWriter(t *testing.T) {
|
||||||
|
reader, writer, err := NewStarPipe(4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := reader.Close(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := writer.Write([]byte("x")); !errors.Is(err, ErrStarBufferClosed) {
|
||||||
|
t.Fatalf("unexpected writer error after reader close: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarPipeNilEndsReportClosedPipe(t *testing.T) {
|
||||||
|
var reader *StarPipeReader
|
||||||
|
var writer *StarPipeWriter
|
||||||
|
|
||||||
|
if _, err := reader.Read(make([]byte, 1)); !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("unexpected nil reader error: %v", err)
|
||||||
|
}
|
||||||
|
if _, err := writer.Write([]byte("x")); !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("unexpected nil writer error: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func BenchmarkQueueBuildMessage64KiB(b *testing.B) {
|
||||||
|
que := NewQueue()
|
||||||
|
payload := make([]byte, 64*1024)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(payload)))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
_ = que.BuildMessage(payload)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkQueueWriteFrame64KiB(b *testing.B) {
|
||||||
|
que := NewQueue()
|
||||||
|
payload := make([]byte, 64*1024)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(payload)))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if err := que.WriteFrame(io.Discard, payload); err != nil {
|
||||||
|
b.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkQueueWriteFrameBuffers64KiB(b *testing.B) {
|
||||||
|
que := NewQueue()
|
||||||
|
payload := make([]byte, 64*1024)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(payload)))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if err := que.WriteFrameBuffers(io.Discard, payload); err != nil {
|
||||||
|
b.Fatalf("WriteFrameBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkQueueWriteFramesBuffers4x64KiB(b *testing.B) {
|
||||||
|
que := NewQueue()
|
||||||
|
payloads := [][]byte{
|
||||||
|
make([]byte, 64*1024),
|
||||||
|
make([]byte, 64*1024),
|
||||||
|
make([]byte, 64*1024),
|
||||||
|
make([]byte, 64*1024),
|
||||||
|
}
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(payloads) * len(payloads[0])))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if err := que.WriteFramesBuffers(io.Discard, payloads...); err != nil {
|
||||||
|
b.Fatalf("WriteFramesBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkQueueParseRestoreHello(b *testing.B) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
if len(frame) == 0 {
|
||||||
|
b.Fatal("BuildMessage returned empty frame")
|
||||||
|
}
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.SetBytes(int64(len(frame)))
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if err := que.ParseMessage(frame, "bench"); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := que.Restore(); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import "encoding/binary"
|
||||||
|
|
||||||
|
// Uint32ToByte 4位uint32转byte。
|
||||||
|
func Uint32ToByte(src uint32) []byte {
|
||||||
|
res := make([]byte, 4)
|
||||||
|
binary.BigEndian.PutUint32(res, src)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
|
// ByteToUint32 byte转4位uint32。
|
||||||
|
func ByteToUint32(src []byte) uint32 {
|
||||||
|
return binary.BigEndian.Uint32(src)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Uint16ToByte 2位uint16转byte。
|
||||||
|
func Uint16ToByte(src uint16) []byte {
|
||||||
|
res := make([]byte, 2)
|
||||||
|
binary.BigEndian.PutUint16(res, src)
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
|
// ByteToUint16 用于byte转uint16。
|
||||||
|
func ByteToUint16(src []byte) uint16 {
|
||||||
|
return binary.BigEndian.Uint16(src)
|
||||||
|
}
|
||||||
|
|
||||||
|
func cloneBytes(src []byte) []byte {
|
||||||
|
if len(src) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return append([]byte(nil), src...)
|
||||||
|
}
|
||||||
+208
@@ -0,0 +1,208 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/binary"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BuildMessage builds one full frame and panics if the payload is too large to
|
||||||
|
// fit in the framing format.
|
||||||
|
//
|
||||||
|
// New code should prefer BuildMessageErr when it needs an explicit error path.
|
||||||
|
func (q *StarQueue) BuildMessage(src []byte) []byte {
|
||||||
|
frame, err := q.BuildMessageErr(src)
|
||||||
|
if err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
return frame
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildMessageErr builds one full frame and returns an explicit error when the
|
||||||
|
// payload exceeds the 32-bit frame length.
|
||||||
|
func (q *StarQueue) BuildMessageErr(src []byte) ([]byte, error) {
|
||||||
|
payload := src
|
||||||
|
if q.Encode && q.EncodeFunc != nil {
|
||||||
|
payload = q.EncodeFunc(payload)
|
||||||
|
}
|
||||||
|
length, err := payloadSizeToUint32(uint64(len(payload)))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
header := q.BuildHeader(length)
|
||||||
|
frame := make([]byte, 0, len(header)+len(payload))
|
||||||
|
frame = append(frame, header...)
|
||||||
|
frame = append(frame, payload...)
|
||||||
|
return frame, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFrame writes one framed payload directly to w without building an
|
||||||
|
// intermediate full-frame slice first.
|
||||||
|
func (q *StarQueue) WriteFrame(w io.Writer, src []byte) error {
|
||||||
|
if w == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
payload := src
|
||||||
|
if q.Encode && q.EncodeFunc != nil {
|
||||||
|
payload = q.EncodeFunc(payload)
|
||||||
|
}
|
||||||
|
length, err := payloadSizeToUint32(uint64(len(payload)))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var header [queHeaderSize]byte
|
||||||
|
writeHeaderBytes(header[:], queHeader{
|
||||||
|
Length: length,
|
||||||
|
Version: queVersionV1,
|
||||||
|
Flags: queSupportedFlags,
|
||||||
|
})
|
||||||
|
if err := writeFull(w, header[:]); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return writeFull(w, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFrameBuffers writes one framed payload using net.Buffers so callers can
|
||||||
|
// opt into gather writes when the underlying writer supports it well.
|
||||||
|
func (q *StarQueue) WriteFrameBuffers(w io.Writer, src []byte) error {
|
||||||
|
if w == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
payload := src
|
||||||
|
if q.Encode && q.EncodeFunc != nil {
|
||||||
|
payload = q.EncodeFunc(payload)
|
||||||
|
}
|
||||||
|
length, err := payloadSizeToUint32(uint64(len(payload)))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var header [queHeaderSize]byte
|
||||||
|
writeHeaderBytes(header[:], queHeader{
|
||||||
|
Length: length,
|
||||||
|
Version: queVersionV1,
|
||||||
|
Flags: queSupportedFlags,
|
||||||
|
})
|
||||||
|
buffers := net.Buffers{header[:], payload}
|
||||||
|
_, err = buffers.WriteTo(w)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// WriteFramesBuffers writes multiple framed payloads using one net.Buffers
|
||||||
|
// batch so callers can reduce write calls on stream transports.
|
||||||
|
func (q *StarQueue) WriteFramesBuffers(w io.Writer, payloads ...[]byte) error {
|
||||||
|
if w == nil {
|
||||||
|
return io.ErrClosedPipe
|
||||||
|
}
|
||||||
|
if len(payloads) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
buffers := make(net.Buffers, 0, len(payloads)*2)
|
||||||
|
headers := make([][queHeaderSize]byte, len(payloads))
|
||||||
|
for i, src := range payloads {
|
||||||
|
payload := src
|
||||||
|
if q.Encode && q.EncodeFunc != nil {
|
||||||
|
payload = q.EncodeFunc(payload)
|
||||||
|
}
|
||||||
|
length, err := payloadSizeToUint32(uint64(len(payload)))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
writeHeaderBytes(headers[i][:], queHeader{
|
||||||
|
Length: length,
|
||||||
|
Version: queVersionV1,
|
||||||
|
Flags: queSupportedFlags,
|
||||||
|
})
|
||||||
|
buffers = append(buffers, headers[i][:], payload)
|
||||||
|
}
|
||||||
|
_, err := buffers.WriteTo(w)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildHeader 生成编码后的Header用于发送。
|
||||||
|
func (q *StarQueue) BuildHeader(length uint32) []byte {
|
||||||
|
return buildHeaderBytes(queHeader{
|
||||||
|
Length: length,
|
||||||
|
Version: queVersionV1,
|
||||||
|
Flags: queSupportedFlags,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildHeaderBytes(header queHeader) []byte {
|
||||||
|
buf := make([]byte, queHeaderSize)
|
||||||
|
writeHeaderBytes(buf, header)
|
||||||
|
return buf
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeHeaderBytes(dst []byte, header queHeader) {
|
||||||
|
if len(dst) < queHeaderSize {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
copy(dst[:queMagicSize], queMagic)
|
||||||
|
binary.BigEndian.PutUint32(dst[queMagicSize:queMagicSize+4], header.Length)
|
||||||
|
dst[12] = header.Version
|
||||||
|
dst[13] = header.Flags
|
||||||
|
}
|
||||||
|
|
||||||
|
func payloadSizeToUint32(size uint64) (uint32, error) {
|
||||||
|
const maxFramePayload = ^uint32(0)
|
||||||
|
if size > uint64(maxFramePayload) {
|
||||||
|
return 0, fmt.Errorf("%w: %d > %d", ErrQueueMessageTooLarge, size, uint64(maxFramePayload))
|
||||||
|
}
|
||||||
|
return uint32(size), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeFull(w io.Writer, data []byte) error {
|
||||||
|
for len(data) > 0 {
|
||||||
|
n, err := w.Write(data)
|
||||||
|
if n > 0 {
|
||||||
|
data = data[n:]
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if n == 0 {
|
||||||
|
return io.ErrNoProgress
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseHeaderBytes(src []byte, maxLength uint32) (queHeader, error) {
|
||||||
|
if len(src) < queHeaderSize {
|
||||||
|
return queHeader{}, ErrQueueDataFormat
|
||||||
|
}
|
||||||
|
if !equalMagic(src[:queMagicSize]) {
|
||||||
|
return queHeader{}, ErrQueueDataFormat
|
||||||
|
}
|
||||||
|
|
||||||
|
header := queHeader{
|
||||||
|
Length: ByteToUint32(src[queMagicSize : queMagicSize+4]),
|
||||||
|
Version: src[12],
|
||||||
|
Flags: src[13],
|
||||||
|
}
|
||||||
|
|
||||||
|
if header.Version != queVersionV1 {
|
||||||
|
return queHeader{}, fmt.Errorf("%w: %d", ErrQueueUnsupportedVersion, header.Version)
|
||||||
|
}
|
||||||
|
if header.Flags != queSupportedFlags {
|
||||||
|
return queHeader{}, fmt.Errorf("%w: %d", ErrQueueUnsupportedFlags, header.Flags)
|
||||||
|
}
|
||||||
|
if maxLength != 0 && header.Length > maxLength {
|
||||||
|
return queHeader{}, fmt.Errorf("%w: %d > %d", ErrQueueMessageTooLarge, header.Length, maxLength)
|
||||||
|
}
|
||||||
|
|
||||||
|
return header, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func equalMagic(src []byte) bool {
|
||||||
|
if len(src) != queMagicSize {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for i, b := range src {
|
||||||
|
if b != queMagic[i] {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
+237
@@ -0,0 +1,237 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"fmt"
|
||||||
|
"reflect"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ParseMessage 用于解析收到的msg信息。
|
||||||
|
func (q *StarQueue) ParseMessage(msg []byte, conn interface{}) error {
|
||||||
|
return q.parseMessage(msg, conn, false, func(payload []byte) error {
|
||||||
|
return q.push2list(MsgQueue{
|
||||||
|
Msg: payload,
|
||||||
|
Conn: conn,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseMessageView parses frames and exposes each payload to fn without
|
||||||
|
// forcing StarQueue to clone it first.
|
||||||
|
//
|
||||||
|
// The provided payload is only valid during the callback. If q uses the legacy
|
||||||
|
// DecodeFunc path, StarQueue still has to allocate a decoded payload first.
|
||||||
|
func (q *StarQueue) ParseMessageView(msg []byte, conn interface{}, fn func(FrameView) error) error {
|
||||||
|
if fn == nil {
|
||||||
|
return ErrQueueFrameHandlerNil
|
||||||
|
}
|
||||||
|
return q.parseMessage(msg, conn, true, func(payload []byte) error {
|
||||||
|
return fn(FrameView{
|
||||||
|
Payload: payload,
|
||||||
|
Conn: conn,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseMessageOwned parses frames and emits owned payload copies to fn without
|
||||||
|
// routing them through RestoreChan.
|
||||||
|
//
|
||||||
|
// Compared with ParseMessage, this keeps StarQueue state handling the same but
|
||||||
|
// lets callers decide how to dispatch parsed messages themselves.
|
||||||
|
func (q *StarQueue) ParseMessageOwned(msg []byte, conn interface{}, fn func(MsgQueue) error) error {
|
||||||
|
if fn == nil {
|
||||||
|
return ErrQueueFrameHandlerNil
|
||||||
|
}
|
||||||
|
frames := make([]MsgQueue, 0, 1)
|
||||||
|
parseErr := q.parseMessage(msg, conn, false, func(payload []byte) error {
|
||||||
|
frames = append(frames, MsgQueue{
|
||||||
|
Msg: payload,
|
||||||
|
Conn: conn,
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
for _, frame := range frames {
|
||||||
|
if err := fn(frame); err != nil {
|
||||||
|
if parseErr != nil {
|
||||||
|
return fmt.Errorf("%v: %w", parseErr, err)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return parseErr
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *StarQueue) parseMessage(msg []byte, conn interface{}, borrowPayload bool, emit func([]byte) error) error {
|
||||||
|
state, err := q.connState(conn)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var firstErr error
|
||||||
|
for {
|
||||||
|
payload, ok, err := q.nextPayload(state, conn, msg, borrowPayload)
|
||||||
|
msg = nil
|
||||||
|
if err != nil && firstErr == nil {
|
||||||
|
firstErr = err
|
||||||
|
}
|
||||||
|
if !ok {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if err := emit(payload); err != nil {
|
||||||
|
if firstErr != nil {
|
||||||
|
return fmt.Errorf("%v: %w", firstErr, err)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return firstErr
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *StarQueue) nextPayload(state *queConnState, conn interface{}, msg []byte, borrowPayload bool) ([]byte, bool, error) {
|
||||||
|
state.mu.Lock()
|
||||||
|
defer state.mu.Unlock()
|
||||||
|
|
||||||
|
if len(msg) != 0 {
|
||||||
|
state.buf = append(state.buf, msg...)
|
||||||
|
}
|
||||||
|
|
||||||
|
var firstErr error
|
||||||
|
for {
|
||||||
|
synced, err := syncFrameStart(&state.buf)
|
||||||
|
if err != nil && firstErr == nil {
|
||||||
|
firstErr = err
|
||||||
|
}
|
||||||
|
if !synced {
|
||||||
|
if len(state.buf) == 0 {
|
||||||
|
q.states.Delete(conn)
|
||||||
|
}
|
||||||
|
return nil, false, firstErr
|
||||||
|
}
|
||||||
|
if len(state.buf) < queHeaderSize {
|
||||||
|
return nil, false, firstErr
|
||||||
|
}
|
||||||
|
|
||||||
|
header, err := parseHeaderBytes(state.buf[:queHeaderSize], q.maxLength)
|
||||||
|
if err != nil {
|
||||||
|
if firstErr == nil {
|
||||||
|
firstErr = err
|
||||||
|
}
|
||||||
|
state.buf = shrinkBuffer(state.buf[1:])
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
frameLen := queHeaderSize + int(header.Length)
|
||||||
|
if len(state.buf) < frameLen {
|
||||||
|
return nil, false, firstErr
|
||||||
|
}
|
||||||
|
|
||||||
|
payload, rest := extractPayload(state.buf, frameLen, borrowPayload && !(q.Encode && q.DecodeFunc != nil))
|
||||||
|
state.buf = rest
|
||||||
|
if q.Encode && q.DecodeFunc != nil {
|
||||||
|
payload = q.DecodeFunc(payload)
|
||||||
|
}
|
||||||
|
if len(state.buf) == 0 {
|
||||||
|
q.states.Delete(conn)
|
||||||
|
}
|
||||||
|
return payload, true, firstErr
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func extractPayload(buf []byte, frameLen int, borrowPayload bool) ([]byte, []byte) {
|
||||||
|
payload := buf[queHeaderSize:frameLen]
|
||||||
|
if !borrowPayload {
|
||||||
|
return cloneBytes(payload), shrinkBuffer(buf[frameLen:])
|
||||||
|
}
|
||||||
|
if frameLen == len(buf) {
|
||||||
|
return payload, nil
|
||||||
|
}
|
||||||
|
return payload, cloneBytes(buf[frameLen:])
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *StarQueue) push2list(msg MsgQueue) error {
|
||||||
|
select {
|
||||||
|
case <-q.ctx.Done():
|
||||||
|
return q.ctx.Err()
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
q.sendMu.RLock()
|
||||||
|
defer q.sendMu.RUnlock()
|
||||||
|
select {
|
||||||
|
case <-q.ctx.Done():
|
||||||
|
return q.ctx.Err()
|
||||||
|
case q.msgPool <- msg:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func validateConnKey(conn interface{}) error {
|
||||||
|
if conn == nil {
|
||||||
|
return ErrQueueConnKeyNil
|
||||||
|
}
|
||||||
|
typ := reflect.TypeOf(conn)
|
||||||
|
if typ != nil && !typ.Comparable() {
|
||||||
|
return ErrQueueConnKeyInvalid
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *StarQueue) connState(conn interface{}) (*queConnState, error) {
|
||||||
|
if err := validateConnKey(conn); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
state, _ := q.states.LoadOrStore(conn, &queConnState{})
|
||||||
|
return state.(*queConnState), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func syncFrameStart(buf *[]byte) (bool, error) {
|
||||||
|
if len(*buf) == 0 {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(*buf) >= queMagicSize && equalMagic((*buf)[:queMagicSize]) {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
idx := bytes.Index(*buf, queMagic)
|
||||||
|
if idx == 0 {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
if idx > 0 {
|
||||||
|
*buf = cloneBytes((*buf)[idx:])
|
||||||
|
return true, ErrQueueDataFormat
|
||||||
|
}
|
||||||
|
|
||||||
|
keep := trailingMagicPrefixLen(*buf)
|
||||||
|
if keep == len(*buf) {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if keep > 0 {
|
||||||
|
*buf = cloneBytes((*buf)[len(*buf)-keep:])
|
||||||
|
return false, ErrQueueDataFormat
|
||||||
|
}
|
||||||
|
|
||||||
|
*buf = (*buf)[:0]
|
||||||
|
return false, ErrQueueDataFormat
|
||||||
|
}
|
||||||
|
|
||||||
|
func trailingMagicPrefixLen(buf []byte) int {
|
||||||
|
max := len(buf)
|
||||||
|
if max > queMagicSize-1 {
|
||||||
|
max = queMagicSize - 1
|
||||||
|
}
|
||||||
|
for keep := max; keep > 0; keep-- {
|
||||||
|
if bytes.Equal(buf[len(buf)-keep:], queMagic[:keep]) {
|
||||||
|
return keep
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
|
||||||
|
func shrinkBuffer(buf []byte) []byte {
|
||||||
|
if len(buf) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if cap(buf) > len(buf)*4 {
|
||||||
|
return cloneBytes(buf)
|
||||||
|
}
|
||||||
|
return buf
|
||||||
|
}
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (q *StarQueue) closedRestoreErr() error {
|
||||||
|
if err := q.ctx.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return os.ErrClosed
|
||||||
|
}
|
||||||
|
|
||||||
|
// Restore blocks until one message is available, the queue is stopped, or the
|
||||||
|
// configured timeout expires.
|
||||||
|
func (q *StarQueue) Restore() (MsgQueue, error) {
|
||||||
|
if q.duration <= 0 {
|
||||||
|
select {
|
||||||
|
case <-q.ctx.Done():
|
||||||
|
return MsgQueue{}, q.ctx.Err()
|
||||||
|
case data, ok := <-q.msgPool:
|
||||||
|
if !ok {
|
||||||
|
return MsgQueue{}, q.closedRestoreErr()
|
||||||
|
}
|
||||||
|
return data, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
timer := time.NewTimer(q.duration)
|
||||||
|
defer timer.Stop()
|
||||||
|
select {
|
||||||
|
case <-q.ctx.Done():
|
||||||
|
return MsgQueue{}, q.ctx.Err()
|
||||||
|
case <-timer.C:
|
||||||
|
return MsgQueue{}, ErrDeadlineExceeded
|
||||||
|
case data, ok := <-q.msgPool:
|
||||||
|
if !ok {
|
||||||
|
return MsgQueue{}, q.closedRestoreErr()
|
||||||
|
}
|
||||||
|
return data, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RestoreOne 获取收到的一个信息。
|
||||||
|
// 兼容性修改。
|
||||||
|
func (q *StarQueue) RestoreOne() (MsgQueue, error) {
|
||||||
|
return q.Restore()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop cancels the queue runtime.
|
||||||
|
//
|
||||||
|
// After Stop returns, Restore unblocks with context.Canceled and RestoreChan is
|
||||||
|
// eventually closed.
|
||||||
|
func (q *StarQueue) Stop() {
|
||||||
|
q.cancel()
|
||||||
|
q.shutdown()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (q *StarQueue) shutdown() {
|
||||||
|
q.stopOnce.Do(func() {
|
||||||
|
q.sendMu.Lock()
|
||||||
|
close(q.msgPool)
|
||||||
|
q.sendMu.Unlock()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// RestoreDuration sets the Restore timeout. A non-positive duration means wait
|
||||||
|
// forever until a message arrives or Stop is called.
|
||||||
|
func (q *StarQueue) RestoreDuration(tm time.Duration) {
|
||||||
|
q.duration = tm
|
||||||
|
}
|
||||||
|
|
||||||
|
// RestoreChan exposes the parsed message stream.
|
||||||
|
//
|
||||||
|
// The returned channel is closed after Stop or when the queue context is
|
||||||
|
// canceled. New code should still prefer Restore when it needs timeout/error
|
||||||
|
// classification instead of a plain stream.
|
||||||
|
func (q *StarQueue) RestoreChan() <-chan MsgQueue {
|
||||||
|
return q.msgPool
|
||||||
|
}
|
||||||
+520
@@ -0,0 +1,520 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestQueueBuildMessageUsesVersionedHeader(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
if len(frame) != queHeaderSize+5 {
|
||||||
|
t.Fatalf("unexpected frame length: got %d want %d", len(frame), queHeaderSize+5)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(frame[:queMagicSize], queMagic) {
|
||||||
|
t.Fatalf("unexpected magic: %v", frame[:queMagicSize])
|
||||||
|
}
|
||||||
|
if got := ByteToUint32(frame[queMagicSize : queMagicSize+4]); got != 5 {
|
||||||
|
t.Fatalf("unexpected payload length: got %d want 5", got)
|
||||||
|
}
|
||||||
|
if frame[12] != queVersionV1 {
|
||||||
|
t.Fatalf("unexpected version: got %d want %d", frame[12], queVersionV1)
|
||||||
|
}
|
||||||
|
if frame[13] != queSupportedFlags {
|
||||||
|
t.Fatalf("unexpected flags: got %d want %d", frame[13], queSupportedFlags)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(frame[queHeaderSize:], []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload: %q", frame[queHeaderSize:])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameMatchesBuildMessage(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
want := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFrame(&buf, []byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFrame mismatch: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameBuffersMatchesBuildMessage(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
want := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFrameBuffers(&buf, []byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrameBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFrameBuffers mismatch: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFramesBuffersMatchesBuildMessage(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
payloads := [][]byte{
|
||||||
|
[]byte("hello"),
|
||||||
|
[]byte("world"),
|
||||||
|
[]byte("batch"),
|
||||||
|
}
|
||||||
|
|
||||||
|
var want []byte
|
||||||
|
for _, payload := range payloads {
|
||||||
|
want = append(want, que.BuildMessage(payload)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFramesBuffers(&buf, payloads...); err != nil {
|
||||||
|
t.Fatalf("WriteFramesBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFramesBuffers mismatch: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameHonorsEncodeFunc(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
que.Encode = true
|
||||||
|
que.EncodeFunc = bytes.ToUpper
|
||||||
|
want := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFrame(&buf, []byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrame failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFrame mismatch with EncodeFunc: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameBuffersHonorsEncodeFunc(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
que.Encode = true
|
||||||
|
que.EncodeFunc = bytes.ToUpper
|
||||||
|
want := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFrameBuffers(&buf, []byte("hello")); err != nil {
|
||||||
|
t.Fatalf("WriteFrameBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFrameBuffers mismatch with EncodeFunc: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFramesBuffersHonorsEncodeFunc(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
que.Encode = true
|
||||||
|
que.EncodeFunc = bytes.ToUpper
|
||||||
|
payloads := [][]byte{
|
||||||
|
[]byte("hello"),
|
||||||
|
[]byte("batch"),
|
||||||
|
}
|
||||||
|
|
||||||
|
var want []byte
|
||||||
|
for _, payload := range payloads {
|
||||||
|
want = append(want, que.BuildMessage(payload)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
if err := que.WriteFramesBuffers(&buf, payloads...); err != nil {
|
||||||
|
t.Fatalf("WriteFramesBuffers failed: %v", err)
|
||||||
|
}
|
||||||
|
if got := buf.Bytes(); !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("WriteFramesBuffers mismatch with EncodeFunc: got %v want %v", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageSplitAcrossCalls(t *testing.T) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
if err := que.ParseMessage(frame[:3], "split"); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial magic: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessage(frame[3:11], "split"); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial header: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessage(frame[11:], "split"); err != nil {
|
||||||
|
t.Fatalf("unexpected error on payload completion: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case data := <-que.RestoreChan():
|
||||||
|
if data.ID != 0 {
|
||||||
|
t.Fatalf("expected deprecated frame ID to stay zero, got %d", data.ID)
|
||||||
|
}
|
||||||
|
if data.Conn != "split" {
|
||||||
|
t.Fatalf("unexpected conn: %#v", data.Conn)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(data.Msg, []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload: %q", data.Msg)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("did not restore parsed frame")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageViewSplitAcrossCalls(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
var got [][]byte
|
||||||
|
|
||||||
|
handler := func(view FrameView) error {
|
||||||
|
got = append(got, cloneBytes(view.Payload))
|
||||||
|
if view.Conn != "split-view" {
|
||||||
|
t.Fatalf("unexpected conn: %#v", view.Conn)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := que.ParseMessageView(frame[:3], "split-view", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial magic: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessageView(frame[3:11], "split-view", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial header: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessageView(frame[11:], "split-view", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on payload completion: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(got) != 1 {
|
||||||
|
t.Fatalf("unexpected frame count: got %d want 1", len(got))
|
||||||
|
}
|
||||||
|
if !bytes.Equal(got[0], []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload: %q", got[0])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageOwnedSplitAcrossCalls(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
var got []MsgQueue
|
||||||
|
|
||||||
|
handler := func(msg MsgQueue) error {
|
||||||
|
got = append(got, MsgQueue{
|
||||||
|
Msg: cloneBytes(msg.Msg),
|
||||||
|
Conn: msg.Conn,
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := que.ParseMessageOwned(frame[:3], "split-owned", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial magic: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessageOwned(frame[3:11], "split-owned", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on partial header: %v", err)
|
||||||
|
}
|
||||||
|
if err := que.ParseMessageOwned(frame[11:], "split-owned", handler); err != nil {
|
||||||
|
t.Fatalf("unexpected error on payload completion: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(got) != 1 {
|
||||||
|
t.Fatalf("unexpected frame count: got %d want 1", len(got))
|
||||||
|
}
|
||||||
|
if got[0].Conn != "split-owned" {
|
||||||
|
t.Fatalf("unexpected conn: %#v", got[0].Conn)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(got[0].Msg, []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload: %q", got[0].Msg)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case msg := <-que.RestoreChan():
|
||||||
|
t.Fatalf("ParseMessageOwned should not use RestoreChan, got %#v", msg)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageSkipsGarbagePrefix(t *testing.T) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
err := que.ParseMessage(append([]byte("junk"), frame...), "garbage")
|
||||||
|
if !errors.Is(err, ErrQueueDataFormat) {
|
||||||
|
t.Fatalf("expected data format error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case data := <-que.RestoreChan():
|
||||||
|
if !bytes.Equal(data.Msg, []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload after resync: %q", data.Msg)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("did not restore frame after skipping garbage")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageViewSkipsGarbagePrefix(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
var got []byte
|
||||||
|
|
||||||
|
err := que.ParseMessageView(append([]byte("junk"), frame...), "garbage-view", func(view FrameView) error {
|
||||||
|
got = cloneBytes(view.Payload)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
if !errors.Is(err, ErrQueueDataFormat) {
|
||||||
|
t.Fatalf("expected data format error, got %v", err)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(got, []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected payload after resync: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageViewNilHandler(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
err := que.ParseMessageView(que.BuildMessage([]byte("hello")), "nil-handler", nil)
|
||||||
|
if !errors.Is(err, ErrQueueFrameHandlerNil) {
|
||||||
|
t.Fatalf("ParseMessageView error = %v, want %v", err, ErrQueueFrameHandlerNil)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageOwnedNilHandler(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
err := que.ParseMessageOwned(que.BuildMessage([]byte("hello")), "nil-handler-owned", nil)
|
||||||
|
if !errors.Is(err, ErrQueueFrameHandlerNil) {
|
||||||
|
t.Fatalf("ParseMessageOwned error = %v, want %v", err, ErrQueueFrameHandlerNil)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameNilWriter(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
err := que.WriteFrame(nil, []byte("hello"))
|
||||||
|
if !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("WriteFrame error = %v, want %v", err, io.ErrClosedPipe)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFrameBuffersNilWriter(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
err := que.WriteFrameBuffers(nil, []byte("hello"))
|
||||||
|
if !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("WriteFrameBuffers error = %v, want %v", err, io.ErrClosedPipe)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueWriteFramesBuffersNilWriter(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
err := que.WriteFramesBuffers(nil, []byte("hello"))
|
||||||
|
if !errors.Is(err, io.ErrClosedPipe) {
|
||||||
|
t.Fatalf("WriteFramesBuffers error = %v, want %v", err, io.ErrClosedPipe)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageRejectsUnsupportedVersion(t *testing.T) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
frame[12] = 2
|
||||||
|
|
||||||
|
err := que.ParseMessage(frame, "version")
|
||||||
|
if !errors.Is(err, ErrQueueUnsupportedVersion) {
|
||||||
|
t.Fatalf("expected unsupported version error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case data := <-que.RestoreChan():
|
||||||
|
t.Fatalf("unexpected restored frame: %#v", data)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageRejectsMessageTooLarge(t *testing.T) {
|
||||||
|
que := NewQueueCtx(nil, 1, 4)
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
|
||||||
|
err := que.ParseMessage(frame, "large")
|
||||||
|
if !errors.Is(err, ErrQueueMessageTooLarge) {
|
||||||
|
t.Fatalf("expected message too large error, got %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case data := <-que.RestoreChan():
|
||||||
|
t.Fatalf("unexpected restored frame: %#v", data)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageRejectsInvalidConnKey(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
err := que.ParseMessage(frame, []byte("not-comparable"))
|
||||||
|
if !errors.Is(err, ErrQueueConnKeyInvalid) {
|
||||||
|
t.Fatalf("expected invalid conn key error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageRejectsNilConnKey(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
frame := que.BuildMessage([]byte("hello"))
|
||||||
|
err := que.ParseMessage(frame, nil)
|
||||||
|
if !errors.Is(err, ErrQueueConnKeyNil) {
|
||||||
|
t.Fatalf("expected nil conn key error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueuePayloadSizeToUint32RejectsOverflow(t *testing.T) {
|
||||||
|
_, err := payloadSizeToUint32(uint64(^uint32(0)) + 1)
|
||||||
|
if !errors.Is(err, ErrQueueMessageTooLarge) {
|
||||||
|
t.Fatalf("expected message too large error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueRestoreDurationZeroWaitsUntilMessage(t *testing.T) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
que.RestoreDuration(0)
|
||||||
|
|
||||||
|
type restoreResult struct {
|
||||||
|
msg MsgQueue
|
||||||
|
err error
|
||||||
|
}
|
||||||
|
resultCh := make(chan restoreResult, 1)
|
||||||
|
go func() {
|
||||||
|
msg, err := que.Restore()
|
||||||
|
resultCh <- restoreResult{msg: msg, err: err}
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
t.Fatalf("Restore returned too early: %#v", result)
|
||||||
|
case <-time.After(50 * time.Millisecond):
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := que.ParseMessage(que.BuildMessage([]byte("hello")), "forever"); err != nil {
|
||||||
|
t.Fatalf("ParseMessage failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
if result.err != nil {
|
||||||
|
t.Fatalf("Restore returned error: %v", result.err)
|
||||||
|
}
|
||||||
|
if result.msg.Conn != "forever" || !bytes.Equal(result.msg.Msg, []byte("hello")) {
|
||||||
|
t.Fatalf("unexpected restore result: %#v", result.msg)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Restore did not return after message arrival")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueRestoreReturnsContextErrorOnStop(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
resultCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
_, err := que.Restore()
|
||||||
|
resultCh <- err
|
||||||
|
}()
|
||||||
|
|
||||||
|
que.Stop()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-resultCh:
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("expected context.Canceled, got %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Restore did not return after Stop")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueRestoreChanClosesOnStop(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
resultCh := make(chan bool, 1)
|
||||||
|
go func() {
|
||||||
|
_, ok := <-que.RestoreChan()
|
||||||
|
resultCh <- ok
|
||||||
|
}()
|
||||||
|
|
||||||
|
que.Stop()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case ok := <-resultCh:
|
||||||
|
if ok {
|
||||||
|
t.Fatal("expected RestoreChan to close after Stop")
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("RestoreChan did not close after Stop")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueRestoreChanClosesOnContextCancel(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
que := NewQueueCtx(ctx, 1, 0)
|
||||||
|
resultCh := make(chan bool, 1)
|
||||||
|
go func() {
|
||||||
|
_, ok := <-que.RestoreChan()
|
||||||
|
resultCh <- ok
|
||||||
|
}()
|
||||||
|
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case ok := <-resultCh:
|
||||||
|
if ok {
|
||||||
|
t.Fatal("expected RestoreChan to close after context cancel")
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("RestoreChan did not close after context cancel")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageReturnsContextErrorWhenStoppedWhilePoolIsFull(t *testing.T) {
|
||||||
|
que := NewQueueWithCount(1)
|
||||||
|
if err := que.ParseMessage(que.BuildMessage([]byte("first")), "full"); err != nil {
|
||||||
|
t.Fatalf("ParseMessage first failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
errCh := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
errCh <- que.ParseMessage(que.BuildMessage([]byte("second")), "full")
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-errCh:
|
||||||
|
t.Fatalf("ParseMessage returned before Stop: %v", err)
|
||||||
|
case <-time.After(50 * time.Millisecond):
|
||||||
|
}
|
||||||
|
|
||||||
|
que.Stop()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-errCh:
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("expected context.Canceled, got %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("ParseMessage did not return after Stop")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestQueueParseMessageViewAllowsReentrantParseOnSameConn(t *testing.T) {
|
||||||
|
que := NewQueue()
|
||||||
|
reentered := false
|
||||||
|
|
||||||
|
err := que.ParseMessageView(que.BuildMessage([]byte("outer")), "reentrant", func(view FrameView) error {
|
||||||
|
if !bytes.Equal(view.Payload, []byte("outer")) {
|
||||||
|
t.Fatalf("unexpected outer payload: %q", view.Payload)
|
||||||
|
}
|
||||||
|
return que.ParseMessageView(que.BuildMessage([]byte("inner")), "reentrant", func(inner FrameView) error {
|
||||||
|
reentered = true
|
||||||
|
if !bytes.Equal(inner.Payload, []byte("inner")) {
|
||||||
|
t.Fatalf("unexpected inner payload: %q", inner.Payload)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ParseMessageView failed: %v", err)
|
||||||
|
}
|
||||||
|
if !reentered {
|
||||||
|
t.Fatal("expected reentrant ParseMessageView to run")
|
||||||
|
}
|
||||||
|
}
|
||||||
+101
@@ -0,0 +1,101 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
var ErrDeadlineExceeded = errors.New("deadline exceeded")
|
||||||
|
var ErrQueueDataFormat = errors.New("data format error")
|
||||||
|
var ErrQueueUnsupportedVersion = errors.New("unsupported frame version")
|
||||||
|
var ErrQueueUnsupportedFlags = errors.New("unsupported frame flags")
|
||||||
|
var ErrQueueMessageTooLarge = errors.New("message too large")
|
||||||
|
var ErrQueueFrameHandlerNil = errors.New("frame handler is nil")
|
||||||
|
var ErrQueueConnKeyInvalid = errors.New("queue conn key must be comparable")
|
||||||
|
var ErrQueueConnKeyNil = errors.New("queue conn key must not be nil")
|
||||||
|
|
||||||
|
const (
|
||||||
|
queMagicSize = 8
|
||||||
|
queHeaderSize = 14
|
||||||
|
queVersionV1 = 1
|
||||||
|
queSupportedFlags = 0
|
||||||
|
)
|
||||||
|
|
||||||
|
// 识别头
|
||||||
|
var queMagic = []byte{11, 27, 19, 96, 12, 25, 02, 20}
|
||||||
|
|
||||||
|
// MsgQueue 为基本的信息单位。
|
||||||
|
type MsgQueue struct {
|
||||||
|
// Deprecated: frame-level IDs are no longer emitted by StarQueue v2 framing.
|
||||||
|
ID uint16
|
||||||
|
|
||||||
|
Msg []byte
|
||||||
|
Conn interface{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FrameView exposes a parsed payload without forcing StarQueue to clone it.
|
||||||
|
//
|
||||||
|
// The payload is only valid during the ParseMessageView callback. Callers must
|
||||||
|
// copy it if they need to keep it after the callback returns.
|
||||||
|
type FrameView struct {
|
||||||
|
Payload []byte
|
||||||
|
Conn interface{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// StarQueue 为流数据中的消息队列分发。
|
||||||
|
type StarQueue struct {
|
||||||
|
maxLength uint32
|
||||||
|
msgPool chan MsgQueue
|
||||||
|
states sync.Map
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
duration time.Duration
|
||||||
|
sendMu sync.RWMutex
|
||||||
|
stopOnce sync.Once
|
||||||
|
|
||||||
|
// Deprecated: new code should keep StarQueue focused on framing only.
|
||||||
|
Encode bool
|
||||||
|
// Deprecated: new code should keep StarQueue focused on framing only.
|
||||||
|
EncodeFunc func([]byte) []byte
|
||||||
|
// Deprecated: new code should keep StarQueue focused on framing only.
|
||||||
|
DecodeFunc func([]byte) []byte
|
||||||
|
}
|
||||||
|
|
||||||
|
type queHeader struct {
|
||||||
|
Length uint32
|
||||||
|
Version uint8
|
||||||
|
Flags uint8
|
||||||
|
}
|
||||||
|
|
||||||
|
type queConnState struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
buf []byte
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewQueueCtx(ctx context.Context, count int64, maxMsgLength uint32) *StarQueue {
|
||||||
|
if count < 0 {
|
||||||
|
panic("stario: negative queue count")
|
||||||
|
}
|
||||||
|
q := &StarQueue{
|
||||||
|
maxLength: maxMsgLength,
|
||||||
|
msgPool: make(chan MsgQueue, count),
|
||||||
|
}
|
||||||
|
if ctx == nil {
|
||||||
|
q.ctx, q.cancel = context.WithCancel(context.Background())
|
||||||
|
} else {
|
||||||
|
q.ctx, q.cancel = context.WithCancel(ctx)
|
||||||
|
}
|
||||||
|
context.AfterFunc(q.ctx, q.shutdown)
|
||||||
|
return q
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewQueueWithCount(count int64) *StarQueue {
|
||||||
|
return NewQueueCtx(nil, count, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewQueue 建立一个新消息队列。
|
||||||
|
func NewQueue() *StarQueue {
|
||||||
|
return NewQueueWithCount(32)
|
||||||
|
}
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import "errors"
|
||||||
|
|
||||||
|
var (
|
||||||
|
ErrSignalInterrupt = errors.New("interrupt")
|
||||||
|
ErrSignalStop = errors.New("stop")
|
||||||
|
ErrSignalQuit = errors.New("quit")
|
||||||
|
)
|
||||||
|
|
||||||
|
type inputSignalError struct {
|
||||||
|
msg string
|
||||||
|
cause error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *inputSignalError) Error() string {
|
||||||
|
return e.msg
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *inputSignalError) Unwrap() error {
|
||||||
|
return e.cause
|
||||||
|
}
|
||||||
|
|
||||||
|
func signalErrorForType(sigtype rune) error {
|
||||||
|
switch sigtype {
|
||||||
|
case 0x03:
|
||||||
|
return &inputSignalError{
|
||||||
|
msg: "SIGNAL SIGINT RECIVED",
|
||||||
|
cause: ErrSignalInterrupt,
|
||||||
|
}
|
||||||
|
case 0x1a:
|
||||||
|
return &inputSignalError{
|
||||||
|
msg: "SIGNAL SIGSTOP RECIVED",
|
||||||
|
cause: ErrSignalStop,
|
||||||
|
}
|
||||||
|
case 0x1c:
|
||||||
|
return &inputSignalError{
|
||||||
|
msg: "SIGNAL SIGQUIT RECIVED",
|
||||||
|
cause: ErrSignalQuit,
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSignalErrorForType(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
sig rune
|
||||||
|
msg string
|
||||||
|
want error
|
||||||
|
}{
|
||||||
|
{name: "interrupt", sig: 0x03, msg: "SIGNAL SIGINT RECIVED", want: ErrSignalInterrupt},
|
||||||
|
{name: "stop", sig: 0x1a, msg: "SIGNAL SIGSTOP RECIVED", want: ErrSignalStop},
|
||||||
|
{name: "quit", sig: 0x1c, msg: "SIGNAL SIGQUIT RECIVED", want: ErrSignalQuit},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
err := signalErrorForType(tc.sig)
|
||||||
|
if err == nil {
|
||||||
|
t.Fatalf("%s: expected non-nil error", tc.name)
|
||||||
|
}
|
||||||
|
if err.Error() != tc.msg {
|
||||||
|
t.Fatalf("%s: unexpected error text: got %q want %q", tc.name, err.Error(), tc.msg)
|
||||||
|
}
|
||||||
|
if !errors.Is(err, tc.want) {
|
||||||
|
t.Fatalf("%s: errors.Is mismatch: got %v want %v", tc.name, err, tc.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSignalErrorForUnknownType(t *testing.T) {
|
||||||
|
if err := signalErrorForType('x'); err != nil {
|
||||||
|
t.Fatalf("expected nil error for unknown signal, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSignalReturnsTypedError(t *testing.T) {
|
||||||
|
err := signal(0x03)
|
||||||
|
if !errors.Is(err, ErrSignalInterrupt) {
|
||||||
|
t.Fatalf("expected interrupt error, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSignalReturnsNilForUnknownType(t *testing.T) {
|
||||||
|
if err := signal('x'); err != nil {
|
||||||
|
t.Fatalf("expected nil error for unknown signal, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
//go:build windows
|
||||||
|
// +build windows
|
||||||
|
|
||||||
|
package stario
|
||||||
|
|
||||||
|
func signal(sigtype rune) error {
|
||||||
|
return signalErrorForType(sigtype)
|
||||||
|
}
|
||||||
@@ -0,0 +1,8 @@
|
|||||||
|
//go:build !windows
|
||||||
|
// +build !windows
|
||||||
|
|
||||||
|
package stario
|
||||||
|
|
||||||
|
func signal(sigtype rune) error {
|
||||||
|
return signalErrorForType(sigtype)
|
||||||
|
}
|
||||||
+123
@@ -0,0 +1,123 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"container/list"
|
||||||
|
"errors"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
var ErrStarRingInvalidCapacity = errors.New("star ring capacity must be greater than zero")
|
||||||
|
|
||||||
|
// StarRing keeps the newest bytes in memory with fixed capacity.
|
||||||
|
// It never blocks writers: when capacity is exceeded, oldest bytes are dropped.
|
||||||
|
type StarRing struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
cap int
|
||||||
|
|
||||||
|
size int
|
||||||
|
chunks list.List // each element: []byte
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewStarRing(capacity int) (*StarRing, error) {
|
||||||
|
if capacity <= 0 {
|
||||||
|
return nil, ErrStarRingInvalidCapacity
|
||||||
|
}
|
||||||
|
return &StarRing{cap: capacity}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) Capacity() int {
|
||||||
|
s.mu.RLock()
|
||||||
|
defer s.mu.RUnlock()
|
||||||
|
return s.cap
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) Len() int {
|
||||||
|
s.mu.RLock()
|
||||||
|
defer s.mu.RUnlock()
|
||||||
|
return s.size
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) Reset() {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
s.size = 0
|
||||||
|
s.chunks.Init()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) Write(p []byte) (int, error) {
|
||||||
|
if len(p) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
payload := append([]byte(nil), p...)
|
||||||
|
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
s.writeLocked(payload)
|
||||||
|
return len(p), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) WriteString(text string) (int, error) {
|
||||||
|
if len(text) == 0 {
|
||||||
|
return 0, nil
|
||||||
|
}
|
||||||
|
return s.Write([]byte(text))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) Snapshot() []byte {
|
||||||
|
s.mu.RLock()
|
||||||
|
defer s.mu.RUnlock()
|
||||||
|
|
||||||
|
out := make([]byte, s.size)
|
||||||
|
pos := 0
|
||||||
|
for node := s.chunks.Front(); node != nil; node = node.Next() {
|
||||||
|
chunk, ok := node.Value.([]byte)
|
||||||
|
if !ok || len(chunk) == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
pos += copy(out[pos:], chunk)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) writeLocked(payload []byte) {
|
||||||
|
if len(payload) >= s.cap {
|
||||||
|
payload = payload[len(payload)-s.cap:]
|
||||||
|
s.chunks.Init()
|
||||||
|
s.chunks.PushBack(payload)
|
||||||
|
s.size = len(payload)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
s.chunks.PushBack(payload)
|
||||||
|
s.size += len(payload)
|
||||||
|
s.trimLocked()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *StarRing) trimLocked() {
|
||||||
|
if s.size <= s.cap {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
overflow := s.size - s.cap
|
||||||
|
for overflow > 0 {
|
||||||
|
front := s.chunks.Front()
|
||||||
|
if front == nil {
|
||||||
|
s.size = 0
|
||||||
|
return
|
||||||
|
}
|
||||||
|
head, ok := front.Value.([]byte)
|
||||||
|
if !ok || len(head) == 0 {
|
||||||
|
s.chunks.Remove(front)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if len(head) <= overflow {
|
||||||
|
s.chunks.Remove(front)
|
||||||
|
s.size -= len(head)
|
||||||
|
overflow -= len(head)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
trimmed := append([]byte(nil), head[overflow:]...)
|
||||||
|
front.Value = trimmed
|
||||||
|
s.size -= overflow
|
||||||
|
overflow = 0
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,104 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewStarRingRejectsInvalidCapacity(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(0)
|
||||||
|
if err != ErrStarRingInvalidCapacity {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if ring != nil {
|
||||||
|
t.Fatal("expected nil ring when capacity is invalid")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarRingWriteAndSnapshot(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(16)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n, err := ring.Write([]byte("abc")); err != nil || n != 3 {
|
||||||
|
t.Fatalf("write failed: n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
if n, err := ring.WriteString("def"); err != nil || n != 3 {
|
||||||
|
t.Fatalf("write string failed: n=%d err=%v", n, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
got := ring.Snapshot()
|
||||||
|
if !bytes.Equal(got, []byte("abcdef")) {
|
||||||
|
t.Fatalf("unexpected snapshot: %q", got)
|
||||||
|
}
|
||||||
|
if ring.Len() != 6 {
|
||||||
|
t.Fatalf("unexpected ring len: %d", ring.Len())
|
||||||
|
}
|
||||||
|
if ring.Capacity() != 16 {
|
||||||
|
t.Fatalf("unexpected ring capacity: %d", ring.Capacity())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarRingOverflowDropsOldestBytesPrecisely(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(10)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, _ = ring.WriteString("abcde")
|
||||||
|
_, _ = ring.WriteString("fghij")
|
||||||
|
_, _ = ring.WriteString("kl")
|
||||||
|
|
||||||
|
got := ring.Snapshot()
|
||||||
|
want := []byte("cdefghijkl")
|
||||||
|
if !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("unexpected snapshot after overflow: got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
if ring.Len() != len(want) {
|
||||||
|
t.Fatalf("unexpected ring len after overflow: %d", ring.Len())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarRingLargeWriteKeepsTail(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(5)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, _ = ring.WriteString("123")
|
||||||
|
_, _ = ring.WriteString("456789")
|
||||||
|
|
||||||
|
got := ring.Snapshot()
|
||||||
|
want := []byte("56789")
|
||||||
|
if !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("unexpected snapshot after large write: got %q want %q", got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarRingSnapshotIsCopy(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(8)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, _ = ring.WriteString("hello")
|
||||||
|
|
||||||
|
first := ring.Snapshot()
|
||||||
|
first[0] = 'H'
|
||||||
|
second := ring.Snapshot()
|
||||||
|
if !bytes.Equal(second, []byte("hello")) {
|
||||||
|
t.Fatalf("snapshot should not expose internal storage: %q", second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStarRingReset(t *testing.T) {
|
||||||
|
ring, err := NewStarRing(8)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, _ = ring.WriteString("hello")
|
||||||
|
ring.Reset()
|
||||||
|
if ring.Len() != 0 {
|
||||||
|
t.Fatalf("expected len=0 after reset, got %d", ring.Len())
|
||||||
|
}
|
||||||
|
if got := ring.Snapshot(); len(got) != 0 {
|
||||||
|
t.Fatalf("expected empty snapshot after reset, got %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,159 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
type waitGroupAddMode uint8
|
||||||
|
|
||||||
|
const (
|
||||||
|
waitGroupAddModeStrict waitGroupAddMode = iota
|
||||||
|
waitGroupAddModeLoose
|
||||||
|
)
|
||||||
|
|
||||||
|
// WaitGroup is a concurrency-limited sync.WaitGroup variant.
|
||||||
|
//
|
||||||
|
// A zero or negative limit means unlimited concurrency. WaitGroup must not be
|
||||||
|
// copied after first use.
|
||||||
|
type WaitGroup struct {
|
||||||
|
wg sync.WaitGroup
|
||||||
|
mu sync.Mutex
|
||||||
|
cond *sync.Cond
|
||||||
|
initOnce sync.Once
|
||||||
|
maxCount int
|
||||||
|
running int
|
||||||
|
addMode waitGroupAddMode
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewWaitGroup creates a WaitGroup with the provided concurrency limit.
|
||||||
|
func NewWaitGroup(maxCount int) WaitGroup {
|
||||||
|
if maxCount < 0 {
|
||||||
|
panic("stario: negative max wait count")
|
||||||
|
}
|
||||||
|
return WaitGroup{
|
||||||
|
maxCount: maxCount,
|
||||||
|
addMode: waitGroupAddModeStrict,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *WaitGroup) init() {
|
||||||
|
w.initOnce.Do(func() {
|
||||||
|
w.cond = sync.NewCond(&w.mu)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add adjusts the running task count.
|
||||||
|
//
|
||||||
|
// Positive deltas may block when the concurrency limit is reached. Negative
|
||||||
|
// deltas release running slots.
|
||||||
|
func (w *WaitGroup) Add(delta int) {
|
||||||
|
w.init()
|
||||||
|
if delta == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if delta < 0 {
|
||||||
|
w.release(-delta)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.acquire(delta)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *WaitGroup) acquire(delta int) {
|
||||||
|
w.mu.Lock()
|
||||||
|
defer w.mu.Unlock()
|
||||||
|
if w.maxCount <= 0 {
|
||||||
|
w.running += delta
|
||||||
|
w.wg.Add(delta)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if delta == 1 {
|
||||||
|
w.wg.Add(1)
|
||||||
|
for w.maxCount > 0 && w.running >= w.maxCount {
|
||||||
|
w.cond.Wait()
|
||||||
|
}
|
||||||
|
w.running++
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if w.running+delta > w.maxCount {
|
||||||
|
if w.addMode == waitGroupAddModeStrict {
|
||||||
|
panic(fmt.Sprintf("stario: WaitGroup.Add(%d) exceeds max limit %d with %d running", delta, w.maxCount, w.running))
|
||||||
|
}
|
||||||
|
w.maxCount = w.running + delta
|
||||||
|
}
|
||||||
|
w.running += delta
|
||||||
|
w.wg.Add(delta)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *WaitGroup) release(delta int) {
|
||||||
|
w.mu.Lock()
|
||||||
|
defer w.mu.Unlock()
|
||||||
|
if delta > w.running {
|
||||||
|
panic(fmt.Sprintf("stario: WaitGroup.Done releases %d tasks but only %d running", delta, w.running))
|
||||||
|
}
|
||||||
|
w.wg.Add(-delta)
|
||||||
|
w.running -= delta
|
||||||
|
w.cond.Broadcast()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Done releases one running task slot.
|
||||||
|
func (w *WaitGroup) Done() {
|
||||||
|
w.Add(-1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Go runs fn in a goroutine while accounting for the concurrency limit.
|
||||||
|
func (w *WaitGroup) Go(fn func()) {
|
||||||
|
w.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer w.Done()
|
||||||
|
fn()
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait blocks until all added work has completed.
|
||||||
|
func (w *WaitGroup) Wait() {
|
||||||
|
w.init()
|
||||||
|
w.wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMaxWaitNum returns the current concurrency limit.
|
||||||
|
func (w *WaitGroup) GetMaxWaitNum() int {
|
||||||
|
w.init()
|
||||||
|
w.mu.Lock()
|
||||||
|
defer w.mu.Unlock()
|
||||||
|
return w.maxCount
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetMaxWaitNum updates the concurrency limit.
|
||||||
|
func (w *WaitGroup) SetMaxWaitNum(num int) {
|
||||||
|
if num < 0 {
|
||||||
|
panic("stario: negative max wait count")
|
||||||
|
}
|
||||||
|
w.init()
|
||||||
|
w.mu.Lock()
|
||||||
|
w.maxCount = num
|
||||||
|
w.mu.Unlock()
|
||||||
|
w.cond.Broadcast()
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetStrictAddMode controls whether Add(n>1) panics or auto-expands the limit
|
||||||
|
// when the requested batch exceeds the current capacity.
|
||||||
|
func (w *WaitGroup) SetStrictAddMode(strict bool) {
|
||||||
|
w.init()
|
||||||
|
w.mu.Lock()
|
||||||
|
if strict {
|
||||||
|
w.addMode = waitGroupAddModeStrict
|
||||||
|
} else {
|
||||||
|
w.addMode = waitGroupAddModeLoose
|
||||||
|
}
|
||||||
|
w.mu.Unlock()
|
||||||
|
w.cond.Broadcast()
|
||||||
|
}
|
||||||
|
|
||||||
|
// StrictAddMode reports whether strict batch-add behavior is enabled.
|
||||||
|
func (w *WaitGroup) StrictAddMode() bool {
|
||||||
|
w.init()
|
||||||
|
w.mu.Lock()
|
||||||
|
defer w.mu.Unlock()
|
||||||
|
return w.addMode == waitGroupAddModeStrict
|
||||||
|
}
|
||||||
+171
@@ -0,0 +1,171 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestWaitGroupAddBlocksAtLimit(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
wg.Add(1)
|
||||||
|
|
||||||
|
unblocked := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
wg.Add(1)
|
||||||
|
close(unblocked)
|
||||||
|
wg.Done()
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-unblocked:
|
||||||
|
t.Fatal("Add(1) should block while running count is at the limit")
|
||||||
|
case <-time.After(30 * time.Millisecond):
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Done()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-unblocked:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Add(1) did not resume after capacity was released")
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupStrictAddPanicsWhenBatchExceedsLimit(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
wg.Add(1)
|
||||||
|
defer wg.Done()
|
||||||
|
|
||||||
|
defer func() {
|
||||||
|
if recover() == nil {
|
||||||
|
t.Fatal("expected Add(2) to panic in strict mode when it exceeds the limit")
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
wg.Add(2)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupLooseAddExpandsLimit(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
wg.SetStrictAddMode(false)
|
||||||
|
|
||||||
|
wg.Add(1)
|
||||||
|
wg.Add(2)
|
||||||
|
|
||||||
|
if got := wg.GetMaxWaitNum(); got != 3 {
|
||||||
|
t.Fatalf("unexpected max count: got %d want 3", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
wg.Done()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupSetMaxWaitNumOverridesValue(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
wg.SetMaxWaitNum(4)
|
||||||
|
if got := wg.GetMaxWaitNum(); got != 4 {
|
||||||
|
t.Fatalf("unexpected max count: got %d want 4", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupSetMaxWaitNumZeroUnblocksWaitingAdd(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
wg.Add(1)
|
||||||
|
|
||||||
|
unblocked := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
wg.Add(1)
|
||||||
|
close(unblocked)
|
||||||
|
wg.Done()
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-unblocked:
|
||||||
|
t.Fatal("Add(1) should still be blocked before removing the limit")
|
||||||
|
case <-time.After(30 * time.Millisecond):
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.SetMaxWaitNum(0)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-unblocked:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Add(1) did not resume after the limit was removed")
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Done()
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupGoWaitsForCompletion(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
done := make(chan struct{})
|
||||||
|
|
||||||
|
wg.Go(func() {
|
||||||
|
close(done)
|
||||||
|
})
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Go(fn) did not run the function")
|
||||||
|
}
|
||||||
|
|
||||||
|
finished := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
wg.Wait()
|
||||||
|
close(finished)
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-finished:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("Wait did not observe Go(fn) completion")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWaitGroupGoRespectsLimit(t *testing.T) {
|
||||||
|
wg := NewWaitGroup(1)
|
||||||
|
releaseFirst := make(chan struct{})
|
||||||
|
firstStarted := make(chan struct{})
|
||||||
|
secondStarted := make(chan struct{})
|
||||||
|
secondReturned := make(chan struct{})
|
||||||
|
|
||||||
|
wg.Go(func() {
|
||||||
|
close(firstStarted)
|
||||||
|
<-releaseFirst
|
||||||
|
})
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-firstStarted:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("first Go(fn) did not start")
|
||||||
|
}
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
wg.Go(func() {
|
||||||
|
close(secondStarted)
|
||||||
|
})
|
||||||
|
close(secondReturned)
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-secondReturned:
|
||||||
|
t.Fatal("second Go(fn) should block while the limit is full")
|
||||||
|
case <-time.After(30 * time.Millisecond):
|
||||||
|
}
|
||||||
|
|
||||||
|
close(releaseFirst)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-secondStarted:
|
||||||
|
case <-time.After(200 * time.Millisecond):
|
||||||
|
t.Fatal("second Go(fn) did not start after capacity was released")
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
}
|
||||||
+73
@@ -0,0 +1,73 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"golang.org/x/term"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ErrTerminalNotTTY reports that terminal-only input was requested from a
|
||||||
|
// non-terminal stdin.
|
||||||
|
var ErrTerminalNotTTY = errors.New("terminal input requires a tty")
|
||||||
|
|
||||||
|
var terminalCheckFunc = term.IsTerminal
|
||||||
|
|
||||||
|
// IsTerminal reports whether os.Stdin is attached to a terminal.
|
||||||
|
func IsTerminal() bool {
|
||||||
|
return IsTerminalFile(os.Stdin)
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsTerminalFD reports whether fd is attached to a terminal.
|
||||||
|
func IsTerminalFD(fd int) bool {
|
||||||
|
return terminalCheckFunc(fd)
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsTerminalFile reports whether file is attached to a terminal.
|
||||||
|
func IsTerminalFile(file *os.File) bool {
|
||||||
|
if file == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return IsTerminalFD(int(file.Fd()))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReadPasswordContext reads one password-style line from the terminal.
|
||||||
|
//
|
||||||
|
// It returns ErrTerminalNotTTY when stdin is not a terminal. If ctx is canceled
|
||||||
|
// while waiting for input, this function returns ctx.Err() without waiting for
|
||||||
|
// the underlying terminal read to finish.
|
||||||
|
func ReadPasswordContext(ctx context.Context, hint string, defaultVal string) (string, error) {
|
||||||
|
return ReadPasswordContextWithMask(ctx, hint, defaultVal, "")
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReadPasswordContextWithMask is like ReadPasswordContext but echoes the given
|
||||||
|
// mask string while the user types.
|
||||||
|
func ReadPasswordContextWithMask(ctx context.Context, hint string, defaultVal string, mask string) (string, error) {
|
||||||
|
if ctx == nil {
|
||||||
|
ctx = context.Background()
|
||||||
|
}
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if !IsTerminal() {
|
||||||
|
return "", ErrTerminalNotTTY
|
||||||
|
}
|
||||||
|
session, err := rawTerminalSessionFactory(hint, true)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
resultCh := make(chan InputMsg, 1)
|
||||||
|
go func() {
|
||||||
|
defer session.Close()
|
||||||
|
resultCh <- rawLineInputSession(session, defaultVal, mask, rawInputSignalReturnError)
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case result := <-resultCh:
|
||||||
|
return result.String()
|
||||||
|
case <-ctx.Done():
|
||||||
|
session.Abort()
|
||||||
|
<-resultCh
|
||||||
|
return "", ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,77 @@
|
|||||||
|
package stario
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func installTerminalCheckStub(t *testing.T, ok bool) {
|
||||||
|
t.Helper()
|
||||||
|
prev := terminalCheckFunc
|
||||||
|
terminalCheckFunc = func(fd int) bool {
|
||||||
|
return ok
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
terminalCheckFunc = prev
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIsTerminalFDUsesStub(t *testing.T) {
|
||||||
|
installTerminalCheckStub(t, true)
|
||||||
|
if !IsTerminalFD(1) {
|
||||||
|
t.Fatal("expected terminal check to return true")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadPasswordContextRejectsNonTTY(t *testing.T) {
|
||||||
|
installTerminalCheckStub(t, false)
|
||||||
|
_, err := ReadPasswordContext(context.Background(), "pwd: ", "")
|
||||||
|
if !errors.Is(err, ErrTerminalNotTTY) {
|
||||||
|
t.Fatalf("expected ErrTerminalNotTTY, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadPasswordContextReadsValue(t *testing.T) {
|
||||||
|
installTerminalCheckStub(t, true)
|
||||||
|
installRawInputStub(t, "secret\n", nil)
|
||||||
|
|
||||||
|
got, err := ReadPasswordContext(context.Background(), "pwd: ", "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ReadPasswordContext returned error: %v", err)
|
||||||
|
}
|
||||||
|
if got != "secret" {
|
||||||
|
t.Fatalf("unexpected password: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReadPasswordContextCanceledWhileWaiting(t *testing.T) {
|
||||||
|
installTerminalCheckStub(t, true)
|
||||||
|
|
||||||
|
prevFactory := rawTerminalSessionFactory
|
||||||
|
pipeReader, pipeWriter := io.Pipe()
|
||||||
|
rawTerminalSessionFactory = func(hint string, printNewline bool) (*rawTerminalSession, error) {
|
||||||
|
return &rawTerminalSession{
|
||||||
|
input: pipeReader,
|
||||||
|
reader: bufio.NewReader(pipeReader),
|
||||||
|
redrawHint: strings.TrimSpace(hint),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
rawTerminalSessionFactory = prevFactory
|
||||||
|
_ = pipeWriter.Close()
|
||||||
|
_ = pipeReader.Close()
|
||||||
|
})
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
_, err := ReadPasswordContext(ctx, "pwd: ", "")
|
||||||
|
if !errors.Is(err, context.DeadlineExceeded) {
|
||||||
|
t.Fatalf("expected deadline exceeded, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
//go:build !windows
|
||||||
|
// +build !windows
|
||||||
|
|
||||||
|
package stario
|
||||||
|
|
||||||
|
import "os"
|
||||||
|
|
||||||
|
func openRawTerminalInput() (*os.File, error) {
|
||||||
|
return os.OpenFile("/dev/tty", os.O_RDWR, 0)
|
||||||
|
}
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
//go:build windows
|
||||||
|
// +build windows
|
||||||
|
|
||||||
|
package stario
|
||||||
|
|
||||||
|
import "os"
|
||||||
|
|
||||||
|
func openRawTerminalInput() (*os.File, error) {
|
||||||
|
return os.OpenFile("CONIN$", os.O_RDWR, 0)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user