Add session keep alive. If client doesn't send a keep alive, after a

specified amount of time the scheduler closes the session
crt-fork · Oct 16, 2017 · 1b0d363 · 1b0d363
1 parent 8769aa3
commit 1b0d363
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 3 deletions.
diff --git a/config/config.go b/config/config.go
@@ -3,6 +3,7 @@ package config
 import (
 	"flag"
 	"fmt"
+	"log"
 	"os"
 	"regexp"
 	"time"
@@ -39,6 +40,8 @@ var SecureCookie *securecookie.SecureCookie
 var GithubClientID, GithubClientSecret string
 var FacebookClientID, FacebookClientSecret string
 var DockerClientID, DockerClientSecret string
+var SessionKeepAlive time.Duration
+var sessionKeepAlive string
 
 type stringslice []string
 
@@ -72,6 +75,7 @@ func ParseFlags() {
 	flag.StringVar(&SSHKeyPath, "ssh_key_path", "", "SSH Private Key to use")
 	flag.StringVar(&CookieHashKey, "cookie-hash-key", "", "Hash key to use to validate cookies")
 	flag.StringVar(&CookieBlockKey, "cookie-block-key", "", "Block key to use to encrypt cookies")
+	flag.StringVar(&sessionKeepAlive, "session-keep-alive", "5m", "Duration for which a session will be kept alive when no more heartbeats arrive")
 
 	flag.StringVar(&GithubClientID, "oauth-github-client-id", "", "Github OAuth Client ID")
 	flag.StringVar(&GithubClientSecret, "oauth-github-client-secret", "", "Github OAuth Client Secret")
@@ -86,6 +90,12 @@ func ParseFlags() {
 
 	SecureCookie = securecookie.New([]byte(CookieHashKey), []byte(CookieBlockKey))
 
+	dur, err := time.ParseDuration(sessionKeepAlive)
+	if err != nil {
+		log.Fatalf("Cannot parse duration of flag [-session-keep-alive]. Got: %v\n", err)
+	}
+	SessionKeepAlive = dur
+
 	registerOAuthProviders()
 }
 

diff --git a/event/event.go b/event/event.go
@@ -15,6 +15,7 @@ var (
 	SESSION_END              = EventType("session end")
 	SESSION_READY            = EventType("session ready")
 	SESSION_BUILDER_OUT      = EventType("session builder out")
+	SESSION_KEEP_ALIVE       = EventType("session keep alive")
 )
 
 type Handler func(sessionId string, args ...interface{})

diff --git a/handlers/ws.go b/handlers/ws.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/googollee/go-socket.io"
 	"github.com/gorilla/mux"
+	"github.com/play-with-docker/play-with-docker/event"
 )
 
 func WS(so socketio.Socket) {
@@ -65,6 +66,10 @@ func WS(so socketio.Socket) {
 		m.Close()
 		core.ClientClose(client)
 	})
+
+	so.On("session keep alive", func() {
+		e.Emit(event.SESSION_KEEP_ALIVE, sessionId)
+	})
 }
 
 func WSError(so socketio.Socket) {

diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"time"
 
+	"github.com/play-with-docker/play-with-docker/config"
 	"github.com/play-with-docker/play-with-docker/event"
 	"github.com/play-with-docker/play-with-docker/pwd"
 	"github.com/play-with-docker/play-with-docker/pwd/types"
@@ -23,8 +24,9 @@ type SchedulerApi interface {
 }
 
 type scheduledSession struct {
-	session *types.Session
-	cancel  context.CancelFunc
+	session   *types.Session
+	keepAlive *time.Timer
+	cancel    context.CancelFunc
 }
 
 type scheduledInstance struct {
@@ -68,6 +70,10 @@ func (s *scheduler) processSession(ctx context.Context, ss *scheduledSession) {
 		// Session has expired. Need to close the session.
 		s.pwd.SessionClose(ss.session)
 		return
+	case <-ss.keepAlive.C:
+		// No keep alive has been received after the defined interval
+		log.Printf("No keep alive has been received for session %s after %s. Closing.\n", ss.session.Id, config.SessionKeepAlive)
+		s.pwd.SessionClose(ss.session)
 	case <-ctx.Done():
 		return
 	}
@@ -141,6 +147,7 @@ func (s *scheduler) scheduleSession(session *types.Session) {
 	s.scheduledSessions[session.Id] = ss
 	ctx, cancel := context.WithCancel(context.Background())
 	ss.cancel = cancel
+	ss.keepAlive = time.NewTimer(config.SessionKeepAlive)
 	go s.processSession(ctx, ss)
 	log.Printf("Scheduled session %s\n", session.Id)
 }
@@ -225,6 +232,18 @@ func (s *scheduler) Start() error {
 		instance := &types.Instance{Name: instanceName}
 		s.unscheduleInstance(instance)
 	})
+	s.event.On(event.SESSION_KEEP_ALIVE, func(sessionId string, args ...interface{}) {
+		log.Printf("Keep alive recevied for session %s\n", sessionId)
+		if _, found := s.scheduledSessions[sessionId]; !found {
+			log.Printf("Session %s was not found. Ignoring.\n", sessionId)
+			return
+		}
+		ss := s.scheduledSessions[sessionId]
+		if ss.keepAlive.Stop() {
+			log.Printf("Keep alive reset for session %s\n", sessionId)
+			ss.keepAlive.Reset(config.SessionKeepAlive)
+		}
+	})
 	s.started = true
 
 	return nil

diff --git a/www/assets/app.js b/www/assets/app.js
@@ -181,6 +181,8 @@
 
         var socket = io({ path: '/sessions/' + sessionId + '/ws' });
 
+
+
         socket.on('instance terminal status', function(name, status) {
             var instance = $scope.idx[name];
             instance.status = status;
@@ -247,8 +249,11 @@
         socket.on('connect_error', function() {
           $scope.connected = false;
         });
-        socket.on('connect', function() {
+        socket.on('connect', function(s) {
           $scope.connected = true;
+  	  setInterval(function() {
+        	socket.emit('session keep alive', 'keep me alive please! ' + new Date());
+	  }, 30*1000);
         });
 
         socket.on('instance stats', function(stats) {