github.com/grailbio/base@v0.0.11/stateio/writer.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package stateio
     6  
     7  import (
     8  	"encoding/binary"
     9  	"io"
    10  	"os"
    11  
    12  	"github.com/grailbio/base/logio"
    13  )
    14  
    15  type syncer interface {
    16  	Sync() error
    17  }
    18  
    19  // Writer writes snapshots and update entries to an underlying log stream.
    20  type Writer struct {
    21  	syncer syncer
    22  	log    *logio.Writer
    23  	epoch  uint64
    24  }
    25  
    26  // NewWriter initializes and returns a new state log writer which
    27  // writes to the stream w, which is positioned at the provided
    28  // offset. The provided epoch must be the current epoch of the log
    29  // file. If the provided io.Writer is also a syncer:
    30  //
    31  //	type Syncer interface {
    32  //		Sync() error
    33  //	}
    34  //
    35  // Then Sync() is called (and errors returned) after each log entry
    36  // has been written.
    37  func NewWriter(w io.Writer, off int64, epoch uint64) *Writer {
    38  	wr := &Writer{log: logio.NewWriter(w, off), epoch: epoch}
    39  	if s, ok := w.(syncer); ok {
    40  		wr.syncer = s
    41  	}
    42  	return wr
    43  }
    44  
    45  // NewFileWriter initializes a state log writer from the provided
    46  // os file.  The file's contents is committed to stable storage
    47  // after each log write.
    48  func NewFileWriter(file *os.File) (*Writer, error) {
    49  	off, err := file.Seek(0, io.SeekEnd)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	_, epoch, _, err := Restore(file, off)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	off, err = file.Seek(0, io.SeekEnd)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	return NewWriter(file, off, epoch), nil
    62  }
    63  
    64  // Snapshot writes a new snapshot to the state log.
    65  // Subsequent updates are based on this snapshot.
    66  func (w *Writer) Snapshot(snap []byte) error {
    67  	off := w.log.Tell()
    68  	entry := make([]byte, len(snap)+9)
    69  	entry[0] = entrySnap
    70  	binary.LittleEndian.PutUint64(entry[1:], w.epoch)
    71  	copy(entry[9:], snap)
    72  	if err := w.log.Append(entry); err != nil {
    73  		return err
    74  	}
    75  	w.epoch = uint64(off)
    76  	return w.sync()
    77  }
    78  
    79  // Update writes a new state update to the log. The update
    80  // refers to the last snapshot written.
    81  func (w *Writer) Update(update []byte) error {
    82  	entry := make([]byte, 9+len(update))
    83  	entry[0] = entryUpdate
    84  	binary.LittleEndian.PutUint64(entry[1:], uint64(w.epoch))
    85  	copy(entry[9:], update)
    86  	if err := w.log.Append(entry); err != nil {
    87  		return err
    88  	}
    89  	return w.sync()
    90  }
    91  
    92  func (w *Writer) sync() error {
    93  	if w.syncer == nil {
    94  		return nil
    95  	}
    96  	return w.syncer.Sync()
    97  }