Repository: geldata/gel
Pull request: Add support for socket activation
Since the idiomatic approach is to have an EdgeDB instance per project
it is easy to end up with dozens of EdgeDB instances running on the user
machine even when idle.  This is problematic from memory consumption
perspective (though a lot less with `on_demand` compiler pool mode).
Thankfully, both systemd and launchd have straightforward support for
socket-activated services.  Combined with auto-shutdown on idle, the
problem of many instances is solved nicely.  This also largely obviates
the need for `--start-conf` on those systems.

Fixes: #3896

Problem statement:
Add support for systemd/launchd/intetd socket activation
edgedb-server should be a good citizen and support common socket-activation protocols.  This, coupled with `--auto-shutdown-after` will largely obviate the need for `instance --start-conf` and will make it so that only actively-used server instances are running on a developer machine.

Hints from issue discussion:

Commit message: Add support for socket activation

Since the idiomatic approach is to have an EdgeDB instance per project
it is easy to end up with dozens of EdgeDB instances running on the user
machine even when idle.  This is problematic from memory consumption
perspective (though a lot less with `on_demand` compiler pool mode).
Thankfully, both systemd and launchd have straightforward support for
socket-activated services.  Combined with auto-shutdown on idle, the
problem of many instances is solved nicely.  This also largely obviates
the need for `--start-conf` on those systems.

Fixes: #3896

Patch to review:
diff --git a/edb/server/args.py b/edb/server/args.py
index 281933cb832..975f37bcba5 100644
--- a/edb/server/args.py
+++ b/edb/server/args.py
@@ -124,8 +124,9 @@ class ServerConfig(NamedTuple):
     default_database_user: Optional[str]
     devmode: bool
     testmode: bool
-    bind_addresses: tuple[str]
+    bind_addresses: list[str]
     port: int
+    activation_socket_names: list[str]
     background: bool
     pidfile_dir: pathlib.Path
     daemon_user: str
@@ -438,6 +439,17 @@ def _writer(status: str) -> None:
     click.option(
         '-P', '--port', type=PortType(), default=None,
         help='port to listen on'),
+    click.option(
+        '--activation-socket-name', type=str, multiple=True,
+        help='The names of the activation sockets to listen on, specify '
+             'multiple times for several sockets.  Should match a Sockets '
+             'entry in service plist on macOS, or FileDescriptorName= in '
+             'systemd .socket unit.  If not specified, defaults to '
+             '"edgedb-server" on macOS and all sockets passed by systemd '
+             'on Linux.  When the server is started via socket activation '
+             '--bind-address, --port, and the corresponding config settings '
+             'are ignored, and the server only listens on the socket passed '
+             'by the system service manager.'),
     click.option(
         '-b', '--background', is_flag=True, help='daemonize'),
     click.option(
@@ -684,6 +696,7 @@ def compiler_options(func):
 
 def parse_args(**kwargs: Any):
     kwargs['bind_addresses'] = kwargs.pop('bind_address')
+    kwargs['activation_socket_names'] = kwargs.pop('activation_socket_name')
 
     if kwargs['echo_runtime_info']:
         warnings.warn(
diff --git a/edb/server/main.py b/edb/server/main.py
index 73ad8956794..2fac99977fb 100644
--- a/edb/server/main.py
+++ b/edb/server/main.py
@@ -29,7 +29,6 @@
 import pathlib
 import resource
 import signal
-import socket
 import sys
 import tempfile
 import uuid
@@ -52,6 +51,7 @@
 from . import defines
 from . import pgconnparams
 from . import pgcluster
+from . import service_manager
 
 
 if TYPE_CHECKING:
@@ -145,25 +145,6 @@ async def _init_cluster(cluster, args: srvargs.ServerConfig) -> bool:
     return need_restart
 
 
-def _sd_notify(message):
-    notify_socket = os.environ.get('NOTIFY_SOCKET')
-    if not notify_socket:
-        return
-
-    if notify_socket[0] == '@':
-        notify_socket = '\0' + notify_socket[1:]
-
-    sd_sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
-
-    try:
-        sd_sock.connect(notify_socket)
-        sd_sock.sendall(message.encode())
-    except Exception as e:
-        logger.info('Could not send systemd notification: %s', e)
-    finally:
-        sd_sock.close()
-
-
 def _init_parsers():
     # Initialize all parsers, rebuilding grammars if
     # necessary.  Do it earlier than later so that we don't
@@ -185,6 +166,12 @@ async def _run_server(
     new_instance: bool,
 ):
 
+    sockets = service_manager.get_activation_listen_sockets(
+        args.activation_socket_names)
+
+    if sockets:
+        logger.info("detected service manager socket activation")
+
     with signalctl.SignalController(signal.SIGINT, signal.SIGTERM) as sc:
         ss = server.Server(
             cluster=cluster,
@@ -196,6 +183,7 @@ async def _run_server(
             compiler_pool_addr=args.compiler_pool_addr,
             nethosts=args.bind_addresses,
             netport=args.port,
+            listen_sockets=tuple(sockets.values()),
             auto_shutdown_after=args.auto_shutdown_after,
             echo_runtime_info=args.echo_runtime_info,
             status_sinks=args.status_sinks,
@@ -240,14 +228,14 @@ async def _run_server(
                 )
 
             # Notify systemd that we've started up.
-            _sd_notify('READY=1')
+            service_manager.sd_notify('READY=1')
 
             try:
                 await sc.wait_for(ss.serve_forever())
             except signalctl.SignalError as e:
                 logger.info('Received signal: %s.', e.signo)
         finally:
-            _sd_notify('STOPPING=1')
+            service_manager.sd_notify('STOPPING=1')
             logger.info('Shutting down.')
             await sc.wait_for(ss.stop())
 
diff --git a/edb/server/server.py b/edb/server/server.py
index 6c333aab6ed..77e00d6e53f 100644
--- a/edb/server/server.py
+++ b/edb/server/server.py
@@ -66,6 +66,9 @@
 
 from . import dbview
 
+if TYPE_CHECKING:
+    import asyncio.base_events
+
 
 ADMIN_PLACEHOLDER = "<edgedb:admin>"
 logger = logging.getLogger('edb.server')
@@ -131,6 +134,7 @@ def __init__(
         nethosts,
         netport,
         new_instance: bool,
+        listen_sockets: tuple[socket.socket, ...] = (),
         testmode: bool = False,
         binary_endpoint_security: srvargs.ServerEndpointSecurityMode = (
             srvargs.ServerEndpointSecurityMode.Tls),
@@ -188,6 +192,11 @@ def __init__(
             defines.MIN_SUGGESTED_CLIENT_POOL_SIZE
         )
 
+        self._listen_sockets = listen_sockets
+        if listen_sockets:
+            nethosts = tuple(s.getsockname()[0] for s in listen_sockets)
+            netport = listen_sockets[0].getsockname()[1]
+
         self._listen_hosts = nethosts
         self._listen_port = netport
 
@@ -896,7 +905,9 @@ async def _restart_servers_new_addr(self, nethosts, netport):
 
         if hosts_to_start:
             new_servers, *_ = await self._start_servers(
-                hosts_to_start, netport, admin
+                hosts_to_start,
+                netport,
+                admin=admin,
             )
             servers.update(new_servers)
         self._servers = servers
@@ -1347,8 +1358,11 @@ async def run_startup_script_and_exit(self):
             await self._destroy_compiler_pool()
 
     async def _start_server(
-        self, host: str, port: int
-    ) -> Optional[asyncio.AbstractServer]:
+        self,
+        host: str,
+        port: int,
+        sock: Optional[socket.socket] = None,
+    ) -> Optional[asyncio.base_events.Server]:
         proto_factory = lambda: protocol.HttpProtocol(
             self,
             self._sslctx,
@@ -1357,15 +1371,22 @@ async def _start_server(
         )
 
         try:
-            return await self.__loop.create_server(
-                proto_factory, host=host, port=port)
+            kwargs: dict[str, Any]
+            if sock is not None:
+                kwargs = {"sock": sock}
+            else:
+                kwargs = {"host": host, "port": port}
+            return await self.__loop.create_server(proto_factory, **kwargs)
         except Exception as e:
             logger.warning(
                 f"could not create listen socket for '{host}:{port}': {e}"
             )
             return None
 
-    async def _start_admin_server(self, port: int) -> asyncio.AbstractServer:
+    async def _start_admin_server(
+        self,
+        port: int,
+    ) -> asyncio.base_events.Server:
         admin_unix_sock_path = os.path.join(
             self._runstate_dir, f'.s.EDGEDB.admin.{port}')
         assert len(admin_unix_sock_path) <= (
@@ -1381,7 +1402,14 @@ async def _start_admin_server(self, port: int) -> asyncio.AbstractServer:
         logger.info('Serving admin on %s', admin_unix_sock_path)
         return admin_unix_srv
 
-    async def _start_servers(self, hosts, port, admin=True):
+    async def _start_servers(
+        self,
+        hosts: tuple[str, ...],
+        port: int,
+        *,
+        admin: bool = True,
+        sockets: tuple[socket.socket, ...] = (),
+    ):
         servers = {}
         if port == 0:
             # Automatic port selection requires us to start servers
@@ -1401,10 +1429,16 @@ async def _start_servers(self, hosts, port, admin=True):
             start_tasks = {}
             try:
                 async with taskgroup.TaskGroup() as g:
-                    for host in hosts:
-                        start_tasks[host] = g.create_task(
-                            self._start_server(host, port)
-                        )
+                    if sockets:
+                        for host, sock in zip(hosts, sockets):
+                            start_tasks[host] = g.create_task(
+                                self._start_server(host, port, sock=sock)
+                            )
+                    else:
+                        for host in hosts:
+                            start_tasks[host] = g.create_task(
+                                self._start_server(host, port)
+                            )
             except Exception:
                 await self._stop_servers([
                     fut.result() for fut in start_tasks.values()
@@ -1547,6 +1581,7 @@ async def start(self):
         self._servers, actual_port, listen_addrs = await self._start_servers(
             await _resolve_interfaces(self._listen_hosts),
             self._listen_port,
+            sockets=self._listen_sockets,
         )
         self._listen_hosts = listen_addrs
         self._listen_port = actual_port
diff --git a/edb/server/service_manager.py b/edb/server/service_manager.py
new file mode 100644
index 00000000000..63c6812edaf
--- /dev/null
+++ b/edb/server/service_manager.py
@@ -0,0 +1,223 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+
+from __future__ import annotations
+from typing import *
+
+import logging
+import os
+import socket
+import sys
+
+
+SD_LISTEN_FDS_START = 3
+
+
+logger = logging.getLogger('edb.server')
+
+
+def _stream_socket_from_fd(fd: int) -> Optional[socket.socket]:
+    try:
+        sock = socket.socket(fileno=fd)
+    except OSError:
+        logger.warning(
+            f"activation file descriptor {fd} is not a socket "
+            f", ignoring"
+        )
+        return None
+
+    if sock.family not in {socket.AF_INET, socket.AF_INET6}:
+        logger.warning(
+            f"activation file descriptor {fd} is not an AF_INET[6] socket "
+            f", ignoring"
+        )
+        return None
+
+    if sock.type != socket.SOCK_STREAM:
+        logger.warning(
+            f"activation file descriptor {fd} is not an SOCK_STREAM "
+            f"socket, ignoring"
+        )
+        return None
+
+    return sock
+
+
+def sd_notify(message: str) -> None:
+    notify_socket = os.environ.get('NOTIFY_SOCKET')
+    if not notify_socket:
+        return
+
+    if notify_socket[0] == '@':
+        notify_socket = '\0' + notify_socket[1:]
+
+    sd_sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+
+    try:
+        sd_sock.connect(notify_socket)
+        sd_sock.sendall(message.encode())
+    except Exception as e:
+        logger.info('Could not send systemd notification: %s', e)
+    finally:
+        sd_sock.close()
+
+
+def sd_get_activation_listen_sockets(
+    names: list[str],
+) -> dict[str, socket.socket]:
+    listen_pid = os.environ.pop("LISTEN_PID", "")
+    listen_fds = os.environ.pop("LISTEN_FDS", "")
+    listen_fdnames = os.environ.pop("LISTEN_FDNAMES", "")
+
+    if not listen_pid or not listen_fds:
+        return {}
+
+    try:
+        expected_pid = int(listen_pid)
+    except ValueError:
+        logger.warning(
+            "the value of LISTEN_PID environment variable "
+            "is not a valid integer, ignoring socket activation data"
+        )
+        return {}
+
+    if expected_pid != os.getpid():
+        logger.warning(
+            "the value of LISTEN_PID does not match the PID of this "
+            "process, ignoring socket activation data"
+        )
+        return {}
+
+    try:
+        num_fds = int(listen_fds)
+    except ValueError:
+        logger.warning(
+            "the value of LISTEN_FDS environment variable "
+            "is not a valid integer, ignoring socket activation data"
+        )
+        return {}
+
+    fd_names = listen_fdnames.split(":")
+    fd_range = range(SD_LISTEN_FDS_START, SD_LISTEN_FDS_START + num_fds)
+    sockets = {}
+
+    for i, fd in enumerate(fd_range):
+        os.set_inheritable(fd, False)
+
+        try:
+            name = fd_names[i]
+        except IndexError:
+            name = f"LISTEN_FD_{fd}"
+
+        if names and name not in names:
+            logger.warning(
+                f"activation file descriptor {name} ({fd}) is not in any "
+                f"of the passed --activation-socket-name= options, ignoring"
+            )
+            continue
+
+        sock = _stream_socket_from_fd(fd)
+        if sock is not None:
+            sockets[name] = sock
+
+    return sockets
+
+
+if sys.platform == "darwin":
+    import ctypes
+    import ctypes.util
+
+    syslib = ctypes.CDLL(ctypes.util.find_library('System'))
+    syslib.launch_activate_socket.argypes = [
+        ctypes.c_char_p,
+        ctypes.POINTER(ctypes.POINTER(ctypes.c_int)),
+        ctypes.POINTER(ctypes.c_size_t),
+    ]
+
+    def _launch_activate_socket(name) -> list[int]:
+        fds = ctypes.POINTER(ctypes.c_int)()
+        num_fds = ctypes.c_size_t()
+        result = syslib.launch_activate_socket(
+            ctypes.c_char_p(name.encode("utf-8")),
+            ctypes.byref(fds),
+            ctypes.byref(num_fds),
+        )
+        if result == 0:
+            return [fds[i] for i in range(num_fds.value)]
+        else:
+            logger.warning(f"launch_activate_socket({name}) returned {result}")
+            return []
+
+    def launchd_get_activation_listen_sockets(
+        names: list[str],
+    ) -> dict[str, socket.socket]:
+        if not names:
+            names = ["edgedb-server"]
+
+        sockets = {}
+
+        for name in names:
+            fds = _launch_activate_socket(name)
+            if len(fds) == 0:
+                logger.warning(f"could not activate socket {name}")
+                continue
+            elif len(fds) != 1:
+                logger.warning(
+                    f"more than one socket returned by "
+                    f"launch_activate_socket({name}), ignoring all but the "
+                    f"first"
+                )
+                continue
+
+            sock = _stream_socket_from_fd(fds[0])
+            if sock is not None:
+                sockets[name] = sock
+
+        return sockets
+
+else:
+    def launchd_get_activation_listen_sockets(
+        names: list[str],
+    ) -> dict[str, socket.socket]:
+        return {}
+
+
+def get_activation_listen_sockets(
+    names: list[str],
+) -> dict[str, socket.socket]:
+    if sys.platform == "darwin":
+        sockets = launchd_get_activation_listen_sockets(names)
+    else:
+        sockets = sd_get_activation_listen_sockets(names)
+
+    port = 0
+
+    for name, sock in list(sockets.items()):
+        this_port = sock.getsockname()[1]
+        if port == 0:
+            port = this_port
+        elif port == this_port:
+            continue
+        else:
+            logger.warning(
+                f"activation sockets are not all on the same TCP port, "
+                f"first socket is at {port} and socket {name!r} is at "
+                f"{this_port}, ignoring the latter")
+            sockets.pop(name)
+
+    return sockets


Code review comments:
diff_hunk: @@ -0,0 +1,223 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2016-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+
+from __future__ import annotations
+from typing import *
+
+import logging
+import os
+import socket
+import sys
+
+
+SD_LISTEN_FDS_START = 3
+
+
+logger = logging.getLogger('edb.server')
+
+
+def _stream_socket_from_fd(fd: int) -> Optional[socket.socket]:
+    try:
+        sock = socket.socket(fileno=fd)
+    except OSError:
+        logger.warning(
+            f"activation file descriptor {fd} is not a socket "
+            f", ignoring"
+        )
+        return None
+
+    if sock.family not in {socket.AF_INET, socket.AF_INET6}:
+        logger.warning(
+            f"activation file descriptor {fd} is not an AF_INET[6] socket "
+            f", ignoring"
+        )
+        return None
+
+    if sock.type != socket.SOCK_STREAM:
+        logger.warning(
+            f"activation file descriptor {fd} is not an SOCK_STREAM "
+            f"socket, ignoring"
+        )
+        return None
+
+    return sock
+
+
+def sd_notify(message: str) -> None:
+    notify_socket = os.environ.get('NOTIFY_SOCKET')
+    if not notify_socket:
+        return
+
+    if notify_socket[0] == '@':
+        notify_socket = '\0' + notify_socket[1:]
+
+    sd_sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+
+    try:
+        sd_sock.connect(notify_socket)
+        sd_sock.sendall(message.encode())
+    except Exception as e:
+        logger.info('Could not send systemd notification: %s', e)
+    finally:
+        sd_sock.close()
path: edb/server/service_manager.py
text: @user1:
nit

```suggestion
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sd_sock:
        try:
            sd_sock.connect(notify_socket)
            sd_sock.sendall(message.encode())
        except Exception as e:
            logger.info('Could not send systemd notification: %s', e)
```