From a6e4779652c24dcef704bead2d42813ee7fd3d71 Mon Sep 17 00:00:00 2001
From: Alexander <zavorotynskiy@pm.me>
Date: Wed, 18 Jan 2023 12:36:34 +0100
Subject: [PATCH] [HTTP] added custom project beacon size (#926)

* feat(backend): added project beacon size select on each sessionStart request
---
 backend/internal/http/router/handlers-web.go |  7 ++-
 backend/internal/http/router/router.go       | 51 +++++++++++++++++++-
 backend/pkg/db/postgres/project.go           |  4 +-
 backend/pkg/db/types/project.go              |  1 +
 4 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/backend/internal/http/router/handlers-web.go b/backend/internal/http/router/handlers-web.go
index 08c4c75c5..7afd184e5 100644
--- a/backend/internal/http/router/handlers-web.go
+++ b/backend/internal/http/router/handlers-web.go
@@ -152,12 +152,15 @@ func (e *Router) startSessionHandlerWeb(w http.ResponseWriter, r *http.Request)
 		}
 	}
 
+	// Save information about session beacon size
+	e.addBeaconSize(tokenData.ID, p.BeaconSize)
+
 	ResponseWithJSON(w, &StartSessionResponse{
 		Token:           e.services.Tokenizer.Compose(*tokenData),
 		UserUUID:        userUUID,
 		SessionID:       strconv.FormatUint(tokenData.ID, 10),
 		ProjectID:       strconv.FormatUint(uint64(p.ProjectID), 10),
-		BeaconSizeLimit: e.cfg.BeaconSizeLimit,
+		BeaconSizeLimit: e.getBeaconSize(tokenData.ID),
 		StartTimestamp:  int64(flakeid.ExtractTimestamp(tokenData.ID)),
 		Delay:           tokenData.Delay,
 	})
@@ -177,7 +180,7 @@ func (e *Router) pushMessagesHandlerWeb(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	bodyBytes, err := e.readBody(w, r, e.cfg.BeaconSizeLimit)
+	bodyBytes, err := e.readBody(w, r, e.getBeaconSize(sessionData.ID))
 	if err != nil {
 		log.Printf("error while reading request body: %s", err)
 		ResponseWithError(w, http.StatusRequestEntityTooLarge, err)
diff --git a/backend/internal/http/router/router.go b/backend/internal/http/router/router.go
index 6d31b7396..964016dfd 100644
--- a/backend/internal/http/router/router.go
+++ b/backend/internal/http/router/router.go
@@ -12,9 +12,15 @@ import (
 	http2 "openreplay/backend/internal/http/services"
 	"openreplay/backend/internal/http/util"
 	"openreplay/backend/pkg/monitoring"
+	"sync"
 	"time"
 )
 
+type BeaconSize struct {
+	size int64
+	time time.Time
+}
+
 type Router struct {
 	router          *mux.Router
 	cfg             *http3.Config
@@ -22,6 +28,8 @@ type Router struct {
 	requestSize     syncfloat64.Histogram
 	requestDuration syncfloat64.Histogram
 	totalRequests   syncfloat64.Counter
+	mutex           *sync.RWMutex
+	beaconSizeCache map[uint64]*BeaconSize // Cache for session's beaconSize
 }
 
 func NewRouter(cfg *http3.Config, services *http2.ServicesBuilder, metrics *monitoring.Metrics) (*Router, error) {
@@ -34,14 +42,53 @@ func NewRouter(cfg *http3.Config, services *http2.ServicesBuilder, metrics *moni
 		return nil, fmt.Errorf("metrics is empty")
 	}
 	e := &Router{
-		cfg:      cfg,
-		services: services,
+		cfg:             cfg,
+		services:        services,
+		mutex:           &sync.RWMutex{},
+		beaconSizeCache: make(map[uint64]*BeaconSize),
 	}
 	e.initMetrics(metrics)
 	e.init()
+	go e.clearBeaconSizes()
 	return e, nil
 }
 
+func (e *Router) addBeaconSize(sessionID uint64, size int64) {
+	if size <= 0 {
+		return
+	}
+	e.mutex.Lock()
+	defer e.mutex.Unlock()
+	e.beaconSizeCache[sessionID] = &BeaconSize{
+		size: size,
+		time: time.Now(),
+	}
+}
+
+func (e *Router) getBeaconSize(sessionID uint64) int64 {
+	e.mutex.RLock()
+	defer e.mutex.RUnlock()
+	if beaconSize, ok := e.beaconSizeCache[sessionID]; ok {
+		beaconSize.time = time.Now()
+		return beaconSize.size
+	}
+	return e.cfg.BeaconSizeLimit
+}
+
+func (e *Router) clearBeaconSizes() {
+	for {
+		time.Sleep(time.Minute * 2)
+		now := time.Now()
+		e.mutex.Lock()
+		for sid, bs := range e.beaconSizeCache {
+			if now.Sub(bs.time) > time.Minute*3 {
+				delete(e.beaconSizeCache, sid)
+			}
+		}
+		e.mutex.Unlock()
+	}
+}
+
 func (e *Router) init() {
 	e.router = mux.NewRouter()
 
diff --git a/backend/pkg/db/postgres/project.go b/backend/pkg/db/postgres/project.go
index f38161885..3239aefb7 100644
--- a/backend/pkg/db/postgres/project.go
+++ b/backend/pkg/db/postgres/project.go
@@ -7,12 +7,12 @@ import (
 func (conn *Conn) GetProjectByKey(projectKey string) (*Project, error) {
 	p := &Project{ProjectKey: projectKey}
 	if err := conn.c.QueryRow(`
-		SELECT max_session_duration, sample_rate, project_id
+		SELECT max_session_duration, sample_rate, project_id, beacon_size
 		FROM projects
 		WHERE project_key=$1 AND active = true
 	`,
 		projectKey,
-	).Scan(&p.MaxSessionDuration, &p.SampleRate, &p.ProjectID); err != nil {
+	).Scan(&p.MaxSessionDuration, &p.SampleRate, &p.ProjectID, &p.BeaconSize); err != nil {
 		return nil, err
 	}
 	return p, nil
diff --git a/backend/pkg/db/types/project.go b/backend/pkg/db/types/project.go
index a5bc2e82c..85ac960ae 100644
--- a/backend/pkg/db/types/project.go
+++ b/backend/pkg/db/types/project.go
@@ -8,6 +8,7 @@ type Project struct {
 	MaxSessionDuration  int64
 	SampleRate          byte
 	SaveRequestPayloads bool
+	BeaconSize          int64
 	Metadata1           *string
 	Metadata2           *string
 	Metadata3           *string