(DAL) Possible fix for memory leak

node/pkg/dal/api/controller.go

@@ -5,11 +5,11 @@ import (
 	"errors"
 	"fmt"
 	"strings"
+	"sync"
 
 	"bisonai.com/orakl/node/pkg/common/types"
 	"bisonai.com/orakl/node/pkg/dal/collector"
 	dalcommon "bisonai.com/orakl/node/pkg/dal/common"
-	"bisonai.com/orakl/node/pkg/dal/utils/stats"
 	"bisonai.com/orakl/node/pkg/utils/request"
 	"github.com/gofiber/contrib/websocket"
 	"github.com/gofiber/fiber/v2"
@@ -19,7 +19,7 @@ import (
 func Setup(ctx context.Context, adminEndpoint string) (*Controller, error) {
 	configs, err := request.Request[[]types.Config](request.WithEndpoint(adminEndpoint + "/config"))
 	if err != nil {
-		log.Error().Err(err).Msg("failed to get configs")
+		log.Error().Str("Player", "controller").Err(err).Msg("failed to get configs")
 		return nil, err
 	}
 
@@ -29,7 +29,7 @@ func Setup(ctx context.Context, adminEndpoint string) (*Controller, error) {
 	}
 	collector, err := collector.NewCollector(ctx, configs)
 	if err != nil {
-		log.Error().Err(err).Msg("failed to create collector")
+		log.Error().Str("Player", "controller").Err(err).Msg("failed to create collector")
 		return nil, err
 	}
 
@@ -46,36 +46,56 @@ func NewController(configs map[string]types.Config, internalCollector *collector
 		register:   make(chan *websocket.Conn),
 		unregister: make(chan *websocket.Conn),
 		broadcast:  make(map[string]chan dalcommon.OutgoingSubmissionData),
+
+		mu: sync.RWMutex{},
 	}
 }
 
 func (c *Controller) Start(ctx context.Context) {
 	go c.Collector.Start(ctx)
-	log.Info().Msg("api collector started")
-	go func() {
-		for {
-			select {
-			case conn := <-c.register:
-				c.mu.Lock()
+	log.Info().Str("Player", "controller").Msg("api collector started")
+	go c.handleConnection(ctx)
+	log.Info().Str("Player", "controller").Msg("connection handler started")
+	go c.startBroadCast()
+}
+
+func (c *Controller) handleConnection(ctx context.Context) {
+	for {
+		select {
+		case conn := <-c.register:
+			c.mu.Lock()
+			if _, ok := c.clients[conn]; !ok {
 				c.clients[conn] = make(map[string]bool)
-				c.mu.Unlock()
-			case conn := <-c.unregister:
-				c.mu.Lock()
+			}
+			c.mu.Unlock()
+		case conn := <-c.unregister:
+			c.mu.Lock()
+			delete(c.clients, conn)
+			conn.Close()
+			c.mu.Unlock()
+		case <-ctx.Done():
+			c.mu.Lock()
+			for conn := range c.clients {
 				delete(c.clients, conn)
 				conn.Close()
-				c.mu.Unlock()
 			}
+			c.mu.Unlock()
+			return
 		}
-	}()
+	}
+}
 
+func (c *Controller) startBroadCast() {
 	for configId, stream := range c.Collector.OutgoingStream {
 		symbol := c.configIdToSymbol(configId)
-		c.broadcast[symbol] = make(chan dalcommon.OutgoingSubmissionData)
+		if symbol == "" {
+			continue
+		}
 		c.broadcast[symbol] = stream
 	}
 
 	for symbol := range c.configs {
-		go c.broadcastDataForSymbol(symbol)
+		c.broadcastDataForSymbol(symbol)
 	}
 }
 
@@ -90,84 +110,20 @@ func (c *Controller) configIdToSymbol(id int32) string {
 
 func (c *Controller) broadcastDataForSymbol(symbol string) {
 	for data := range c.broadcast[symbol] {
-
-		go c.castSubmissionData(&data, &symbol)
+		go c.castSubmissionData(&data, symbol)
 	}
 }
 
-// pass by pointer to reduce memory copy time
-func (c *Controller) castSubmissionData(data *dalcommon.OutgoingSubmissionData, symbol *string) {
+func (c *Controller) castSubmissionData(data *dalcommon.OutgoingSubmissionData, symbol string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	for conn := range c.clients {
-		if _, ok := c.clients[conn][*symbol]; ok {
+		if _, ok := c.clients[conn][symbol]; ok {
 			if err := conn.WriteJSON(*data); err != nil {
-				log.Error().Err(err).Msg("failed to write message")
+				log.Error().Str("Player", "controller").Err(err).Msg("failed to write message")
 				delete(c.clients, conn)
-				conn.Close()
-			}
-		}
-	}
-}
-
-func HandleWebsocket(conn *websocket.Conn) {
-	c, ok := conn.Locals("apiController").(*Controller)
-	if !ok {
-		log.Error().Msg("api controller not found")
-		return
-	}
-
-	ctx, ok := conn.Locals("context").(*context.Context)
-	if !ok {
-		log.Error().Msg("ctx not found")
-		return
-	}
-
-	c.register <- conn
-	apiKey := conn.Headers("X-Api-Key")
-
-	id, err := stats.InsertWebsocketConnection(*ctx, apiKey)
-	if err != nil {
-		log.Error().Err(err).Msg("failed to insert websocket connection")
-		return
-	}
-	log.Info().Int32("id", id).Msg("inserted websocket connection")
-
-	defer func() {
-		c.unregister <- conn
-		conn.Close()
-		err = stats.UpdateWebsocketConnection(*ctx, id)
-		if err != nil {
-			log.Error().Err(err).Msg("failed to update websocket connection")
-			return
-		}
-		log.Info().Int32("id", id).Msg("updated websocket connection")
-	}()
-
-	for {
-		var msg Subscription
-		if err = conn.ReadJSON(&msg); err != nil {
-			log.Error().Err(err).Msg("failed to read message")
-			return
-		}
-
-		if msg.Method == "SUBSCRIBE" {
-			c.mu.Lock()
-			if c.clients[conn] == nil {
-				c.clients[conn] = make(map[string]bool)
+				conn.CloseHandler()(websocket.CloseInternalServerErr, "failed to write message")
 			}
-			for _, param := range msg.Params {
-				symbol := strings.TrimPrefix(param, "submission@")
-				if _, ok := c.configs[symbol]; !ok {
-					continue
-				}
-				c.clients[conn][symbol] = true
-				err = stats.InsertWebsocketSubscription(*ctx, id, param)
-				if err != nil {
-					log.Error().Err(err).Msg("failed to insert websocket subscription")
-				}
-			}
-			c.mu.Unlock()
 		}
 	}
 }

node/pkg/dal/api/websocketcontroller.go

@@ -0,0 +1,97 @@
+package api
+
+import (
+	"context"
+	"strings"
+
+	"bisonai.com/orakl/node/pkg/dal/utils/stats"
+	"github.com/gofiber/contrib/websocket"
+	"github.com/rs/zerolog/log"
+)
+
+func HandleWebsocket(conn *websocket.Conn) {
+	c, ok := conn.Locals("apiController").(*Controller)
+	if !ok {
+		log.Error().Str("Player", "controller").Msg("api controller not found")
+		return
+	}
+
+	closeHandler := conn.CloseHandler()
+	conn.SetCloseHandler(func(code int, text string) error {
+		log.Info().Str("Player", "controller").Int("code", code).Str("text", text).Msg("close handler")
+		c.unregister <- conn
+		return closeHandler(code, text)
+	})
+
+	ctxPointer, ok := conn.Locals("context").(*context.Context)
+	if !ok {
+		log.Error().Str("Player", "controller").Str("RemoteAddr", conn.RemoteAddr().String()).Msg("ctx not found")
+		return
+	}
+
+	ctx := *ctxPointer
+	apiKey := conn.Headers("X-Api-Key")
+
+	c.register <- conn
+
+	id, err := stats.InsertWebsocketConnection(ctx, apiKey)
+	if err != nil {
+		log.Error().Str("Player", "controller").Str("RemoteAddr", conn.RemoteAddr().String()).Err(err).Msg("failed to insert websocket connection")
+		return
+	}
+	log.Info().Str("Player", "controller").Int32("id", id).Msg("inserted websocket connection")
+
+	defer func() {
+		conn.CloseHandler()(websocket.CloseNormalClosure, "normal closure")
+		err = stats.UpdateWebsocketConnection(ctx, id)
+		if err != nil {
+			log.Error().Str("Player", "controller").Str("RemoteAddr", conn.RemoteAddr().String()).Err(err).Msg("failed to update websocket connection")
+			return
+		}
+		log.Info().Str("Player", "controller").Int32("id", id).Msg("updated websocket connection")
+	}()
+
+	for {
+		if err := handleMessage(ctx, conn, c, id); err != nil {
+			log.Error().Str("Player", "controller").Str("RemoteAddr", conn.RemoteAddr().String()).Err(err).Msg("failed to handle message")
+			return
+		}
+	}
+}
+
+func handleMessage(ctx context.Context, conn *websocket.Conn, c *Controller, id int32) error {
+	var msg Subscription
+	if err := conn.ReadJSON(&msg); err != nil {
+		return err
+	}
+
+	if msg.Method == "SUBSCRIBE" {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		if c.clients[conn] == nil {
+			c.clients[conn] = make(map[string]bool)
+		}
+		addedSymbols := make([]string, 0, len(msg.Params))
+		for _, param := range msg.Params {
+			symbol := strings.TrimPrefix(param, "submission@")
+			if _, ok := c.configs[symbol]; !ok {
+				continue
+			}
+
+			if _, ok := c.clients[conn][symbol]; ok {
+				continue
+			}
+
+			c.clients[conn][symbol] = true
+			addedSymbols = append(addedSymbols, symbol)
+		}
+		defer func() {
+			err := stats.InsertWebsocketSubscriptions(ctx, id, addedSymbols)
+			if err != nil {
+				log.Error().Str("Player", "controller").Err(err).Msg("failed to insert websocket subscriptions")
+			}
+		}()
+	}
+
+	return nil
+}

node/pkg/dal/k6/script.js

@@ -0,0 +1,84 @@
+import { randomIntBetween } from "https://jslib.k6.io/k6-utils/1.2.0/index.js";
+import { check } from "k6";
+
+import ws from "k6/ws";
+const sessionDuration = randomIntBetween(120000, 240000);
+const msg =
+  '{"method":"SUBSCRIBE","params":["submission@ADA-KRW","submission@AKT-KRW","submission@AAVE-KRW","submission@ADA-USDT","submission@ATOM-USDT","submission@APT-KRW","submission@ASTR-KRW","submission@AUCTION-KRW","submission@AVAX-USDT","submission@ARB-KRW","submission@AVAX-KRW","submission@BCH-KRW","submission@BLUR-KRW","submission@AXS-KRW","submission@BLAST-KRW","submission@BNB-USDT","submission@BORA-KRW","submission@BTC-KRW","submission@BTC-USDT","submission@BSV-KRW","submission@CHF-USD","submission@BTG-KRW","submission@BTT-KRW","submission@BONK-KRW","submission@CHZ-KRW","submission@CTC-KRW","submission@DAI-USDT","submission@DOGE-USDT","submission@DOGE-KRW","submission@DOT-KRW","submission@DOT-USDT","submission@ENS-KRW","submission@EOS-KRW","submission@ETH-USDT","submission@ETC-KRW","submission@FET-KRW","submission@FTM-USDT","submission@FLOW-KRW","submission@GAS-KRW","submission@GBP-USD","submission@GLM-KRW","submission@ETH-KRW","submission@GRT-KRW","submission@HPO-KRW","submission@HBAR-KRW","submission@EUR-USD","submission@IQ-KRW","submission@JOY-USDT","submission@IMX-KRW","submission@JPY-USD","submission@KLAY-KRW","submission@KLAY-USDT","submission@KNC-KRW","submission@KRW-USD","submission@LINK-KRW","submission@KSP-KRW","submission@LTC-USDT","submission@MBL-KRW","submission@MATIC-KRW","submission@MBX-KRW","submission@MED-KRW","submission@MATIC-USDT","submission@MLK-KRW","submission@MINA-KRW","submission@MNR-KRW","submission@NEAR-KRW","submission@MTL-KRW","submission@ONDO-KRW","submission@ONG-KRW","submission@PEPE-KRW","submission@SAND-KRW","submission@PYTH-KRW","submission@PAXG-USDT","submission@SEI-KRW","submission@SHIB-KRW","submission@SNT-KRW","submission@PER-KLAY","submission@SHIB-USDT","submission@PEPE-USDT","submission@SOL-USDT","submission@SOL-KRW","submission@STPT-KRW","submission@STG-KRW","submission@STX-KRW","submission@STRK-KRW","submission@TRX-KRW","submission@SUI-KRW","submission@TRX-USDT","submission@USDT-KRW","submission@WEMIX-KRW","submission@USDT-USD","submission@USDC-USDT","submission@WEMIX-USDT","submission@WLD-KRW","submission@WAVES-KRW","submission@XEC-KRW","submission@XLM-KRW","submission@XRP-KRW","submission@XRP-USDT","submission@ZK-KRW","submission@ZETA-KRW","submission@ZRO-KRW","submission@UNI-USDT"]}';
+export const options = {
+  // A number specifying the number of VUs to run concurrently.
+  vus: 120,
+  // A string specifying the total duration of the test run.
+  iterations: 240,
+
+  // The following section contains configuration options for execution of this
+  // test script in Grafana Cloud.
+  //
+  // See https://grafana.com/docs/grafana-cloud/k6/get-started/run-cloud-tests-from-the-cli/
+  // to learn about authoring and running k6 test scripts in Grafana k6 Cloud.
+  //
+  // cloud: {
+  //   // The ID of the project to which the test is assigned in the k6 Cloud UI.
+  //   // By default tests are executed in default project.
+  //   projectID: "",
+  //   // The name of the test in the k6 Cloud UI.
+  //   // Test runs with the same name will be grouped.
+  //   name: "script.js"
+  // },
+
+  // Uncomment this section to enable the use of Browser API in your tests.
+  //
+  // See https://grafana.com/docs/k6/latest/using-k6-browser/running-browser-tests/ to learn more
+  // about using Browser API in your test scripts.
+  //
+  // scenarios: {
+  //   // The scenario name appears in the result summary, tags, and so on.
+  //   // You can give the scenario any name, as long as each name in the script is unique.
+  //   ui: {
+  //     // Executor is a mandatory parameter for browser-based tests.
+  //     // Shared iterations in this case tells k6 to reuse VUs to execute iterations.
+  //     //
+  //     // See https://grafana.com/docs/k6/latest/using-k6/scenarios/executors/ for other executor types.
+  //     executor: 'shared-iterations',
+  //     options: {
+  //       browser: {
+  //         // This is a mandatory parameter that instructs k6 to launch and
+  //         // connect to a chromium-based browser, and use it to run UI-based
+  //         // tests.
+  //         type: 'chromium',
+  //       },
+  //     },
+  //   },
+  // }
+};
+
+// The function that defines VU logic.
+//
+// See https://grafana.com/docs/k6/latest/examples/get-started-with-k6/ to learn more
+// about authoring k6 scripts.
+//
+export default function () {
+  const url = "ws://dal.baobab.orakl.network/ws";
+  const params = { headers: { "X-API-Key": "" } };
+
+  const res = ws.connect(url, params, function (socket) {
+    socket.on("open", function open() {
+      console.log("connected");
+      socket.send(msg);
+    });
+
+    socket.on("ping", function () {
+      console.log("PING!");
+    });
+
+    socket.on("message", (data) => console.log("Message received: ", data));
+    socket.on("close", () => console.log("disconnected"));
+
+    socket.setTimeout(function () {
+      console.log(`Closing the socket forcefully 3s after graceful LEAVE`);
+      socket.close();
+    }, sessionDuration + 15000);
+  });
+
+  check(res, { "status is 101": (r) => r && r.status === 101 });
+}