From 67c8c575ccb9e962df521435cc3827435b5e16dd Mon Sep 17 00:00:00 2001
From: Ilya Raykker <ilya.th.kapitonov@gmail.com>
Date: Tue, 26 Nov 2024 16:31:24 +0400
Subject: [PATCH] Correctly implement WasmSnapshotter

---
 x/compute/internal/keeper/wasm_snapshotter.go | 69 +++++--------------
 1 file changed, 17 insertions(+), 52 deletions(-)

diff --git a/x/compute/internal/keeper/wasm_snapshotter.go b/x/compute/internal/keeper/wasm_snapshotter.go
index 20701a542..8e181a6f5 100644
--- a/x/compute/internal/keeper/wasm_snapshotter.go
+++ b/x/compute/internal/keeper/wasm_snapshotter.go
@@ -13,31 +13,15 @@ import (
 	storetypes "cosmossdk.io/store/types"
 	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	protoio "github.com/cosmos/gogoproto/io"
 	"github.com/scrtlabs/SecretNetwork/x/compute/internal/types"
 )
 
 /*
 API to implement:
 
-// Snapshotter is something that can create and restore snapshots, consisting of streamed binary
-// chunks - all of which must be read from the channel and closed. If an unsupported format is
-// given, it must return ErrUnknownFormat (possibly wrapped with fmt.Errorf).
-type Snapshotter interface {
-	// Snapshot writes snapshot items into the protobuf writer.
-	Snapshot(height uint64, protoWriter protoio.Writer) error
-
-	// Restore restores a state snapshot from the protobuf items read from the reader.
-	// If the ready channel is non-nil, it returns a ready signal (by being closed) once the
-	// restorer is ready to accept chunks.
-	Restore(height uint64, format uint32, protoReader protoio.Reader) (SnapshotItem, error)
-}
-
 // ExtensionSnapshotter is an extension Snapshotter that is appended to the snapshot stream.
 // ExtensionSnapshotter has an unique name and manages it's own internal formats.
 type ExtensionSnapshotter interface {
-	Snapshotter
-
 	// SnapshotName returns the name of snapshotter, it should be unique in the manager.
 	SnapshotName() string
 
@@ -48,6 +32,13 @@ type ExtensionSnapshotter interface {
 
 	// SupportedFormats returns a list of formats it can restore from.
 	SupportedFormats() []uint32
+
+	// SnapshotExtension writes extension payloads into the underlying protobuf stream.
+	SnapshotExtension(height uint64, payloadWriter ExtensionPayloadWriter) error
+
+	// RestoreExtension restores an extension state snapshot,
+	// the payload reader returns `io.EOF` when reached the extension boundaries.
+	RestoreExtension(height uint64, format uint32, payloadReader ExtensionPayloadReader) error
 }
 */
 
@@ -81,7 +72,7 @@ func (ws *WasmSnapshotter) SupportedFormats() []uint32 {
 	return []uint32{1}
 }
 
-func (ws *WasmSnapshotter) Snapshot(height uint64, protoWriter protoio.Writer) error {
+func (ws *WasmSnapshotter) SnapshotExtension(height uint64, payloadWriter snapshottypes.ExtensionPayloadWriter) error {
 	// TODO: This seems more correct (historical info), but kills my tests
 	// Since codeIDs and wasm are immutible, it is never wrong to return new wasm data than the
 	// user requests
@@ -113,7 +104,7 @@ func (ws *WasmSnapshotter) Snapshot(height uint64, protoWriter protoio.Writer) e
 			return true
 		}
 
-		err = snapshottypes.WriteExtensionPayload(protoWriter, wasmBytes)
+		err = payloadWriter(wasmBytes)
 		if err != nil {
 			rerr = err
 			return true
@@ -125,54 +116,28 @@ func (ws *WasmSnapshotter) Snapshot(height uint64, protoWriter protoio.Writer) e
 	return rerr
 }
 
-func (ws *WasmSnapshotter) Restore(
-	height uint64, format uint32, protoReader protoio.Reader, //nolint:all
-) (snapshottypes.SnapshotItem, error) {
+func (ws *WasmSnapshotter) RestoreExtension(
+	height uint64, format uint32, payloadReader snapshottypes.ExtensionPayloadReader, //nolint:all
+) error {
 	if format != 1 {
-		return snapshottypes.SnapshotItem{}, snapshottypes.ErrUnknownFormat
+		return snapshottypes.ErrUnknownFormat
 	}
 
 	for {
-		item := snapshottypes.SnapshotItem{}
-		err := protoReader.ReadMsg(&item)
+		wasmBytes, err := payloadReader()
 		if err == io.EOF {
-			return snapshottypes.SnapshotItem{}, nil
+			return nil
 		} else if err != nil {
-			return snapshottypes.SnapshotItem{}, errorsmod.Wrap(err, "invalid protobuf message")
-		}
-
-		// if it is not another ExtensionPayload message, then it is not for us.
-		// we should return it an let the manager handle this one
-		payload := item.GetExtensionPayload()
-		if payload == nil {
-			return item, nil
+			return errorsmod.Wrap(err, "invalid protobuf message")
 		}
 
-		wasmBytes := payload.Payload
-
 		codeHash := sha256.Sum256(wasmBytes)
 		wasmFileName := hex.EncodeToString(codeHash[:])
 		wasmFilePath := filepath.Join(ws.wasmDirectory, wasmFileName)
 
 		err = os.WriteFile(wasmFilePath, wasmBytes, 0o600 /* -rw------- */)
 		if err != nil {
-			return snapshottypes.SnapshotItem{}, errorsmod.Wrapf(err, "failed to write wasm file '%v' to disk", wasmFilePath)
+			return errorsmod.Wrapf(err, "failed to write wasm file '%v' to disk", wasmFilePath)
 		}
 	}
 }
-
-func (ws *WasmSnapshotter) PruneSnapshotHeight(_ int64) {
-	panic("not implemented")
-}
-
-func (ws *WasmSnapshotter) SetSnapshotInterval(_ uint64) {
-	panic("not implemented")
-}
-
-func (ws *WasmSnapshotter) RestoreExtension(_ uint64, _ uint32, _ snapshottypes.ExtensionPayloadReader) error {
-	panic("not implemented")
-}
-
-func (ws *WasmSnapshotter) SnapshotExtension(_ uint64, _ snapshottypes.ExtensionPayloadWriter) error {
-	panic("not implemented")
-}