From 63a58e1b1214834f4ccf2027ac0ee3e5f7f4f564 Mon Sep 17 00:00:00 2001
From: kbearXD <s199180626@gmail.com>
Date: Tue, 3 Sep 2024 15:34:55 +0800
Subject: [PATCH] FIX: fix memory leak

---
 pkg/bbgo/marketdatastore.go      | 18 ++++++++++---
 pkg/bbgo/marketdatastore_test.go | 45 ++++++++++++++++++++++++++++++++
 pkg/types/kline_test.go          |  3 ++-
 3 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 pkg/bbgo/marketdatastore_test.go

diff --git a/pkg/bbgo/marketdatastore.go b/pkg/bbgo/marketdatastore.go
index dd9eaf0992..f506e8a89a 100644
--- a/pkg/bbgo/marketdatastore.go
+++ b/pkg/bbgo/marketdatastore.go
@@ -5,6 +5,8 @@ import "github.com/c9s/bbgo/pkg/types"
 const MaxNumOfKLines = 5_000
 const MaxNumOfKLinesTruncate = 100
 
+const CapacityOfKLineWindowLimit = 5_000
+
 // MarketDataStore receives and maintain the public market data of a single symbol
 //go:generate callbackgen -type MarketDataStore
 type MarketDataStore struct {
@@ -57,10 +59,20 @@ func (store *MarketDataStore) AddKLine(k types.KLine) {
 	}
 	window.Add(k)
 
-	if len(*window) > MaxNumOfKLines {
-		*window = (*window)[MaxNumOfKLinesTruncate-1:]
-	}
+	truncateKLineWindowIfNeeded(window)
 
 	store.EmitKLineClosed(k)
 	store.EmitKLineWindowUpdate(k.Interval, *window)
 }
+
+func truncateKLineWindowIfNeeded(window *types.KLineWindow) {
+	lenOfWindow := len(*window)
+	capOfWindow := cap(*window)
+
+	if lenOfWindow == capOfWindow && capOfWindow > CapacityOfKLineWindowLimit {
+		size := CapacityOfKLineWindowLimit / 2
+		start := lenOfWindow - size
+		copy(*window, (*window)[start:])
+		*window = (*window)[:size]
+	}
+}
diff --git a/pkg/bbgo/marketdatastore_test.go b/pkg/bbgo/marketdatastore_test.go
new file mode 100644
index 0000000000..e6bdb40b52
--- /dev/null
+++ b/pkg/bbgo/marketdatastore_test.go
@@ -0,0 +1,45 @@
+package bbgo
+
+import (
+	"testing"
+
+	"github.com/c9s/bbgo/pkg/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMarketDataStore_AddKLineAndTruncateWindow(t *testing.T) {
+	store := NewMarketDataStore("BTCUSD")
+
+	interval := types.Interval1s
+
+	var maxCap int = 0
+	capFixed := false
+
+	var gid uint64 = 0
+	// insert 1.5 * CapacityOfKLineWindowLimit KLine into window
+	for ; gid < CapacityOfKLineWindowLimit+(CapacityOfKLineWindowLimit/2); gid++ {
+		store.AddKLine(types.KLine{
+			Interval: interval,
+			GID:      gid,
+		})
+
+		// if the capacity is > CapacityOfKLineWindowLimit, the capacity should be fixed. We use this if expression to verify it then.
+		if !capFixed && cap(*store.KLineWindows[interval]) > CapacityOfKLineWindowLimit {
+			maxCap = cap(*store.KLineWindows[interval])
+			capFixed = true
+		}
+	}
+
+	window := store.KLineWindows[interval]
+
+	// make sure the capacity is fixed
+	assert.Equal(t, maxCap, cap(*window))
+
+	// after truncate, it will remain (CapacityOfKLineWindowLimit / 2) KLine in the window
+	// so the first GIC will be the maxCap - (CapacityOfKLineWindowLimit / 2)
+	truncatedGID := uint64(maxCap - (CapacityOfKLineWindowLimit / 2))
+	for _, kline := range *window {
+		assert.Equal(t, truncatedGID, kline.GID)
+		truncatedGID++
+	}
+}
diff --git a/pkg/types/kline_test.go b/pkg/types/kline_test.go
index f125f39d9c..cea9c5adbe 100644
--- a/pkg/types/kline_test.go
+++ b/pkg/types/kline_test.go
@@ -2,8 +2,9 @@ package types
 
 import (
 	"encoding/json"
-	"github.com/stretchr/testify/assert"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestKLineWindow_Tail(t *testing.T) {