From 0211a0bd3dc7bf4e0005029327fb80a4f4d90319 Mon Sep 17 00:00:00 2001
From: Jeffery Walsh <cyberhorsey@gmail.com>
Date: Tue, 26 Sep 2023 10:54:18 -0700
Subject: [PATCH] capacity needs to be taken before generating proof

---
 prover/capacity_manager/capacity_manager.go   | 18 ++++++---
 .../capacity_manager/capacity_manager_test.go |  5 ++-
 prover/prover.go                              | 37 ++++++++++++-------
 3 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/prover/capacity_manager/capacity_manager.go b/prover/capacity_manager/capacity_manager.go
index 78164f1cf..28daff6ff 100644
--- a/prover/capacity_manager/capacity_manager.go
+++ b/prover/capacity_manager/capacity_manager.go
@@ -8,13 +8,14 @@ import (
 
 // CapacityManager manages the prover capacity concurrent-safely.
 type CapacityManager struct {
-	capacity uint64
-	mutex    sync.RWMutex
+	capacity    uint64
+	maxCapacity uint64
+	mutex       sync.RWMutex
 }
 
 // New creates a new CapacityManager instance.
 func New(capacity uint64) *CapacityManager {
-	return &CapacityManager{capacity: capacity}
+	return &CapacityManager{capacity: capacity, maxCapacity: capacity}
 }
 
 // ReadCapacity reads the current capacity.
@@ -27,16 +28,21 @@ func (m *CapacityManager) ReadCapacity() uint64 {
 	return m.capacity
 }
 
-// ReleaseCapacity releases one capacity.
-func (m *CapacityManager) ReleaseOneCapacity() uint64 {
+// ReleaseOneCapacity releases one capacity.
+func (m *CapacityManager) ReleaseOneCapacity() (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
+	}
+
 	m.capacity += 1
 
 	log.Info("Released capacity", "capacityAfterRelease", m.capacity)
 
-	return m.capacity
+	return m.capacity, true
 }
 
 // TakeOneCapacity takes one capacity.
diff --git a/prover/capacity_manager/capacity_manager_test.go b/prover/capacity_manager/capacity_manager_test.go
index 1c1213e50..be9ff36d1 100644
--- a/prover/capacity_manager/capacity_manager_test.go
+++ b/prover/capacity_manager/capacity_manager_test.go
@@ -22,7 +22,10 @@ func (s *CapacityManagerTestSuite) TestReadCapacity() {
 }
 
 func (s *CapacityManagerTestSuite) TestReleaseOneCapacity() {
-	s.Equal(testCapacity+1, s.m.ReleaseOneCapacity())
+	capacity, released := s.m.ReleaseOneCapacity()
+	s.Equal(true, released)
+
+	s.Equal(testCapacity+1, capacity)
 	s.Equal(testCapacity+1, s.m.ReadCapacity())
 }
 
diff --git a/prover/prover.go b/prover/prover.go
index 138895651..de37d3287 100644
--- a/prover/prover.go
+++ b/prover/prover.go
@@ -601,6 +601,12 @@ func (p *Prover) onBlockProposed(
 			metrics.ProverProofsAssigned.Inc(1)
 		}
 
+		if !p.cfg.OracleProver {
+			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
+				return errNoCapacity
+			}
+		}
+
 		ctx, cancelCtx := context.WithCancel(ctx)
 		p.currentBlocksBeingProvenMutex.Lock()
 		p.currentBlocksBeingProven[event.BlockId.Uint64()] = cancelFunc(func() {
@@ -611,12 +617,6 @@ func (p *Prover) onBlockProposed(
 		})
 		p.currentBlocksBeingProvenMutex.Unlock()
 
-		if !p.cfg.OracleProver {
-			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
-				return errNoCapacity
-			}
-		}
-
 		return p.validProofSubmitter.RequestProof(ctx, event)
 	}
 
@@ -656,7 +656,10 @@ func (p *Prover) submitProofOp(ctx context.Context, proofWithHeader *proofProduc
 		defer func() {
 			<-p.submitProofConcurrencyGuard
 			if !p.cfg.OracleProver {
-				p.capacityManager.ReleaseOneCapacity()
+				_, released := p.capacityManager.ReleaseOneCapacity()
+				if !released {
+					log.Error("unable to release capacity")
+				}
 			}
 		}()
 
@@ -866,10 +869,15 @@ func (p *Prover) cancelProof(ctx context.Context, blockID uint64) {
 	defer p.currentBlocksBeingProvenMutex.Unlock()
 
 	if cancel, ok := p.currentBlocksBeingProven[blockID]; ok {
+		log.Info("cancelling proof", "blockID", blockID)
+
 		cancel()
 		delete(p.currentBlocksBeingProven, blockID)
 		if !p.cfg.OracleProver {
-			p.capacityManager.ReleaseOneCapacity()
+			capacity, released := p.capacityManager.ReleaseOneCapacity()
+			if !released {
+				log.Error("unable to release capacity while cancelling proof", "capacity", capacity)
+			}
 		}
 	}
 }
@@ -999,6 +1007,13 @@ func (p *Prover) requestProofForBlockId(blockId *big.Int, l1Height *big.Int) err
 			return nil
 		}
 
+		// make sure to takea capacity before requesting proof
+		if !p.cfg.OracleProver {
+			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
+				return errNoCapacity
+			}
+		}
+
 		ctx, cancelCtx := context.WithCancel(ctx)
 		p.currentBlocksBeingProvenMutex.Lock()
 		p.currentBlocksBeingProven[event.BlockId.Uint64()] = cancelFunc(func() {
@@ -1015,12 +1030,6 @@ func (p *Prover) requestProofForBlockId(blockId *big.Int, l1Height *big.Int) err
 			return err
 		}
 
-		if !p.cfg.OracleProver {
-			if _, ok := p.capacityManager.TakeOneCapacity(); !ok {
-				return errNoCapacity
-			}
-		}
-
 		return nil
 	}