# HG changeset patch
# User Gregory Szorc <gregory.szorc@gmail.com>
# Date 2018-02-24 20:07:21
# Node ID e7411fb7ba7ff616e9657efebfa75478042d0ca2
# Parent  c3df20906689676de549f81ce9c62c43212c3e34

wireprotoserver: ability to run an SSH server until an event is set

It seems useful to be able to start an SSH protocol server that
won't run forever and won't call sys.exit() when it stops. This
could be used to facilitate intra-process testing of the SSH
protocol, for example.

We teach the server function to loop until a threading.Event is set
and invent a new API to run the server until an event is set. It also
won't sys.exit() afterwards.

There aren't many callers of serve_forever(). So we could refactor
them relatively easily. But I was lazy.

threading.Event might be a bit heavyweight. An alternative would be
a list whose only elements is changed. We can't use a simple scalar
value like a bool or int because those types are immutable. Events
are what you use in systems programming for this use case, so the
use of threading.Event seems justified.

Differential Revision: https://phab.mercurial-scm.org/D2461

diff --git a/mercurial/wireprotoserver.py b/mercurial/wireprotoserver.py
--- a/mercurial/wireprotoserver.py
+++ b/mercurial/wireprotoserver.py
@@ -9,6 +9,7 @@ from __future__ import absolute_import
 import contextlib
 import struct
 import sys
+import threading
 
 from .i18n import _
 from . import (
@@ -373,7 +374,7 @@ class sshv1protocolhandler(wireprototype
 class sshv2protocolhandler(sshv1protocolhandler):
     """Protocol handler for version 2 of the SSH protocol."""
 
-def _runsshserver(ui, repo, fin, fout):
+def _runsshserver(ui, repo, fin, fout, ev):
     # This function operates like a state machine of sorts. The following
     # states are defined:
     #
@@ -430,7 +431,7 @@ def _runsshserver(ui, repo, fin, fout):
     proto = sshv1protocolhandler(ui, fin, fout)
     protoswitched = False
 
-    while True:
+    while not ev.is_set():
         if state == 'protov1-serving':
             # Commands are issued on new lines.
             request = fin.readline()[:-1]
@@ -601,5 +602,9 @@ class sshserver(object):
         util.setbinary(self._fout)
 
     def serve_forever(self):
-        _runsshserver(self._ui, self._repo, self._fin, self._fout)
+        self.serveuntil(threading.Event())
         sys.exit(0)
+
+    def serveuntil(self, ev):
+        """Serve until a threading.Event is set."""
+        _runsshserver(self._ui, self._repo, self._fin, self._fout, ev)