refactor(coordinator): refactor CollectProofs

Author: 0xmountaintop

coordinator/manager.go

@@ -196,7 +196,7 @@ func (m *Manager) restorePrevSessions() {
 			}
 			m.sessions[sess.info.ID] = sess
 			log.Info("Coordinator restart reload sessions", "ID", sess.info.ID, "sess", sess.info)
-			go m.CollectProofs(sess.info.ID, sess)
+			go m.CollectProofs(sess)
 		}
 	}
 }
@@ -308,7 +308,7 @@ func (m *Manager) handleZkProof(pk string, msg *message.ProofDetail) error {
 }
 
 // CollectProofs collects proofs corresponding to a proof generation session.
-func (m *Manager) CollectProofs(id string, sess *session) {
+func (m *Manager) CollectProofs(sess *session) {
 	timer := time.NewTimer(time.Duration(m.cfg.CollectionTime) * time.Minute)
 
 	for {
@@ -318,7 +318,7 @@ func (m *Manager) CollectProofs(id string, sess *session) {
 
 			// Ensure proper clean-up of resources.
 			defer func() {
-				delete(m.sessions, id)
+				delete(m.sessions, sess.info.ID)
 				m.mu.Unlock()
 			}()
 
@@ -335,13 +335,13 @@ func (m *Manager) CollectProofs(id string, sess *session) {
 				// record failed session.
 				errMsg := "proof generation session ended without receiving any valid proofs"
 				m.addFailedSession(sess, errMsg)
-				log.Warn(errMsg, "session id", id)
+				log.Warn(errMsg, "session id", sess.info.ID)
 				// Set status as skipped.
 				// Note that this is only a workaround for testnet here.
 				// TODO: In real cases we should reset to orm.ProvingTaskUnassigned
 				// so as to re-distribute the task in the future
-				if err := m.orm.UpdateProvingStatus(id, orm.ProvingTaskFailed); err != nil {
-					log.Error("fail to reset task_status as Unassigned", "id", id, "err", err)
+				if err := m.orm.UpdateProvingStatus(sess.info.ID, orm.ProvingTaskFailed); err != nil {
+					log.Error("fail to reset task_status as Unassigned", "id", sess.info.I, "err", err)
 				}
 				return
 			}
@@ -456,7 +456,7 @@ func (m *Manager) StartProofGenerationSession(task *orm.BlockBatch) (success boo
 	m.mu.Lock()
 	m.sessions[task.ID] = s
 	m.mu.Unlock()
-	go m.CollectProofs(task.ID, s)
+	go m.CollectProofs(s)
 
 	return true
 }