Fix thread exhaustion in SSH daemon

5 files changed, 29 insertions, 14 deletions

diff --git a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
index 4d64cfb7..7c512907 100644
--- a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
+++ b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
@@ -41,9 +41,9 @@ import com.gitblit.IStoredSettings;
 import com.gitblit.Keys;
 import com.gitblit.manager.IGitblit;
 import com.gitblit.transport.ssh.commands.SshCommandFactory;
-import com.gitblit.utils.IdGenerator;
 import com.gitblit.utils.JnaUtils;
 import com.gitblit.utils.StringUtils;
+import com.gitblit.utils.WorkQueue;
 import com.google.common.io.Files;
 
 /**
@@ -76,8 +76,9 @@ public class SshDaemon {
 	 * Construct the Gitblit SSH daemon.
 	 *
 	 * @param gitblit
+	 * @param workQueue
 	 */
-	public SshDaemon(IGitblit gitblit, IdGenerator idGenerator) {
+	public SshDaemon(IGitblit gitblit, WorkQueue workQueue) {
 		this.gitblit = gitblit;
 
 		IStoredSettings settings = gitblit.getSettings();
@@ -126,7 +127,7 @@ public class SshDaemon {
 		sshd.setSessionFactory(new SshServerSessionFactory());
 		sshd.setFileSystemFactory(new DisabledFilesystemFactory());
 		sshd.setTcpipForwardingFilter(new NonForwardingFilter());
-		sshd.setCommandFactory(new SshCommandFactory(gitblit, idGenerator));
+		sshd.setCommandFactory(new SshCommandFactory(gitblit, workQueue));
 		sshd.setShellFactory(new WelcomeShell(settings));
 
 		// Set the server id.  This can be queried with:
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
index d996ea9a..ab2756d0 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
@@ -33,12 +33,13 @@ import org.apache.sshd.server.Environment;
 import org.apache.sshd.server.ExitCallback;
 import org.apache.sshd.server.SessionAware;
 import org.apache.sshd.server.session.ServerSession;
+import org.kohsuke.args4j.Argument;
 import org.kohsuke.args4j.CmdLineException;
+import org.kohsuke.args4j.Option;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.gitblit.Keys;
-import com.gitblit.utils.IdGenerator;
 import com.gitblit.utils.StringUtils;
 import com.gitblit.utils.WorkQueue;
 import com.gitblit.utils.WorkQueue.CancelableRunnable;
@@ -80,13 +81,10 @@ public abstract class BaseCommand implements Command, SessionAware {
 	/** The task, as scheduled on a worker thread. */
 	private final AtomicReference<Future<?>> task;
 
-	private final WorkQueue.Executor executor;
+	private WorkQueue workQueue;
 
 	public BaseCommand() {
 		task = Atomics.newReference();
-		IdGenerator gen = new IdGenerator();
-		WorkQueue w = new WorkQueue(gen);
-		this.executor = w.getDefaultQueue();
 	}
 
 	@Override
@@ -97,6 +95,10 @@ public abstract class BaseCommand implements Command, SessionAware {
 	@Override
 	public void destroy() {
 		log.debug("destroying " + getClass().getName());
+		Future<?> future = task.getAndSet(null);
+		if (future != null && !future.isDone()) {
+			future.cancel(true);
+		}
 		session = null;
 		ctx = null;
 	}
@@ -110,12 +112,21 @@ public abstract class BaseCommand implements Command, SessionAware {
 
 	protected void provideStateTo(final BaseCommand cmd) {
 		cmd.setContext(ctx);
+		cmd.setWorkQueue(workQueue);
 		cmd.setInputStream(in);
 		cmd.setOutputStream(out);
 		cmd.setErrorStream(err);
 		cmd.setExitCallback(exit);
 	}
 
+	public WorkQueue getWorkQueue() {
+		return workQueue;
+	}
+
+	public void setWorkQueue(WorkQueue workQueue) {
+		this.workQueue = workQueue;
+	}
+
 	public void setContext(SshCommandContext ctx) {
 		this.ctx = ctx;
 	}
@@ -467,7 +478,7 @@ public abstract class BaseCommand implements Command, SessionAware {
 	 */
 	protected void startThread(final CommandRunnable thunk) {
 		final TaskThunk tt = new TaskThunk(thunk);
-		task.set(executor.submit(tt));
+		task.set(workQueue.getDefaultQueue().submit(tt));
 	}
 
 	/** Thrown from {@link CommandRunnable#run()} with client message and code. */
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
index 86b3369d..d17a4ebf 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
@@ -154,6 +154,7 @@ public abstract class DispatchCommand extends BaseCommand implements ExtensionPo
 
 		try {
 			dispatcher.setContext(getContext());
+			dispatcher.setWorkQueue(getWorkQueue());
 			dispatcher.setup();
 			if (dispatcher.commands.isEmpty() && dispatcher.dispatchers.isEmpty()) {
 				log.debug(MessageFormat.format("excluding empty dispatcher {0} for {1}",
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/RootDispatcher.java b/src/main/java/com/gitblit/transport/ssh/commands/RootDispatcher.java
index 0bf6d510..e41ee19f 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/RootDispatcher.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/RootDispatcher.java
@@ -26,6 +26,7 @@ import com.gitblit.manager.IGitblit;
 import com.gitblit.transport.ssh.SshDaemonClient;
 import com.gitblit.transport.ssh.git.GitDispatcher;
 import com.gitblit.transport.ssh.keys.KeysDispatcher;
+import com.gitblit.utils.WorkQueue;
 
 /**
  * The root dispatcher is the dispatch command that handles registering all
@@ -37,9 +38,10 @@ class RootDispatcher extends DispatchCommand {
 
 	private Logger log = LoggerFactory.getLogger(getClass());
 
-	public RootDispatcher(IGitblit gitblit, SshDaemonClient client, String cmdLine) {
+	public RootDispatcher(IGitblit gitblit, SshDaemonClient client, String cmdLine, WorkQueue workQueue) {
 		super();
 		setContext(new SshCommandContext(gitblit, client, cmdLine));
+		setWorkQueue(workQueue);
 
 		register(VersionCommand.class);
 		register(GitDispatcher.class);
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/SshCommandFactory.java b/src/main/java/com/gitblit/transport/ssh/commands/SshCommandFactory.java
index 599d94b3..fa4b9167 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/SshCommandFactory.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/SshCommandFactory.java
@@ -40,7 +40,6 @@ import org.slf4j.LoggerFactory;
 import com.gitblit.Keys;
 import com.gitblit.manager.IGitblit;
 import com.gitblit.transport.ssh.SshDaemonClient;
-import com.gitblit.utils.IdGenerator;
 import com.gitblit.utils.WorkQueue;
 import com.google.common.util.concurrent.Atomics;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
@@ -48,15 +47,16 @@ import com.google.common.util.concurrent.ThreadFactoryBuilder;
 public class SshCommandFactory implements CommandFactory {
 	private static final Logger logger = LoggerFactory.getLogger(SshCommandFactory.class);
 
+	private final WorkQueue workQueue;
 	private final IGitblit gitblit;
 	private final ScheduledExecutorService startExecutor;
 	private final ExecutorService destroyExecutor;
 
-	public SshCommandFactory(IGitblit gitblit, IdGenerator idGenerator) {
+	public SshCommandFactory(IGitblit gitblit, WorkQueue workQueue) {
 		this.gitblit = gitblit;
+		this.workQueue = workQueue;
 
 		int threads = gitblit.getSettings().getInteger(Keys.git.sshCommandStartThreads, 2);
-		WorkQueue workQueue = new WorkQueue(idGenerator);
 		startExecutor = workQueue.createQueue(threads, "SshCommandStart");
 		destroyExecutor = Executors.newSingleThreadExecutor(
 				new ThreadFactoryBuilder()
@@ -70,7 +70,7 @@ public class SshCommandFactory implements CommandFactory {
 	}
 
 	public RootDispatcher createRootDispatcher(SshDaemonClient client, String commandLine) {
-		return new RootDispatcher(gitblit, client, commandLine);
+		return new RootDispatcher(gitblit, client, commandLine, workQueue);
 	}
 
 	@Override