Commilitia-Drop/internal/relay/ring.go

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
}