review

Author: sergei-bronnikov

internal/message/consumer.go

@@ -1,6 +1,8 @@
 package message
 
 import (
+	"sync"
+
 	"github.com/bricks-cloud/bricksllm/internal/event"
 	"github.com/bricks-cloud/bricksllm/internal/key"
 	"go.uber.org/zap"
@@ -12,6 +14,7 @@ type Consumer struct {
 	log                 *zap.Logger
 	numOfEventConsumers int
 	handle              func(Message) error
+	wg                  sync.WaitGroup
 }
 
 type recorder interface {
@@ -33,12 +36,25 @@ func NewConsumer(mc <-chan Message, log *zap.Logger, num int, handle func(Messag
 
 func (c *Consumer) StartEventMessageConsumers() {
 	for i := 0; i < c.numOfEventConsumers; i++ {
+		c.wg.Add(1)
 		go func() {
+			defer c.wg.Done()
+
 			for {
 				select {
 				case <-c.done:
-					c.log.Info("event message consumer stoped...")
-					return
+					for {
+						select {
+						case m := <-c.messageChan:
+							err := c.handle(m)
+							if err != nil {
+								continue
+							}
+						default:
+							c.log.Info("event message consumer stoped...")
+							return
+						}
+					}
 
 				case m := <-c.messageChan:
 					err := c.handle(m)
@@ -56,5 +72,6 @@ func (c *Consumer) StartEventMessageConsumers() {
 func (c *Consumer) Stop() {
 	c.log.Info("shutting down consumer...")
 
-	c.done <- true
+	close(c.done)
+	c.wg.Wait()
 }

internal/server/web/proxy/proxy.go

@@ -91,13 +91,13 @@ func NewProxyServer(log *zap.Logger, mode, privacyMode string, c cache, m KeyMan
 	router.Use(getTimeoutMiddleware(timeout))
 	router.Use(getMiddleware(cpm, rm, pm, a, prod, private, log, pub, "proxy", ac, uac, http.Client{}, scanner, cd, um, removeAgentHeaders))
 
+	baseTransport := http.DefaultTransport.(*http.Transport).Clone()
+	baseTransport.MaxIdleConns = 200
+	baseTransport.MaxIdleConnsPerHost = 100
+	baseTransport.IdleConnTimeout = 90 * time.Second
+	baseTransport.TLSHandshakeTimeout = 10 * time.Second
 	client := http.Client{
-		Transport: &http.Transport{
-			MaxIdleConns:        200,
-			MaxIdleConnsPerHost: 100,
-			IdleConnTimeout:     90 * time.Second,
-			TLSHandshakeTimeout: 10 * time.Second,
-		},
+		Transport: baseTransport,
 	}
 
 	// health check