From b3b3afad4ba6a67e701b269de9278cfb736c3c36 Mon Sep 17 00:00:00 2001
From: Cassandra Heart <cassandra@quilibrium.com>
Date: Wed, 23 Oct 2024 01:34:39 -0500
Subject: [PATCH] qol: use separate channel for worker

---
 node/consensus/data/data_clock_consensus_engine.go | 10 ++++++++++
 node/consensus/data/pre_midnight_proof_worker.go   | 10 ++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/node/consensus/data/data_clock_consensus_engine.go b/node/consensus/data/data_clock_consensus_engine.go
index ab62a5d..51b4ef9 100644
--- a/node/consensus/data/data_clock_consensus_engine.go
+++ b/node/consensus/data/data_clock_consensus_engine.go
@@ -317,6 +317,16 @@ func (e *DataClockConsensusEngine) Start() <-chan error {
 		); err != nil {
 			panic(err)
 		}
+
+		if e.dataTimeReel.GetFrameProverTries()[0].Contains(e.provingKeyAddress) {
+			if err := e.pubSub.StartDirectChannelListener(
+				e.pubSub.GetPeerID(),
+				"worker",
+				server,
+			); err != nil {
+				panic(err)
+			}
+		}
 	}()
 
 	e.state = consensus.EngineStateCollecting
diff --git a/node/consensus/data/pre_midnight_proof_worker.go b/node/consensus/data/pre_midnight_proof_worker.go
index ce605aa..349fba0 100644
--- a/node/consensus/data/pre_midnight_proof_worker.go
+++ b/node/consensus/data/pre_midnight_proof_worker.go
@@ -81,7 +81,7 @@ outer:
 			continue
 		}
 
-		cc, err := e.pubSub.GetDirectChannel([]byte(peerId), "")
+		cc, err := e.pubSub.GetDirectChannel([]byte(peerId), "worker")
 		if err != nil {
 			e.logger.Info(
 				"could not establish direct channel, waiting...",
@@ -90,7 +90,6 @@ outer:
 			time.Sleep(10 * time.Second)
 			continue
 		}
-		defer cc.Close()
 
 		client := protobufs.NewDataServiceClient(cc)
 
@@ -107,6 +106,7 @@ outer:
 				zap.Error(err),
 			)
 			time.Sleep(10 * time.Second)
+			cc.Close()
 			continue
 		}
 
@@ -123,6 +123,7 @@ outer:
 
 		if status.Increment == 0 && !bytes.Equal(status.Address, make([]byte, 32)) {
 			e.logger.Info("already completed pre-midnight mint")
+			cc.Close()
 			return
 		}
 
@@ -150,6 +151,7 @@ outer:
 					zap.String("peer_id", peer.ID(e.pubSub.GetPeerID()).String()),
 					zap.Int("increment", i),
 				)
+				cc.Close()
 				return
 			}
 
@@ -163,6 +165,7 @@ outer:
 				}
 				sig, err := e.pubSub.SignMessage(payload)
 				if err != nil {
+					cc.Close()
 					panic(err)
 				}
 
@@ -185,6 +188,7 @@ outer:
 						zap.Error(err),
 					)
 					time.Sleep(10 * time.Second)
+					cc.Close()
 					continue outer
 				}
 
@@ -197,10 +201,12 @@ outer:
 
 				if i == 0 {
 					e.logger.Info("pre-midnight proofs submitted, returning")
+					cc.Close()
 					return
 				}
 			}
 		}
+		cc.Close()
 	}
 }