Add resource constraints to arbos storage and precompiles for resource constraints

arbos/arbosState/arbosstate.go

@@ -25,6 +25,7 @@ import (
 	"github.com/offchainlabs/nitro/arbos/arbostypes"
 	"github.com/offchainlabs/nitro/arbos/blockhash"
 	"github.com/offchainlabs/nitro/arbos/burn"
+	"github.com/offchainlabs/nitro/arbos/constraints"
 	"github.com/offchainlabs/nitro/arbos/features"
 	"github.com/offchainlabs/nitro/arbos/l1pricing"
 	"github.com/offchainlabs/nitro/arbos/l2pricing"
@@ -49,6 +50,7 @@ type ArbosState struct {
 	networkFeeAccount      storage.StorageBackedAddress
 	l1PricingState         *l1pricing.L1PricingState
 	l2PricingState         *l2pricing.L2PricingState
+	resourceConstraints    *constraints.StorageResourceConstraints
 	retryableState         *retryables.RetryableState
 	addressTable           *addressTable.AddressTable
 	chainOwners            *addressSet.AddressSet
@@ -79,13 +81,16 @@ func OpenArbosState(stateDB vm.StateDB, burner burn.Burner) (*ArbosState, error)
 	if arbosVersion == 0 {
 		return nil, ErrUninitializedArbOS
 	}
+	constraintsBytes := backingStorage.OpenStorageBackedBytes(constraintsSubspace)
+
 	return &ArbosState{
 		arbosVersion:           arbosVersion,
 		upgradeVersion:         backingStorage.OpenStorageBackedUint64(uint64(upgradeVersionOffset)),
 		upgradeTimestamp:       backingStorage.OpenStorageBackedUint64(uint64(upgradeTimestampOffset)),
 		networkFeeAccount:      backingStorage.OpenStorageBackedAddress(uint64(networkFeeAccountOffset)),
 		l1PricingState:         l1pricing.OpenL1PricingState(backingStorage.OpenCachedSubStorage(l1PricingSubspace), arbosVersion),
 		l2PricingState:         l2pricing.OpenL2PricingState(backingStorage.OpenCachedSubStorage(l2PricingSubspace)),
+		resourceConstraints:    constraints.NewStorageResourceConstraints(&constraintsBytes),
 		retryableState:         retryables.OpenRetryableState(backingStorage.OpenCachedSubStorage(retryablesSubspace), stateDB),
 		addressTable:           addressTable.Open(backingStorage.OpenCachedSubStorage(addressTableSubspace)),
 		chainOwners:            addressSet.OpenAddressSet(backingStorage.OpenCachedSubStorage(chainOwnerSubspace)),
@@ -187,6 +192,7 @@ var (
 	programsSubspace         SubspaceID = []byte{8}
 	featuresSubspace         SubspaceID = []byte{9}
 	nativeTokenOwnerSubspace SubspaceID = []byte{10}
+	constraintsSubspace      SubspaceID = []byte{11}
 )
 
 var PrecompileMinArbOSVersions = make(map[common.Address]uint64)
@@ -524,6 +530,10 @@ func (state *ArbosState) Blockhashes() *blockhash.Blockhashes {
 	return state.blockhashes
 }
 
+func (state *ArbosState) ResourceConstraints() *constraints.StorageResourceConstraints {
+	return state.resourceConstraints
+}
+
 func (state *ArbosState) NetworkFeeAccount() (common.Address, error) {
 	return state.networkFeeAccount.Get()
 }

arbos/constraints/constraints.go

@@ -16,56 +16,101 @@ import (
 // Going over the target for this given period will increase the gas price.
 type PeriodSecs uint32
 
-// ResourceSet is a set of resources.
-type ResourceSet uint32
+// ResourceWeight is a multiplier applied to a resource’s gas usage when computing backlog.
+type ResourceWeight uint64
 
-// EmptyResourceSet creates a new set.
-func EmptyResourceSet() ResourceSet {
-	return ResourceSet(0)
+// WeightedResourceSet tracks resource weights for constraint calculation.
+type WeightedResourceSet struct {
+	weights [multigas.NumResourceKind]ResourceWeight
+}
+
+// ResourceSet represents presence of each resource kind.
+type ResourceSet struct {
+	kinds [multigas.NumResourceKind]bool
 }
 
-// WithResources adds the list of resources to the set.
-func (s ResourceSet) WithResources(resources ...multigas.ResourceKind) ResourceSet {
-	for _, resource := range resources {
-		s = s | (1 << resource)
+const MaxResourceWeight = 1_000_000 // 1e6
+
+// NewWeightedResourceSet creates a new weighted set with all weights initialized to zero.
+func NewWeightedResourceSet() WeightedResourceSet {
+	return WeightedResourceSet{
+		weights: [multigas.NumResourceKind]ResourceWeight{},
 	}
+}
+
+// WithResource sets the weight for a single resource.
+func (s WeightedResourceSet) WithResource(resource multigas.ResourceKind, weight ResourceWeight) WeightedResourceSet {
+	s.weights[resource] = weight
 	return s
 }
 
-// HasResource returns whether the given resource is in the set.
-func (s ResourceSet) HasResource(resource multigas.ResourceKind) bool {
-	return (s & (1 << resource)) != 0
+// HasResource returns true if the resource has a non-zero weight in the set.
+func (s WeightedResourceSet) HasResource(resource multigas.ResourceKind) bool {
+	return s.weights[resource] != 0
+}
+
+// WithoutWeights returns resources set without weights
+func (s WeightedResourceSet) WithoutWeights() ResourceSet {
+	var rs ResourceSet
+	for i, weight := range s.weights {
+		if weight > 0 {
+			rs.kinds[i] = true
+		}
+	}
+	return rs
+}
+
+// EmptyResourceSet creates a new empty resource set.
+func EmptyResourceSet() ResourceSet {
+	return ResourceSet{
+		kinds: [multigas.NumResourceKind]bool{},
+	}
+}
+
+// WithResource returns a copy of the set with the given resource weight updated.
+func (s ResourceSet) WithResource(resource multigas.ResourceKind) ResourceSet {
+	s.kinds[resource] = true
+	return s
 }
 
-// GetResources returns the list of resources in the set.
-func (s ResourceSet) GetResources() []multigas.ResourceKind {
-	var resources []multigas.ResourceKind
-	for resource := range multigas.NumResourceKind {
-		if s.HasResource(resource) {
-			resources = append(resources, resource)
+// All returns all resources with non-zero weights.
+func (s WeightedResourceSet) All() iter.Seq2[multigas.ResourceKind, ResourceWeight] {
+	return func(yield func(multigas.ResourceKind, ResourceWeight) bool) {
+		for i, weight := range s.weights {
+			if weight != 0 {
+				//nolint:gosec // G115: Safe conversion, s.weights length is multigas.NumResourceKind
+				resource := multigas.ResourceKind(i)
+				if !yield(resource, weight) {
+					break
+				}
+			}
 		}
 	}
-	return resources
 }
 
 // ResourceConstraint defines the max gas target per second for the given period for a single resource.
 type ResourceConstraint struct {
-	Resources    ResourceSet
+	Resources    WeightedResourceSet
 	Period       PeriodSecs
 	TargetPerSec uint64
 	Backlog      uint64
 }
 
-// AddToBacklog increases the constraint backlog given the multi-dimensional gas used.
+// AddToBacklog increases the constraint backlog given the multi-dimensional gas used multiplied by their weights.
 func (c *ResourceConstraint) AddToBacklog(gasUsed multigas.MultiGas) {
-	for _, resource := range c.Resources.GetResources() {
-		c.Backlog = arbmath.SaturatingUAdd(c.Backlog, gasUsed.Get(resource))
+	for resource, weight := range c.Resources.All() {
+		if weight == 0 {
+			continue
+		}
+		weightedGas := arbmath.SaturatingUMul(gasUsed.Get(resource), uint64(weight))
+		c.Backlog = arbmath.SaturatingUAdd(c.Backlog, weightedGas)
 	}
 }
 
 // RemoveFromBacklog decreases the backlog by its target given the amount of time passed.
 func (c *ResourceConstraint) RemoveFromBacklog(timeElapsed uint64) {
-	c.Backlog = arbmath.SaturatingUSub(c.Backlog, timeElapsed*c.TargetPerSec)
+	amount := arbmath.SaturatingUMul(timeElapsed, c.TargetPerSec)
+	c.Backlog = arbmath.SaturatingUSub(c.Backlog, amount)
 }
 
 // constraintKey identifies a resource constraint. There can be only one constraint given the
@@ -99,10 +144,10 @@ func NewResourceConstraints() *ResourceConstraints {
 // Set adds or updates the given resource constraint.
 // The set of resources and the period are the key that defines the constraint.
 func (rc *ResourceConstraints) Set(
-	resources ResourceSet, periodSecs PeriodSecs, targetPerSec uint64,
+	resources WeightedResourceSet, periodSecs PeriodSecs, targetPerSec uint64,
 ) {
 	key := constraintKey{
-		resources: resources,
+		resources: resources.WithoutWeights(),
 		period:    periodSecs,
 	}
 	constraint := &ResourceConstraint{

arbos/constraints/constraints_test.go

@@ -12,38 +12,53 @@ import (
 	"github.com/ethereum/go-ethereum/arbitrum/multigas"
 )
 
-func TestResourceSetWithResources(t *testing.T) {
-	s := EmptyResourceSet()
-	resources := []multigas.ResourceKind{
-		multigas.ResourceKindComputation,
-		multigas.ResourceKindStorageAccess,
+func TestResourceSetWithResource(t *testing.T) {
+	s := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1).
+		WithResource(multigas.ResourceKindStorageAccess, 2)
+	for resource, weight := range s.All() {
+		require.True(t, s.HasResource(resource))
+		require.Equal(t, weight, s.weights[resource])
 	}
-	s = s.WithResources(resources...)
-	for _, r := range resources {
-		require.True(t, s.HasResource(r))
-	}
-}
-
-func TestResourceSetHasResource(t *testing.T) {
-	s := EmptyResourceSet()
-	require.False(t, s.HasResource(multigas.ResourceKindComputation))
-	s = s.WithResources(multigas.ResourceKindComputation)
-	require.True(t, s.HasResource(multigas.ResourceKindComputation))
 }
 
 func TestResourceSetGetResources(t *testing.T) {
-	s := EmptyResourceSet()
-	resources := []multigas.ResourceKind{
-		multigas.ResourceKindComputation,
-		multigas.ResourceKindStorageAccess,
+	s := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1).
+		WithResource(multigas.ResourceKindStorageAccess, 1)
+
+	expected := map[multigas.ResourceKind]ResourceWeight{
+		multigas.ResourceKindComputation:   1,
+		multigas.ResourceKindStorageAccess: 1,
+	}
+
+	actual := make(map[multigas.ResourceKind]ResourceWeight)
+	for resource, weight := range s.All() {
+		actual[resource] = weight
 	}
-	s = s.WithResources(resources...)
-	retrieved := s.GetResources()
-	require.Equal(t, resources, retrieved)
+
+	require.Equal(t, expected, actual)
+}
+
+func TestOverrideResourceWeights(t *testing.T) {
+	s := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1).
+		WithResource(multigas.ResourceKindStorageAccess, 2)
+	require.Equal(t, ResourceWeight(1), s.weights[multigas.ResourceKindComputation])
+	require.Equal(t, ResourceWeight(2), s.weights[multigas.ResourceKindStorageAccess])
+	used := s.WithoutWeights()
+
+	s = s.WithResource(multigas.ResourceKindComputation, 3).
+		WithResource(multigas.ResourceKindStorageAccess, 10)
+	require.Equal(t, ResourceWeight(3), s.weights[multigas.ResourceKindComputation])
+	require.Equal(t, ResourceWeight(10), s.weights[multigas.ResourceKindStorageAccess])
+	require.Equal(t, used, s.WithoutWeights())
 }
 
 func TestAddToBacklog(t *testing.T) {
-	resources := EmptyResourceSet().WithResources(multigas.ResourceKindComputation, multigas.ResourceKindStorageAccess)
+	resources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1).
+		WithResource(multigas.ResourceKindStorageAccess, 1)
 	c := &ResourceConstraint{
 		Resources: resources,
 		Backlog:   0,
@@ -64,6 +79,25 @@ func TestAddToBacklog(t *testing.T) {
 	require.Equal(t, c.Backlog, uint64(math.MaxUint64))
 }
 
+func TestAddToBacklogWithWeights(t *testing.T) {
+	resources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 2).
+		WithResource(multigas.ResourceKindStorageAccess, 3)
+	c := &ResourceConstraint{
+		Resources: resources,
+		Backlog:   0,
+	}
+
+	gasUsed := multigas.MultiGasFromPairs(
+		multigas.Pair{Kind: multigas.ResourceKindComputation, Amount: 10},    // 10 * 2 = 20
+		multigas.Pair{Kind: multigas.ResourceKindStorageAccess, Amount: 20},  // 20 * 3 = 60
+		multigas.Pair{Kind: multigas.ResourceKindStorageGrowth, Amount: 100}, // ignored
+	)
+
+	c.AddToBacklog(gasUsed)
+	require.Equal(t, uint64(80), c.Backlog) // 20 + 60 = 80
+}
+
 func TestRemoveFromBacklog(t *testing.T) {
 	c := &ResourceConstraint{
 		Backlog:      1000,
@@ -93,13 +127,14 @@ func TestNewResourceConstraints(t *testing.T) {
 
 func TestSetResourceConstraints(t *testing.T) {
 	rc := NewResourceConstraints()
-	resources := EmptyResourceSet().WithResources(multigas.ResourceKindComputation)
+	resources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1)
 	periodSecs := PeriodSecs(10)
 	targetPerSec := uint64(100)
 
 	rc.Set(resources, periodSecs, targetPerSec)
 
-	constraint := rc.Get(resources, periodSecs)
+	constraint := rc.Get(resources.WithoutWeights(), periodSecs)
 	require.NotNil(t, constraint)
 	require.Equal(t, resources, constraint.Resources)
 	require.Equal(t, periodSecs, constraint.Period)
@@ -108,53 +143,58 @@ func TestSetResourceConstraints(t *testing.T) {
 
 func TestGetResourceConstraints(t *testing.T) {
 	rc := NewResourceConstraints()
-	resources := EmptyResourceSet().WithResources(multigas.ResourceKindComputation)
+	resources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1)
 	periodSecs := PeriodSecs(10)
 	targetPerSec := uint64(100)
 
 	rc.Set(resources, periodSecs, targetPerSec)
 
 	// Test getting an existing constraint
-	constraint := rc.Get(resources, periodSecs)
+	constraint := rc.Get(resources.WithoutWeights(), periodSecs)
 	require.NotNil(t, constraint)
 	require.Equal(t, resources, constraint.Resources)
 	require.Equal(t, periodSecs, constraint.Period)
 	require.Equal(t, targetPerSec, constraint.TargetPerSec)
 	require.Equal(t, uint64(0), constraint.Backlog)
 
 	// Test getting a non-existent constraint
-	nonExistentResources := EmptyResourceSet().WithResources(multigas.ResourceKindStorageAccess)
-	constraint = rc.Get(nonExistentResources, periodSecs)
+	nonExistentResources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindStorageAccess, 1)
+	constraint = rc.Get(nonExistentResources.WithoutWeights(), periodSecs)
 	require.Nil(t, constraint)
 }
 
 func TestClearResourceConstraints(t *testing.T) {
 	rc := NewResourceConstraints()
-	resources := EmptyResourceSet().WithResources(multigas.ResourceKindComputation)
+	resources := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1)
 	periodSecs := PeriodSecs(10)
 	targetPerSec := uint64(100)
 
 	rc.Set(resources, periodSecs, targetPerSec)
 
 	// Ensure the constraint was set
-	constraint := rc.Get(resources, periodSecs)
+	constraint := rc.Get(resources.WithoutWeights(), periodSecs)
 	require.NotNil(t, constraint)
 
 	// Clear the constraint
-	rc.Clear(resources, periodSecs)
+	rc.Clear(resources.WithoutWeights(), periodSecs)
 
 	// Ensure the constraint is gone
-	constraint = rc.Get(resources, periodSecs)
+	constraint = rc.Get(resources.WithoutWeights(), periodSecs)
 	require.Nil(t, constraint)
 }
 
 func TestAllResourceConstraints(t *testing.T) {
 	rc := NewResourceConstraints()
-	resources1 := EmptyResourceSet().WithResources(multigas.ResourceKindComputation)
+	resources1 := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindComputation, 1)
 	periodSecs1 := PeriodSecs(10)
 	targetPerSec1 := uint64(100)
 
-	resources2 := EmptyResourceSet().WithResources(multigas.ResourceKindStorageAccess)
+	resources2 := NewWeightedResourceSet().
+		WithResource(multigas.ResourceKindStorageAccess, 1)
 	periodSecs2 := PeriodSecs(20)
 	targetPerSec2 := uint64(200)