Add SetSubnetOwner to Chain interface (#2031)

vms/platformvm/config/execution_config.go

@@ -17,6 +17,7 @@ var DefaultExecutionConfig = ExecutionConfig{
 	ChainCacheSize:               2048,
 	ChainDBCacheSize:             2048,
 	BlockIDCacheSize:             8192,
+	FxOwnerCacheSize:             4 * units.MiB,
 	ChecksumsEnabled:             false,
 }
 
@@ -29,6 +30,7 @@ type ExecutionConfig struct {
 	ChainCacheSize               int  `json:"chain-cache-size"`
 	ChainDBCacheSize             int  `json:"chain-db-cache-size"`
 	BlockIDCacheSize             int  `json:"block-id-cache-size"`
+	FxOwnerCacheSize             int  `json:"fx-owner-cache-size"`
 	ChecksumsEnabled             bool `json:"checksums-enabled"`
 }
 

vms/platformvm/config/execution_config_test.go

@@ -46,6 +46,7 @@ func TestExecutionConfigUnmarshal(t *testing.T) {
 			"chain-cache-size": 6,
 			"chain-db-cache-size": 7,
 			"block-id-cache-size": 8,
+			"fx-owner-cache-size": 9,
 			"checksums-enabled": true
 		}`)
 		ec, err := GetExecutionConfig(b)
@@ -58,6 +59,7 @@ func TestExecutionConfigUnmarshal(t *testing.T) {
 			ChainCacheSize:               6,
 			ChainDBCacheSize:             7,
 			BlockIDCacheSize:             8,
+			FxOwnerCacheSize:             9,
 			ChecksumsEnabled:             true,
 		}
 		require.Equal(expected, ec)

vms/platformvm/state/diff.go

@@ -72,6 +72,7 @@ func NewDiff(
 		parentID:      parentID,
 		stateVersions: stateVersions,
 		timestamp:     parentState.GetTimestamp(),
+		subnetOwners:  make(map[ids.ID]fx.Owner),
 	}, nil
 }
 
@@ -293,16 +294,6 @@ func (d *diff) AddSubnet(createSubnetTx *txs.Tx) {
 	if d.cachedSubnets != nil {
 		d.cachedSubnets = append(d.cachedSubnets, createSubnetTx)
 	}
-
-	castTx := createSubnetTx.Unsigned.(*txs.CreateSubnetTx)
-	subnetID := createSubnetTx.ID()
-	if d.subnetOwners == nil {
-		d.subnetOwners = map[ids.ID]fx.Owner{
-			subnetID: castTx.Owner,
-		}
-	} else {
-		d.subnetOwners[subnetID] = castTx.Owner
-	}
 }
 
 func (d *diff) GetSubnetOwner(subnetID ids.ID) (fx.Owner, error) {
@@ -319,6 +310,10 @@ func (d *diff) GetSubnetOwner(subnetID ids.ID) (fx.Owner, error) {
 	return parentState.GetSubnetOwner(subnetID)
 }
 
+func (d *diff) SetSubnetOwner(subnetID ids.ID, owner fx.Owner) {
+	d.subnetOwners[subnetID] = owner
+}
+
 func (d *diff) GetSubnetTransformation(subnetID ids.ID) (*txs.Tx, error) {
 	tx, exists := d.transformedSubnets[subnetID]
 	if exists {
@@ -562,5 +557,8 @@ func (d *diff) Apply(baseState State) error {
 			baseState.DeleteUTXO(utxoID)
 		}
 	}
+	for subnetID, owner := range d.subnetOwners {
+		baseState.SetSubnetOwner(subnetID, owner)
+	}
 	return nil
 }

vms/platformvm/state/diff_test.go

@@ -516,3 +516,65 @@ func assertChainsEqual(t *testing.T, expected, actual Chain) {
 		}
 	}
 }
+
+func TestDiffSubnetOwner(t *testing.T) {
+	require := require.New(t)
+	ctrl := gomock.NewController(t)
+
+	state, _ := newInitializedState(require)
+
+	states := NewMockVersions(ctrl)
+	lastAcceptedID := ids.GenerateTestID()
+	states.EXPECT().GetState(lastAcceptedID).Return(state, true).AnyTimes()
+
+	var (
+		owner1 = fx.NewMockOwner(ctrl)
+		owner2 = fx.NewMockOwner(ctrl)
+
+		createSubnetTx = &txs.Tx{
+			Unsigned: &txs.CreateSubnetTx{
+				BaseTx: txs.BaseTx{},
+				Owner:  owner1,
+			},
+		}
+
+		subnetID = createSubnetTx.ID()
+	)
+
+	// Create subnet on base state
+	owner, err := state.GetSubnetOwner(subnetID)
+	require.ErrorIs(err, database.ErrNotFound)
+	require.Nil(owner)
+
+	state.AddSubnet(createSubnetTx)
+	state.SetSubnetOwner(subnetID, owner1)
+
+	owner, err = state.GetSubnetOwner(subnetID)
+	require.NoError(err)
+	require.Equal(owner1, owner)
+
+	// Create diff and verify that subnet owner returns correctly
+	d, err := NewDiff(lastAcceptedID, states)
+	require.NoError(err)
+
+	owner, err = d.GetSubnetOwner(subnetID)
+	require.NoError(err)
+	require.Equal(owner1, owner)
+
+	// Transferring subnet ownership on diff should be reflected on diff not state
+	d.SetSubnetOwner(subnetID, owner2)
+	owner, err = d.GetSubnetOwner(subnetID)
+	require.NoError(err)
+	require.Equal(owner2, owner)
+
+	owner, err = state.GetSubnetOwner(subnetID)
+	require.NoError(err)
+	require.Equal(owner1, owner)
+
+	// State should reflect new subnet owner after diff is applied.
+	require.NoError(d.Apply(state))
+
+	owner, err = state.GetSubnetOwner(subnetID)
+	require.NoError(err)
+	require.Equal(owner2, owner)
+}

vms/platformvm/state/state.go

@@ -58,7 +58,6 @@ const (
 var (
 	_ State = (*state)(nil)
 
-	ErrCantFindSubnet               = errors.New("couldn't find subnet")
 	errMissingValidatorSet          = errors.New("missing validator set")
 	errValidatorSetAlreadyPopulated = errors.New("validator set already populated")
 	errDuplicateValidatorSet        = errors.New("duplicate validator set")
@@ -81,6 +80,7 @@ var (
 	rewardUTXOsPrefix                   = []byte("rewardUTXOs")
 	utxoPrefix                          = []byte("utxo")
 	subnetPrefix                        = []byte("subnet")
+	subnetOwnerPrefix                   = []byte("subnetOwner")
 	transformedSubnetPrefix             = []byte("transformedSubnet")
 	supplyPrefix                        = []byte("supply")
 	chainPrefix                         = []byte("chain")
@@ -115,6 +115,7 @@ type Chain interface {
 	AddSubnet(createSubnetTx *txs.Tx)
 
 	GetSubnetOwner(subnetID ids.ID) (fx.Owner, error)
+	SetSubnetOwner(subnetID ids.ID, owner fx.Owner)
 
 	GetSubnetTransformation(subnetID ids.ID) (*txs.Tx, error)
 	AddSubnetTransformation(transformSubnetTx *txs.Tx)
@@ -274,6 +275,8 @@ type stateBlk struct {
  * |-. subnets
  * | '-. list
  * |   '-- txID -> nil
+ * |-. subnetOwners
+ * | '-. subnetID -> owner
  * |-. chains
  * | '-. subnetID
  * |   '-. list
@@ -352,6 +355,11 @@ type state struct {
 	subnetBaseDB  database.Database
 	subnetDB      linkeddb.LinkedDB
 
+	// Subnet ID --> Owner of the subnet
+	subnetOwners     map[ids.ID]fx.Owner
+	subnetOwnerCache cache.Cacher[ids.ID, fxOwnerAndSize] // cache of subnetID -> owner if the entry is nil, it is not in the database
+	subnetOwnerDB    database.Database
+
 	transformedSubnets     map[ids.ID]*txs.Tx            // map of subnetID -> transformSubnetTx
 	transformedSubnetCache cache.Cacher[ids.ID, *txs.Tx] // cache of subnetID -> transformSubnetTx if the entry is nil, it is not in the database
 	transformedSubnetDB    database.Database
@@ -420,6 +428,11 @@ type txAndStatus struct {
 	status status.Status
 }
 
+type fxOwnerAndSize struct {
+	owner fx.Owner
+	size  int
+}
+
 func txSize(_ ids.ID, tx *txs.Tx) int {
 	if tx == nil {
 		return ids.IDLen + constants.PointerOverhead
@@ -573,6 +586,18 @@ func newState(
 
 	subnetBaseDB := prefixdb.New(subnetPrefix, baseDB)
 
+	subnetOwnerDB := prefixdb.New(subnetOwnerPrefix, baseDB)
+	subnetOwnerCache, err := metercacher.New[ids.ID, fxOwnerAndSize](
+		"subnet_owner_cache",
+		metricsReg,
+		cache.NewSizedLRU[ids.ID, fxOwnerAndSize](execCfg.FxOwnerCacheSize, func(_ ids.ID, f fxOwnerAndSize) int {
+			return ids.IDLen + f.size
+		}),
+	)
+	if err != nil {
+		return nil, err
+	}
+
 	transformedSubnetCache, err := metercacher.New(
 		"transformed_subnet_cache",
 		metricsReg,
@@ -669,6 +694,10 @@ func newState(
 		subnetBaseDB: subnetBaseDB,
 		subnetDB:     linkeddb.NewDefault(subnetBaseDB),
 
+		subnetOwners:     make(map[ids.ID]fx.Owner),
+		subnetOwnerDB:    subnetOwnerDB,
+		subnetOwnerCache: subnetOwnerCache,
+
 		transformedSubnets:     make(map[ids.ID]*txs.Tx),
 		transformedSubnetCache: transformedSubnetCache,
 		transformedSubnetDB:    prefixdb.New(transformedSubnetPrefix, baseDB),
@@ -819,24 +848,54 @@ func (s *state) AddSubnet(createSubnetTx *txs.Tx) {
 }
 
 func (s *state) GetSubnetOwner(subnetID ids.ID) (fx.Owner, error) {
+	if owner, exists := s.subnetOwners[subnetID]; exists {
+		return owner, nil
+	}
+
+	if ownerAndSize, cached := s.subnetOwnerCache.Get(subnetID); cached {
+		if ownerAndSize.owner == nil {
+			return nil, database.ErrNotFound
+		}
+		return ownerAndSize.owner, nil
+	}
+
+	ownerBytes, err := s.subnetOwnerDB.Get(subnetID[:])
+	if err == nil {
+		var owner fx.Owner
+		if _, err := block.GenesisCodec.Unmarshal(ownerBytes, &owner); err != nil {
+			return nil, err
+		}
+		s.subnetOwnerCache.Put(subnetID, fxOwnerAndSize{
+			owner: owner,
+			size:  len(ownerBytes),
+		})
+		return owner, nil
+	}
+	if err != database.ErrNotFound {
+		return nil, err
+	}
+
 	subnetIntf, _, err := s.GetTx(subnetID)
 	if err != nil {
-		return nil, fmt.Errorf(
-			"%w %q: %w",
-			ErrCantFindSubnet,
-			subnetID,
-			err,
-		)
+		if err == database.ErrNotFound {
+			s.subnetOwnerCache.Put(subnetID, fxOwnerAndSize{})
+		}
+		return nil, err
 	}
 
 	subnet, ok := subnetIntf.Unsigned.(*txs.CreateSubnetTx)
 	if !ok {
 		return nil, fmt.Errorf("%q %w", subnetID, errIsNotSubnet)
 	}
 
+	s.SetSubnetOwner(subnetID, subnet.Owner)
 	return subnet.Owner, nil
 }
 
+func (s *state) SetSubnetOwner(subnetID ids.ID, owner fx.Owner) {
+	s.subnetOwners[subnetID] = owner
+}
+
 func (s *state) GetSubnetTransformation(subnetID ids.ID) (*txs.Tx, error) {
 	if tx, exists := s.transformedSubnets[subnetID]; exists {
 		return tx, nil
@@ -1683,6 +1742,7 @@ func (s *state) write(updateValidators bool, height uint64) error {
 		s.writeRewardUTXOs(),
 		s.writeUTXOs(),
 		s.writeSubnets(),
+		s.writeSubnetOwners(),
 		s.writeTransformedSubnets(),
 		s.writeSubnetSupplies(),
 		s.writeChains(),
@@ -2273,6 +2333,29 @@ func (s *state) writeSubnets() error {
 	return nil
 }
 
+func (s *state) writeSubnetOwners() error {
+	for subnetID, owner := range s.subnetOwners {
+		subnetID := subnetID
+		owner := owner
+		delete(s.subnetOwners, subnetID)
+
+		ownerBytes, err := block.GenesisCodec.Marshal(block.Version, &owner)
+		if err != nil {
+			return fmt.Errorf("failed to marshal subnet owner: %w", err)
+		}
+
+		s.subnetOwnerCache.Put(subnetID, fxOwnerAndSize{
+			owner: owner,
+			size:  len(ownerBytes),
+		})
+
+		if err := s.subnetOwnerDB.Put(subnetID[:], ownerBytes); err != nil {
+			return fmt.Errorf("failed to write subnet owner: %w", err)
+		}
+	}
+	return nil
+}
+
 func (s *state) writeTransformedSubnets() error {
 	for subnetID, tx := range s.transformedSubnets {
 		txID := tx.ID()