3 files changed, 35 insertions, 6 deletions

diff --git a/nbd.c b/nbd.c
index eb72f4a6e5..4922d38ea3 100644
--- a/nbd.c
+++ b/nbd.c
@@ -89,6 +89,7 @@ struct NBDRequest {
 };
 
 struct NBDExport {
+    int refcount;
     BlockDriverState *bs;
     off_t dev_offset;
     off_t size;
@@ -664,6 +665,7 @@ void nbd_client_put(NBDClient *client)
         qemu_set_fd_handler2(client->sock, NULL, NULL, NULL, NULL);
         close(client->sock);
         client->sock = -1;
+        nbd_export_put(client->exp);
         g_free(client);
     }
 }
@@ -722,6 +724,7 @@ NBDExport *nbd_export_new(BlockDriverState *bs, off_t dev_offset,
 {
     NBDExport *exp = g_malloc0(sizeof(NBDExport));
     QSIMPLEQ_INIT(&exp->requests);
+    exp->refcount = 1;
     exp->bs = bs;
     exp->dev_offset = dev_offset;
     exp->nbdflags = nbdflags;
@@ -731,14 +734,34 @@ NBDExport *nbd_export_new(BlockDriverState *bs, off_t dev_offset,
 
 void nbd_export_close(NBDExport *exp)
 {
-    while (!QSIMPLEQ_EMPTY(&exp->requests)) {
-        NBDRequest *first = QSIMPLEQ_FIRST(&exp->requests);
-        QSIMPLEQ_REMOVE_HEAD(&exp->requests, entry);
-        qemu_vfree(first->data);
-        g_free(first);
+    assert(exp->refcount == 1);
+
+    /* stub */
+}
+
+void nbd_export_get(NBDExport *exp)
+{
+    assert(exp->refcount > 0);
+    exp->refcount++;
+}
+
+void nbd_export_put(NBDExport *exp)
+{
+    assert(exp->refcount > 0);
+    if (exp->refcount == 1) {
+        nbd_export_close(exp);
     }
 
-    g_free(exp);
+    if (--exp->refcount == 0) {
+        while (!QSIMPLEQ_EMPTY(&exp->requests)) {
+            NBDRequest *first = QSIMPLEQ_FIRST(&exp->requests);
+            QSIMPLEQ_REMOVE_HEAD(&exp->requests, entry);
+            qemu_vfree(first->data);
+            g_free(first);
+        }
+
+        g_free(exp);
+    }
 }
 
 static int nbd_can_read(void *opaque);
@@ -1011,5 +1034,7 @@ NBDClient *nbd_client_new(NBDExport *exp, int csock,
     client->close = close;
     qemu_co_mutex_init(&client->send_lock);
     qemu_set_fd_handler2(csock, nbd_can_read, nbd_read, NULL, client);
+
+    nbd_export_get(exp);
     return client;
 }
diff --git a/nbd.h b/nbd.h
index 8b84a50ed4..86921cd048 100644
--- a/nbd.h
+++ b/nbd.h
@@ -81,6 +81,8 @@ typedef struct NBDClient NBDClient;
 NBDExport *nbd_export_new(BlockDriverState *bs, off_t dev_offset,
                           off_t size, uint32_t nbdflags);
 void nbd_export_close(NBDExport *exp);
+void nbd_export_get(NBDExport *exp);
+void nbd_export_put(NBDExport *exp);
 
 NBDClient *nbd_client_new(NBDExport *exp, int csock,
                           void (*close)(NBDClient *));
diff --git a/qemu-nbd.c b/qemu-nbd.c
index 23392e0d3c..2a2cba3f73 100644
--- a/qemu-nbd.c
+++ b/qemu-nbd.c
@@ -586,7 +586,9 @@ int main(int argc, char **argv)
     } while (!sigterm_reported && (persistent || !nbd_started || nb_fds > 0));
 
     nbd_export_close(exp);
+    nbd_export_put(exp);
     bdrv_close(bs);
+
     if (sockpath) {
         unlink(sockpath);
     }