libbtrfsutil: add subvolume iterator helpers

This is how we can implement stuff like `btrfs subvol list`. Rather than
producing the entire list upfront, the iterator approach uses less
memory in the common case where the whole list is not stored (O(max
subvolume path length)). It supports both pre-order traversal (useful
for, e.g, recursive snapshot) and post-order traversal (useful for
recursive delete).

Signed-off-by: Omar Sandoval <osandov@fb.com>
Signed-off-by: David Sterba <dsterba@suse.com>

4 files changed, 276 insertions, 0 deletions

diff --git a/libbtrfsutil/python/btrfsutilpy.h b/libbtrfsutil/python/btrfsutilpy.h
index a56e41fc..79ac2eec 100644
--- a/libbtrfsutil/python/btrfsutilpy.h
+++ b/libbtrfsutil/python/btrfsutilpy.h
@@ -37,6 +37,7 @@ typedef struct {
 extern PyTypeObject BtrfsUtilError_type;
 extern PyStructSequence_Desc SubvolumeInfo_desc;
 extern PyTypeObject SubvolumeInfo_type;
+extern PyTypeObject SubvolumeIterator_type;
 extern PyTypeObject QgroupInherit_type;
 
 /*
diff --git a/libbtrfsutil/python/module.c b/libbtrfsutil/python/module.c
index 91c4b9be..9a237142 100644
--- a/libbtrfsutil/python/module.c
+++ b/libbtrfsutil/python/module.c
@@ -239,6 +239,10 @@ PyInit_btrfsutil(void)
 	if (PyStructSequence_InitType2(&SubvolumeInfo_type, &SubvolumeInfo_desc) < 0)
 		return NULL;
 
+	SubvolumeIterator_type.tp_new = PyType_GenericNew;
+	if (PyType_Ready(&SubvolumeIterator_type) < 0)
+		return NULL;
+
 	QgroupInherit_type.tp_new = PyType_GenericNew;
 	if (PyType_Ready(&QgroupInherit_type) < 0)
 		return NULL;
@@ -254,6 +258,10 @@ PyInit_btrfsutil(void)
 	Py_INCREF(&SubvolumeInfo_type);
 	PyModule_AddObject(m, "SubvolumeInfo", (PyObject *)&SubvolumeInfo_type);
 
+	Py_INCREF(&SubvolumeIterator_type);
+	PyModule_AddObject(m, "SubvolumeIterator",
+			   (PyObject *)&SubvolumeIterator_type);
+
 	Py_INCREF(&QgroupInherit_type);
 	PyModule_AddObject(m, "QgroupInherit",
 			   (PyObject *)&QgroupInherit_type);
diff --git a/libbtrfsutil/python/subvolume.c b/libbtrfsutil/python/subvolume.c
index fa3ec4a7..6c384583 100644
--- a/libbtrfsutil/python/subvolume.c
+++ b/libbtrfsutil/python/subvolume.c
@@ -348,3 +348,214 @@ PyObject *create_subvolume(PyObject *self, PyObject *args, PyObject *kwds)
 	else
 		Py_RETURN_NONE;
 }
+
+typedef struct {
+	PyObject_HEAD
+	struct btrfs_util_subvolume_iterator *iter;
+	bool info;
+} SubvolumeIterator;
+
+static void SubvolumeIterator_dealloc(SubvolumeIterator *self)
+{
+	btrfs_util_destroy_subvolume_iterator(self->iter);
+	Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+static PyObject *SubvolumeIterator_next(SubvolumeIterator *self)
+{
+	enum btrfs_util_error err;
+	PyObject *ret, *tmp;
+	char *path;
+
+	if (!self->iter) {
+		PyErr_SetString(PyExc_ValueError,
+				"operation on closed iterator");
+		return NULL;
+	}
+
+	if (self->info) {
+		struct btrfs_util_subvolume_info subvol;
+
+		err = btrfs_util_subvolume_iterator_next_info(self->iter, &path,
+							      &subvol);
+		if (err == BTRFS_UTIL_ERROR_STOP_ITERATION) {
+			PyErr_SetNone(PyExc_StopIteration);
+			return NULL;
+		} else if (err) {
+			SetFromBtrfsUtilError(err);
+			return NULL;
+		}
+
+		tmp = subvolume_info_to_object(&subvol);
+	} else {
+		uint64_t id;
+
+		err = btrfs_util_subvolume_iterator_next(self->iter, &path, &id);
+		if (err == BTRFS_UTIL_ERROR_STOP_ITERATION) {
+			PyErr_SetNone(PyExc_StopIteration);
+			return NULL;
+		} else if (err) {
+			SetFromBtrfsUtilError(err);
+			return NULL;
+		}
+
+		tmp = PyLong_FromUnsignedLongLong(id);
+
+	}
+	if (tmp) {
+		ret = Py_BuildValue("O&O", PyUnicode_DecodeFSDefault, path,
+				    tmp);
+		Py_DECREF(tmp);
+		free(path);
+	} else {
+		ret = NULL;
+	}
+	return ret;
+}
+
+static int SubvolumeIterator_init(SubvolumeIterator *self, PyObject *args,
+				  PyObject *kwds)
+{
+	static char *keywords[] = {"path", "top", "info", "post_order", NULL};
+	struct path_arg path = {.allow_fd = true};
+	enum btrfs_util_error err;
+	unsigned long long top = 5;
+	int info = 0;
+	int post_order = 0;
+	int flags = 0;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|Kpp:SubvolumeIterator",
+					 keywords, &path_converter, &path, &top,
+					 &info, &post_order))
+		return -1;
+
+	if (post_order)
+		flags |= BTRFS_UTIL_SUBVOLUME_ITERATOR_POST_ORDER;
+
+	if (path.path) {
+		err = btrfs_util_create_subvolume_iterator(path.path, top,
+							   flags, &self->iter);
+	} else {
+		err = btrfs_util_create_subvolume_iterator_fd(path.fd, top,
+							      flags,
+							      &self->iter);
+	}
+	if (err) {
+		SetFromBtrfsUtilErrorWithPath(err, &path);
+		path_cleanup(&path);
+		return -1;
+	}
+
+	self->info = info;
+
+	return 0;
+}
+
+static PyObject *SubvolumeIterator_close(SubvolumeIterator *self)
+{
+	if (self->iter) {
+		btrfs_util_destroy_subvolume_iterator(self->iter);
+		self->iter = NULL;
+	}
+	Py_RETURN_NONE;
+}
+
+static PyObject *SubvolumeIterator_fileno(SubvolumeIterator *self)
+{
+	if (!self->iter) {
+		PyErr_SetString(PyExc_ValueError,
+				"operation on closed iterator");
+		return NULL;
+	}
+	return PyLong_FromLong(btrfs_util_subvolume_iterator_fd(self->iter));
+}
+
+static PyObject *SubvolumeIterator_enter(SubvolumeIterator *self)
+{
+	Py_INCREF((PyObject *)self);
+	return (PyObject *)self;
+}
+
+static PyObject *SubvolumeIterator_exit(SubvolumeIterator *self, PyObject *args,
+				       PyObject *kwds)
+{
+	static char *keywords[] = {"exc_type", "exc_value", "traceback", NULL};
+	PyObject *exc_type, *exc_value, *traceback;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOO:__exit__", keywords,
+					 &exc_type, &exc_value, &traceback))
+		return NULL;
+
+	return SubvolumeIterator_close(self);
+}
+
+#define SubvolumeIterator_DOC	\
+	 "SubvolumeIterator(path, top=0, info=False, post_order=False) -> new subvolume iterator\n\n"	\
+	 "Create a new iterator that produces tuples of (path, ID) representing\n"	\
+	 "subvolumes on a filesystem.\n\n"						\
+	 "Arguments:\n"									\
+	 "path -- string, bytes, path-like object, or open file descriptor in\n"	\
+	 "filesystem to list\n"								\
+	 "top -- if not zero, instead of only listing subvolumes beneath the\n"		\
+	 "given path, list subvolumes beneath the subvolume with this ID; passing\n"	\
+	 "BTRFS_FS_TREE_OBJECTID (5) here lists all subvolumes. The subvolumes\n"	\
+	 "are listed relative to the subvolume with this ID.\n"				\
+	 "info -- bool indicating the iterator should yield SubvolumeInfo instead of\n"	\
+	 "the subvolume ID\n"								\
+	 "post_order -- bool indicating whether to yield parent subvolumes before\n"	\
+	 "child subvolumes (e.g., 'foo/bar' before 'foo')"
+
+static PyMethodDef SubvolumeIterator_methods[] = {
+	{"close", (PyCFunction)SubvolumeIterator_close,
+	 METH_NOARGS,
+	 "close()\n\n"
+	 "Close this iterator."},
+	{"fileno", (PyCFunction)SubvolumeIterator_fileno,
+	 METH_NOARGS,
+	 "fileno() -> int\n\n"
+	 "Get the file descriptor associated with this iterator."},
+	{"__enter__", (PyCFunction)SubvolumeIterator_enter,
+	 METH_NOARGS, ""},
+	{"__exit__", (PyCFunction)SubvolumeIterator_exit,
+	 METH_VARARGS | METH_KEYWORDS, ""},
+	{},
+};
+
+PyTypeObject SubvolumeIterator_type = {
+	PyVarObject_HEAD_INIT(NULL, 0)
+	"btrfsutil.SubvolumeIterator",		/* tp_name */
+	sizeof(SubvolumeIterator),		/* tp_basicsize */
+	0,					/* tp_itemsize */
+	(destructor)SubvolumeIterator_dealloc,	/* tp_dealloc */
+	NULL,					/* tp_print */
+	NULL,					/* tp_getattr */
+	NULL,					/* tp_setattr */
+	NULL,					/* tp_as_async */
+	NULL,					/* tp_repr */
+	NULL,					/* tp_as_number */
+	NULL,					/* tp_as_sequence */
+	NULL,					/* tp_as_mapping */
+	NULL,					/* tp_hash  */
+	NULL,					/* tp_call */
+	NULL,					/* tp_str */
+	NULL,					/* tp_getattro */
+	NULL,					/* tp_setattro */
+	NULL,					/* tp_as_buffer */
+	Py_TPFLAGS_DEFAULT,			/* tp_flags */
+	SubvolumeIterator_DOC,			/* tp_doc */
+	NULL,					/* tp_traverse */
+	NULL,					/* tp_clear */
+	NULL,					/* tp_richcompare */
+	0,					/* tp_weaklistoffset */
+	PyObject_SelfIter,			/* tp_iter */
+	(iternextfunc)SubvolumeIterator_next,	/* tp_iternext */
+	SubvolumeIterator_methods,		/* tp_methods */
+	NULL,					/* tp_members */
+	NULL,					/* tp_getset */
+	NULL,					/* tp_base */
+	NULL,					/* tp_dict */
+	NULL,					/* tp_descr_get */
+	NULL,					/* tp_descr_set */
+	0,					/* tp_dictoffset */
+	(initproc)SubvolumeIterator_init,	/* tp_init */
+};
diff --git a/libbtrfsutil/python/tests/test_subvolume.py b/libbtrfsutil/python/tests/test_subvolume.py
index 937a4397..b43beca7 100644
--- a/libbtrfsutil/python/tests/test_subvolume.py
+++ b/libbtrfsutil/python/tests/test_subvolume.py
@@ -214,3 +214,59 @@ class TestSubvolume(BtrfsTestCase):
         wstatus = os.waitpid(pid, 0)[1]
         self.assertTrue(os.WIFEXITED(wstatus))
         self.assertEqual(os.WEXITSTATUS(wstatus), 0)
+
+    def test_subvolume_iterator(self):
+        pwd = os.getcwd()
+        try:
+            os.chdir(self.mountpoint)
+            btrfsutil.create_subvolume('foo')
+
+            path, subvol = next(btrfsutil.SubvolumeIterator('.', info=True))
+            self.assertEqual(path, 'foo')
+            self.assertIsInstance(subvol, btrfsutil.SubvolumeInfo)
+            self.assertEqual(subvol.id, 256)
+            self.assertEqual(subvol.parent_id, 5)
+
+            btrfsutil.create_subvolume('foo/bar')
+            btrfsutil.create_subvolume('foo/bar/baz')
+
+            subvols = [
+                ('foo', 256),
+                ('foo/bar', 257),
+                ('foo/bar/baz', 258),
+            ]
+
+            for arg in self.path_or_fd('.'):
+                with self.subTest(type=type(arg)):
+                    self.assertEqual(list(btrfsutil.SubvolumeIterator(arg)), subvols)
+            self.assertEqual(list(btrfsutil.SubvolumeIterator('.', top=0)), subvols)
+
+            self.assertEqual(list(btrfsutil.SubvolumeIterator('.', post_order=True)),
+                             [('foo/bar/baz', 258),
+                              ('foo/bar', 257),
+                              ('foo', 256)])
+
+            subvols = [
+                ('bar', 257),
+                ('bar/baz', 258),
+            ]
+
+            self.assertEqual(list(btrfsutil.SubvolumeIterator('.', top=256)), subvols)
+            self.assertEqual(list(btrfsutil.SubvolumeIterator('foo', top=0)), subvols)
+
+            os.rename('foo/bar/baz', 'baz')
+            self.assertEqual(sorted(btrfsutil.SubvolumeIterator('.')),
+                             [('baz', 258),
+                              ('foo', 256),
+                              ('foo/bar', 257)])
+
+            with btrfsutil.SubvolumeIterator('.') as it:
+                self.assertGreaterEqual(it.fileno(), 0)
+                it.close()
+                with self.assertRaises(ValueError):
+                    next(iter(it))
+                with self.assertRaises(ValueError):
+                    it.fileno()
+                it.close()
+        finally:
+            os.chdir(pwd)