Merge pull request #3 from sks/feature/serverMap

[Fix] Using sync.Map for concurrency in server ID storage

Author: strowk

pkg/sse/options.go

@@ -9,3 +9,23 @@ type KeepAliveOption struct {
 func (o KeepAliveOption) apply(t *sseTransport) {
 	t.keepAliveInterval = o.Interval
 }
+
+type PortOption struct {
+	Port int
+}
+
+func (o PortOption) apply(t *sseTransport) {
+	t.port = o.Port
+}
+
+func WithPort(port int) SSETransportOption {
+	return PortOption{
+		Port: port,
+	}
+}
+
+func WithKeepAliveInterval(interval time.Duration) SSETransportOption {
+	return KeepAliveOption{
+		Interval: interval,
+	}
+}

pkg/sse/sse_transport.go

@@ -3,7 +3,10 @@ package sse
 import (
 	"context"
 	"io"
+	"net"
 	"net/http"
+	"strconv"
+	"sync"
 	"time"
 
 	foxyevent "github.com/strowk/foxy-contexts/pkg/foxy_event"
@@ -41,8 +44,8 @@ func NewTransport(options ...SSETransportOption) server.Transport {
 type sseTransport struct {
 	keepAliveInterval time.Duration
 	e                 *echo.Echo
-
-	sessionManager *session.SessionManager
+	port              int
+	sessionManager    *session.SessionManager
 }
 
 func newResponseEvent(res jsonrpc2.JsonRpcResponse) (*Event, error) {
@@ -67,14 +70,14 @@ func (s *sseTransport) Run(
 
 	// e.Use(middleware.Logger())
 
-	servers := map[uuid.UUID]server.Server{}
+	servers := sync.Map{}
 
 	postEndpoint := "/message"
 
 	e.GET("/sse", func(c echo.Context) error {
 		sessionId := uuid.New()
 		srv := server.NewServer(capabilities, serverInfo, options...)
-		servers[sessionId] = srv
+		servers.Store(sessionId, srv)
 		w := c.Response()
 		w.Header().Set("Content-Type", "text/event-stream")
 		w.Header().Set("Cache-Control", "no-cache")
@@ -103,10 +106,10 @@ func (s *sseTransport) Run(
 				// Protocol does not seem to have a way to notify client about server initiated shutdown
 				// so we would just close the connection to allow server to shutdown and client to reconnect
 				// to, hopefully, a new server instance started by orchestrator
-				delete(servers, sessionId)
+				servers.Delete(sessionId)
 				return nil
 			case <-c.Request().Context().Done():
-				delete(servers, sessionId)
+				servers.Delete(sessionId)
 				srv.GetLogger().LogEvent(foxyevent.SSEClientDisconnected{ClientIP: c.RealIP()})
 				return nil
 			case res := <-srv.GetResponses():
@@ -141,7 +144,7 @@ func (s *sseTransport) Run(
 			return c.String(http.StatusBadRequest, "sessionId is not a valid UUID")
 		}
 
-		r, ok := servers[parsedSessionId]
+		r, ok := servers.Load(parsedSessionId)
 		if !ok {
 			return c.String(http.StatusNotFound, "session not found")
 		}
@@ -156,11 +159,11 @@ func (s *sseTransport) Run(
 			return c.String(http.StatusNotFound, "failed to resolve session")
 		}
 
-		r.Handle(ctx, b)
+		r.(server.Server).Handle(ctx, b)
 		return c.JSON(http.StatusAccepted, "Accepted")
 	})
 
-	return e.Start("127.0.0.1:1323")
+	return e.Start(net.JoinHostPort("127.0.0.1", strconv.Itoa(s.port)))
 }
 
 func (s *sseTransport) Shutdown(ctx context.Context) error {