From aafca835986ad3fe3c72956ddd41a69a34852513 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Faruk=20Irmak?= <omerfirmak@gmail.com>
Date: Fri, 21 Jun 2024 12:39:43 +0300
Subject: [PATCH] fix: free rust traces that never make it to CCC (#844)

---
 params/version.go                               |  2 +-
 rollup/circuitcapacitychecker/impl.go           |  4 ++++
 rollup/circuitcapacitychecker/libzkp/libzkp.h   |  3 ++-
 rollup/circuitcapacitychecker/libzkp/src/lib.rs |  8 ++++++++
 rollup/circuitcapacitychecker/mock.go           |  3 +++
 rollup/pipeline/pipeline.go                     | 14 +++++++++++++-
 6 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/params/version.go b/params/version.go
index 87e50a0db9c5..481197f291ec 100644
--- a/params/version.go
+++ b/params/version.go
@@ -24,7 +24,7 @@ import (
 const (
 	VersionMajor = 5         // Major version component of the current release
 	VersionMinor = 4         // Minor version component of the current release
-	VersionPatch = 6         // Patch version component of the current release
+	VersionPatch = 7         // Patch version component of the current release
 	VersionMeta  = "mainnet" // Version metadata to append to the version string
 )
 
diff --git a/rollup/circuitcapacitychecker/impl.go b/rollup/circuitcapacitychecker/impl.go
index fe10a02b0604..91e1111ad5f8 100644
--- a/rollup/circuitcapacitychecker/impl.go
+++ b/rollup/circuitcapacitychecker/impl.go
@@ -223,3 +223,7 @@ func MakeRustTrace(trace *types.BlockTrace, buffer *bytes.Buffer) unsafe.Pointer
 
 	return C.parse_json_to_rust_trace(tracesStr)
 }
+
+func FreeRustTrace(ptr unsafe.Pointer) {
+	C.free_rust_trace(ptr)
+}
diff --git a/rollup/circuitcapacitychecker/libzkp/libzkp.h b/rollup/circuitcapacitychecker/libzkp/libzkp.h
index d15363f3e753..9862b86de03d 100644
--- a/rollup/circuitcapacitychecker/libzkp/libzkp.h
+++ b/rollup/circuitcapacitychecker/libzkp/libzkp.h
@@ -9,4 +9,5 @@ char* apply_block(uint64_t id, void* block_trace);
 char* get_tx_num(uint64_t id);
 char* set_light_mode(uint64_t id, bool light_mode);
 void free_c_chars(char* ptr);
-void* parse_json_to_rust_trace(char* trace_json_ptr);
\ No newline at end of file
+void* parse_json_to_rust_trace(char* trace_json_ptr);
+void free_rust_trace(void* trace_ptr);
diff --git a/rollup/circuitcapacitychecker/libzkp/src/lib.rs b/rollup/circuitcapacitychecker/libzkp/src/lib.rs
index 0b93c89319ac..d8eee0013a17 100644
--- a/rollup/circuitcapacitychecker/libzkp/src/lib.rs
+++ b/rollup/circuitcapacitychecker/libzkp/src/lib.rs
@@ -271,6 +271,7 @@ pub mod utils {
     use std::ffi::{CStr, CString};
     use std::os::raw::c_char;
     use std::str::Utf8Error;
+    use prover::BlockTrace;
 
     /// # Safety
     #[no_mangle]
@@ -283,6 +284,13 @@ pub mod utils {
         let _ = CString::from_raw(ptr);
     }
 
+    /// # Safety
+    #[no_mangle]
+    pub unsafe extern "C" fn free_rust_trace(trace_ptr: *mut BlockTrace) {
+        let _ = Box::from_raw(trace_ptr);
+    }
+
+
     #[allow(dead_code)]
     pub(crate) fn c_char_to_str(c: *const c_char) -> Result<&'static str, Utf8Error> {
         let cstr = unsafe { CStr::from_ptr(c) };
diff --git a/rollup/circuitcapacitychecker/mock.go b/rollup/circuitcapacitychecker/mock.go
index 4f89464124dc..dc603bfb3515 100644
--- a/rollup/circuitcapacitychecker/mock.go
+++ b/rollup/circuitcapacitychecker/mock.go
@@ -96,3 +96,6 @@ func MakeRustTrace(trace *types.BlockTrace, buffer *bytes.Buffer) unsafe.Pointer
 	goTraces[unsafe.Pointer(rustTrace)] = trace
 	return unsafe.Pointer(rustTrace)
 }
+
+func FreeRustTrace(ptr unsafe.Pointer) {
+}
diff --git a/rollup/pipeline/pipeline.go b/rollup/pipeline/pipeline.go
index 787b176e00e1..e09eb5d4f728 100644
--- a/rollup/pipeline/pipeline.go
+++ b/rollup/pipeline/pipeline.go
@@ -368,7 +368,10 @@ func (p *Pipeline) encodeStage(traces <-chan *BlockCandidate) <-chan *BlockCandi
 				encodeTimer.UpdateSince(encodeStart)
 
 				stallStart := time.Now()
-				sendCancellable(downstreamCh, trace, p.ctx.Done())
+				if sendCancellable(downstreamCh, trace, p.ctx.Done()) {
+					// failed to send the trace downstream, free it here.
+					circuitcapacitychecker.FreeRustTrace(trace.RustTrace)
+				}
 				encodeStallTimer.UpdateSince(stallStart)
 			case <-p.ctx.Done():
 				return
@@ -395,6 +398,15 @@ func (p *Pipeline) cccStage(candidates <-chan *BlockCandidate, deadline time.Tim
 			close(resultCh)
 			deadlineTimer.Stop()
 			lifetimeTimer.UpdateSince(p.start)
+			// consume candidates and free all rust traces
+			for candidate := range candidates {
+				if candidate == nil {
+					break
+				}
+				if candidate.RustTrace != nil {
+					circuitcapacitychecker.FreeRustTrace(candidate.RustTrace)
+				}
+			}
 			p.wg.Done()
 		}()
 		for {