package relay import ( "context" "errors" "io" "sync" ) // ErrClosed is returned by Read/Write after Abort. var ErrClosed = errors.New("relay: ring aborted") // Ring is a bounded byte buffer with single-producer / single-consumer // blocking semantics. Plan §M5: 单 session 64 MiB. // // Internally it's a flat slice with head/tail wrap-around, guarded by a // sync.Cond. Both Read and Write accept ctx so handlers can abort fast on // client disconnect — the watchdog goroutine broadcasts the cond when ctx // fires so blocked Wait calls re-check ctx.Err(). type Ring struct { mu sync.Mutex cond *sync.Cond buf []byte head int tail int used int cap int closed bool // sender called CloseWriter; reader drains to EOF aborted bool // both sides torn down } func newRing(capacity int) *Ring { r := &Ring{ buf: make([]byte, capacity), cap: capacity, } r.cond = sync.NewCond(&r.mu) return r } func (r *Ring) Write(ctx context.Context, p []byte) (int, error) { if len(p) == 0 { return 0, nil } stop := r.watchCtx(ctx) defer stop() r.mu.Lock() defer r.mu.Unlock() written := 0 for written < len(p) { for r.used == r.cap && !r.aborted && !r.closed && ctxOK(ctx) { r.cond.Wait() } if !ctxOK(ctx) { return written, ctx.Err() } if r.aborted { return written, ErrClosed } if r.closed { return written, ErrClosed } free := r.cap - r.used toWrite := len(p) - written if toWrite > free { toWrite = free } end := r.tail + toWrite if end <= r.cap { copy(r.buf[r.tail:end], p[written:written+toWrite]) } else { first := r.cap - r.tail copy(r.buf[r.tail:], p[written:written+first]) copy(r.buf[:toWrite-first], p[written+first:written+toWrite]) } r.tail = (r.tail + toWrite) % r.cap r.used += toWrite written += toWrite r.cond.Broadcast() } return written, nil } func (r *Ring) Read(ctx context.Context, p []byte) (int, error) { if len(p) == 0 { return 0, nil } stop := r.watchCtx(ctx) defer stop() r.mu.Lock() defer r.mu.Unlock() for r.used == 0 && !r.closed && !r.aborted && ctxOK(ctx) { r.cond.Wait() } if !ctxOK(ctx) { return 0, ctx.Err() } if r.used == 0 { if r.aborted { return 0, ErrClosed } return 0, io.EOF } n := r.used if n > len(p) { n = len(p) } end := r.head + n if end <= r.cap { copy(p[:n], r.buf[r.head:end]) } else { first := r.cap - r.head copy(p[:first], r.buf[r.head:]) copy(p[first:n], r.buf[:n-first]) } r.head = (r.head + n) % r.cap r.used -= n r.cond.Broadcast() return n, nil } // CloseWriter signals end-of-stream from the sender side. Reader returns // io.EOF after draining whatever is still buffered. func (r *Ring) CloseWriter() { r.mu.Lock() r.closed = true r.cond.Broadcast() r.mu.Unlock() } // Abort tears the ring down hard. Subsequent Read/Write return ErrClosed. func (r *Ring) Abort() { r.mu.Lock() r.aborted = true r.cond.Broadcast() r.mu.Unlock() } func (r *Ring) watchCtx(ctx context.Context) func() { if ctx == nil || ctx.Err() != nil { return func() {} } stop := make(chan struct{}) go func() { select { case <-ctx.Done(): r.mu.Lock() r.cond.Broadcast() r.mu.Unlock() case <-stop: } }() return func() { close(stop) } } func ctxOK(ctx context.Context) bool { if ctx == nil { return true } return ctx.Err() == nil }