Something went wrong. Try again.
forked from tailscale/tailscale
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187// Copyright (c) Tailscale Inc & contributors// SPDX-License-Identifier: BSD-3-Clause
package ioqueue
import ( "io" "math/bits" "sync"
"tailscale.com/types/bools")
// Statically verify that VolatileBuffer implements Buffer.var _ Buffer = (*VolatileBuffer)(nil)
// VolatileBuffer is an in-memory implementation of [Buffer].// The zero value is an empty buffer ready for use.//// Unread data is not durable: it is lost when the process exits.// Close the write side with [VolatileBuffer.CloseWrite] to signal// that no further bytes will be produced; subsequent [Buffer.Read]// and [Buffer.Peek] calls return [io.EOF] once the buffer is drained,// and outstanding [Buffer.WaitUntil] waiters are unblocked.type VolatileBuffer struct { mu sync.Mutex
// buf may retain a consumed prefix after reads/discards. // Unread data is always the trailing writeOffset-readOffset bytes. buf []byte readOffset int64 writeOffset int64 peakLen int64 // decayed by compactLocked
closedWrite bool waiters offsetWaiters}
// Len reports the size of the buffer,// which is the number of written, but unread bytes.func (b *VolatileBuffer) Len() int64 { b.mu.Lock() defer b.mu.Unlock() return b.lenLocked()}
func (b *VolatileBuffer) lenLocked() int64 { return b.writeOffset - b.readOffset}
// WriteOffset is the total number of bytes written.func (b *VolatileBuffer) WriteOffset() int64 { b.mu.Lock() defer b.mu.Unlock() return b.writeOffset}
// ReadOffset is the total number of bytes read.func (b *VolatileBuffer) ReadOffset() int64 { b.mu.Lock() defer b.mu.Unlock() return b.readOffset}
// dataLocked returns the unread portion of buf.func (b *VolatileBuffer) dataLocked() []byte { n := int(b.writeOffset - b.readOffset) return b.buf[len(b.buf)-n:]}
// compactLocked reclaims backing storage after a consume,// by sliding data forward or allocating a smaller capacity buffer.func (b *VolatileBuffer) compactLocked() { data := b.dataLocked() // If already read data is greater than unread data, then compact. if len(b.buf) > 2*len(data) { b.peakLen = max(b.peakLen, b.lenLocked()) // If total capacity is 4x greater than the peak Len, // then allocate smaller capacity, otherwise slide data to front. if int64(cap(b.buf)/4) > b.peakLen && cap(b.buf) > 4<<10 { newCap := 2 << bits.Len(uint(b.peakLen)-1) // double the peak Len, but round up to power-of-2 newCap = max(newCap, 4<<10) // minimum buffer capacity to shrink to b.buf = append(make([]byte, 0, newCap), data...) // allocate smaller buffer } else { b.buf = b.buf[:copy(b.buf, data)] // slide data to the front } b.peakLen = b.peakLen * 7 / 8 // reduce peakLen by 12.5% }}
// Write writes data to the end of the buffer,// atomically incrementing Len and WriteOffset// by the amount of bytes written.// Write does not block.// After [CloseWrite], Write returns an error.func (b *VolatileBuffer) Write(p []byte) (int, error) { b.mu.Lock() defer b.mu.Unlock() if b.closedWrite { return 0, wrapError("write", errClosed) } b.buf = append(b.buf, p...) b.writeOffset += int64(len(p)) b.peakLen = max(b.peakLen, b.lenLocked()) b.waiters.notify(b.writeOffset, b.closedWrite) return len(p), nil}
// Read reads data from the front of the buffer, atomically decrementing// Len and incrementing ReadOffset by the amount of bytes read.// Rather than blocking, it returns [ErrEmpty] when the buffer is empty.// After [CloseWrite], it returns [io.EOF] once drained.func (b *VolatileBuffer) Read(p []byte) (int, error) { b.mu.Lock() defer b.mu.Unlock() data := b.dataLocked() if len(data) == 0 { return 0, wrapError("read", bools.IfElse(b.closedWrite, io.EOF, ErrEmpty)) } n := copy(p, data) b.readOffset += int64(n) b.compactLocked() return n, nil}
// Peek copies data from the front of the buffer into p// without affecting Len or ReadOffset.// It reports the current ReadOffset and the number of bytes copied.// If no bytes are available, it reports [ErrEmpty],// or [io.EOF] after [VolatileBuffer.CloseWrite] once drained.func (b *VolatileBuffer) Peek(p []byte) (readOffset int64, n int64, err error) { b.mu.Lock() defer b.mu.Unlock() readOffset = b.readOffset data := b.dataLocked() if len(data) == 0 { return readOffset, 0, wrapError("peek", bools.IfElse(b.closedWrite, io.EOF, ErrEmpty)) } return readOffset, int64(copy(p, data)), nil}
// DiscardUntil discards bytes from the front of the buffer by// incrementing ReadOffset to match the specified readOffset.// See [Buffer.DiscardUntil] for full semantics.func (b *VolatileBuffer) DiscardUntil(readOffset int64) (n int64, err error) { b.mu.Lock() defer b.mu.Unlock() if readOffset <= b.readOffset { return 0, nil } if readOffset > b.writeOffset { err = wrapError("discard", bools.IfElse(b.closedWrite, io.EOF, ErrEmpty)) } readOffset = min(readOffset, b.writeOffset) n = readOffset - b.readOffset b.readOffset = readOffset b.compactLocked() return n, err}
// WaitUntil returns a channel that is closed when WriteOffset exceeds// the specified writeOffset, or when the write side is closed via// [CloseWrite] (so waiters do not block indefinitely after close).//// Multiple waiters for the same writeOffset share one channel.func (b *VolatileBuffer) WaitUntil(writeOffset int64) <-chan struct{} { b.mu.Lock() defer b.mu.Unlock() return b.waiters.wait(b.writeOffset, writeOffset, b.closedWrite)}
// CloseWrite closes the write side of the buffer.// After CloseWrite, [Write] fails, [Read]/[Peek] return [io.EOF]// once all buffered data has been consumed, and all [WaitUntil]// waiters are unblocked.// CloseWrite is idempotent and returns an error if already closed.func (b *VolatileBuffer) CloseWrite() error { b.mu.Lock() defer b.mu.Unlock() if b.closedWrite { return wrapError("close", errClosed) } b.closedWrite = true b.waiters.notify(b.writeOffset, b.closedWrite) return nil}