Add type hints to synmark. (#16421)

1 files changed, 42 insertions, 26 deletions

diff --git a/synmark/suites/logging.py b/synmark/suites/logging.py
index 04e5b29dc9..e160443643 100644
--- a/synmark/suites/logging.py
+++ b/synmark/suites/logging.py
@@ -11,14 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import logging
+import logging.config
 import warnings
 from io import StringIO
+from typing import Optional
 from unittest.mock import Mock
 
 from pyperf import perf_counter
 
+from twisted.internet.address import IPv4Address, IPv6Address
 from twisted.internet.defer import Deferred
 from twisted.internet.protocol import ServerFactory
 from twisted.logger import LogBeginner, LogPublisher
@@ -26,45 +28,53 @@ from twisted.protocols.basic import LineOnlyReceiver
 
 from synapse.config.logger import _setup_stdlib_logging
 from synapse.logging import RemoteHandler
+from synapse.synapse_rust import reset_logging_config
+from synapse.types import ISynapseReactor
 from synapse.util import Clock
 
 
 class LineCounter(LineOnlyReceiver):
     delimiter = b"\n"
+    count = 0
 
-    def __init__(self, *args, **kwargs):
-        self.count = 0
-        super().__init__(*args, **kwargs)
-
-    def lineReceived(self, line):
+    def lineReceived(self, line: bytes) -> None:
         self.count += 1
 
+        assert isinstance(self.factory, Factory)
+
         if self.count >= self.factory.wait_for and self.factory.on_done:
             on_done = self.factory.on_done
             self.factory.on_done = None
             on_done.callback(True)
 
 
-async def main(reactor, loops):
+class Factory(ServerFactory):
+    protocol = LineCounter
+    wait_for: int
+    on_done: Optional[Deferred]
+
+
+async def main(reactor: ISynapseReactor, loops: int) -> float:
     """
     Benchmark how long it takes to send `loops` messages.
     """
-    servers = []
 
-    def protocol():
-        p = LineCounter()
-        servers.append(p)
-        return p
-
-    logger_factory = ServerFactory.forProtocol(protocol)
+    logger_factory = Factory()
     logger_factory.wait_for = loops
     logger_factory.on_done = Deferred()
-    port = reactor.listenTCP(0, logger_factory, interface="127.0.0.1")
+    port = reactor.listenTCP(0, logger_factory, backlog=50, interface="127.0.0.1")
 
     # A fake homeserver config.
     class Config:
-        server_name = "synmark-" + str(loops)
-        no_redirect_stdio = True
+        class server:
+            server_name = "synmark-" + str(loops)
+
+        # This odd construct is to avoid mypy thinking that logging escapes the
+        # scope of Config.
+        class _logging:
+            no_redirect_stdio = True
+
+        logging = _logging
 
     hs_config = Config()
 
@@ -78,28 +88,34 @@ async def main(reactor, loops):
         publisher, errors, mock_sys, warnings, initialBufferSize=loops
     )
 
+    address = port.getHost()
+    assert isinstance(address, (IPv4Address, IPv6Address))
     log_config = {
         "version": 1,
-        "loggers": {"synapse": {"level": "DEBUG", "handlers": ["tersejson"]}},
+        "loggers": {"synapse": {"level": "DEBUG", "handlers": ["remote"]}},
         "formatters": {"tersejson": {"class": "synapse.logging.TerseJsonFormatter"}},
         "handlers": {
-            "tersejson": {
+            "remote": {
                 "class": "synapse.logging.RemoteHandler",
-                "host": "127.0.0.1",
-                "port": port.getHost().port,
+                "formatter": "tersejson",
+                "host": address.host,
+                "port": address.port,
                 "maximum_buffer": 100,
-                "_reactor": reactor,
             }
         },
     }
 
-    logger = logging.getLogger("synapse.logging.test_terse_json")
+    logger = logging.getLogger("synapse")
     _setup_stdlib_logging(
-        hs_config,
-        log_config,
+        hs_config,  # type: ignore[arg-type]
+        None,
         logBeginner=beginner,
     )
 
+    # Force a new logging config without having to load it from a file.
+    logging.config.dictConfig(log_config)
+    reset_logging_config()
+
     # Wait for it to connect...
     for handler in logging.getLogger("synapse").handlers:
         if isinstance(handler, RemoteHandler):
@@ -107,7 +123,7 @@ async def main(reactor, loops):
     else:
         raise RuntimeError("Improperly configured: no RemoteHandler found.")
 
-    await handler._service.whenConnected()
+    await handler._service.whenConnected(failAfterFailures=10)
 
     start = perf_counter()