from nullptr.models import * from os.path import isfile, dirname, isdir import os from os.path import basename import json from .util import * from time import time import pickle from struct import unpack, pack from functools import partial from io import BytesIO from copy import copy class StorePickler(pickle.Pickler): def persistent_id(self, obj): return "STORE" if type(obj) == Store else None class StoreUnpickler(pickle.Unpickler): def __init__(self, stream, store): self.store = store super().__init__(stream) def persistent_load(self, pers_id): if pers_id == "STORE": return self.store raise pickle.UnpicklingError("I don know the persid!") CHUNK_MAGIC = b'ChNkcHnK' class ChunkHeader: def __init__(self): self.magic = CHUNK_MAGIC self.offset = 0 self.in_use = True self.size = 0 self.used = 0 @classmethod def parse(cls, fil): offset = fil.tell() d = fil.read(24) if len(d) < 24: return None o = cls() o.offset = offset o.magic, d, o.used = unpack('<8sQQ', d) o.size = d & 0x7fffffffffffffff o.in_use = d & 0x8000000000000000 != 0 if o.magic != CHUNK_MAGIC: raise ValueError(f"Invalid chunk magic: {o.magic}") # print(o) return o def write(self, f): d = self.size if self.in_use: d |= 1 << 63 d = pack('<8sQQ', self.magic, d, self.used) f.write(d) def __repr__(self): return f'chunk {self.in_use} {self.size} {self.used}' def f(self, detail=1): if detail == 1: return f'chunk {self.offset} {self.used}/{self.size}' else: r = f'Stored at: {self.offset}\n' slack = self.size - self.used r += f'Used: {self.used}/{self.size} (slack {slack})' return r class Store: def __init__(self, data_file, verbose=False): self.init_models() self.data_file = data_file self.data_dir = os.path.dirname(data_file) self.fil = open_file(data_file) self.data = {m: {} for m in self.models} self.system_members = {} self.dirty_objects = set() self.cleanup_interval = 600 self.last_cleanup = 0 self.slack = 0.1 self.slack_min = 64 self.slack_max = 1024 self.verbose = verbose self.load() def p(self, m): if not self.verbose: return print(m) def f(self, detail): return f'Store {self.data_file}' def close(self): self.flush() self.fil.close() def init_models(self): self.models = all_subclasses(Base) self.extensions = {c.ext(): c for c in self.models} self.model_names = {c.__name__: c for c in self.models} def dirty(self, obj): self.dirty_objects.add(obj) def dump_object(self, obj): buf = BytesIO() p = StorePickler(buf) p.dump(obj) return buf.getvalue() def load_object(self, data, offset): buf = BytesIO(data) p = StoreUnpickler(buf, self) obj = p.load() x = self.get(type(obj), obj.symbol) if x is not None and x in self.dirty_objects: self.dirty_objects.remove(obj) obj._file_offset = offset self.hold(obj) def load(self): cnt = 0 start_time = time() total = 0 free = 0 self.fil.seek(0) offset = 0 while (hdr := ChunkHeader.parse(self.fil)): # self.p(hdr) total += hdr.size if not hdr.in_use: # print(f"skip {hdr.size} {self.fil.tell()}") self.fil.seek(hdr.size, 1) free += hdr.size else: data = self.fil.read(hdr.used) self.load_object(data, offset) # print(f"pad {hdr.size - hdr.used}") self.fil.seek(hdr.size - hdr.used, 1) cnt += 1 offset = self.fil.tell() dur = time() - start_time # just in case any temp objects were created self.dirty_objects = set() self.p(f'Loaded {cnt} objects in {dur:.2f} seconds') self.p(f'Fragmented space: {free} / {total} bytes') def allocate_chunk(self, sz): used = sz slack = sz * self.slack slack = min(slack, self.slack_max) slack = max(slack, self.slack_min) sz += int(slack) self.fil.seek(0, 2) offset = self.fil.tell() h = ChunkHeader() h.size = sz h.used = used h.offset = offset h.write(self.fil) return offset, h def get_header(self, obj): if obj._file_offset is None: return None self.fil.seek(obj._file_offset) hdr = ChunkHeader.parse(self.fil) return hdr def purge(self, obj): if obj._file_offset is not None: self.fil.seek(obj._file_offset) hdr = ChunkHeader.parse(self.fil) hdr.in_use = False self.fil.seek(obj._file_offset) hdr.write(self.fil) if type(obj) in self.data and obj.symbol in self.data[type(obj)]: del self.data[type(obj)][obj.symbol] self.remove_from_members(obj) if obj in self.dirty_objects: self.dirty_objects.remove(obj) obj._file_offset = None def store(self, obj): data = self.dump_object(obj) osize = len(data) # is there an existing chunk for this obj? if obj._file_offset is not None: # read chunk hdr self.fil.seek(obj._file_offset) hdr = ChunkHeader.parse(self.fil) csize = hdr.size # if the chunk is too small if csize < osize: # free the chunk hdr.in_use = False # force a new chunk obj._file_offset = None else: # if it is big enough, update the used field hdr.used = osize self.fil.seek(hdr.offset) hdr.write(self.fil) if obj._file_offset is None: obj._file_offset, hdr = self.allocate_chunk(osize) # print(type(obj).__name__, hdr) self.fil.write(data) slack = b'\x00' * (hdr.size - hdr.used) self.fil.write(slack) def remove_from_members(self, obj): if type(obj).__name__ in ['Waypoint','Marketplace', 'Jumpgate', 'Survey']: system_str = obj.system.symbol if system_str not in self.system_members: return self.system_members[system_str].remove(obj) def hold(self, obj): typ = type(obj) symbol = obj.symbol obj.store = self self.data[typ][symbol] = obj if type(obj).__name__ in ['Waypoint','Marketplace', 'Jumpgate', 'Survey']: system_str = obj.system.symbol if system_str not in self.system_members: self.system_members[system_str] = set() self.system_members[system_str].add(obj) def create(self, typ, symbol): obj = typ(symbol, self) obj.created() self.hold(obj) self.dirty(obj) return obj def get(self, typ, symbol, create=False): if type(typ) == str and typ in self.model_names: typ = self.model_names[typ] symbol = symbol.upper() if typ not in self.data: return None if symbol not in self.data[typ]: if create: return self.create(typ, symbol) else: return None return self.data[typ][symbol] def getter(self, typ, create=False): if type(typ) == str and typ in self.model_names: typ = self.model_names[typ] return partial(self.get, typ, create=create) def update(self, typ, data, symbol=None): if type(typ) == str and typ in self.model_names: typ = self.model_names[typ] if symbol is None: symbol = mg(data, typ.identifier) obj = self.get(typ, symbol, True) obj.update(data) return obj def update_list(self, typ, lst): return [self.update(typ, d) for d in lst] def all(self, typ): if type(typ) == str and typ in self.model_names: typ = self.model_names[typ] for m in self.data[typ].values(): if m.is_expired(): self.dirty(m) continue yield m def all_members(self, system, typ=None): if type(typ) == str and typ in self.model_names: typ = self.model_names[typ] if type(system) == System: system = system.symbol if system not in self.system_members: return garbage = set() for m in self.system_members[system]: if m.is_expired(): self.dirty(m) garbage.add(m) continue if typ is None or type(m) == typ: yield m for m in garbage: self.system_members[system].remove(m) def cleanup(self): self.last_cleanup = time() start_time = time() expired = list() for t in self.data: for o in self.data[t].values(): if o.is_expired(): expired.append(o) for o in expired: self.purge(o) dur = time() - start_time # self.p(f'cleaned {len(expired)} in {dur:.03f} seconds') def flush(self): self.cleanup() it = 0 start_time = time() for obj in copy(self.dirty_objects): it += 1 if obj.symbol not in self.data[type(obj)] or self.data[type(obj)][obj.symbol] != obj: # print(f"Dirty object not in data {type(obj)} {obj.symbol} {obj}") continue self.store(obj) self.fil.flush() self.dirty_objects = set() dur = time() - start_time #self.p(f'flush done {it} items {dur:.2f}') def defrag(self): self.flush() nm = self.fil.name self.fil.close() bakfile = nm+'.bak' if os.path.isfile(bakfile): os.remove(bakfile) os.rename(nm, nm + '.bak') self.fil = open_file(nm) for t in self.data: for o in self.data[t].values(): o._file_offset = None self.store(o)