Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/cas/storage/size_differentiated.py: 100.00%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

40 statements  

1# Copyright (C) 2020 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16""" 

17SizeDifferentiatedStorage 

18========================= 

19 

20A storage provider which passes requests to other storage providers 

21based on the size of the blob being requested. 

22 

23""" 

24 

25import io 

26import logging 

27from typing import List, MutableMapping, NamedTuple, Union 

28 

29from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest 

30from .storage_abc import StorageABC 

31 

32 

33StorageListType = List[MutableMapping[str, Union[int, StorageABC]]] 

34 

35 

36class _SizeLimitedStorage(NamedTuple): 

37 max_size: int 

38 storage: StorageABC 

39 

40 

41class SizeDifferentiatedStorage(StorageABC): 

42 

43 def __init__(self, storages: StorageListType, fallback: StorageABC): 

44 self._logger = logging.getLogger(__name__) 

45 self._fallback_storage = fallback 

46 self._storages = [_SizeLimitedStorage(**storage) for storage in storages] 

47 self._storages.sort(key=lambda storage: storage.max_size) 

48 

49 def _storage_from_digest(self, digest: Digest) -> StorageABC: 

50 for storage in self._storages: 

51 if digest.size_bytes < storage.max_size: 

52 return storage.storage 

53 # If the blob is too big for any of the size-limited storages, 

54 # put it in the fallback. 

55 return self._fallback_storage 

56 

57 def has_blob(self, digest: Digest) -> bool: 

58 self._logger.debug(f"Checking for blob: [{digest}]") 

59 storage = self._storage_from_digest(digest) 

60 return storage.has_blob(digest) 

61 

62 def get_blob(self, digest: Digest) -> io.IOBase: 

63 self._logger.debug(f"Getting blob: [{digest}]") 

64 storage = self._storage_from_digest(digest) 

65 return storage.get_blob(digest) 

66 

67 def delete_blob(self, digest: Digest): 

68 self._logger.debug(f"Deleting blob: [{digest}]") 

69 storage = self._storage_from_digest(digest) 

70 storage.delete_blob(digest) 

71 

72 def begin_write(self, digest: Digest) -> io.IOBase: 

73 storage = self._storage_from_digest(digest) 

74 return storage.begin_write(digest) 

75 

76 def commit_write(self, digest: Digest, write_session: io.IOBase): 

77 self._logger.debug(f"Writing blob: [{digest}]") 

78 storage = self._storage_from_digest(digest) 

79 storage.commit_write(digest, write_session)