github.com/badrootd/nibiru-cometbft@v0.37.5-0.20240307173500-2a75559eee9b/test/e2e/app/snapshots.go (about)

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