tye capacity to a block iD

prover/capacity_manager/capacity_manager.go

@@ -8,56 +8,56 @@ import (
 
 // CapacityManager manages the prover capacity concurrent-safely.
 type CapacityManager struct {
-	capacity    uint64
+	capacity    map[uint64]bool
 	maxCapacity uint64
 	mutex       sync.RWMutex
 }
 
 // New creates a new CapacityManager instance.
 func New(capacity uint64) *CapacityManager {
-	return &CapacityManager{capacity: capacity, maxCapacity: capacity}
+	return &CapacityManager{capacity: make(map[uint64]bool), maxCapacity: capacity}
 }
 
 // ReadCapacity reads the current capacity.
 func (m *CapacityManager) ReadCapacity() uint64 {
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 
-	log.Info("Reading capacity", "capacity", m.capacity)
+	log.Info("Reading capacity", "capacity", len(m.capacity))
 
-	return m.capacity
+	return uint64(len(m.capacity))
 }
 
 // ReleaseOneCapacity releases one capacity.
-func (m *CapacityManager) ReleaseOneCapacity() (uint64, bool) {
+func (m *CapacityManager) ReleaseOneCapacity(blockID uint64) (uint64, bool) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if m.capacity+1 > m.maxCapacity {
-		log.Info("Can not release capacity", "currentCapacity", m.capacity, "maxCapacity", m.maxCapacity)
-		return m.capacity, false
+	if _, ok := m.capacity[blockID]; !ok {
+		log.Info("Can not release capacity", "blockID", blockID, "currentCapacity", m.capacity, "maxCapacity", m.maxCapacity)
+		return uint64(len(m.capacity)), false
 	}
 
-	m.capacity += 1
+	delete(m.capacity, blockID)
 
-	log.Info("Released capacity", "capacityAfterRelease", m.capacity)
+	log.Info("Released capacity", "capacityAfterRelease", len(m.capacity))
 
-	return m.capacity, true
+	return uint64(len(m.capacity)), true
 }
 
 // TakeOneCapacity takes one capacity.
-func (m *CapacityManager) TakeOneCapacity() (uint64, bool) {
+func (m *CapacityManager) TakeOneCapacity(blockID uint64) (uint64, bool) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	if m.capacity == 0 {
-		log.Info("Could not take one capacity", "capacity", m.capacity)
+	if len(m.capacity) == int(m.maxCapacity) {
+		log.Info("Could not take one capacity", "capacity", len(m.capacity))
 		return 0, false
 	}
 
-	m.capacity -= 1
+	m.capacity[blockID] = true
 
-	log.Info("Took one capacity", "capacityAfterTaking", m.capacity)
+	log.Info("Took one capacity", "blockID", blockID, "capacityAfterTaking", m.capacity)
 
-	return m.capacity, true
+	return uint64(len(m.capacity)), true
 }

prover/capacity_manager/capacity_manager_test.go

@@ -22,15 +22,25 @@ func (s *CapacityManagerTestSuite) TestReadCapacity() {
 }
 
 func (s *CapacityManagerTestSuite) TestReleaseOneCapacity() {
-	capacity, released := s.m.ReleaseOneCapacity()
+	var blockID uint64 = 1
+	capacity, released := s.m.ReleaseOneCapacity(blockID)
+	s.Equal(false, released)
+
+	capacity, ok := s.m.TakeOneCapacity(blockID)
+
+	s.Equal(true, ok)
+
+	capacity, released = s.m.ReleaseOneCapacity(blockID)
 	s.Equal(true, released)
 
 	s.Equal(testCapacity+1, capacity)
 	s.Equal(testCapacity+1, s.m.ReadCapacity())
 }
 
 func (s *CapacityManagerTestSuite) TestTakeOneCapacity() {
-	capacity, ok := s.m.TakeOneCapacity()
+	var blockID uint64 = 1
+
+	capacity, ok := s.m.TakeOneCapacity(blockID)
 	s.True(ok)
 	s.Equal(testCapacity-1, capacity)
 	s.Equal(testCapacity-1, s.m.ReadCapacity())

prover/prover.go

@@ -604,7 +604,7 @@ func (p *Prover) onBlockProposed(
 		}
 
 		if !p.cfg.OracleProver {
-			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
+			if _, ok := p.capacityManager.TakeOneCapacity(event.BlockId.Uint64()); !ok {
 				return errNoCapacity
 			}
 		}
@@ -658,7 +658,7 @@ func (p *Prover) submitProofOp(ctx context.Context, proofWithHeader *proofProduc
 		defer func() {
 			<-p.submitProofConcurrencyGuard
 			if !p.cfg.OracleProver {
-				_, released := p.capacityManager.ReleaseOneCapacity()
+				_, released := p.capacityManager.ReleaseOneCapacity(proofWithHeader.Meta.Id)
 				if !released {
 					log.Error("unable to release capacity")
 				}
@@ -876,9 +876,12 @@ func (p *Prover) cancelProof(ctx context.Context, blockID uint64) {
 		cancel()
 		delete(p.currentBlocksBeingProven, blockID)
 		if !p.cfg.OracleProver {
-			capacity, released := p.capacityManager.ReleaseOneCapacity()
+			capacity, released := p.capacityManager.ReleaseOneCapacity(blockID)
 			if !released {
-				log.Error("unable to release capacity while cancelling proof", "capacity", capacity)
+				log.Error("unable to release capacity while cancelling proof",
+					"capacity", capacity,
+					"blockID", blockID,
+				)
 			}
 		}
 	}
@@ -1011,7 +1014,7 @@ func (p *Prover) requestProofForBlockId(blockId *big.Int, l1Height *big.Int) err
 
 		// make sure to takea capacity before requesting proof
 		if !p.cfg.OracleProver {
-			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
+			if _, ok := p.capacityManager.TakeOneCapacity(event.BlockId.Uint64()); !ok {
 				return errNoCapacity
 			}
 		}