Fix workspace connection tracker and heartbeats

Author: janekbaraniewski

pkg/daemon/workspace/network/connection_tracker.go

@@ -1,27 +1,36 @@
 package network
 
-import "sync"
+import (
+	"sync"
+
+	"github.com/loft-sh/log"
+)
 
 // ConnTracker is a simple connection counter used by several services.
 type ConnTracker struct {
 	mu    sync.Mutex
 	count int
+
+	logger log.Logger
 }
 
-func (c *ConnTracker) Add() {
+func (c *ConnTracker) Add(serviceName string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	c.count++
+	c.logger.Debugf("%s: Added new connection, connection count %d\n", serviceName, c.count)
 }
 
-func (c *ConnTracker) Remove() {
+func (c *ConnTracker) Remove(serviceName string) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	c.count--
+	c.logger.Debugf("%s: Removed connection, connection count %d\n", serviceName, c.count)
 }
 
-func (c *ConnTracker) Count() int {
+func (c *ConnTracker) Count(serviceName string) int {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	c.logger.Debugf("%s: Get connection count %d\n", serviceName, c.count)
 	return c.count
 }

pkg/daemon/workspace/network/heartbeat.go

@@ -37,19 +37,24 @@ func NewHeartbeatService(config *WorkspaceServerConfig, tsServer *tsnet.Server,
 
 // Start begins the heartbeat loop.
 func (s *HeartbeatService) Start(ctx context.Context) {
+	s.log.Info("HeartbeatService: Start")
 	transport := &http.Transport{DialContext: s.tsServer.Dial}
 	client := &http.Client{Transport: transport, Timeout: 10 * time.Second}
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 	for {
 		select {
 		case <-ctx.Done():
+			s.log.Info("HeartbeatService: Exit")
 			return
 		case <-ticker.C:
-			if s.tracker.Count() > 0 {
+			s.log.Debugf("HeartbeatService: checking connection count")
+			if s.tracker.Count("HeartbeatService") > 0 {
 				if err := s.sendHeartbeat(ctx, client); err != nil {
 					s.log.Errorf("HeartbeatService: failed to send heartbeat: %v", err)
 				}
+			} else {
+				s.log.Debugf("HeartbeatService: No active connections, skipping heartbeat.")
 			}
 		}
 	}
@@ -58,25 +63,33 @@ func (s *HeartbeatService) Start(ctx context.Context) {
 func (s *HeartbeatService) sendHeartbeat(ctx context.Context, client *http.Client) error {
 	hbCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
+
 	discoveredRunner, err := discoverRunner(hbCtx, s.lc, s.log)
 	if err != nil {
+		s.log.Errorf("HeartbeatService: failed to discover runner: %v", err)
 		return fmt.Errorf("failed to discover runner: %w", err)
 	}
+
 	heartbeatURL := fmt.Sprintf("http://%s.ts.loft/devpod/%s/%s/heartbeat", discoveredRunner, s.projectName, s.workspaceName)
-	s.log.Infof("HeartbeatService: sending heartbeat to %s, active connections: %d", heartbeatURL, s.tracker.Count())
+	s.log.Infof("HeartbeatService: sending heartbeat to %s, active connections: %d", heartbeatURL, s.tracker.Count("HeartbeatService"))
 	req, err := http.NewRequestWithContext(hbCtx, "GET", heartbeatURL, nil)
 	if err != nil {
+		s.log.Errorf("HeartbeatService: failed to create request for %s: %v", heartbeatURL, err)
 		return fmt.Errorf("failed to create request for %s: %w", heartbeatURL, err)
 	}
 	req.Header.Set("Authorization", "Bearer "+s.config.AccessKey)
 	resp, err := client.Do(req)
 	if err != nil {
+		s.log.Errorf("HeartbeatService: request to %s failed: %v", heartbeatURL, err)
 		return fmt.Errorf("request to %s failed: %w", heartbeatURL, err)
 	}
 	defer resp.Body.Close()
+
 	if resp.StatusCode != http.StatusOK {
+		s.log.Errorf("HeartbeatService: received non-OK response from %s - Status: %d", heartbeatURL, resp.StatusCode)
 		return fmt.Errorf("received response from %s - Status: %d", heartbeatURL, resp.StatusCode)
 	}
-	s.log.Infof("HeartbeatService: received response from %s - Status: %d", heartbeatURL, resp.StatusCode)
+
+	s.log.Debugf("HeartbeatService: received response from %s - Status: %d", heartbeatURL, resp.StatusCode)
 	return nil
 }

pkg/daemon/workspace/network/http_proxy.go

@@ -29,7 +29,6 @@ func NewHttpProxyHandler(tsSrv *tsnet.Server, logger log.Logger) *HttpProxyHandl
 	}
 }
 
-// tsDialer is a helper within HttpProxyHandler for dialing over tsnet.
 func (h *HttpProxyHandler) tsDialer(ctx context.Context, addr string) (net.Conn, error) {
 	h.log.Debugf("HttpProxyHandler: Dialing target %s via tsnet", addr)
 	conn, err := h.tsServer.Dial(ctx, "tcp", addr)

pkg/daemon/workspace/network/netmap.go

@@ -31,21 +31,19 @@ func NewNetmapWatcherService(rootDir string, lc *tailscale.LocalClient, log log.
 
 // Start begins watching the netmap.
 func (s *NetmapWatcherService) Start(ctx context.Context) {
-	go func() {
-		lastUpdate := time.Now()
-		if err := ts.WatchNetmap(ctx, s.lc, func(netMap *netmap.NetworkMap) {
-			if time.Since(lastUpdate) < netMapCooldown {
-				return
-			}
-			lastUpdate = time.Now()
-			nm, err := json.Marshal(netMap)
-			if err != nil {
-				s.log.Errorf("NetmapWatcherService: failed to marshal netmap: %v", err)
-			} else {
-				_ = os.WriteFile(filepath.Join(s.rootDir, "netmap.json"), nm, 0644)
-			}
-		}); err != nil {
-			s.log.Errorf("NetmapWatcherService: failed to watch netmap: %v", err)
+	lastUpdate := time.Now()
+	if err := ts.WatchNetmap(ctx, s.lc, func(netMap *netmap.NetworkMap) {
+		if time.Since(lastUpdate) < netMapCooldown {
+			return
 		}
-	}()
+		lastUpdate = time.Now()
+		nm, err := json.Marshal(netMap)
+		if err != nil {
+			s.log.Errorf("NetmapWatcherService: failed to marshal netmap: %v", err)
+		} else {
+			_ = os.WriteFile(filepath.Join(s.rootDir, "netmap.json"), nm, 0644)
+		}
+	}); err != nil {
+		s.log.Errorf("NetmapWatcherService: failed to watch netmap: %v", err)
+	}
 }

pkg/daemon/workspace/network/port_forward.go

@@ -58,8 +58,8 @@ func (s *HTTPPortForwardService) Start(ctx context.Context) {
 }
 
 func (s *HTTPPortForwardService) portForwardHandler(w http.ResponseWriter, r *http.Request) {
-	s.tracker.Add()
-	defer s.tracker.Remove()
+	s.tracker.Add("PortForward")
+	defer s.tracker.Remove("PortForward")
 	s.log.Debugf("HTTPPortForwardService: received request")
 
 	targetPort := r.Header.Get("X-Loft-Forward-Port")

pkg/daemon/workspace/network/server.go

@@ -46,9 +46,11 @@ type WorkspaceServer struct {
 // NewWorkspaceServer creates a new WorkspaceServer.
 func NewWorkspaceServer(config *WorkspaceServerConfig, logger log.Logger) *WorkspaceServer {
 	return &WorkspaceServer{
-		config:      config,
-		log:         logger,
-		connTracker: &ConnTracker{},
+		config: config,
+		log:    logger,
+		connTracker: &ConnTracker{
+			logger: logger,
+		},
 	}
 }
 
@@ -72,7 +74,6 @@ func (s *WorkspaceServer) Start(ctx context.Context) error {
 	}
 	s.sshSvc.Start(ctx)
 
-	// Create and start the HTTP port forward service.
 	s.httpProxySvc, err = NewHTTPPortForwardService(s.network, s.connTracker, s.log)
 	if err != nil {
 		return err
@@ -92,15 +93,15 @@ func (s *WorkspaceServer) Start(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	s.netProxySvc.Start(ctx)
+	go s.netProxySvc.Start(ctx)
 
 	// Start the heartbeat service.
 	s.heartbeatSvc = NewHeartbeatService(s.config, s.network, lc, projectName, workspaceName, s.connTracker, s.log)
 	go s.heartbeatSvc.Start(ctx)
 
 	// Start netmap watcher.
 	s.netmapWatcher = NewNetmapWatcherService(s.config.RootDir, lc, s.log)
-	s.netmapWatcher.Start(ctx)
+	go s.netmapWatcher.Start(ctx)
 
 	// Wait until the context is canceled.
 	<-ctx.Done()

pkg/daemon/workspace/network/ssh.go

@@ -59,8 +59,8 @@ func (s *SSHService) acceptLoop(ctx context.Context) {
 }
 
 func (s *SSHService) handleConnection(conn net.Conn) {
-	s.tracker.Add()
-	defer s.tracker.Remove()
+	s.tracker.Add("SSHService")
+	defer s.tracker.Remove("SSHService")
 	defer conn.Close()
 
 	localAddr := fmt.Sprintf("127.0.0.1:%d", sshServer.DefaultUserPort)
@@ -71,10 +71,23 @@ func (s *SSHService) handleConnection(conn net.Conn) {
 	}
 	defer backendConn.Close()
 
+	// We need to wait for copying to finish before the function returns and Remove is called.
+	errChan := make(chan error, 2)
+
 	go func() {
-		_, _ = io.Copy(backendConn, conn)
+		_, err := io.Copy(backendConn, conn)
+		errChan <- err
 	}()
-	_, _ = io.Copy(conn, backendConn)
+
+	go func() {
+		_, err := io.Copy(conn, backendConn)
+		errChan <- err
+	}()
+
+	// Wait for one side of the connection to close or error
+	<-errChan
+	// Optionally wait for the second one too, or just proceed to cleanup
+	// <-errChan
 }
 
 // Stop stops the SSH server by closing its listener.