From 4e8d10cb442723b569170e8595fba197cfad443d Mon Sep 17 00:00:00 2001 From: Lex Berezhny Date: Fri, 6 Aug 2021 10:44:57 -0400 Subject: [PATCH] disk space manager and status API --- lbry/blob/disk_space_manager.py | 21 ++++++++++++++++ lbry/conf.py | 1 + lbry/extras/daemon/components.py | 25 +++++++++++++++++++ lbry/extras/daemon/daemon.py | 2 +- .../datanetwork/test_file_commands.py | 12 +++++++++ tests/unit/blob/test_disk_space_manager.py | 23 +++++++++++++++++ 6 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 lbry/blob/disk_space_manager.py create mode 100644 tests/unit/blob/test_disk_space_manager.py diff --git a/lbry/blob/disk_space_manager.py b/lbry/blob/disk_space_manager.py new file mode 100644 index 000000000..aadb78b79 --- /dev/null +++ b/lbry/blob/disk_space_manager.py @@ -0,0 +1,21 @@ +import os + + +class DiskSpaceManager: + + def __init__(self, config): + self.config = config + + @property + def space_used_bytes(self): + used = 0 + data_dir = self.config.data_dir + for item in os.listdir(data_dir): + blob_path = os.path.join(data_dir, item) + if os.path.isfile(blob_path): + used += os.path.getsize(blob_path) + return used + + @property + def space_used_mb(self): + return int(self.space_used_bytes/1024.0/1024.0) diff --git a/lbry/conf.py b/lbry/conf.py index 4250fe579..480be87ab 100644 --- a/lbry/conf.py +++ b/lbry/conf.py @@ -634,6 +634,7 @@ class Config(CLIConfig): # blob announcement and download save_blobs = Toggle("Save encrypted blob files for hosting, otherwise download blobs to memory only.", True) + blob_storage_limit = Integer("Disk space in MB to be allocated for blob storage. 0 = no limit", 0) blob_lru_cache_size = Integer( "LRU cache size for decrypted downloaded blobs used to minimize re-downloading the same blobs when " "replying to a range request. Set to 0 to disable.", 32 diff --git a/lbry/extras/daemon/components.py b/lbry/extras/daemon/components.py index 822022f32..fbaa598cb 100644 --- a/lbry/extras/daemon/components.py +++ b/lbry/extras/daemon/components.py @@ -15,6 +15,7 @@ from lbry.dht.node import Node from lbry.dht.peer import is_valid_public_ipv4 from lbry.dht.blob_announcer import BlobAnnouncer from lbry.blob.blob_manager import BlobManager +from lbry.blob.disk_space_manager import DiskSpaceManager from lbry.blob_exchange.server import BlobServer from lbry.stream.stream_manager import StreamManager from lbry.file.file_manager import FileManager @@ -40,6 +41,7 @@ WALLET_SERVER_PAYMENTS_COMPONENT = "wallet_server_payments" DHT_COMPONENT = "dht" HASH_ANNOUNCER_COMPONENT = "hash_announcer" FILE_MANAGER_COMPONENT = "file_manager" +DISK_SPACE_COMPONENT = "disk_space" PEER_PROTOCOL_SERVER_COMPONENT = "peer_protocol_server" UPNP_COMPONENT = "upnp" EXCHANGE_RATE_MANAGER_COMPONENT = "exchange_rate_manager" @@ -375,6 +377,29 @@ class FileManagerComponent(Component): self.file_manager.stop() +class DiskSpaceComponent(Component): + component_name = DISK_SPACE_COMPONENT + + def __init__(self, component_manager): + super().__init__(component_manager) + self.disk_space_manager = DiskSpaceManager(self.conf) + + @property + def component(self) -> typing.Optional[DiskSpaceManager]: + return self.disk_space_manager + + async def get_status(self): + return { + 'used': str(self.disk_space_manager.space_used_mb), + } + + async def start(self): + pass + + async def stop(self): + pass + + class TorrentComponent(Component): component_name = LIBTORRENT_COMPONENT diff --git a/lbry/extras/daemon/daemon.py b/lbry/extras/daemon/daemon.py index 4443c209f..bf33c0541 100644 --- a/lbry/extras/daemon/daemon.py +++ b/lbry/extras/daemon/daemon.py @@ -41,7 +41,7 @@ from lbry.error import ( from lbry.extras import system_info from lbry.extras.daemon import analytics from lbry.extras.daemon.components import WALLET_COMPONENT, DATABASE_COMPONENT, DHT_COMPONENT, BLOB_COMPONENT -from lbry.extras.daemon.components import FILE_MANAGER_COMPONENT +from lbry.extras.daemon.components import FILE_MANAGER_COMPONENT, DISK_SPACE_COMPONENT from lbry.extras.daemon.components import EXCHANGE_RATE_MANAGER_COMPONENT, UPNP_COMPONENT from lbry.extras.daemon.componentmanager import RequiredCondition from lbry.extras.daemon.componentmanager import ComponentManager diff --git a/tests/integration/datanetwork/test_file_commands.py b/tests/integration/datanetwork/test_file_commands.py index ee65c755b..7d472f314 100644 --- a/tests/integration/datanetwork/test_file_commands.py +++ b/tests/integration/datanetwork/test_file_commands.py @@ -511,3 +511,15 @@ class FileCommands(CommandTestCase): await tx.sign([self.account]) await self.broadcast(tx) await self.confirm_tx(tx.id) + + +class DiskSpaceManagement(CommandTestCase): + + async def test_file_management(self): + status = await self.daemon.jsonrpc_status() + self.assertIn('disk_space', status) + self.assertEqual(status['disk_space']['used'], '0') + await self.stream_create('foo', '0.01', data=('0' * 3 * 1024 * 1024).encode()) + status = await self.daemon.jsonrpc_status() + self.assertIn('disk_space', status) + self.assertEqual(status['disk_space']['used'], '3') diff --git a/tests/unit/blob/test_disk_space_manager.py b/tests/unit/blob/test_disk_space_manager.py new file mode 100644 index 000000000..75aa07c4c --- /dev/null +++ b/tests/unit/blob/test_disk_space_manager.py @@ -0,0 +1,23 @@ +import os +import unittest +import tempfile + +import lbry.wallet +from lbry.conf import Config +from lbry.blob.disk_space_manager import DiskSpaceManager + + +class ConfigurationTests(unittest.TestCase): + + def test_space_calculation(self): + with tempfile.TemporaryDirectory() as temp_dir: + config = Config( + data_dir=temp_dir, + wallet_dir=temp_dir, + config=os.path.join(temp_dir, 'settings.yml') + ) + dsm = DiskSpaceManager(config) + self.assertEqual(0, dsm.space_used_mb) + with open(os.path.join(config.data_dir, '3mb-file'), 'w') as blob: + blob.write('0' * 3 * 1024 * 1024) + self.assertEqual(3, dsm.space_used_mb)