Add timeout option

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,
@@ -119,7 +127,8 @@ func New(command []string, options *Options) (*App, error) {
 
 		titleTemplate: titleTemplate,
 
-		onceMutex: umutex.New(),
+		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 {
-		if app.connections >= app.options.MaxConnection {
+	app.stopTimer()
+
+	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{

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)
-		}
 	}()
 }