storj.io/uplink@v1.13.0/private/storage/streams/buffer/buffer.go (about)

     1  // Copyright (C) 2023 Storj Labs, Inc.
     2  // See LICENSE for copying information.
     3  
     4  package buffer
     5  
     6  import (
     7  	"io"
     8  	"sync"
     9  )
    10  
    11  // Buffer lets one write to some Backend and lazily create readers from it.
    12  type Buffer struct {
    13  	mu      sync.Mutex
    14  	cursor  *Cursor
    15  	backend Backend
    16  	wrote   int64
    17  }
    18  
    19  // New constructs a Buffer using the underlying Backend and allowing
    20  // the writer to write writeAhead extra bytes in front of the most advanced reader.
    21  func New(backend Backend, writeAhead int64) *Buffer {
    22  	return &Buffer{
    23  		cursor:  NewCursor(writeAhead),
    24  		backend: backend,
    25  	}
    26  }
    27  
    28  // DoneReading signals to the Buffer that no more Read calls are coming and
    29  // that Write calls should return the provided error. The first call that
    30  // causes both DoneReading and DoneWriting to have been called closes the
    31  // underlying backend.
    32  func (w *Buffer) DoneReading(err error) {
    33  	if w.cursor.DoneReading(err) {
    34  		_ = w.backend.Close()
    35  	}
    36  }
    37  
    38  // DoneWriting signals to the Buffer that no more Write calls are coming and
    39  // that Read calls should return the provided error. The first call that
    40  // causes both DoneReading and DoneWriting to have been called closes the
    41  // underlying backend.
    42  func (w *Buffer) DoneWriting(err error) {
    43  	if w.cursor.DoneWriting(err) {
    44  		_ = w.backend.Close()
    45  	}
    46  }
    47  
    48  // Write appends the bytes in p to the Buffer. It blocks until the furthest
    49  // advanced reader is less than the write ahead amount.
    50  func (w *Buffer) Write(p []byte) (n int, err error) {
    51  	w.mu.Lock()
    52  	defer w.mu.Unlock()
    53  
    54  	for len(p) > 0 {
    55  		m, ok, err := w.cursor.WaitWrite(w.wrote + int64(len(p)))
    56  		if err != nil {
    57  			return n, err
    58  		} else if !ok {
    59  			return n, nil
    60  		}
    61  
    62  		var nn int
    63  		nn, err = w.backend.Write(p[:m-w.wrote])
    64  		n += nn
    65  		p = p[nn:]
    66  		w.wrote += int64(nn)
    67  		w.cursor.WroteTo(w.wrote)
    68  
    69  		if err != nil {
    70  			w.cursor.DoneWriting(err)
    71  			return n, err
    72  		}
    73  	}
    74  	return n, nil
    75  }
    76  
    77  // bufferReader is the concrete implementation returned from the Buffer.Reader method.
    78  type bufferReader struct {
    79  	mu     sync.Mutex
    80  	cursor *Cursor
    81  	buffer Backend
    82  	read   int64
    83  }
    84  
    85  // Reader returns a fresh io.Reader that can be used to read all of the previously
    86  // and future written bytes to the Buffer.
    87  func (w *Buffer) Reader() io.Reader {
    88  	return &bufferReader{
    89  		cursor: w.cursor,
    90  		buffer: w.backend,
    91  	}
    92  }
    93  
    94  // Read reads bytes into p, waiting for writes to happen if needed. It returns
    95  // 0, io.EOF when the end of the data has been reached.
    96  func (w *bufferReader) Read(p []byte) (n int, err error) {
    97  	w.mu.Lock()
    98  	defer w.mu.Unlock()
    99  
   100  	m, ok, err := w.cursor.WaitRead(w.read + int64(len(p)))
   101  	if err != nil {
   102  		return 0, err
   103  	} else if m == w.read {
   104  		return 0, io.EOF
   105  	}
   106  
   107  	n, err = w.buffer.ReadAt(p[:m-w.read], w.read)
   108  	w.read += int64(n)
   109  	w.cursor.ReadTo(w.read)
   110  
   111  	if err != nil {
   112  		return n, err
   113  	} else if !ok {
   114  		return n, io.EOF
   115  	}
   116  	return n, nil
   117  }