[RFC PATCH 08/22] nbd: Add max-connections to nbd-server-start

[Kevin Wolf <kwolf@redhat.com>]

This is a QMP equivalent of qemu-nbd's --share option, limiting the
maximum number of clients that can attach at the same time.

Signed-off-by: Kevin Wolf <kwolf@redhat.com>
---
 qapi/block-export.json         | 10 ++++++++--
 include/block/nbd.h            |  3 ++-
 block/monitor/block-hmp-cmds.c |  2 +-
 blockdev-nbd.c                 | 33 ++++++++++++++++++++++++++-------
 qemu-storage-daemon.c          |  4 ++--
 5 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/qapi/block-export.json b/qapi/block-export.json
index 40369814b4..1fdc55c53a 100644
--- a/qapi/block-export.json
+++ b/qapi/block-export.json
@@ -14,6 +14,8 @@
 #             is only resolved at time of use, so can be deleted and
 #             recreated on the fly while the NBD server is active.
 #             If missing, it will default to denying access (since 4.0).
+# @max-connections: The maximum number of connections to allow at the same
+#                   time, 0 for unlimited. (since 5.2; default: 0)
 #
 # Keep this type consistent with the nbd-server-start arguments. The only
 # intended difference is using SocketAddress instead of SocketAddressLegacy.
@@ -23,7 +25,8 @@
 { 'struct': 'NbdServerOptions',
   'data': { 'addr': 'SocketAddress',
             '*tls-creds': 'str',
-            '*tls-authz': 'str'} }
+            '*tls-authz': 'str',
+            '*max-connections': 'uint32' } }
 
 ##
 # @nbd-server-start:
@@ -40,6 +43,8 @@
 #             is only resolved at time of use, so can be deleted and
 #             recreated on the fly while the NBD server is active.
 #             If missing, it will default to denying access (since 4.0).
+# @max-connections: The maximum number of connections to allow at the same
+#                   time, 0 for unlimited. (since 5.2; default: 0)
 #
 # Returns: error if the server is already running.
 #
@@ -51,7 +56,8 @@
 { 'command': 'nbd-server-start',
   'data': { 'addr': 'SocketAddressLegacy',
             '*tls-creds': 'str',
-            '*tls-authz': 'str'} }
+            '*tls-authz': 'str',
+            '*max-connections': 'uint32' } }
 
 ##
 # @BlockExportOptionsNbd:
diff --git a/include/block/nbd.h b/include/block/nbd.h
index ffca3be78f..6fc1f05ef4 100644
--- a/include/block/nbd.h
+++ b/include/block/nbd.h
@@ -352,7 +352,8 @@ void nbd_client_get(NBDClient *client);
 void nbd_client_put(NBDClient *client);
 
 void nbd_server_start(SocketAddress *addr, const char *tls_creds,
-                      const char *tls_authz, Error **errp);
+                      const char *tls_authz, uint32_t max_connections,
+                      Error **errp);
 void nbd_server_start_options(NbdServerOptions *arg, Error **errp);
 
 /* nbd_read
diff --git a/block/monitor/block-hmp-cmds.c b/block/monitor/block-hmp-cmds.c
index 56bc83ac97..a651954e16 100644
--- a/block/monitor/block-hmp-cmds.c
+++ b/block/monitor/block-hmp-cmds.c
@@ -411,7 +411,7 @@ void hmp_nbd_server_start(Monitor *mon, const QDict *qdict)
         goto exit;
     }
 
-    nbd_server_start(addr, NULL, NULL, &local_err);
+    nbd_server_start(addr, NULL, NULL, 0, &local_err);
     qapi_free_SocketAddress(addr);
     if (local_err != NULL) {
         goto exit;
diff --git a/blockdev-nbd.c b/blockdev-nbd.c
index 019c37c0bc..28159a92b2 100644
--- a/blockdev-nbd.c
+++ b/blockdev-nbd.c
@@ -23,23 +23,41 @@ typedef struct NBDServerData {
     QIONetListener *listener;
     QCryptoTLSCreds *tlscreds;
     char *tlsauthz;
+    uint32_t max_connections;
+    uint32_t connections;
 } NBDServerData;
 
 static NBDServerData *nbd_server;
 
+static void nbd_update_server_watch(NBDServerData *s);
+
 static void nbd_blockdev_client_closed(NBDClient *client, bool ignored)
 {
     nbd_client_put(client);
+    assert(nbd_server->connections > 0);
+    nbd_server->connections--;
+    nbd_update_server_watch(nbd_server);
 }
 
 static void nbd_accept(QIONetListener *listener, QIOChannelSocket *cioc,
                        gpointer opaque)
 {
+    nbd_server->connections++;
+    nbd_update_server_watch(nbd_server);
+
     qio_channel_set_name(QIO_CHANNEL(cioc), "nbd-server");
     nbd_client_new(cioc, nbd_server->tlscreds, nbd_server->tlsauthz,
                    nbd_blockdev_client_closed);
 }
 
+static void nbd_update_server_watch(NBDServerData *s)
+{
+    if (!s->max_connections || s->connections < s->max_connections) {
+        qio_net_listener_set_client_func(s->listener, nbd_accept, NULL, NULL);
+    } else {
+        qio_net_listener_set_client_func(s->listener, NULL, NULL, NULL);
+    }
+}
 
 static void nbd_server_free(NBDServerData *server)
 {
@@ -88,7 +106,8 @@ static QCryptoTLSCreds *nbd_get_tls_creds(const char *id, Error **errp)
 
 
 void nbd_server_start(SocketAddress *addr, const char *tls_creds,
-                      const char *tls_authz, Error **errp)
+                      const char *tls_authz, uint32_t max_connections,
+                      Error **errp)
 {
     if (nbd_server) {
         error_setg(errp, "NBD server already running");
@@ -96,6 +115,7 @@ void nbd_server_start(SocketAddress *addr, const char *tls_creds,
     }
 
     nbd_server = g_new0(NBDServerData, 1);
+    nbd_server->max_connections = max_connections;
     nbd_server->listener = qio_net_listener_new();
 
     qio_net_listener_set_name(nbd_server->listener,
@@ -120,10 +140,7 @@ void nbd_server_start(SocketAddress *addr, const char *tls_creds,
 
     nbd_server->tlsauthz = g_strdup(tls_authz);
 
-    qio_net_listener_set_client_func(nbd_server->listener,
-                                     nbd_accept,
-                                     NULL,
-                                     NULL);
+    nbd_update_server_watch(nbd_server);
 
     return;
 
@@ -134,17 +151,19 @@ void nbd_server_start(SocketAddress *addr, const char *tls_creds,
 
 void nbd_server_start_options(NbdServerOptions *arg, Error **errp)
 {
-    nbd_server_start(arg->addr, arg->tls_creds, arg->tls_authz, errp);
+    nbd_server_start(arg->addr, arg->tls_creds, arg->tls_authz,
+                     arg->max_connections, errp);
 }
 
 void qmp_nbd_server_start(SocketAddressLegacy *addr,
                           bool has_tls_creds, const char *tls_creds,
                           bool has_tls_authz, const char *tls_authz,
+                          bool has_max_connections, uint32_t max_connections,
                           Error **errp)
 {
     SocketAddress *addr_flat = socket_address_flatten(addr);
 
-    nbd_server_start(addr_flat, tls_creds, tls_authz, errp);
+    nbd_server_start(addr_flat, tls_creds, tls_authz, max_connections, errp);
     qapi_free_SocketAddress(addr_flat);
 }
 
diff --git a/qemu-storage-daemon.c b/qemu-storage-daemon.c
index b6f678d3ab..0fcab6ed2d 100644
--- a/qemu-storage-daemon.c
+++ b/qemu-storage-daemon.c
@@ -101,9 +101,9 @@ static void help(void)
 "                         configure a QMP monitor\n"
 "\n"
 "  --nbd-server addr.type=inet,addr.host=<host>,addr.port=<port>\n"
-"               [,tls-creds=<id>][,tls-authz=<id>]\n"
+"               [,tls-creds=<id>][,tls-authz=<id>][,max-connections=<n>]\n"
 "  --nbd-server addr.type=unix,addr.path=<path>\n"
-"               [,tls-creds=<id>][,tls-authz=<id>]\n"
+"               [,tls-creds=<id>][,tls-authz=<id>][,max-connections=<n>]\n"
 "                         start an NBD server for exporting block nodes\n"
 "\n"
 "  --object help          list object types that can be added\n"
-- 
2.25.4