From 6fefaba097eaacbc2a81dc44cea568baa5e6afc9 Mon Sep 17 00:00:00 2001
From: Rafael da Fonseca <rafael.fonseca@wildlifestudios.com>
Date: Sat, 30 Sep 2023 19:54:49 +0100
Subject: [PATCH] Add optional session drain

---
 app.go           | 7 +++++++
 config/config.go | 1 +
 2 files changed, 8 insertions(+)

diff --git a/app.go b/app.go
index a440b732..cd3393ff 100644
--- a/app.go
+++ b/app.go
@@ -434,6 +434,13 @@ func Start() {
 		logger.Log.Warn("the app will shutdown in a few seconds")
 	case s := <-sg:
 		logger.Log.Warn("got signal: ", s, ", shutting down...")
+		if app.config.GetBool("pitaya.session.drain") {
+			logger.Log.Info("Session drain is enabled, draining all sessions before shutting down")
+			for session.SessionCount > 0 {
+				logger.Log.Infof("Waiting for all sessions to finish: %d sessions remaining...", session.SessionCount)
+				time.Sleep(5 * time.Second)
+			}
+		}
 		close(app.dieChan)
 	}
 
diff --git a/config/config.go b/config/config.go
index f4a50430..7845e83d 100644
--- a/config/config.go
+++ b/config/config.go
@@ -114,6 +114,7 @@ func (c *Config) fillDefaultValues() {
 		"pitaya.conn.ratelimiting.interval":                "1s",
 		"pitaya.conn.ratelimiting.forcedisable":            false,
 		"pitaya.session.unique":                            true,
+		"pitaya.session.drain":                             false,
 		"pitaya.worker.concurrency":                        1,
 		"pitaya.worker.redis.pool":                         "10",
 		"pitaya.worker.redis.url":                          "localhost:6379",