chore(container): fix rm bbn mount, add pprof, fix OOM with config (#3)

* fix del, add pprof cfg

* disable iavl cache so we don't go OOM on bbn node

* fix lint

container/container.go

@@ -3,13 +3,16 @@ package container
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"github.com/babylonlabs-io/babylon-benchmark/lib"
+	"github.com/docker/docker/api/types"
 	"regexp"
 	"strconv"
 	"time"
 
 	bbn "github.com/babylonlabs-io/babylon/types"
+	rawDc "github.com/docker/docker/client"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 
@@ -32,28 +35,36 @@ var errRegex = regexp.MustCompile(`(E|e)rror`)
 // Manager is a wrapper around all Docker instances, and the Docker API.
 // It provides utilities to run and interact with all Docker containers
 type Manager struct {
-	cfg       ImageConfig
-	pool      *dockertest.Pool
-	resources map[string]*dockertest.Resource
+	cfg         ImageConfig
+	pool        *dockertest.Pool
+	resources   map[string]*dockertest.Resource
+	rawDcClient *rawDc.Client
 }
 
 // NewManager creates a new Manager instance and initializes
 // all Docker specific utilities. Returns an error if initialization fails.
-func NewManager() (docker *Manager, err error) {
+func NewManager() (mgr *Manager, err error) {
 	imgCfg, err := NewImageConfig()
 	if err != nil {
 		return nil, err
 	}
-	docker = &Manager{
+	mgr = &Manager{
 		cfg:       *imgCfg,
 		resources: make(map[string]*dockertest.Resource),
 	}
-	docker.pool, err = dockertest.NewPool("")
+	mgr.pool, err = dockertest.NewPool("")
 	if err != nil {
 		return nil, err
 	}
 
-	return docker, nil
+	dc, err := rawDc.NewClientWithOpts(rawDc.FromEnv)
+	if err != nil {
+		return nil, err
+	}
+
+	mgr.rawDcClient = dc
+
+	return mgr, nil
 }
 
 func (m *Manager) ExecBitcoindCliCmd(ctx context.Context, command []string) (bytes.Buffer, bytes.Buffer, error) {
@@ -196,7 +207,11 @@ func (m *Manager) RunBabylondResource(
 				"--btc-confirmation-depth=2 --additional-sender-account --btc-network=regtest "+
 				"--min-staking-time-blocks=200 --min-staking-amount-sat=10000 "+
 				"--epoch-interval=%d --slashing-pk-script=%s --btc-base-header=%s "+
-				"--covenant-quorum=1 --covenant-pks=%s && chmod -R 777 /home && babylond start --home=/home/node0/babylond",
+				"--covenant-quorum=1 --covenant-pks=%s && "+
+				"chmod -R 777 /home && "+
+				"sed -i -e 's/iavl-cache-size = 781250/iavl-cache-size = 0/' /home/node0/babylond/config/app.toml && "+ // disable the cache otherwise we go OOM
+				"sed -i -e 's/iavl-disable-fastnode = false/iavl-disable-fastnode = true/' /home/node0/babylond/config/app.toml && "+
+				"babylond start --home=/home/node0/babylond --rpc.pprof_laddr=0.0.0.0:6060",
 			epochInterval, slashingPkScript, baseHeaderHex, bbn.NewBIP340PubKeyFromBTCPK(CovenantPubKey).MarshalHex()),
 	}
 
@@ -215,14 +230,15 @@ func (m *Manager) RunBabylondResource(
 			ExposedPorts: []string{
 				"9090/tcp", // only expose what we need
 				"26657/tcp",
+				"6060/tcp",
 			},
 			Cmd: cmd,
 		},
 		func(config *docker.HostConfig) {
 			config.PortBindings = map[docker.Port][]docker.PortBinding{
 				"9090/tcp":  {{HostIP: "", HostPort: "9090"}},
 				"26657/tcp": {{HostIP: "", HostPort: "26657"}},
-				"8080/tcp":  {{HostIP: "", HostPort: "8080"}},
+				"6060/tcp":  {{HostIP: "", HostPort: "6060"}},
 			}
 		},
 		noRestart,
@@ -236,6 +252,26 @@ func (m *Manager) RunBabylondResource(
 	return resource, nil
 }
 
+func (m *Manager) MemoryUsage(ctx context.Context, containerName string) (uint64, error) {
+	containerId := m.resources[containerName].Container.ID
+
+	res, err := m.rawDcClient.ContainerStats(ctx, containerId, false)
+	if err != nil {
+		return 0, err
+	}
+
+	defer res.Body.Close()
+
+	// Decode stats JSON
+	var containerStats types.StatsJSON //nolint:staticcheck
+	decoder := json.NewDecoder(res.Body)
+	if err := decoder.Decode(&containerStats); err != nil {
+		return 0, err
+	}
+
+	return containerStats.MemoryStats.Usage, nil
+}
+
 // ClearResources removes all outstanding Docker resources created by the Manager.
 func (m *Manager) ClearResources() error {
 	for _, resource := range m.resources {

harness/app.go

@@ -101,7 +101,7 @@ func startHarness(ctx context.Context, cfg config.Config) error {
 		}
 	}
 
-	go printStatsForever(ctx, stopChan, cfg)
+	go printStatsForever(ctx, tm, stopChan, cfg)
 
 	covenantSender, err := NewSenderWithBabylonClient(ctx, "covenant", tm.Config.Babylon.RPCAddr, tm.Config.Babylon.GRPCAddr)
 	if err != nil {
@@ -128,7 +128,7 @@ func startHarness(ctx context.Context, cfg config.Config) error {
 	return nil
 }
 
-func printStatsForever(ctx context.Context, stopChan chan struct{}, cfg config.Config) {
+func printStatsForever(ctx context.Context, tm *TestManager, stopChan chan struct{}, cfg config.Config) {
 	t := time.NewTicker(5 * time.Second)
 	defer t.Stop()
 
@@ -146,7 +146,13 @@ func printStatsForever(ctx context.Context, stopChan chan struct{}, cfg config.C
 				close(stopChan)
 			}
 
-			fmt.Printf("📄 Delegations sent: %d\n", atomic.LoadInt32(&delegationsSentCounter))
+			mem, err := tm.manger.MemoryUsage(ctx, "babylond")
+			if err != nil {
+				fmt.Printf("err getting memory usage for bbn node %v\n", err)
+			}
+			now := time.Now()
+			fmt.Printf("📄 Delegations sent: %d, ts: %s, mem: %d MB\n",
+				atomic.LoadInt32(&delegationsSentCounter), now.Format(time.UnixDate), mem/1e6)
 			prevSent = atomic.LoadInt32(&delegationsSentCounter)
 		case <-ctx.Done():
 			return

harness/manager.go

@@ -188,7 +188,7 @@ func (tm *TestManager) Stop() {
 
 	tm.BitcoindHandler.Stop()
 
-	if tm.benchConfig.BabylonPath != "" {
+	if tm.benchConfig.BabylonPath == "" {
 		cleanupDir(tm.babylonDir) // don't cleanup babylon if user specified a path
 	}
 }