From 676daa1782f0130a6832f9ac248f03704567d71c Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 18 Jan 2022 10:08:13 +0300
Subject: [PATCH] [#215] container: remove old estimations when adding new ones

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 container/container_contract.go | 55 ++++++++++++++++++++++++++++++---
 tests/container_test.go         | 18 +++++++++--
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/container/container_contract.go b/container/container_contract.go
index 1a9ff96..f132876 100644
--- a/container/container_contract.go
+++ b/container/container_contract.go
@@ -3,6 +3,7 @@ package container
 import (
 	"github.com/nspcc-dev/neo-go/pkg/interop"
 	"github.com/nspcc-dev/neo-go/pkg/interop/contract"
+	"github.com/nspcc-dev/neo-go/pkg/interop/convert"
 	"github.com/nspcc-dev/neo-go/pkg/interop/iterator"
 	"github.com/nspcc-dev/neo-go/pkg/interop/native/crypto"
 	"github.com/nspcc-dev/neo-go/pkg/interop/native/management"
@@ -59,10 +60,15 @@ const (
 	// V2 format
 	containerIDSize = 32 // SHA256 size
 
-	estimateKeyPrefix   = "cnr"
-	estimatePostfixSize = 10
-	// CleanupDelta contains number last epochs for which container estimations are present.
+	singleEstimatePrefix = "est"
+	estimateKeyPrefix    = "cnr"
+	estimatePostfixSize  = 10
+	// CleanupDelta contains number of last epochs for which container estimations are present.
 	CleanupDelta = 3
+	// TotalCleanupDelta contains number of epochs after which estimation
+	// will be removed by epoch tick cleanup if any node didn't updated
+	// container size and/or container was removed. Must be greater than CleanupDelta.
+	TotalCleanupDelta = CleanupDelta + 1
 
 	// NotFoundError is returned if container is missing.
 	NotFoundError = "container does not exist"
@@ -88,6 +94,20 @@ func _deploy(data interface{}, isUpdate bool) {
 		args := data.([]interface{})
 		common.CheckVersion(args[len(args)-1].(int))
 		storage.Delete(ctx, common.LegacyOwnerKey)
+
+		// Migrate container estimation keys.
+		it := storage.Find(ctx, []byte(estimateKeyPrefix), storage.DeserializeValues)
+		for iterator.Next(it) {
+			kv := iterator.Value(it).(struct {
+				key   []byte
+				value estimation
+			})
+
+			end := len(kv.key) - containerIDSize - estimatePostfixSize
+			rawEpoch := kv.key[len(estimateKeyPrefix):end]
+			cid := kv.key[end : len(kv.key)-estimatePostfixSize]
+			updateEstimations(ctx, convert.ToInteger(rawEpoch), cid, kv.value.from, true)
+		}
 		return
 	}
 
@@ -503,6 +523,7 @@ func PutContainerSize(epoch int, cid []byte, usedSize int, pubKey interop.Public
 	}
 
 	storage.Put(ctx, key, std.Serialize(s))
+	updateEstimations(ctx, epoch, cid, pubKey, false)
 
 	runtime.Log("saved container size estimation")
 }
@@ -772,6 +793,32 @@ func isStorageNode(ctx storage.Context, key interop.PublicKey) bool {
 	return false
 }
 
+func updateEstimations(ctx storage.Context, epoch int, cid []byte, pub interop.PublicKey, isUpdate bool) {
+	h := crypto.Ripemd160(pub)
+	estKey := append([]byte(singleEstimatePrefix), cid...)
+	estKey = append(estKey, h...)
+
+	var newEpochs []int
+	rawList := storage.Get(ctx, estKey).([]byte)
+
+	if rawList != nil {
+		epochs := std.Deserialize(rawList).([]int)
+		for _, oldEpoch := range epochs {
+			if !isUpdate && epoch-oldEpoch > CleanupDelta {
+				key := append([]byte(estimateKeyPrefix), convert.ToBytes(oldEpoch)...)
+				key = append(key, cid...)
+				key = append(key, h[:estimatePostfixSize]...)
+				storage.Delete(ctx, key)
+			} else {
+				newEpochs = append(newEpochs, oldEpoch)
+			}
+		}
+	}
+
+	newEpochs = append(newEpochs, epoch)
+	common.SetSerialized(ctx, estKey, newEpochs)
+}
+
 func cleanupContainers(ctx storage.Context, epoch int) {
 	it := storage.Find(ctx, []byte(estimateKeyPrefix), storage.KeysOnly)
 	for iterator.Next(it) {
@@ -781,7 +828,7 @@ func cleanupContainers(ctx storage.Context, epoch int) {
 
 		var n interface{} = nbytes
 
-		if epoch-n.(int) > CleanupDelta {
+		if epoch-n.(int) > TotalCleanupDelta {
 			storage.Delete(ctx, k)
 		}
 	}
diff --git a/tests/container_test.go b/tests/container_test.go
index 15f5b71..520a725 100644
--- a/tests/container_test.go
+++ b/tests/container_test.go
@@ -308,9 +308,23 @@ func TestContainerSizeEstimation(t *testing.T) {
 		checkEstimations(t, c, 3, cnt, estimation{nodes[2].pub, 888})
 	}
 
-	cNm.Invoke(t, stackitem.Null{}, "newEpoch", int64(2+container.CleanupDelta+1))
-	checkEstimations(t, c, 2, cnt)
+	epoch := int64(2 + container.CleanupDelta + 1)
+	cNm.Invoke(t, stackitem.Null{}, "newEpoch", epoch)
+	checkEstimations(t, c, 2, cnt, estimations...) // not yet removed
 	checkEstimations(t, c, 3, cnt, estimation{nodes[2].pub, 888})
+
+	c.WithSigners(nodes[1].signer).Invoke(t, stackitem.Null{}, "putContainerSize",
+		epoch, cnt.id[:], int64(999), nodes[1].pub)
+
+	checkEstimations(t, c, 2, cnt, estimations[:1]...)
+	checkEstimations(t, c, epoch, cnt, estimation{nodes[1].pub, int64(999)})
+
+	// Estimation from node 0 should be cleaned during epoch tick.
+	for i := int64(1); i <= container.TotalCleanupDelta-container.CleanupDelta; i++ {
+		cNm.Invoke(t, stackitem.Null{}, "newEpoch", epoch+i)
+	}
+	checkEstimations(t, c, 2, cnt)
+	checkEstimations(t, c, epoch, cnt, estimation{nodes[1].pub, int64(999)})
 }
 
 type estimation struct {