113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
import os
|
|
import asyncio
|
|
import logging
|
|
from threading import Thread
|
|
from multiprocessing import Queue, Event
|
|
from concurrent import futures
|
|
|
|
from lbry.wallet.stream import StreamController
|
|
from lbry.db import Database
|
|
|
|
from .lbrycrd import Lbrycrd
|
|
from . import worker
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class ProgressMonitorThread(Thread):
|
|
|
|
STOP = 'stop'
|
|
FORMAT = '{l_bar}{bar}| {n_fmt:>6}/{total_fmt:>7} [{elapsed}<{remaining:>5}, {rate_fmt:>15}]'
|
|
|
|
def __init__(self, state: dict, queue: Queue, stream_controller: StreamController):
|
|
super().__init__()
|
|
self.state = state
|
|
self.queue = queue
|
|
self.stream_controller = stream_controller
|
|
self.loop = asyncio.get_event_loop()
|
|
|
|
def run(self):
|
|
asyncio.set_event_loop(self.loop)
|
|
while True:
|
|
msg = self.queue.get()
|
|
if msg == self.STOP:
|
|
return
|
|
self.stream_controller.add(msg)
|
|
|
|
def shutdown(self):
|
|
self.queue.put(self.STOP)
|
|
self.join()
|
|
|
|
def __enter__(self):
|
|
self.start()
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.shutdown()
|
|
|
|
|
|
class BlockchainSync:
|
|
|
|
def __init__(self, chain: Lbrycrd, db: Database, use_process_pool=False):
|
|
self.chain = chain
|
|
self.db = db
|
|
self.use_process_pool = use_process_pool
|
|
self._on_progress_controller = StreamController()
|
|
self.on_progress = self._on_progress_controller.stream
|
|
|
|
def get_worker_pool(self, queue, full_stop) -> futures.Executor:
|
|
args = dict(
|
|
initializer=worker.initializer,
|
|
initargs=(self.chain.data_dir, self.chain.regtest, self.db.db_path, queue, full_stop)
|
|
)
|
|
if not self.use_process_pool:
|
|
return futures.ThreadPoolExecutor(max_workers=1, **args)
|
|
return futures.ProcessPoolExecutor(max_workers=max(os.cpu_count()-1, 4), **args)
|
|
|
|
def get_progress_monitor(self, state, queue) -> ProgressMonitorThread:
|
|
return ProgressMonitorThread(state, queue, self._on_progress_controller)
|
|
|
|
async def load_blocks(self):
|
|
jobs = []
|
|
queue, full_stop = Queue(), Event()
|
|
executor = self.get_worker_pool(queue, full_stop)
|
|
files = list(await self.chain.get_block_files_not_synced())
|
|
state = {
|
|
file.file_number: {
|
|
'status': worker.PENDING,
|
|
'done_txs': 0,
|
|
'total_txs': file.txs,
|
|
'done_blocks': 0,
|
|
'total_blocks': file.blocks,
|
|
} for file in files
|
|
}
|
|
progress = self.get_progress_monitor(state, queue)
|
|
progress.start()
|
|
|
|
def cancel_all_the_things():
|
|
for job in jobs:
|
|
job.cancel()
|
|
full_stop.set()
|
|
for job in jobs:
|
|
exception = job.exception()
|
|
if exception is not None:
|
|
raise exception
|
|
|
|
try:
|
|
|
|
for file in files:
|
|
jobs.append(executor.submit(worker.process_block_file, file.file_number))
|
|
|
|
done, not_done = await asyncio.get_event_loop().run_in_executor(
|
|
None, futures.wait, jobs, None, futures.FIRST_EXCEPTION
|
|
)
|
|
if not_done:
|
|
cancel_all_the_things()
|
|
|
|
except asyncio.CancelledError:
|
|
cancel_all_the_things()
|
|
raise
|
|
|
|
finally:
|
|
progress.shutdown()
|
|
executor.shutdown()
|