CNS-352: Reset CU at month end

common/fixation_entry_index.go

@@ -40,10 +40,11 @@ func (fs FixationStore) removeEntryIndex(ctx sdk.Context, safeIndex string) {
 	store.Delete(types.KeyPrefix(fs.createEntryIndexKey(safeIndex)))
 }
 
-// GetAllEntryIndex returns all Entry indices
-func (fs FixationStore) GetAllEntryIndices(ctx sdk.Context) []string {
+// GetAllEntryIndexWithPrefix returns all Entry indices with a given prefix
+func (fs FixationStore) GetAllEntryIndicesWithPrefix(ctx sdk.Context, prefix string) []string {
 	store := fs.getEntryIndexStore(ctx)
-	iterator := sdk.KVStorePrefixIterator(store, []byte{})
+	entryPrefix := types.KeyPrefix(fs.createEntryIndexKey(prefix))
+	iterator := sdk.KVStorePrefixIterator(store, entryPrefix)
 	defer iterator.Close()
 
 	// iterate over the store's values and save the indices in a list
@@ -57,6 +58,11 @@ func (fs FixationStore) GetAllEntryIndices(ctx sdk.Context) []string {
 	return indexList
 }
 
+// GetAllEntryIndex returns all Entry indices
+func (fs FixationStore) GetAllEntryIndices(ctx sdk.Context) []string {
+	return fs.GetAllEntryIndicesWithPrefix(ctx, "")
+}
+
 func (fs FixationStore) createEntryIndexStoreKey() string {
 	return types.EntryIndexKey + fs.prefix
 }

common/fixation_entry_test.go

@@ -134,6 +134,9 @@ func testWithTemplate(t *testing.T, playbook []template, countObj int, countVS i
 		case "getall":
 			indexList := vs[play.store].GetAllEntryIndices(ctx)
 			require.Equal(t, int(play.count), len(indexList), what)
+		case "getallprefix":
+			indexList := vs[play.store].GetAllEntryIndicesWithPrefix(ctx, index)
+			require.Equal(t, int(play.count), len(indexList), what)
 		}
 	}
 }
@@ -316,3 +319,23 @@ func TestEntriesSort(t *testing.T) {
 
 	testWithTemplate(t, playbook, 3, 1)
 }
+
+func TestGetAllEntries(t *testing.T) {
+	block0 := int64(10)
+
+	playbook := []template{
+		{ op: "append", name: "entry #1", index: "prefix1_a", count: block0, coin: 0 },
+		{ op: "append", name: "entry #1", index: "prefix1_b", count: block0, coin: 1 },
+		{ op: "append", name: "entry #1", index: "prefix1_c", count: block0, coin: 2 },
+		{ op: "append", name: "entry #1", index: "prefix2_a", count: block0, coin: 3 },
+		{ op: "append", name: "entry #1", index: "prefix2_b", count: block0, coin: 4 },
+		{ op: "append", name: "entry #1", index: "prefix3_a", count: block0, coin: 5 },
+		{ op: "getall", name: "to check all indices", count: 6 },
+		{ op: "getallprefix", name: "to check all indices with prefix", index: "prefix", count: 6 },
+		{ op: "getallprefix", name: "to check indices with prefix1", index: "prefix1", count: 3 },
+		{ op: "getallprefix", name: "to check indices with prefix2", index: "prefix2", count: 2 },
+		{ op: "getallprefix", name: "to check indices with prefix3", index: "prefix3", count: 1 },
+	}
+
+	testWithTemplate(t, playbook, 6, 1)
+}

x/pairing/keeper/pairing_subscription_test.go

@@ -15,32 +15,55 @@ import (
 func TestGetPairingForSubscription(t *testing.T) {
 	ts := setupForPaymentTest(t)
 	var balance int64 = 10000
-	consumer := common.CreateNewAccount(ts.ctx, *ts.keepers, balance)
+
+	consumer := common.CreateNewAccount(ts.ctx, *ts.keepers, balance).Addr.String()
 	vrfpk_stub := "testvrfpk"
-	_, err := ts.servers.SubscriptionServer.Buy(ts.ctx, &subtypes.MsgBuy{Creator: consumer.Addr.String(), Consumer: consumer.Addr.String(), Index: ts.plan.Index, Duration: 1, Vrfpk: vrfpk_stub})
+	msgBuy := &subtypes.MsgBuy{
+		Creator:  consumer,
+		Consumer: consumer,
+		Index:    ts.plan.Index,
+		Duration: 1,
+		Vrfpk:    vrfpk_stub,
+	}
+	_, err := ts.servers.SubscriptionServer.Buy(ts.ctx, msgBuy)
 	require.Nil(t, err)
 
 	ts.ctx = testkeeper.AdvanceEpoch(ts.ctx, ts.keepers)
+	ctx := sdk.UnwrapSDKContext(ts.ctx)
 
-	pairingReq := types.QueryGetPairingRequest{ChainID: ts.spec.Index, Client: consumer.Addr.String()}
+	pairingReq := types.QueryGetPairingRequest{
+		ChainID: ts.spec.Index,
+		Client:  consumer,
+	}
 	pairing, err := ts.keepers.Pairing.GetPairing(ts.ctx, &pairingReq)
 	require.Nil(t, err)
 
-	verifyPairingQuery := &types.QueryVerifyPairingRequest{ChainID: ts.spec.Index, Client: consumer.Addr.String(), Provider: pairing.Providers[0].Address, Block: uint64(sdk.UnwrapSDKContext(ts.ctx).BlockHeight())}
+	verifyPairingQuery := &types.QueryVerifyPairingRequest{
+		ChainID:  ts.spec.Index,
+		Client:   consumer,
+		Provider: pairing.Providers[0].Address,
+		Block:    uint64(ctx.BlockHeight()),
+	}
 	vefiry, err := ts.keepers.Pairing.VerifyPairing(ts.ctx, verifyPairingQuery)
 	require.Nil(t, err)
 	require.True(t, vefiry.Valid)
 
-	project, vrfpk, err := ts.keepers.Projects.GetProjectForDeveloper(sdk.UnwrapSDKContext(ts.ctx), consumer.Addr.String(), uint64(sdk.UnwrapSDKContext(ts.ctx).BlockHeight()))
+	project, vrfpk, err := ts.keepers.Projects.GetProjectForDeveloper(ctx, consumer, uint64(ctx.BlockHeight()))
 	require.Nil(t, err)
 	require.Equal(t, vrfpk, vrfpk_stub)
-	err = ts.keepers.Projects.DeleteProject(sdk.UnwrapSDKContext(ts.ctx), project.Index)
+
+	err = ts.keepers.Projects.DeleteProject(ctx, project.Index)
 	require.Nil(t, err)
 
 	_, err = ts.keepers.Pairing.GetPairing(ts.ctx, &pairingReq)
 	require.NotNil(t, err)
 
-	verifyPairingQuery = &types.QueryVerifyPairingRequest{ChainID: ts.spec.Index, Client: consumer.Addr.String(), Provider: pairing.Providers[0].Address, Block: uint64(sdk.UnwrapSDKContext(ts.ctx).BlockHeight())}
+	verifyPairingQuery = &types.QueryVerifyPairingRequest{
+		ChainID:  ts.spec.Index,
+		Client:   consumer,
+		Provider: pairing.Providers[0].Address,
+		Block:    uint64(ctx.BlockHeight()),
+	}
 	vefiry, err = ts.keepers.Pairing.VerifyPairing(ts.ctx, verifyPairingQuery)
 	require.NotNil(t, err)
 	require.False(t, vefiry.Valid)

x/projects/keeper/creation.go

@@ -15,7 +15,7 @@ func (k Keeper) CreateAdminProject(ctx sdk.Context, subscriptionAddress string,
 
 // add a new project to the subscription
 func (k Keeper) CreateProject(ctx sdk.Context, subscriptionAddress string, projectName string, adminAddress string, enable bool, totalCU uint64, cuPerEpoch uint64, providers uint64, geolocation uint64, vrfpk string) error {
-	project := types.CreateProject(subscriptionAddress, projectName)
+	project := types.NewProject(subscriptionAddress, projectName)
 	var emptyProject types.Project
 
 	blockHeight := uint64(ctx.BlockHeight())
@@ -38,6 +38,7 @@ func (k Keeper) CreateProject(ctx sdk.Context, subscriptionAddress string, proje
 	}
 
 	project.Enabled = enable
+
 	return k.projectsFS.AppendEntry(ctx, project.Index, blockHeight, &project)
 }
 
@@ -57,8 +58,19 @@ func (k Keeper) RegisterDeveloperKey(ctx sdk.Context, developerKey string, proje
 	return nil
 }
 
+// Snapshot all projects of a given subscription
+func (k Keeper) SnapshotSubscriptionProjects(ctx sdk.Context, subscriptionAddr string) {
+	projects := k.projectsFS.GetAllEntryIndicesWithPrefix(ctx, subscriptionAddr)
+	for _, projectID := range projects {
+		err := k.snapshotProject(ctx, projectID)
+		if err != nil {
+			panic(err)
+		}
+	}
+}
+
 // snapshot project, create a snapshot of a project and reset the cu
-func (k Keeper) SnapshotProject(ctx sdk.Context, projectID string) error {
+func (k Keeper) snapshotProject(ctx sdk.Context, projectID string) error {
 	var project types.Project
 	err, found := k.projectsFS.FindEntry(ctx, projectID, uint64(ctx.BlockHeight()), &project)
 	if err != nil || !found {

x/projects/types/keys.go

@@ -19,7 +19,7 @@ const (
 	// prefix for the projects fixation store
 	ProjectsFixationPrefix = "prj-fs"
 
-	// prefix for the projects fixation store
+	// prefix for the developer keys fixation store
 	DeveloperKeysFixationPrefix = "dev-fs"
 )
 

x/projects/types/project.go

@@ -10,7 +10,7 @@ func ProjectIndex(subscriptionAddress string, projectName string) string {
 	return subscriptionAddress + "-" + projectName
 }
 
-func CreateProject(subscriptionAddress string, projectName string) Project {
+func NewProject(subscriptionAddress string, projectName string) Project {
 	return Project{
 		Index:        ProjectIndex(subscriptionAddress, projectName),
 		Subscription: subscriptionAddress,

x/subscription/keeper/epoch_start.go

@@ -58,8 +58,11 @@ func (k Keeper) EpochStart(ctx sdk.Context) {
 			date = nextMonth(date)
 			sub.MonthExpiryTime = uint64(date.Unix())
 
-			// reset CU allowance for this coming month
+			// reset subscription CU allowance for this coming month
 			sub.MonthCuLeft = sub.MonthCuTotal
+
+			// reset projects' CU allowance for this coming month
+			k.projectsKeeper.SnapshotSubscriptionProjects(ctx, sub.Consumer)
 		} else {
 			// duration ended, but don't delete yet - keep around for another
 			// EpochsToSave epochs before removing, to allow for payments for

x/subscription/keeper/subscription.go

@@ -251,3 +251,23 @@ func (k Keeper) CreateSubscription(
 
 	return nil
 }
+
+func (k Keeper) ChargeSubscription(ctx sdk.Context, consumer string, cuAmount uint64) error {
+	sub, found := k.GetSubscription(ctx, consumer)
+	if !found {
+		details := map[string]string{"consumer": consumer}
+		return utils.LavaError(ctx, k.Logger(ctx), "ChargeSubscription", details, "invalid subscription")
+	}
+
+	// TODO: if subscription is exhausted, should we push back (and the provider
+	// may not be paid?
+
+	if sub.MonthCuLeft < cuAmount {
+		sub.MonthCuLeft = 0
+	} else {
+		sub.MonthCuLeft -= cuAmount
+	}
+
+	k.SetSubscription(ctx, sub)
+	return nil
+}

x/subscription/keeper/subscription_test.go

@@ -119,7 +119,10 @@ func (ts *testStruct) expireSubscription(sub types.Subscription) types.Subscript
 
 func setupTestStruct(t *testing.T, numPlans int) testStruct {
 	_, keepers, _ctx := keepertest.InitAllKeepers(t)
+
+	_ctx = keepertest.AdvanceEpoch(_ctx, keepers)
 	ctx := sdk.UnwrapSDKContext(_ctx)
+
 	plans := createNPlans(&keepers.Plans, ctx, numPlans)
 
 	ts := testStruct{
@@ -278,7 +281,7 @@ func TestRenewSubscription(t *testing.T) {
 	sub, found := keeper.GetSubscription(ts.ctx, creator)
 	require.True(t, found)
 
-	// fast-forward two months
+	// fast-forward three months
 	sub = ts.expireSubscription(sub)
 	sub = ts.expireSubscription(sub)
 	sub = ts.expireSubscription(sub)
@@ -317,6 +320,69 @@ func TestSubscriptionAdminProject(t *testing.T) {
 	require.Nil(t, err)
 }
 
+func TestMonthlyRechargeCU(t *testing.T) {
+	ts := setupTestStruct(t, 1)
+	keeper := ts.keepers.Subscription
+	projectKeeper := ts.keepers.Projects
+
+	account := common.CreateNewAccount(ts._ctx, *ts.keepers, 10000)
+	creator := account.Addr.String()
+
+	err := keeper.CreateSubscription(ts.ctx, creator, creator, ts.plans[0].Index, 2, "")
+	require.Nil(t, err)
+
+	block1 := uint64(ts.ctx.BlockHeight())
+
+	sub, found := keeper.GetSubscription(ts.ctx, creator)
+	require.True(t, found)
+
+	ts._ctx = keepertest.AdvanceEpoch(ts._ctx, ts.keepers)
+	ts.ctx = sdk.UnwrapSDKContext(ts._ctx)
+
+	// use the subscription and the project
+	keeper.ChargeSubscription(ts.ctx, creator, 1000)
+	require.Equal(t, sub.PrevCuLeft, sub.MonthCuTotal-1000)
+	err = projectKeeper.AddComputeUnitsToProject(ts.ctx, creator, block1, 1000)
+	require.Nil(t, err)
+
+	// verify that project used the CU
+	proj, _, err := projectKeeper.GetProjectForDeveloper(ts.ctx, creator, block1)
+	require.Nil(t, err)
+	require.Equal(t, uint64(1000), proj.UsedCu)
+
+	block2 := uint64(ts.ctx.BlockHeight())
+
+	// force fixation entry (by adding project key)
+	projKey := []projectstypes.ProjectKey{
+		{
+			Key: common.CreateNewAccount(ts._ctx, *ts.keepers, 10000).Addr.String(),
+			Types: []projectstypes.ProjectKey_KEY_TYPE{projectstypes.ProjectKey_ADMIN},
+		},
+	}
+	projectKeeper.AddKeysToProject(ts.ctx, projectstypes.ADMIN_PROJECT_NAME, creator, projKey)
+
+	// fast-forward one months
+	sub = ts.expireSubscription(sub)
+	require.Equal(t, uint64(1), sub.DurationLeft)
+
+	block3 := uint64(ts.ctx.BlockHeight())
+
+	// check that subscription and project have renewed CUs, and that the
+	// project created a snapshot for last month
+	sub, found = keeper.GetSubscription(ts.ctx, creator)
+	require.True(t, found)
+	require.Equal(t, sub.MonthCuLeft, sub.MonthCuTotal)
+	proj, _, err = projectKeeper.GetProjectForDeveloper(ts.ctx, creator, block1)
+	require.Nil(t, err)
+	require.Equal(t, uint64(1000), proj.UsedCu)
+	proj, _, err = projectKeeper.GetProjectForDeveloper(ts.ctx, creator, block2)
+	require.Nil(t, err)
+	require.Equal(t, uint64(1000), proj.UsedCu)
+	proj, _, err = projectKeeper.GetProjectForDeveloper(ts.ctx, creator, block3)
+	require.Nil(t, err)
+	require.Equal(t, uint64(0), proj.UsedCu)
+}
+
 func TestExpiryTime(t *testing.T) {
 	ts := setupTestStruct(t, 1)
 	keeper := ts.keepers.Subscription
@@ -376,8 +442,10 @@ func TestExpiryTime(t *testing.T) {
 			require.Equal(t, tt.months, sub.DurationTotal)
 
 			keeper.RemoveSubscription(ts.ctx, creator)
+
 			// TODO: remove when RemoveSubscriptions properly removes projects
-			ts.keepers.Projects.DeleteProject(ts.ctx, projectstypes.ProjectIndex(creator, "default"))
+			projectID := projectstypes.ProjectIndex(creator, projectstypes.ADMIN_PROJECT_NAME)
+			ts.keepers.Projects.DeleteProject(ts.ctx, projectID)
 		})
 	}
 }
@@ -421,8 +489,10 @@ func TestPrice(t *testing.T) {
 			require.Equal(t, balance.Amount.Int64(), int64(10000-tt.cost))
 
 			keeper.RemoveSubscription(ts.ctx, creator)
+
 			// TODO: remove when RemoveSubscriptions properly removes projects
-			ts.keepers.Projects.DeleteProject(ts.ctx, projectstypes.ProjectIndex(creator, "default"))
+			projectID := projectstypes.ProjectIndex(creator, projectstypes.ADMIN_PROJECT_NAME)
+			ts.keepers.Projects.DeleteProject(ts.ctx, projectID)
 		})
 	}
 }

x/subscription/types/expected_keepers.go

@@ -28,6 +28,7 @@ type EpochstorageKeeper interface {
 type ProjectsKeeper interface {
 	CreateAdminProject(ctx sdk.Context, subscriptionAddress string, totalCU uint64, cuPerEpoch uint64, providers uint64, vrfpk string) error
 	DeleteProject(ctx sdk.Context, index string) error
+	SnapshotSubscriptionProjects(ctx sdk.Context, subscriptionAddr string)
 	// Methods imported from projectskeeper should be defined here
 }