fix: Handle zombie sessions

Author: viccon

server/handlers.go

@@ -1,8 +1,6 @@
 package server
 
 import (
-	"strconv"
-
 	"github.com/creativecreature/pulse"
 )
 
@@ -16,10 +14,9 @@ func (s *Server) FocusGained(event pulse.Event, reply *string) {
 		"os", event.OS,
 	)
 
-	// Lock the mutex to prevent race conditions with the heartbeat check.
+	s.checkHeartbeat()
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
-
 	s.lastHeartbeat = s.clock.GetTime()
 
 	// The FocusGained event will be triggered when I switch back to an active
@@ -82,9 +79,10 @@ func (s *Server) OpenFile(event pulse.Event, reply *string) {
 		return
 	}
 
-	// Lock the mutex to prevent race conditions with the heartbeat check.
+	s.checkHeartbeat()
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
+	s.lastHeartbeat = s.clock.GetTime()
 
 	// The editor could have been inactive, while focused, for 10 minutes.
 	// That would end the session, and we could get a OpenFile event without
@@ -96,7 +94,6 @@ func (s *Server) OpenFile(event pulse.Event, reply *string) {
 	}
 
 	// If a new file was opened it means that the session is still active.
-	s.lastHeartbeat = s.clock.GetTime()
 	gitFile, gitFileErr := s.fileReader.GitFile(event.Path)
 	if gitFileErr != nil {
 		s.log.Debug("Failed to get git file.",
@@ -122,9 +119,10 @@ func (s *Server) SendHeartbeat(event pulse.Event, reply *string) {
 		"os", event.OS,
 	)
 
-	// Lock the mutex to prevent race conditions with the heartbeat check.
+	s.checkHeartbeat()
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
+	s.lastHeartbeat = s.clock.GetTime()
 
 	// This is to handle the case where the server would have ended the clients
 	// session due to inactivity. When a session ends it is written to disk and
@@ -147,8 +145,6 @@ func (s *Server) SendHeartbeat(event pulse.Event, reply *string) {
 		s.setActiveBuffer(gitFile)
 	}
 
-	// Update the time for the last heartbeat.
-	s.lastHeartbeat = s.clock.GetTime()
 	*reply = "Successfully sent heartbeat"
 }
 
@@ -160,7 +156,7 @@ func (s *Server) EndSession(event pulse.Event, reply *string) {
 		"os", event.OS,
 	)
 
-	// Lock the mutex to prevent race conditions with the heartbeat check.
+	s.checkHeartbeat()
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
@@ -188,29 +184,3 @@ func (s *Server) EndSession(event pulse.Event, reply *string) {
 	s.saveActiveSession()
 	*reply = "The session was ended successfully."
 }
-
-// CheckHeartbeat is used to check if the session has been inactive for more than
-// ten minutes. If that is the case, the session will be terminated and saved to disk.
-func (s *Server) CheckHeartbeat() {
-	s.log.Debug("Checking heartbeat.",
-		"active_editor_id", s.activeEditorID,
-		"last_heartbeat", s.lastHeartbeat,
-		"time_now", s.clock.GetTime(),
-	)
-	if s.activeEditorID == "" {
-		return
-	}
-
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
-	if s.lastHeartbeat+HeartbeatTTL.Milliseconds() < s.clock.GetTime() {
-		s.log.Info(
-			"Ending all active sessions due to inactivity",
-			"last_heartbeat", strconv.FormatInt(s.lastHeartbeat, 10),
-			"current_time", strconv.FormatInt(s.clock.GetTime(), 10),
-		)
-		s.saveAllSessions()
-		s.activeEditorID = ""
-	}
-}

server/heartbeat.go

@@ -0,0 +1,60 @@
+package server
+
+import (
+	"strconv"
+	"time"
+)
+
+const (
+	HeartbeatTTL      = time.Minute * 10
+	heartbeatInterval = time.Second * 10
+)
+
+// CheckHeartbeat is used to check if the session has been inactive for more than
+// ten minutes. If that is the case, the session will be terminated and saved to disk.
+func (s *Server) checkHeartbeat() {
+	s.log.Debug("Checking heartbeat.",
+		"active_editor_id", s.activeEditorID,
+		"last_heartbeat", s.lastHeartbeat,
+		"time_now", s.clock.GetTime(),
+	)
+	if s.activeEditorID == "" {
+		return
+	}
+
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if s.lastHeartbeat+HeartbeatTTL.Milliseconds() < s.clock.GetTime() {
+		s.log.Info(
+			"Ending all active sessions due to inactivity",
+			"last_heartbeat", strconv.FormatInt(s.lastHeartbeat, 10),
+			"current_time", strconv.FormatInt(s.clock.GetTime(), 10),
+			"end_time", strconv.FormatInt(s.lastHeartbeat+int64(HeartbeatTTL), 10),
+		)
+
+		// The machine may have entered sleep mode, potentially stopping the heartbeat
+		// check from executing at its scheduled interval. To mitigate this, the session
+		// will be terminated based on the time of the last recorded heartbeat plus the
+		// TTL. This prevents the creation of inaccurately long sessions.
+		s.saveAllSessions(s.lastHeartbeat + int64(HeartbeatTTL/time.Millisecond))
+		s.activeEditorID = ""
+	}
+}
+
+// startHeartbeatChecks runs in a separate goroutine and makes sure
+// that no session is allowed to be idle for more than 10 minutes.
+func (s *Server) startHeartbeatChecks() {
+	go func() {
+		ticker, stopTicker := s.clock.NewTicker(heartbeatInterval)
+		defer stopTicker()
+		for {
+			select {
+			case <-s.stopHeartbeatChecks:
+				return
+			case <-ticker:
+				s.checkHeartbeat()
+			}
+		}
+	}()
+}

server/server.go

@@ -18,39 +18,36 @@ import (
 	"github.com/creativecreature/pulse/proxy"
 )
 
-const (
-	HeartbeatTTL      = time.Minute * 10
-	heartbeatInterval = time.Second * 45
-)
-
 type Server struct {
-	name           string
-	activeEditorID string
-	activeSessions map[string]*pulse.CodingSession
-	lastHeartbeat  int64
-	clock          pulse.Clock
-	fileReader     FileReader
-	log            *log.Logger
-	mutex          sync.Mutex
-	storage        pulse.TemporaryStorage
+	name                string
+	activeEditorID      string
+	activeSessions      map[string]*pulse.CodingSession
+	lastHeartbeat       int64
+	stopHeartbeatChecks chan struct{}
+	clock               pulse.Clock
+	fileReader          FileReader
+	log                 *log.Logger
+	mutex               sync.Mutex
+	storage             pulse.TemporaryStorage
 }
 
 // New creates a new server.
 func New(serverName string, opts ...Option) (*Server, error) {
-	a := &Server{
-		name:           serverName,
-		activeSessions: make(map[string]*pulse.CodingSession),
-		clock:          pulse.NewClock(),
-		fileReader:     filereader.New(),
+	s := &Server{
+		name:                serverName,
+		activeSessions:      make(map[string]*pulse.CodingSession),
+		clock:               pulse.NewClock(),
+		stopHeartbeatChecks: make(chan struct{}),
+		fileReader:          filereader.New(),
 	}
 	for _, opt := range opts {
-		err := opt(a)
+		err := opt(s)
 		if err != nil {
 			return &Server{}, err
 		}
 	}
-
-	return a, nil
+	s.startHeartbeatChecks()
+	return s, nil
 }
 
 // createSession creates a new session and sets it as the current session.
@@ -81,8 +78,7 @@ func (s *Server) setActiveBuffer(gitFile pulse.GitFile) {
 	)
 }
 
-func (s *Server) saveAllSessions() {
-	now := s.clock.GetTime()
+func (s *Server) saveAllSessions(endedAt int64) {
 	s.log.Debug("Saving all sessions.")
 
 	for _, session := range s.activeSessions {
@@ -95,7 +91,7 @@ func (s *Server) saveAllSessions() {
 			return
 		}
 
-		finishedSession := session.End(now)
+		finishedSession := session.End(endedAt)
 		err := s.storage.Write(finishedSession)
 		if err != nil {
 			s.log.Error(err)
@@ -177,20 +173,6 @@ func (s *Server) startServer(port string) (*http.Server, error) {
 	return httpServer, nil
 }
 
-// HeartbeatCheck runs a heartbeat ticker that ensures that
-// the current session is not idle for more than ten minutes.
-func (s *Server) HeartbeatCheck() func() {
-	s.log.Info("Starting the heartbeat checks.")
-	ticker, stop := s.clock.NewTicker(heartbeatInterval)
-	go func() {
-		for range ticker {
-			s.CheckHeartbeat()
-		}
-	}()
-
-	return stop
-}
-
 // Start starts the server on the given port.
 func (s *Server) Start(port string) error {
 	s.log.Info("Starting up...")
@@ -199,9 +181,6 @@ func (s *Server) Start(port string) error {
 		return err
 	}
 
-	// Start the ECG. It will end inactive sessions.
-	stopHeartbeat := s.HeartbeatCheck()
-
 	// Catch shutdown signals from the OS
 	quit := make(chan os.Signal, 1)
 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
@@ -211,15 +190,15 @@ func (s *Server) Start(port string) error {
 	s.log.Info("Received shutdown signal", "signal", sig.String())
 
 	// Stop the heartbeat checks and shutdown the http server.
-	stopHeartbeat()
+	s.stopHeartbeatChecks <- struct{}{}
 	err = httpServer.Shutdown(context.Background())
 	if err != nil {
 		s.log.Error(err)
 		return err
 	}
 
 	// Save the all sessions before shutting down.
-	s.saveAllSessions()
+	s.saveAllSessions(s.clock.GetTime())
 	s.log.Info("Shutting down.")
 
 	return nil

server/server_test.go

@@ -398,8 +398,6 @@ func TestNoActivityShouldEndSession(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	s.HeartbeatCheck()
-
 	// Open an initial VIM window.
 	s.OpenFile(pulse.Event{
 		EditorID: "123",
@@ -445,11 +443,17 @@ func TestNoActivityShouldEndSession(t *testing.T) {
 	if storedSessions[0].DurationMs != 100 {
 		t.Errorf("expected the sessions duration to be 100; got %d", storedSessions[0].DurationMs)
 	}
+	if storedSessions[0].Files[0].DurationMs != 100 {
+		t.Errorf("expected the file duration to be 100; got %d", storedSessions[0].Files[0].DurationMs)
+	}
 
-	// The second session should last for 11 minutes. 10 for the
-	// heartbeat to expire, and 1 for the heartbeat to trigger.
-	dur := int64(11 * time.Minute / time.Millisecond)
+	// The second session should have been terminated by the
+	// heartbeat check after 10 minutes of inactivity.
+	dur := int64(10 * time.Minute / time.Millisecond)
 	if storedSessions[1].DurationMs != dur {
 		t.Errorf("expected the sessions duration to be %d; got %d", dur, storedSessions[1].DurationMs)
 	}
+	if storedSessions[1].Files[0].DurationMs != dur {
+		t.Errorf("expected the file duration to be %d; got %d", dur, storedSessions[1].Files[0].DurationMs)
+	}
 }