buildgrid 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

buildgrid/server/cas/storage/redis_fmb_cache.py

@@ -0,0 +1,220 @@
+# Copyright (C) 2026 Bloomberg LP
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#  <http://www.apache.org/licenses/LICENSE-2.0>
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+A storage provider that uses redis to maintain a cache of existence for some underlying storage.
+"""
+
+from datetime import timedelta
+from typing import Callable, IO, TypeVar
+
+import redis
+
+from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest
+from buildgrid._protos.google.rpc import code_pb2
+from buildgrid._protos.google.rpc.status_pb2 import Status
+from buildgrid.server.cas.storage.storage_abc import StorageABC
+from buildgrid.server.decorators import timed
+from buildgrid.server.enums import BlobExistence
+from buildgrid.server.logging import buildgrid_logger
+from buildgrid.server.metrics_names import METRIC
+from buildgrid.server.redis.provider import RedisProvider
+
+LOGGER = buildgrid_logger(__name__)
+
+
+T = TypeVar("T")
+
+
+class RedisFMBCache(StorageABC):
+    TYPE = "RedisFMBCache"
+
+    def __init__(self, redis: RedisProvider, storage: StorageABC, ttl_secs: int, prefix: str | None = None) -> None:
+        self._redis = redis
+        self._storage = storage
+        self._prefix = "C"
+        if prefix == "C":
+            LOGGER.error("Prefix 'C' is reserved as the default prefix and cannot be used")
+            raise ValueError("Prefix 'C' is reserved as the default prefix and cannot be used")
+        elif prefix:
+            self._prefix = prefix
+        self._ttl = timedelta(seconds=ttl_secs)
+
+    def start(self) -> None:
+        self._storage.start()
+
+    def stop(self) -> None:
+        self._storage.stop()
+
+    def _construct_key(self, digest: Digest) -> str:
+        """Helper to get the redis key name for a particular digest"""
+        # The tag prefix serves to distinguish between our keys and
+        # actual blobs if the same redis is used for both index and storage
+        return f"{self._prefix}:{digest.hash}_{digest.size_bytes}"
+
+    def _deconstruct_key(self, keystr: str) -> Digest | None:
+        """Helper to attempt to recover a Digest from a redis key"""
+        try:
+            tag, rest = keystr.split(":", 1)
+            if tag != self._prefix:
+                return None
+            hash, size_bytes = rest.rsplit("_", 1)
+            return Digest(hash=hash, size_bytes=int(size_bytes))
+        except ValueError:
+            return None
+
+    def _safe_redis_ro(self, func: Callable[["redis.Redis[bytes]"], T]) -> T | None:
+        try:
+            return self._redis.execute_ro(func)
+        except Exception:
+            return None
+
+    def _safe_redis_rw(self, func: Callable[["redis.Redis[bytes]"], T]) -> T | None:
+        try:
+            return self._redis.execute_rw(func)
+        except Exception:
+            return None
+
+    @timed(METRIC.STORAGE.STAT_DURATION, type=TYPE)
+    def has_blob(self, digest: Digest) -> bool:
+        blob_present = self._safe_redis_ro(lambda r: r.get(self._construct_key(digest)))
+
+        # If we don't have a key cached for this digest, check the underlying storage
+        if blob_present is None:
+            storage_has_blob = self._storage.has_blob(digest)
+            value = BlobExistence.EXISTS.value if storage_has_blob else BlobExistence.DOES_NOT_EXIST.value
+            # Cache the result
+            self._safe_redis_rw(lambda r: r.set(self._construct_key(digest), value, ex=self._ttl))
+            return storage_has_blob
+
+        return bool(int(blob_present))
+
+    @timed(METRIC.STORAGE.READ_DURATION, type=TYPE)
+    def get_blob(self, digest: Digest) -> IO[bytes] | None:
+        if blob := self._storage.get_blob(digest):
+            self._safe_redis_rw(lambda r: r.set(self._construct_key(digest), BlobExistence.EXISTS.value, ex=self._ttl))
+            return blob
+
+        self._safe_redis_rw(
+            lambda r: r.set(self._construct_key(digest), BlobExistence.DOES_NOT_EXIST.value, ex=self._ttl)
+        )
+        return None
+
+    @timed(METRIC.STORAGE.DELETE_DURATION, type=TYPE)
+    def delete_blob(self, digest: Digest) -> None:
+        self._safe_redis_rw(lambda r: r.delete(self._construct_key(digest)))
+        self._storage.delete_blob(digest)
+
+        # Delete the key again, in case we raced with an insertion and deleted the
+        # storage out from under it. This lets us err on the side of having an uncached
+        # existence check rather than an incorrect cache.
+        self._safe_redis_rw(lambda r: r.delete(self._construct_key(digest)))
+
+    @timed(METRIC.STORAGE.WRITE_DURATION, type=TYPE)
+    def commit_write(self, digest: Digest, write_session: IO[bytes]) -> None:
+        self._storage.commit_write(digest, write_session)
+        self._safe_redis_rw(lambda r: r.set(self._construct_key(digest), BlobExistence.EXISTS.value, ex=self._ttl))
+
+    @timed(METRIC.STORAGE.BULK_DELETE_DURATION, type=TYPE)
+    def bulk_delete(self, digests: list[Digest]) -> list[str]:
+        # Delete any digests in the list that we have cached
+        keys = [self._construct_key(digest) for digest in digests]
+        self._safe_redis_rw(lambda r: r.delete(*keys))
+
+        failed_deletes = self._storage.bulk_delete(digests)
+
+        # Redo the deletion in case we raced an insert. Worst-case this leaves us with an
+        # uncached key
+        self._safe_redis_rw(lambda r: r.delete(*keys))
+        return failed_deletes
+
+    @timed(METRIC.STORAGE.BULK_STAT_DURATION, type=TYPE)
+    def missing_blobs(self, digests: list[Digest]) -> list[Digest]:
+        def check_existence(r: "redis.Redis[bytes]") -> list[bytes]:
+            pipe = r.pipeline(transaction=False)
+            for digest in digests:
+                pipe.get(self._construct_key(digest))
+            return pipe.execute()
+
+        results = self._safe_redis_ro(check_existence)
+        missing_digests = []
+        uncached_digests = []
+
+        # Map digests to pipeline results, then work out the subsets which are uncached
+        # (needing to be checked in the underlying storage) and missing (cached as not present)
+        if results is not None:
+            for digest, result in zip(digests, results):
+                if result == BlobExistence.DOES_NOT_EXIST.value:
+                    missing_digests.append(digest)
+                elif result is None:
+                    uncached_digests.append(digest)
+
+        # If the redis call failed, just fall back to the underlying storage for all digests
+        else:
+            uncached_digests = digests
+
+        if uncached_digests:
+            missing_digests.extend(self._storage.missing_blobs(uncached_digests))
+
+        def populate_cache(r: "redis.Redis[bytes]") -> None:
+            pipe = r.pipeline()
+            for digest in uncached_digests:
+                key = self._construct_key(digest)
+                if digest in missing_digests:
+                    pipe.set(key, BlobExistence.DOES_NOT_EXIST.value, ex=self._ttl)
+                else:
+                    pipe.set(key, BlobExistence.EXISTS.value, ex=self._ttl)
+            pipe.execute()
+
+        self._safe_redis_rw(populate_cache)
+        return missing_digests
+
+    @timed(METRIC.STORAGE.BULK_WRITE_DURATION, type=TYPE)
+    def bulk_update_blobs(self, blobs: list[tuple[Digest, bytes]]) -> list[Status]:
+        result_map: dict[str, Status] = {}
+        missing_blob_pairs: list[tuple[Digest, bytes]] = []
+        missing_blobs = self.missing_blobs([digest for digest, _ in blobs])
+        for digest, blob in blobs:
+            if digest not in missing_blobs:
+                result_map[digest.hash] = Status(code=code_pb2.OK)
+            else:
+                missing_blob_pairs.append((digest, blob))
+        results = self._storage.bulk_update_blobs(missing_blob_pairs)
+
+        def cache_existence(r: "redis.Redis[bytes]") -> None:
+            pipe = r.pipeline()
+            for digest, result in zip(missing_blobs, results):
+                result_map[digest.hash] = result
+                if result.code == code_pb2.OK:
+                    key = self._construct_key(digest)
+                    pipe.set(key, BlobExistence.EXISTS.value, ex=self._ttl)
+            pipe.execute()
+
+        self._safe_redis_rw(cache_existence)
+        return [result_map[digest.hash] for digest, _ in blobs]
+
+    @timed(METRIC.STORAGE.BULK_READ_DURATION, type=TYPE)
+    def bulk_read_blobs(self, digests: list[Digest]) -> dict[str, bytes]:
+        fetched_digests = self._storage.bulk_read_blobs(digests)
+
+        fetched_digest_hashes = set(digest_hash for (digest_hash, _) in fetched_digests.items())
+        digests_not_in_storage: list[Digest] = []
+        for expected_digest in digests:
+            if expected_digest.hash not in fetched_digest_hashes:
+                digests_not_in_storage.append(expected_digest)
+
+        if digests_not_in_storage:
+            self._safe_redis_rw(lambda r: r.delete(*[self._construct_key(digest) for digest in digests_not_in_storage]))
+
+        return fetched_digests

buildgrid/server/enums.py

@@ -172,3 +172,8 @@ class JobAssignmentStrategy(Enum):
     CAPACITY = "capacity"
     PROACTIVE = "proactive"
     PREEMPTION = "preemption"
+
+
+class BlobExistence(Enum):
+    EXISTS = b"1"
+    DOES_NOT_EXIST = b"0"

buildgrid/server/metrics_names.py

@@ -64,8 +64,6 @@ class METRIC:
             SIZE_CALCULATION_DURATION       = "storage.sql_index.size_calculation.duration.ms"
             DELETE_N_BYTES_DURATION         = "storage.sql_index.delete_n_bytes.duration.ms"
             BULK_DELETE_INDEX_DURATION      = "storage.sql_index.bulk_delete_index.duration.ms"
-            MARK_DELETED_DURATION           = "storage.sql_index.mark_deleted.duration.ms"
-            PREMARKED_DELETED_COUNT         = "storage.sql_index.premarked_deleted.count"
 
         class REPLICATED:
             REQUIRED_REPLICATION_COUNT      = "storage.replicated.required_replication.count"