github.com/ari-anchor/sei-tendermint@v0.0.0-20230519144642-dc826b7b56bb/internal/dbsync/snapshot.go (about)

     1  package dbsync
     2  
     3  import (
     4  	"crypto/md5"
     5  	"fmt"
     6  	"io"
     7  	"io/ioutil"
     8  	"os"
     9  	"path"
    10  	"sync"
    11  
    12  	"github.com/ari-anchor/sei-tendermint/config"
    13  	dstypes "github.com/ari-anchor/sei-tendermint/proto/tendermint/dbsync"
    14  )
    15  
    16  func Snapshot(height uint64, dbsyncConfig config.DBSyncConfig, baseConfig config.BaseConfig) error {
    17  	src := path.Join(baseConfig.DBDir(), ApplicationDBSubdirectory)
    18  	wasmSrc := path.Join(baseConfig.RootDir, WasmDirectory)
    19  	dst := path.Join(dbsyncConfig.SnapshotDirectory, fmt.Sprintf("%s%d", HeightSubdirectoryPrefix, height))
    20  	os.RemoveAll(dst)
    21  	err := os.MkdirAll(dst, os.ModePerm)
    22  	if err != nil {
    23  		return fmt.Errorf("error creating directory %s - %s", dst, err)
    24  	}
    25  	var fds []os.FileInfo
    26  	if fds, err = ioutil.ReadDir(src); err != nil {
    27  		return err
    28  	}
    29  	wasmNames := map[string]struct{}{}
    30  	if wasmFds, _ := ioutil.ReadDir(wasmSrc); wasmFds != nil {
    31  		fds = append(fds, wasmFds...)
    32  		for _, fd := range wasmFds {
    33  			wasmNames[fd.Name()] = struct{}{}
    34  		}
    35  	}
    36  
    37  	assignments := make([][]os.FileInfo, dbsyncConfig.SnapshotWorkerCount)
    38  
    39  	for i, fd := range fds {
    40  		assignments[i%dbsyncConfig.SnapshotWorkerCount] = append(assignments[i%dbsyncConfig.SnapshotWorkerCount], fd)
    41  	}
    42  
    43  	metadata := dstypes.MetadataResponse{
    44  		Height:      height,
    45  		Filenames:   []string{},
    46  		Md5Checksum: [][]byte{},
    47  	}
    48  	metadataMtx := &sync.Mutex{}
    49  
    50  	wg := sync.WaitGroup{}
    51  	for i := 0; i < dbsyncConfig.SnapshotWorkerCount; i++ {
    52  		wg.Add(1)
    53  		assignment := assignments[i]
    54  		go func() {
    55  			for _, fd := range assignment {
    56  				var srcfp, dstfp string
    57  				if _, ok := wasmNames[fd.Name()]; ok {
    58  					srcfp = path.Join(wasmSrc, fd.Name())
    59  					dstfp = path.Join(dst, fd.Name()) + WasmSuffix
    60  				} else {
    61  					srcfp = path.Join(src, fd.Name())
    62  					dstfp = path.Join(dst, fd.Name())
    63  				}
    64  
    65  				var srcfd *os.File
    66  				var dstfd *os.File
    67  				if srcfd, err = os.Open(srcfp); err != nil {
    68  					panic(err)
    69  				}
    70  
    71  				if dstfd, err = os.Create(dstfp); err != nil {
    72  					srcfd.Close()
    73  					panic(err)
    74  				}
    75  
    76  				if _, err = io.Copy(dstfd, srcfd); err != nil {
    77  					srcfd.Close()
    78  					dstfd.Close()
    79  					panic(err)
    80  				}
    81  
    82  				filename := fd.Name()
    83  				if _, ok := wasmNames[fd.Name()]; ok {
    84  					filename += WasmSuffix
    85  				}
    86  
    87  				bz, err := ioutil.ReadFile(path.Join(dst, filename))
    88  				if err != nil {
    89  					panic(err)
    90  				}
    91  				sum := md5.Sum(bz)
    92  
    93  				metadataMtx.Lock()
    94  				metadata.Filenames = append(metadata.Filenames, filename)
    95  				metadata.Md5Checksum = append(metadata.Md5Checksum, sum[:])
    96  
    97  				metadataMtx.Unlock()
    98  				srcfd.Close()
    99  				dstfd.Close()
   100  			}
   101  			wg.Done()
   102  		}()
   103  	}
   104  	wg.Wait()
   105  
   106  	metadataBz, err := metadata.Marshal()
   107  	if err != nil {
   108  		return err
   109  	}
   110  
   111  	metadataFile, err := os.Create(path.Join(dst, MetadataFilename))
   112  	if err != nil {
   113  		return err
   114  	}
   115  	defer metadataFile.Close()
   116  	_, err = metadataFile.Write(metadataBz)
   117  	if err != nil {
   118  		return err
   119  	}
   120  
   121  	heightFile, err := os.Create(path.Join(dbsyncConfig.SnapshotDirectory, MetadataHeightFilename))
   122  	if err != nil {
   123  		return err
   124  	}
   125  	defer heightFile.Close()
   126  	_, err = heightFile.Write([]byte(fmt.Sprintf("%d", height)))
   127  	return err
   128  }