github.com/lazyledger/lazyledger-core@v0.35.0-dev.0.20210613111200-4c651f053571/test/e2e/app/snapshots.go (about)

     1  // nolint: gosec
     2  package main
     3  
     4  import (
     5  	"crypto/sha256"
     6  	"encoding/json"
     7  	"errors"
     8  	"fmt"
     9  	"io/ioutil"
    10  	"math"
    11  	"os"
    12  	"path/filepath"
    13  	"sync"
    14  
    15  	abci "github.com/lazyledger/lazyledger-core/abci/types"
    16  )
    17  
    18  const (
    19  	snapshotChunkSize = 1e6
    20  )
    21  
    22  // SnapshotStore stores state sync snapshots. Snapshots are stored simply as
    23  // JSON files, and chunks are generated on-the-fly by splitting the JSON data
    24  // into fixed-size chunks.
    25  type SnapshotStore struct {
    26  	sync.RWMutex
    27  	dir      string
    28  	metadata []abci.Snapshot
    29  }
    30  
    31  // NewSnapshotStore creates a new snapshot store.
    32  func NewSnapshotStore(dir string) (*SnapshotStore, error) {
    33  	store := &SnapshotStore{dir: dir}
    34  	if err := os.MkdirAll(dir, 0755); err != nil {
    35  		return nil, err
    36  	}
    37  	if err := store.loadMetadata(); err != nil {
    38  		return nil, err
    39  	}
    40  	return store, nil
    41  }
    42  
    43  // loadMetadata loads snapshot metadata. Does not take out locks, since it's
    44  // called internally on construction.
    45  func (s *SnapshotStore) loadMetadata() error {
    46  	file := filepath.Join(s.dir, "metadata.json")
    47  	metadata := []abci.Snapshot{}
    48  
    49  	bz, err := ioutil.ReadFile(file)
    50  	switch {
    51  	case errors.Is(err, os.ErrNotExist):
    52  	case err != nil:
    53  		return fmt.Errorf("failed to load snapshot metadata from %q: %w", file, err)
    54  	}
    55  	if len(bz) != 0 {
    56  		err = json.Unmarshal(bz, &metadata)
    57  		if err != nil {
    58  			return fmt.Errorf("invalid snapshot data in %q: %w", file, err)
    59  		}
    60  	}
    61  	s.metadata = metadata
    62  	return nil
    63  }
    64  
    65  // saveMetadata saves snapshot metadata. Does not take out locks, since it's
    66  // called internally from e.g. Create().
    67  func (s *SnapshotStore) saveMetadata() error {
    68  	bz, err := json.Marshal(s.metadata)
    69  	if err != nil {
    70  		return err
    71  	}
    72  
    73  	// save the file to a new file and move it to make saving atomic.
    74  	newFile := filepath.Join(s.dir, "metadata.json.new")
    75  	file := filepath.Join(s.dir, "metadata.json")
    76  	err = ioutil.WriteFile(newFile, bz, 0644) // nolint: gosec
    77  	if err != nil {
    78  		return err
    79  	}
    80  	return os.Rename(newFile, file)
    81  }
    82  
    83  // Create creates a snapshot of the given application state's key/value pairs.
    84  func (s *SnapshotStore) Create(state *State) (abci.Snapshot, error) {
    85  	s.Lock()
    86  	defer s.Unlock()
    87  	bz, err := state.Export()
    88  	if err != nil {
    89  		return abci.Snapshot{}, err
    90  	}
    91  	hash := sha256.Sum256(bz)
    92  	snapshot := abci.Snapshot{
    93  		Height: state.Height,
    94  		Format: 1,
    95  		Hash:   hash[:],
    96  		Chunks: byteChunks(bz),
    97  	}
    98  	err = ioutil.WriteFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", state.Height)), bz, 0644)
    99  	if err != nil {
   100  		return abci.Snapshot{}, err
   101  	}
   102  	s.metadata = append(s.metadata, snapshot)
   103  	err = s.saveMetadata()
   104  	if err != nil {
   105  		return abci.Snapshot{}, err
   106  	}
   107  	return snapshot, nil
   108  }
   109  
   110  // List lists available snapshots.
   111  func (s *SnapshotStore) List() ([]*abci.Snapshot, error) {
   112  	s.RLock()
   113  	defer s.RUnlock()
   114  	snapshots := []*abci.Snapshot{}
   115  	for _, snapshot := range s.metadata {
   116  		s := snapshot // copy to avoid pointer to range variable
   117  		snapshots = append(snapshots, &s)
   118  	}
   119  	return snapshots, nil
   120  }
   121  
   122  // LoadChunk loads a snapshot chunk.
   123  func (s *SnapshotStore) LoadChunk(height uint64, format uint32, chunk uint32) ([]byte, error) {
   124  	s.RLock()
   125  	defer s.RUnlock()
   126  	for _, snapshot := range s.metadata {
   127  		if snapshot.Height == height && snapshot.Format == format {
   128  			bz, err := ioutil.ReadFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", height)))
   129  			if err != nil {
   130  				return nil, err
   131  			}
   132  			return byteChunk(bz, chunk), nil
   133  		}
   134  	}
   135  	return nil, nil
   136  }
   137  
   138  // byteChunk returns the chunk at a given index from the full byte slice.
   139  func byteChunk(bz []byte, index uint32) []byte {
   140  	start := int(index * snapshotChunkSize)
   141  	end := int((index + 1) * snapshotChunkSize)
   142  	switch {
   143  	case start >= len(bz):
   144  		return nil
   145  	case end >= len(bz):
   146  		return bz[start:]
   147  	default:
   148  		return bz[start:end]
   149  	}
   150  }
   151  
   152  // byteChunks calculates the number of chunks in the byte slice.
   153  func byteChunks(bz []byte) uint32 {
   154  	return uint32(math.Ceil(float64(len(bz)) / snapshotChunkSize))
   155  }