Implement discovery loop in python

examples/discover-loop.py

@@ -16,46 +16,54 @@ def disc_supp_str(dlp_supp_opts):
     }
     return [txt for msk, txt in d.items() if dlp_supp_opts & msk]
 
-root = nvme.root()
-host = nvme.host(root)
-
-subsysnqn = nvme.NVME_DISC_SUBSYS_NAME
-transport = 'tcp'
-traddr = '127.0.0.1'
-trsvcid = '4420'
+def discover(host, ctrl, iteration):
+    # Only 8 levels of indirection are supported
+    if iteration > 8:
+        return
 
-ctrl = nvme.ctrl(root, subsysnqn=subsysnqn, transport=transport, traddr=traddr, trsvcid=trsvcid)
+    try:
+        ctrl.connect(host)
+    except Exception as e:
+        print(f'Failed to connect: {e}')
+        return
 
-try:
-    ctrl.connect(host)
-except Exception as e:
-    sys.exit(f'Failed to connect: {e}')
+    print(f'{ctrl.name} connected to {ctrl.subsystem}')
 
-print(f'{ctrl.name} connected to subsys {ctrl.subsystem}')
+    slp = ctrl.supported_log_pages()
+    try:
+        dlp_supp_opts = slp[nvme.NVME_LOG_LID_DISCOVER] >> 16
+    except (TypeError, IndexError):
+        dlp_supp_opts = 0
 
-slp = ctrl.supported_log_pages()
+    print(f"LID {nvme.NVME_LOG_LID_DISCOVER}h (Discovery), supports: {disc_supp_str(dlp_supp_opts)}")
 
-try:
-    dlp_supp_opts = slp[nvme.NVME_LOG_LID_DISCOVER] >> 16
-except (TypeError, IndexError):
-    dlp_supp_opts = 0
+    try:
+        lsp = nvme.NVMF_LOG_DISC_LSP_PLEO if dlp_supp_opts & nvme.NVMF_LOG_DISC_LID_PLEOS else 0
+        disc_log = ctrl.discover(lsp=lsp)
+    except Exception as e:
+        print(f'Failed to discover: {e}')
+        return
 
-print(f"LID {nvme.NVME_LOG_LID_DISCOVER}h (Discovery), supports: {disc_supp_str(dlp_supp_opts)}")
+    for dlpe in disc_log:
+        if dlpe['subtype'] == 'nvme':
+            print(f'{iteration}: {dlpe["subtype"]} {dlpe["subnqn"]}')
+            continue
+        if dlpe['subtype'] == 'discovery' and dlpe['subnqn'] == nvme.NVME_DISC_SUBSYS_NAME:
+            continue
+        print(f'{iteration}: {dlpe["subtype"]} {dlpe["subnqn"]}')
+        with nvme.ctrl(root, subsysnqn=dlpe['subnqn'], transport=dlpe['trtype'], traddr=dlpe['traddr'], trsvcid=dlpe['trsvcid']) as new_ctrl:
+            discover(host, new_ctrl, iteration + 1)
 
-try:
-    lsp = nvme.NVMF_LOG_DISC_LSP_PLEO if dlp_supp_opts & nvme.NVMF_LOG_DISC_LID_PLEOS else 0
-    disc_log = ctrl.discover(lsp=lsp)
-except Exception as e:
-    print(f'Failed to discover: {e}')
-    disc_log = []
+root = nvme.root()
+host = nvme.host(root)
 
-for dlpe in disc_log:
-    print(f'log entry {dlpe["portid"]}: {dlpe["subtype"]} {dlpe["subnqn"]}')
+subsysnqn = nvme.NVME_DISC_SUBSYS_NAME
+transport = 'tcp'
+traddr = '127.0.0.1'
+trsvcid = '4420'
 
-try:
-    ctrl.disconnect()
-except Exception as e:
-    sys.exit(f'Failed to disconnect: {e}')
+with nvme.ctrl(root, subsysnqn=subsysnqn, transport=transport, traddr=traddr, trsvcid=trsvcid) as ctrl:
+    discover(host, ctrl, 0)
 
 for s in host.subsystems():
     for c in s.controllers():

libnvme/nvme.i

@@ -496,6 +496,12 @@ struct nvme_ns {
 	~nvme_root() {
 		nvme_free_tree($self);
 	}
+	struct nvme_root* __enter__() {
+		return $self;
+	}
+	struct nvme_root* __exit__(PyObject *type, PyObject *value, PyObject *traceback) {
+		return $self;
+	}
 	void log_level(const char *level) {
 		int log_level = DEFAULT_LOGLEVEL;
 		if (!strcmp(level, "debug")) log_level = LOG_DEBUG;
@@ -559,6 +565,12 @@ struct nvme_ns {
 	~nvme_host() {
 		nvme_free_host($self);
 	}
+	struct nvme_host* __enter__() {
+		return $self;
+	}
+	struct nvme_host* __exit__(PyObject *type, PyObject *value, PyObject *traceback) {
+		return $self;
+	}
 	%feature("autodoc", SET_SYMNAME_DOCSTRING) set_symname;
 	void set_symname(const char *hostsymname) {
 		nvme_host_set_hostsymname($self, hostsymname);
@@ -599,6 +611,12 @@ struct nvme_ns {
 	~nvme_subsystem() {
 		nvme_free_subsystem($self);
 	}
+	struct nvme_subsystem* __enter__() {
+		return $self;
+	}
+	struct nvme_subsystem* __exit__(PyObject *type, PyObject *value, PyObject *traceback) {
+		return $self;
+	}
 	PyObject *__str__() {
 		return PyUnicode_FromFormat("nvme.subsystem(%s,%s)", STR_OR_NONE($self->name), STR_OR_NONE($self->subsysnqn));
 	}
@@ -664,6 +682,14 @@ struct nvme_ns {
 	~nvme_ctrl() {
 		nvme_free_ctrl($self);
 	}
+	struct nvme_ctrl* __enter__() {
+		return $self;
+	}
+	struct nvme_ctrl* __exit__(PyObject *type, PyObject *value, PyObject *traceback) {
+		if (nvme_ctrl_get_name($self))
+			nvme_disconnect_ctrl($self);
+		return $self;
+	}
 
 	%pythoncode %{
 	def discovery_ctrl_set(self, discovery: bool):
@@ -962,6 +988,12 @@ struct nvme_ns {
 	~nvme_ns() {
 		nvme_free_ns($self);
 	}
+	struct nvme_ns* __enter__() {
+		return $self;
+	}
+	struct nvme_ns* __exit__(PyObject *type, PyObject *value, PyObject *traceback) {
+		return $self;
+	}
 	PyObject *__str__() {
 		return PyUnicode_FromFormat("nvme.ns(%u)", $self->nsid);
 	}