Add timeout option

Iwasaki Yudai
Commit 8c9433ff authored Jan 09, 2017 by Iwasaki Yudai
Showing with 59 additions and 15 deletions
.gotty
README.md
app/app.go
app/client_context.go
main.go
--- a/.gotty
+++ b/.gotty
@@ -52,7 +52,10 @@
 // [int] Interval time to try reconnection (seconds)
 //       To enable reconnection, set `true` to `enable_reconnect`
-// reconnect_time = false
+// reconnect_time = 10
+// [int] Timeout seconds for waiting a client (0 to disable)
+// timeout = 60
 // [int] Maximum connection to gotty, 0(default) means no limit.
 // max_connection = 0

--- a/README.md
+++ b/README.md
@@ -64,6 +64,7 @@ By default, GoTTY starts a web server at port 8080. Open the URL on your web bro
 --title-format "GoTTY - {{ .Command }} ({{ .Hostname }})"    Title format of browser window [$GOTTY_TITLE_FORMAT]
 --reconnect                                                  Enable reconnection [$GOTTY_RECONNECT]
 --reconnect-time "10"                                        Time to reconnect [$GOTTY_RECONNECT_TIME]
+--timeout "0"                                                Timeout seconds for waiting a client (0 to disable) [$GOTTY_TIMEOUT]
 --once                                                       Accept only one client and exit on disconnection [$GOTTY_ONCE]
 --permit-arguments                                           Permit clients to send command line arguments in URL (e.g. http://example.com:8080/?arg=AAA&arg=BBB) [$GOTTY_PERMIT_ARGUMENTS]
 --close-signal "1"                                           Signal sent to the command process when gotty close it (default: SIGHUP) [$GOTTY_CLOSE_SIGNAL]

--- a/app/app.go
+++ b/app/app.go
@@ -18,7 +18,9 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"text/template"
+	"time"
 	"github.com/braintree/manners"
 	"github.com/elazarl/go-bindata-assetfs"
@@ -43,8 +45,11 @@ type App struct {
 	titleTemplate *template.Template
 	onceMutex *umutex.UnblockingMutex
+	timer     *time.Timer
-	connections int
+	// clientContext writes concurrently
+	// Use atomic operations.
+	connections *int64
 }
 type Options struct {
@@ -66,6 +71,7 @@ type Options struct {
 	ReconnectTime       int                    `hcl:"reconnect_time"`
 	MaxConnection       int                    `hcl:"max_connection"`
 	Once                bool                   `hcl:"once"`
+	Timeout             int                    `hcl:"timeout"`
 	PermitArguments     bool                   `hcl:"permit_arguments"`
 	CloseSignal         int                    `hcl:"close_signal"`
 	Preferences         HtermPrefernces        `hcl:"preferences"`
@@ -107,6 +113,8 @@ func New(command []string, options *Options) (*App, error) {
 		return nil, errors.New("Title format string syntax error")
 	}
+	connections := int64(0)
 	return &App{
 		command: command,
 		options: options,
@@ -120,6 +128,7 @@ func New(command []string, options *Options) (*App, error) {
 		titleTemplate: titleTemplate,
 		onceMutex:   umutex.New(),
+		connections: &connections,
 	}, nil
 }
@@ -235,6 +244,14 @@ func (app *App) Run() error {
 		server,
 	)
+	if app.options.Timeout > 0 {
+		app.timer = time.NewTimer(time.Duration(app.options.Timeout) * time.Second)
+		go func() {
+			<-app.timer.C
+			app.Exit()
+		}()
+	}
 	if app.options.EnableTLS {
 		crtFile := ExpandHomeDir(app.options.TLSCrtFile)
 		keyFile := ExpandHomeDir(app.options.TLSKeyFile)
@@ -281,9 +298,24 @@ func (app *App) makeServer(addr string, handler *http.Handler) (*http.Server, er
 	return server, nil
 }
+func (app *App) stopTimer() {
+	if app.options.Timeout > 0 {
+		app.timer.Stop()
+	}
+}
+func (app *App) restartTimer() {
+	if app.options.Timeout > 0 {
+		app.timer.Reset(time.Duration(app.options.Timeout) * time.Second)
+	}
+}
 func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
-	if app.options.MaxConnection != 0 {
+	app.stopTimer()
-		if app.connections >= app.options.MaxConnection {
+	connections := atomic.AddInt64(app.connections, 1)
+	if int64(app.options.MaxConnection) != 0 {
+		if connections >= int64(app.options.MaxConnection) {
 			log.Printf("Reached max connection: %d", app.options.MaxConnection)
 			return
 		}
@@ -357,13 +389,12 @@ func (app *App) handleWS(w http.ResponseWriter, r *http.Request) {
 		return
 	}
-	app.connections++
 	if app.options.MaxConnection != 0 {
 		log.Printf("Command is running for client %s with PID %d (args=%q), connections: %d/%d",
-			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), app.connections, app.options.MaxConnection)
+			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), connections, app.options.MaxConnection)
 	} else {
 		log.Printf("Command is running for client %s with PID %d (args=%q), connections: %d",
-			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), app.connections)
+			r.RemoteAddr, cmd.Process.Pid, strings.Join(argv, " "), connections)
 	}
 	context := &clientContext{

--- a/app/client_context.go
+++ b/app/client_context.go
@@ -10,6 +10,7 @@ import (
 	"os/exec"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"unsafe"
@@ -69,6 +70,21 @@ func (context *clientContext) goHandleClient() {
 	go func() {
 		defer context.app.server.FinishRoutine()
+		defer func() {
+			connections := atomic.AddInt64(context.app.connections, -1)
+			if context.app.options.MaxConnection != 0 {
+				log.Printf("Connection closed: %s, connections: %d/%d",
+					context.request.RemoteAddr, connections, context.app.options.MaxConnection)
+			} else {
+				log.Printf("Connection closed: %s, connections: %d",
+					context.request.RemoteAddr, connections)
+			}
+			if connections == 0 {
+				context.app.restartTimer()
+			}
+		}()
 		<-exit
 		context.pty.Close()
@@ -79,14 +95,6 @@ func (context *clientContext) goHandleClient() {
 		context.command.Wait()
 		context.connection.Close()
-		context.app.connections--
-		if context.app.options.MaxConnection != 0 {
-			log.Printf("Connection closed: %s, connections: %d/%d",
-				context.request.RemoteAddr, context.app.connections, context.app.options.MaxConnection)
-		} else {
-			log.Printf("Connection closed: %s, connections: %d",
-				context.request.RemoteAddr, context.app.connections)
-		}
 	}()
 }

--- a/main.go
+++ b/main.go
@@ -33,6 +33,7 @@ func main() {
 		flag{"title-format", "", "Title format of browser window"},
 		flag{"reconnect", "", "Enable reconnection"},
 		flag{"reconnect-time", "", "Time to reconnect"},
+		flag{"timeout", "", "Timeout seconds for waiting a client (0 to disable)"},
 		flag{"max-connection", "", "Maximum connection to gotty, 0(default) means no limit"},
 		flag{"once", "", "Accept only one client and exit on disconnection"},
 		flag{"permit-arguments", "", "Permit clients to send command line arguments in URL (e.g. http://example.com:8080/?arg=AAA&arg=BBB)"},