usage of session proxy inside orchestrator

Author: sitole

packages/orchestrator/cmd/mock-sandbox/mock.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"log"
 	"os"
 	"os/signal"
@@ -104,6 +105,7 @@ func mockSandbox(
 	buildId,
 	sandboxId string,
 	dns *dns.DNS,
+	proxy *proxy.SessionProxy,
 	keepAlive time.Duration,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
@@ -128,6 +130,7 @@ func mockSandbox(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{

packages/orchestrator/cmd/mock-snapshot/mock.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"log"
 	"os"
 	"os/signal"
@@ -113,6 +114,7 @@ func mockSnapshot(
 	buildId,
 	sandboxId string,
 	dns *dns.DNS,
+	proxy *proxy.SessionProxy,
 	keepAlive time.Duration,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
@@ -137,6 +139,7 @@ func mockSnapshot(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{
@@ -234,14 +237,15 @@ func mockSnapshot(
 		return fmt.Errorf("failed to add snapshot to template cache: %w", err)
 	}
 
-	fmt.Println("Add snapshot to template cache time: ", time.Since(snapshotTime).Milliseconds())
+	fmt.Println("AddSandbox snapshot to template cache time: ", time.Since(snapshotTime).Milliseconds())
 
 	start = time.Now()
 
 	sbx, cleanup2, err := sandbox.NewSandbox(
 		childCtx,
 		tracer,
 		dns,
+		proxy,
 		networkPool,
 		templateCache,
 		&orchestrator.SandboxConfig{

packages/orchestrator/internal/sandbox/sandbox_linux.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"net/http"
 	"os"
 	"sync/atomic"
@@ -86,6 +87,7 @@ func NewSandbox(
 	ctx context.Context,
 	tracer trace.Tracer,
 	dns *dns.DNS,
+	proxy *proxy.SessionProxy,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
 	config *orchestrator.SandboxConfig,
@@ -314,11 +316,13 @@ func NewSandbox(
 	sbx.StartedAt = time.Now()
 
 	dns.Add(config.SandboxId, ips.HostIP())
+	proxy.AddSandbox(config.SandboxId, ips.HostIP())
 
 	telemetry.ReportEvent(childCtx, "added DNS record", attribute.String("ip", ips.HostIP()), attribute.String("hostname", config.SandboxId))
 
 	cleanup.Add(func() error {
 		dns.Remove(config.SandboxId, ips.HostIP())
+		proxy.RemoveSandbox(config.SandboxId)
 
 		return nil
 	})

packages/orchestrator/internal/sandbox/sandbox_other.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"net/http"
 	"sync/atomic"
 	"time"
@@ -74,6 +75,7 @@ func NewSandbox(
 	ctx context.Context,
 	tracer trace.Tracer,
 	dns *dns.DNS,
+	proxy *proxy.SessionProxy,
 	networkPool *network.Pool,
 	templateCache *template.Cache,
 	config *orchestrator.SandboxConfig,

packages/orchestrator/internal/server/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"math"
 	"net"
 	"os"
@@ -38,6 +39,7 @@ type server struct {
 	orchestrator.UnimplementedSandboxServiceServer
 	sandboxes       *smap.Map[*sandbox.Sandbox]
 	dns             *dns.DNS
+	proxy           *proxy.SessionProxy
 	tracer          trace.Tracer
 	networkPool     *network.Pool
 	templateCache   *template.Cache
@@ -54,6 +56,7 @@ type Service struct {
 	server   *server
 	grpc     *grpc.Server
 	dns      *dns.DNS
+	proxy    *proxy.SessionProxy
 	port     uint16
 	shutdown struct {
 		once sync.Once
@@ -68,7 +71,7 @@ type Service struct {
 	useClickhouseMetrics string
 }
 
-func New(ctx context.Context, port uint, clientID string) (*Service, error) {
+func New(ctx context.Context, port uint, clientID string, proxy *proxy.SessionProxy) (*Service, error) {
 	if port > math.MaxUint16 {
 		return nil, fmt.Errorf("%d is larger than maximum possible port %d", port, math.MaxInt16)
 	}
@@ -92,6 +95,7 @@ func New(ctx context.Context, port uint, clientID string) (*Service, error) {
 	// BLOCK: initialize services
 	{
 		srv.dns = dns.New()
+		srv.proxy = proxy
 
 		opts := []logging.Option{
 			logging.WithLogOnEvents(logging.StartCall, logging.PayloadReceived, logging.PayloadSent, logging.FinishCall),
@@ -143,6 +147,7 @@ func New(ctx context.Context, port uint, clientID string) (*Service, error) {
 		srv.server = &server{
 			tracer:               otel.Tracer(ServiceName),
 			dns:                  srv.dns,
+			proxy:                srv.proxy,
 			sandboxes:            smap.New[*sandbox.Sandbox](),
 			networkPool:          networkPool,
 			templateCache:        templateCache,

packages/orchestrator/internal/server/sandboxes.go

@@ -49,6 +49,7 @@ func (s *server) Create(ctxConn context.Context, req *orchestrator.SandboxCreate
 		childCtx,
 		s.tracer,
 		s.dns,
+		s.proxy,
 		s.networkPool,
 		s.templateCache,
 		req.Sandbox,
@@ -182,6 +183,7 @@ func (s *server) Delete(ctxConn context.Context, in *orchestrator.SandboxDeleteR
 
 	// Don't allow connecting to the sandbox anymore.
 	s.dns.Remove(in.SandboxId, sbx.Slot.HostIP())
+	s.proxy.RemoveSandbox(in.SandboxId)
 
 	// Remove the sandbox from the cache to prevent loading it again in API during the time the instance is stopping.
 	// Old comment:

packages/orchestrator/main.go

@@ -5,12 +5,14 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"log"
 	"os"
 	"os/signal"
 	"sync"
 	"sync/atomic"
 	"syscall"
+	"time"
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -24,8 +26,9 @@ import (
 )
 
 const (
-	defaultPort = 5008
-	ServiceName = "orchestrator"
+	defaultPort      = 5008
+	defaultProxyPort = 5007
+	ServiceName      = "orchestrator"
 )
 
 var commitSHA string
@@ -38,8 +41,10 @@ func main() {
 	defer sigCancel()
 
 	var port uint
-
 	flag.UintVar(&port, "port", defaultPort, "orchestrator server port")
+
+	var proxyPort uint
+	flag.UintVar(&proxyPort, "proxy-port", defaultProxyPort, "orchestrator proxy port")
 	flag.Parse()
 
 	wg := &sync.WaitGroup{}
@@ -97,8 +102,9 @@ func main() {
 	log.Println("Starting orchestrator", "commit", commitSHA)
 
 	clientID := consul.GetClientID()
+	sessionProxy := proxy.New(proxyPort)
 
-	srv, err := server.New(ctx, port, clientID)
+	srv, err := server.New(ctx, port, clientID, sessionProxy)
 	if err != nil {
 		zap.L().Fatal("failed to create server", zap.Error(err))
 	}
@@ -149,6 +155,32 @@ func main() {
 		}
 	}()
 
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		errChan := make(chan error, 1)
+		go func() {
+			err := sessionProxy.Start()
+			errChan <- err
+		}()
+
+		select {
+		case <-ctx.Done():
+		case <-errChan:
+			if err != nil {
+				zap.L().Error("session proxy failed", zap.Error(err))
+				exitCode.Add(1)
+				cancel()
+			}
+		}
+
+		// close session proxy, wait 5 seconds until all connections are closed
+		shutdownCtx, shutdownCtxCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer shutdownCtxCancel()
+		defer sessionProxy.Shutdown(shutdownCtx)
+	}()
+
 	wg.Wait()
 
 	os.Exit(int(exitCode.Load()))