github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/snapshot/snapshot.go (about)

     1  // snapshot manages the interactions between Nomad and Raft in order to take
     2  // and restore snapshots for disaster recovery. The internal format of a
     3  // snapshot is simply a tar file, as described in archive.go.
     4  package snapshot
     5  
     6  import (
     7  	"compress/gzip"
     8  	"crypto/sha256"
     9  	"encoding/base64"
    10  	"fmt"
    11  	"io"
    12  	"io/ioutil"
    13  	"os"
    14  
    15  	"github.com/hashicorp/go-hclog"
    16  	"github.com/hashicorp/raft"
    17  )
    18  
    19  // Snapshot is a structure that holds state about a temporary file that is used
    20  // to hold a snapshot. By using an intermediate file we avoid holding everything
    21  // in memory.
    22  type Snapshot struct {
    23  	file     *os.File
    24  	index    uint64
    25  	checksum string
    26  }
    27  
    28  // New takes a state snapshot of the given Raft instance into a temporary file
    29  // and returns an object that gives access to the file as an io.Reader. You must
    30  // arrange to call Close() on the returned object or else you will leak a
    31  // temporary file.
    32  func New(logger hclog.Logger, r *raft.Raft) (*Snapshot, error) {
    33  	// Take the snapshot.
    34  	future := r.Snapshot()
    35  	if err := future.Error(); err != nil {
    36  		return nil, fmt.Errorf("Raft error when taking snapshot: %v", err)
    37  	}
    38  
    39  	// Open up the snapshot.
    40  	metadata, snap, err := future.Open()
    41  	if err != nil {
    42  		return nil, fmt.Errorf("failed to open snapshot: %v:", err)
    43  	}
    44  	defer func() {
    45  		if err := snap.Close(); err != nil {
    46  			logger.Error("Failed to close Raft snapshot", "error", err)
    47  		}
    48  	}()
    49  
    50  	// Make a scratch file to receive the contents so that we don't buffer
    51  	// everything in memory. This gets deleted in Close() since we keep it
    52  	// around for re-reading.
    53  	archive, err := ioutil.TempFile("", "snapshot")
    54  	if err != nil {
    55  		return nil, fmt.Errorf("failed to create snapshot file: %v", err)
    56  	}
    57  
    58  	// If anything goes wrong after this point, we will attempt to clean up
    59  	// the temp file. The happy path will disarm this.
    60  	var keep bool
    61  	defer func() {
    62  		if keep {
    63  			return
    64  		}
    65  
    66  		if err := os.Remove(archive.Name()); err != nil {
    67  			logger.Error("Failed to clean up temp snapshot", "error", err)
    68  		}
    69  	}()
    70  
    71  	hash := sha256.New()
    72  	out := io.MultiWriter(hash, archive)
    73  
    74  	// Wrap the file writer in a gzip compressor.
    75  	compressor := gzip.NewWriter(out)
    76  
    77  	// Write the archive.
    78  	if err := write(compressor, metadata, snap); err != nil {
    79  		return nil, fmt.Errorf("failed to write snapshot file: %v", err)
    80  	}
    81  
    82  	// Finish the compressed stream.
    83  	if err := compressor.Close(); err != nil {
    84  		return nil, fmt.Errorf("failed to compress snapshot file: %v", err)
    85  	}
    86  
    87  	// Sync the compressed file and rewind it so it's ready to be streamed
    88  	// out by the caller.
    89  	if err := archive.Sync(); err != nil {
    90  		return nil, fmt.Errorf("failed to sync snapshot: %v", err)
    91  	}
    92  	if _, err := archive.Seek(0, 0); err != nil {
    93  		return nil, fmt.Errorf("failed to rewind snapshot: %v", err)
    94  	}
    95  
    96  	checksum := "sha-256=" + base64.StdEncoding.EncodeToString(hash.Sum(nil))
    97  
    98  	keep = true
    99  	return &Snapshot{archive, metadata.Index, checksum}, nil
   100  }
   101  
   102  // Index returns the index of the snapshot. This is safe to call on a nil
   103  // snapshot, it will just return 0.
   104  func (s *Snapshot) Index() uint64 {
   105  	if s == nil {
   106  		return 0
   107  	}
   108  	return s.index
   109  }
   110  
   111  func (s *Snapshot) Checksum() string {
   112  	if s == nil {
   113  		return ""
   114  	}
   115  	return s.checksum
   116  }
   117  
   118  // Read passes through to the underlying snapshot file. This is safe to call on
   119  // a nil snapshot, it will just return an EOF.
   120  func (s *Snapshot) Read(p []byte) (n int, err error) {
   121  	if s == nil {
   122  		return 0, io.EOF
   123  	}
   124  	return s.file.Read(p)
   125  }
   126  
   127  // Close closes the snapshot and removes any temporary storage associated with
   128  // it. You must arrange to call this whenever NewSnapshot() has been called
   129  // successfully. This is safe to call on a nil snapshot.
   130  func (s *Snapshot) Close() error {
   131  	if s == nil {
   132  		return nil
   133  	}
   134  
   135  	if err := s.file.Close(); err != nil {
   136  		return err
   137  	}
   138  	return os.Remove(s.file.Name())
   139  }
   140  
   141  type Discard struct {
   142  	io.Writer
   143  }
   144  
   145  func (dc Discard) Close() error { return nil }
   146  
   147  // Verify takes the snapshot from the reader and verifies its contents.
   148  func Verify(in io.Reader) (*raft.SnapshotMeta, error) {
   149  	return CopySnapshot(in, Discard{Writer: io.Discard})
   150  }
   151  
   152  // CopySnapshot copies the snapshot content from snapshot archive to dest.
   153  // It will close the destination once complete.
   154  func CopySnapshot(in io.Reader, dest io.WriteCloser) (*raft.SnapshotMeta, error) {
   155  	defer dest.Close()
   156  
   157  	// Wrap the reader in a gzip decompressor.
   158  	decomp, err := gzip.NewReader(in)
   159  	if err != nil {
   160  		return nil, fmt.Errorf("failed to decompress snapshot: %v", err)
   161  	}
   162  	defer decomp.Close()
   163  
   164  	// Read the archive, throwing away the snapshot data.
   165  	var metadata raft.SnapshotMeta
   166  	if err := read(decomp, &metadata, dest); err != nil {
   167  		return nil, fmt.Errorf("failed to read snapshot file: %v", err)
   168  	}
   169  
   170  	if err := concludeGzipRead(decomp); err != nil {
   171  		return nil, err
   172  	}
   173  
   174  	return &metadata, nil
   175  }
   176  
   177  // concludeGzipRead should be invoked after you think you've consumed all of
   178  // the data from the gzip stream. It will error if the stream was corrupt.
   179  //
   180  // The docs for gzip.Reader say: "Clients should treat data returned by Read as
   181  // tentative until they receive the io.EOF marking the end of the data."
   182  func concludeGzipRead(decomp *gzip.Reader) error {
   183  	extra, err := ioutil.ReadAll(decomp) // ReadAll consumes the EOF
   184  	if err != nil {
   185  		return err
   186  	} else if len(extra) != 0 {
   187  		return fmt.Errorf("%d unread uncompressed bytes remain", len(extra))
   188  	}
   189  	return nil
   190  }
   191  
   192  type readWrapper struct {
   193  	in io.Reader
   194  	c  int
   195  }
   196  
   197  func (r *readWrapper) Read(b []byte) (int, error) {
   198  	n, err := r.in.Read(b)
   199  	r.c += n
   200  	if err != nil && err != io.EOF {
   201  		return n, fmt.Errorf("failed to read after %v: %v", r.c, err)
   202  	}
   203  	return n, err
   204  }
   205  
   206  // Restore takes the snapshot from the reader and attempts to apply it to the
   207  // given Raft instance.
   208  func Restore(logger hclog.Logger, in io.Reader, r *raft.Raft) error {
   209  	// Wrap the reader in a gzip decompressor.
   210  	decomp, err := gzip.NewReader(&readWrapper{in, 0})
   211  	if err != nil {
   212  		return fmt.Errorf("failed to decompress snapshot: %v", err)
   213  	}
   214  	defer func() {
   215  		if err := decomp.Close(); err != nil {
   216  			logger.Error("Failed to close snapshot decompressor", "error", err)
   217  		}
   218  	}()
   219  
   220  	// Make a scratch file to receive the contents of the snapshot data so
   221  	// we can avoid buffering in memory.
   222  	snap, err := ioutil.TempFile("", "snapshot")
   223  	if err != nil {
   224  		return fmt.Errorf("failed to create temp snapshot file: %v", err)
   225  	}
   226  	defer func() {
   227  		if err := snap.Close(); err != nil {
   228  			logger.Error("Failed to close temp snapshot", "error", err)
   229  		}
   230  		if err := os.Remove(snap.Name()); err != nil {
   231  			logger.Error("Failed to clean up temp snapshot", "error", err)
   232  		}
   233  	}()
   234  
   235  	// Read the archive.
   236  	var metadata raft.SnapshotMeta
   237  	if err := read(decomp, &metadata, snap); err != nil {
   238  		return fmt.Errorf("failed to read snapshot file: %v", err)
   239  	}
   240  
   241  	if err := concludeGzipRead(decomp); err != nil {
   242  		return err
   243  	}
   244  
   245  	// Sync and rewind the file so it's ready to be read again.
   246  	if err := snap.Sync(); err != nil {
   247  		return fmt.Errorf("failed to sync temp snapshot: %v", err)
   248  	}
   249  	if _, err := snap.Seek(0, 0); err != nil {
   250  		return fmt.Errorf("failed to rewind temp snapshot: %v", err)
   251  	}
   252  
   253  	// Feed the snapshot into Raft.
   254  	if err := r.Restore(&metadata, snap, 0); err != nil {
   255  		return fmt.Errorf("Raft error when restoring snapshot: %v", err)
   256  	}
   257  
   258  	return nil
   259  }