From 0eb6324fcc6abccd97130bae45ad70e17d6b8e65 Mon Sep 17 00:00:00 2001
From: John Mulligan <jmulligan@redhat.com>
Date: Wed, 1 May 2024 10:55:14 -0400
Subject: [PATCH] mgr/smb: cache new and stored keys separately in the staging
 area

We would like a quick way of determining what is new. Add methods to the
staging area class to make it possible to do so using (cached) keys.

Signed-off-by: John Mulligan <jmulligan@redhat.com>
---
 src/pybind/mgr/smb/handler.py | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/pybind/mgr/smb/handler.py b/src/pybind/mgr/smb/handler.py
index 9931fe0a68835..1bf96557f6b4b 100644
--- a/src/pybind/mgr/smb/handler.py
+++ b/src/pybind/mgr/smb/handler.py
@@ -190,10 +190,11 @@ class _Staging:
         self.destination_store = store
         self.incoming: Dict[EntryKey, SMBResource] = {}
         self.deleted: Dict[EntryKey, SMBResource] = {}
-        self._keycache: Set[EntryKey] = set()
+        self._store_keycache: Set[EntryKey] = set()
+        self._virt_keycache: Set[EntryKey] = set()
 
     def stage(self, resource: SMBResource) -> None:
-        self._keycache = set()
+        self._virt_keycache = set()
         ekey = resource_key(resource)
         if resource.intent == Intent.REMOVED:
             self.deleted[ekey] = resource
@@ -201,19 +202,21 @@ class _Staging:
             self.deleted.pop(ekey, None)
             self.incoming[ekey] = resource
 
-    def _virtual_keys(self) -> Iterator[EntryKey]:
-        new = set(self.incoming.keys())
-        for ekey in self.destination_store:
-            if ekey in self.deleted:
-                continue
-            yield ekey
-            new.discard(ekey)
-        for ekey in new:
-            yield ekey
+    def _virtual_keys(self) -> Collection[EntryKey]:
+        if self._virt_keycache:
+            return self._virt_keycache
+        self._virt_keycache = set(self._store_keys()) - set(
+            self.deleted
+        ) | set(self.incoming)
+        return self._virt_keycache
+
+    def _store_keys(self) -> Collection[EntryKey]:
+        if not self._store_keycache:
+            self._store_keycache = set(self.destination_store)
+        return self._store_keycache
 
     def __iter__(self) -> Iterator[EntryKey]:
-        self._keycache = set(self._virtual_keys())
-        return iter(self._keycache)
+        return iter(self._virtual_keys())
 
     def namespaces(self) -> Collection[str]:
         return {k[0] for k in self}
@@ -221,6 +224,10 @@ class _Staging:
     def contents(self, ns: str) -> Collection[str]:
         return {kname for kns, kname in self if kns == ns}
 
+    def is_new(self, resource: SMBResource) -> bool:
+        ekey = resource_key(resource)
+        return ekey not in self._store_keys()
+
     def get_cluster(self, cluster_id: str) -> resources.Cluster:
         ekey = (str(ClusterEntry.namespace), cluster_id)
         if ekey in self.incoming:
-- 
2.39.5