github.com/onflow/flow-go@v0.33.17/utils/debug/registerCache.go (about)

     1  package debug
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/hex"
     6  	"encoding/json"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  
    11  	"github.com/onflow/flow-go/model/flow"
    12  )
    13  
    14  type registerCache interface {
    15  	Get(owner, key string) (value []byte, found bool)
    16  	Set(owner, key string, value []byte)
    17  	Persist() error
    18  }
    19  
    20  type memRegisterCache struct {
    21  	data map[string]flow.RegisterValue
    22  }
    23  
    24  func newMemRegisterCache() *memRegisterCache {
    25  	return &memRegisterCache{data: make(map[string]flow.RegisterValue)}
    26  
    27  }
    28  func (c *memRegisterCache) Get(owner, key string) ([]byte, bool) {
    29  	v, found := c.data[owner+"~"+key]
    30  	return v, found
    31  }
    32  
    33  func (c *memRegisterCache) Set(owner, key string, value []byte) {
    34  	c.data[owner+"~"+key] = value
    35  }
    36  func (c *memRegisterCache) Persist() error {
    37  	// No-op
    38  	return nil
    39  }
    40  
    41  type fileRegisterCache struct {
    42  	filePath string
    43  	data     map[string]flow.RegisterEntry
    44  }
    45  
    46  func newFileRegisterCache(filePath string) *fileRegisterCache {
    47  	cache := &fileRegisterCache{filePath: filePath}
    48  	data := make(map[string]flow.RegisterEntry)
    49  
    50  	if _, err := os.Stat(filePath); err == nil {
    51  		f, err := os.Open(filePath)
    52  		if err != nil {
    53  			fmt.Printf("error opening file: %v\n", err)
    54  			os.Exit(1)
    55  		}
    56  		defer f.Close()
    57  		r := bufio.NewReader(f)
    58  		var s string
    59  		for {
    60  			s, err = r.ReadString('\n')
    61  			if err != nil && err != io.EOF {
    62  				break
    63  			}
    64  			if len(s) > 0 {
    65  				var d flow.RegisterEntry
    66  				if err := json.Unmarshal([]byte(s), &d); err != nil {
    67  					panic(err)
    68  				}
    69  				owner, err := hex.DecodeString(d.Key.Owner)
    70  				if err != nil {
    71  					panic(err)
    72  				}
    73  				keyCopy, err := hex.DecodeString(d.Key.Key)
    74  				if err != nil {
    75  					panic(err)
    76  				}
    77  				data[string(owner)+"~"+string(keyCopy)] = d
    78  			}
    79  			if err != nil {
    80  				break
    81  			}
    82  		}
    83  	}
    84  
    85  	cache.data = data
    86  	return cache
    87  }
    88  
    89  func (f *fileRegisterCache) Get(owner, key string) ([]byte, bool) {
    90  	v, found := f.data[owner+"~"+key]
    91  	if found {
    92  		return v.Value, found
    93  	}
    94  	return nil, found
    95  }
    96  
    97  func (f *fileRegisterCache) Set(owner, key string, value []byte) {
    98  	valueCopy := make([]byte, len(value))
    99  	copy(valueCopy, value)
   100  	ownerAddr := flow.BytesToAddress([]byte(owner))
   101  	fmt.Println(ownerAddr.Hex(), hex.EncodeToString([]byte(key)), len(value))
   102  	f.data[owner+"~"+key] = flow.RegisterEntry{
   103  		Key:   flow.NewRegisterID(ownerAddr, hex.EncodeToString([]byte(key))),
   104  		Value: flow.RegisterValue(valueCopy),
   105  	}
   106  }
   107  
   108  func (c *fileRegisterCache) Persist() error {
   109  	f, err := os.OpenFile(c.filePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0755)
   110  	if err != nil {
   111  		return err
   112  	}
   113  	defer f.Close()
   114  	w := bufio.NewWriter(f)
   115  	for _, v := range c.data {
   116  		fltV, err := json.Marshal(v)
   117  		if err != nil {
   118  			return err
   119  		}
   120  		_, err = w.WriteString(string(fltV))
   121  		if err != nil {
   122  			return err
   123  		}
   124  		err = w.WriteByte('\n')
   125  		if err != nil {
   126  			return err
   127  		}
   128  	}
   129  	err = w.Flush()
   130  	if err != nil {
   131  		return err
   132  	}
   133  	return nil
   134  }