<commit_before># cython: profile=True

import os, logging, warnings
import struct, mmap, platform
import datetime, csv
from collections import Counter

from sparsehash cimport SparseHashMap
from utils cimport compactify_deck, canonicalize, State, Location
from utils import hands5_factor_rev, hands5_factor, expand_deck, Handle, load, dump


total_scores = 40
bytes_per_entry = 2
prob_format = '>H'
prob_range = 65534
usage_filename = 'usage.json'
usage_total_key = 'total'
usage_memory_key = 'memory'

dense_ext = 'dat'
sparse_ext = 'shash'

config_filename = 'config.json'
history_filename = 'history.csv'

memory_usage_threshold = 701000000

history_report_period = 10000
small_report_period = 100000
large_report_period = 1000000
save_period = 50000000


def posix_fallocate(path, size):
    with open(path, 'wb') as f:
        os.posix_fallocate(f.fileno(), 0, size)

def win_fallocate(path, size):
    # I haven't tested this.
    with open(path, 'wb') as f:
        f.seek(size - 1)
        f.write(b'\x00')

def posix_fadvise_willneed(fileno):
    os.posix_fadvise(fileno, 0, 0, os.POSIX_FADV_WILLNEED)

def posix_fadvise_dontneed(fileno):
    os.posix_fadvise(fileno, 0, 0, os.POSIX_FADV_DONTNEED)

def win_fadvise_willneed(fileno):
    pass

def win_fadvise_dontneed(fileno):
    pass

if platform.system() == 'Linux':
    fallocate = posix_fallocate
    fadvise_willneed = posix_fadvise_willneed
    fadvise_dontneed = posix_fadvise_dontneed
elif platform.system() == 'Windows':
    logging.warning("Windows file interaction functions are untested.")
    fallocate = win_fallocate
    fadvise_willneed = win_fadvise_willneed
    fadvise_dontneed = win_fadvise_dontneed


cdef class Storage:
    cdef str usage_path, history_path, config_path, curr_path, last_path
    cdef list state_dirs, storage_path, memory_storage, history
    cdef bytes curr_value
    cdef dict storage_handles, config
    cdef object storage_usage
    cdef int curr_action
    cdef long curr_offset
    def __init__(self, state_dirs):
        if len(state_dirs) == 1:
            warnings.warn("It is inefficient to use only one storage drive. Add another one.")

        self.state_dirs = state_dirs
        self.storage_handles = {}

        self.usage_path = os.path.join(state_dirs[0], usage_filename)
        if os.path.exists(self.usage_path):
            with open(self.usage_path, 'r') as f:
                self.storage_usage = Counter(load(f))
        else:
            self.storage_usage = Counter()

        self.config_path = os.path.join(state_dirs[0], config_filename)
        if os.path.exists(self.config_path):
            with open(self.config_path, 'r') as f:
                self.config = load(f)
        else:
            self.config = {
                'in_memory': [True] * total_scores
            }

        self.storage_path = [
            os.path.join(self.state_dirs[score % len(state_dirs)],
                         '%02d.%s' % (score, sparse_ext if self.config['in_memory'][score] else dense_ext))
            for score in range(total_scores)
        ]

        self.history_path = os.path.join(state_dirs[0], history_filename)
        self.history = []
        self.register_history('startup')

        # In case the generator is aborted, I want to do at least something to preserve consistency.
        # When I/O is performed, this variable stores the current I/O action.
        # There are 2 types of actions available: writing values in existing files and creating new files.
        # -1 = nothing
        # 0 = writing value
        # 1 = creating file
        # 2 = transferring data from memory storage to disk
        # I'm missing Haskell's algebraic data types.
        self.curr_action = -1

        self.curr_value = None
        self.curr_offset = -1
        self.curr_path = None

        # Before accessing each file I call os.posix_fadvise to load the whole file in memory.
        # Since I don't want to do that every time I store or retrieve data, I save the last
        # accessed file and only call os.posix_fadvise when it changes.
        self.last_path = None

        # This is the storage for the data that resides in memory.
        self.memory_storage = [None] * total_scores

        logging.info("Currently using %d entries in the table" % self.storage_usage[usage_total_key])
        self.report_breakdown()

    cdef long _get_offset(self, State cstate) except -1:
        cdef long hand_offset, deck_offset, offset
        hand_offset = hands5_factor_rev[cstate.hand]
        deck_offset = compactify_deck(cstate.hand, cstate.deck)

        # Multiplication by 2**19. Note: without the parentheses
        # this is interpreted as `offset = hand_offset << (19 + deck_offset)`
        offset = (hand_offset << 19) + deck_offset

        return offset

    cdef _unpack_offset(self, offset):
        # bottom 19 bits: compactified deck
        # the rest: hands5_factor index
        hand_idx = offset >> 19
        compact_deck = offset & ((1 << 19) - 1)
        hand = hands5_factor[hand_idx]
        deck = expand_deck(hand, compact_deck)
        return hand, deck

    cdef int _state_is_in_memory(self, State cstate) except -1:
        return self.config['in_memory'][cstate.score]

    cdef int _do_prefetch(self, Location loc) except -1:
        # TODO: this is ugly, but is there a better way?
        return loc.path == self.storage_path[0]

    cdef _reset_storage(self, loc):
        # assumes that the path exists but the file doesn't

        file_size = loc.size * bytes_per_entry
        logging.info("Creating a new table at %s for %d entries (%d bytes)." % (loc.path, loc.size, file_size))

        self.curr_action = 1
        self.curr_path = loc.path

        fallocate(loc.path, file_size)

        self.curr_action = -1
        self.curr_path = None

        logging.info("Done.")

    cdef _ensure_usable(self, path):
        directory = os.path.dirname(path)
        if not os.path.exists(directory):
            os.makedirs(directory)
        elif not os.path.isdir(directory):
            raise RuntimeError("%s exists but is not a directory" % directory)

    cdef _ensure_initialized(self, Location loc):
        if loc.in_memory:
            if self.memory_storage[loc.idx] is None:
                self._ensure_usable(loc.path)

                if os.path.exists(loc.path):
                    logging.debug("Loading %s from disk to memory..." % loc.path)
                    # Loading from on-disk storage
                    with open(loc.path, 'rb') as f:
                        self.memory_storage[loc.idx] = SparseHashMap(f)
                    assert self.storage_usage[loc.path] == len(self.memory_storage[loc.idx])
                else:
                    self.memory_storage[loc.idx] = SparseHashMap()
        else:

            if loc.path not in self.storage_handles:
                self._ensure_usable(loc.path)

                if not os.path.exists(loc.path):
                    self._reset_storage(loc)

                assert loc.path not in self.storage_handles

                fileobj = open(loc.path, 'r+b')
                mmapobj = mmap.mmap(fileobj.fileno(), loc.size * bytes_per_entry)

                self.storage_handles[loc.path] = Handle(fileobj, mmapobj)

            if loc.path != self.last_path and self._do_prefetch(loc):
                # os.posix_fadvise tells the kernel to prefetch the whole file in memory.
                # API reference: https://docs.python.org/3/library/os.html#os.posix_fadvise
                # According to http://linux.die.net/man/2/posix_fadvise , len=0 (3rd argument)
                # indicates the intention to access the whole file.
                if self.last_path is not None:
                    logging.debug("Telling the kernel that %s will be used instead of %s..." % (loc.path, self.last_path))<commit_msg>Added automatic recovery of usage information<commit_after># cython: profile=True

import os, logging, warnings
import struct, mmap, platform
import datetime, csv
from collections import Counter

from sparsehash cimport SparseHashMap
from utils cimport compactify_deck, canonicalize, State, Location
from utils import hands5_factor_rev, hands5_factor, expand_deck, Handle, load, dump


total_scores = 40
bytes_per_entry = 2
prob_format = '>H'
prob_range = 65534
usage_filename = 'usage.json'
usage_total_key = 'total'
usage_memory_key = 'memory'

dense_ext = 'dat'
sparse_ext = 'shash'

config_filename = 'config.json'
history_filename = 'history.csv'

memory_usage_threshold = 701000000

history_report_period = 10000
small_report_period = 100000
large_report_period = 1000000
save_period = 50000000


def posix_fallocate(path, size):
    with open(path, 'wb') as f:
        os.posix_fallocate(f.fileno(), 0, size)

def win_fallocate(path, size):
    # I haven't tested this.
    with open(path, 'wb') as f:
        f.seek(size - 1)
        f.write(b'\x00')

def posix_fadvise_willneed(fileno):
    os.posix_fadvise(fileno, 0, 0, os.POSIX_FADV_WILLNEED)

def posix_fadvise_dontneed(fileno):
    os.posix_fadvise(fileno, 0, 0, os.POSIX_FADV_DONTNEED)

def win_fadvise_willneed(fileno):
    pass

def win_fadvise_dontneed(fileno):
    pass

if platform.system() == 'Linux':
    fallocate = posix_fallocate
    fadvise_willneed = posix_fadvise_willneed
    fadvise_dontneed = posix_fadvise_dontneed
elif platform.system() == 'Windows':
    logging.warning("Windows file interaction functions are untested.")
    fallocate = win_fallocate
    fadvise_willneed = win_fadvise_willneed
    fadvise_dontneed = win_fadvise_dontneed


cdef class Storage:
    cdef str usage_path, history_path, config_path, curr_path, last_path
    cdef list state_dirs, storage_path, memory_storage, history
    cdef bytes curr_value
    cdef dict storage_handles, config
    cdef object storage_usage
    cdef int curr_action
    cdef long curr_offset
    def __init__(self, state_dirs):
        if len(state_dirs) == 1:
            warnings.warn("It is inefficient to use only one storage drive. Add another one.")

        self.state_dirs = state_dirs
        self.storage_handles = {}

        self.usage_path = os.path.join(state_dirs[0], usage_filename)
        if os.path.exists(self.usage_path):
            with open(self.usage_path, 'r') as f:
                self.storage_usage = Counter(load(f))
        else:
            self.storage_usage = Counter()

        self.config_path = os.path.join(state_dirs[0], config_filename)
        if os.path.exists(self.config_path):
            with open(self.config_path, 'r') as f:
                self.config = load(f)
        else:
            self.config = {
                'in_memory': [True] * total_scores
            }

        self.storage_path = [
            os.path.join(self.state_dirs[score % len(state_dirs)],
                         '%02d.%s' % (score, sparse_ext if self.config['in_memory'][score] else dense_ext))
            for score in range(total_scores)
        ]

        self.history_path = os.path.join(state_dirs[0], history_filename)
        self.history = []
        self.register_history('startup')

        # In case the generator is aborted, I want to do at least something to preserve consistency.
        # When I/O is performed, this variable stores the current I/O action.
        # There are 2 types of actions available: writing values in existing files and creating new files.
        # -1 = nothing
        # 0 = writing value
        # 1 = creating file
        # 2 = transferring data from memory storage to disk
        # I'm missing Haskell's algebraic data types.
        self.curr_action = -1

        self.curr_value = None
        self.curr_offset = -1
        self.curr_path = None

        # Before accessing each file I call os.posix_fadvise to load the whole file in memory.
        # Since I don't want to do that every time I store or retrieve data, I save the last
        # accessed file and only call os.posix_fadvise when it changes.
        self.last_path = None

        # This is the storage for the data that resides in memory.
        self.memory_storage = [None] * total_scores

        logging.info("Currently using %d entries in the table" % self.storage_usage[usage_total_key])
        self.report_breakdown()

    cdef long _get_offset(self, State cstate) except -1:
        cdef long hand_offset, deck_offset, offset
        hand_offset = hands5_factor_rev[cstate.hand]
        deck_offset = compactify_deck(cstate.hand, cstate.deck)

        # Multiplication by 2**19. Note: without the parentheses
        # this is interpreted as `offset = hand_offset << (19 + deck_offset)`
        offset = (hand_offset << 19) + deck_offset

        return offset

    cdef _unpack_offset(self, offset):
        # bottom 19 bits: compactified deck
        # the rest: hands5_factor index
        hand_idx = offset >> 19
        compact_deck = offset & ((1 << 19) - 1)
        hand = hands5_factor[hand_idx]
        deck = expand_deck(hand, compact_deck)
        return hand, deck

    cdef int _state_is_in_memory(self, State cstate) except -1:
        return self.config['in_memory'][cstate.score]

    cdef int _do_prefetch(self, Location loc) except -1:
        # TODO: this is ugly, but is there a better way?
        return loc.path == self.storage_path[0]

    cdef _reset_storage(self, loc):
        # assumes that the path exists but the file doesn't

        file_size = loc.size * bytes_per_entry
        logging.info("Creating a new table at %s for %d entries (%d bytes)." % (loc.path, loc.size, file_size))

        self.curr_action = 1
        self.curr_path = loc.path

        fallocate(loc.path, file_size)

        self.curr_action = -1
        self.curr_path = None

        logging.info("Done.")

    cdef _ensure_usable(self, path):
        directory = os.path.dirname(path)
        if not os.path.exists(directory):
            os.makedirs(directory)
        elif not os.path.isdir(directory):
            raise RuntimeError("%s exists but is not a directory" % directory)

    cdef _ensure_initialized(self, Location loc):
        if loc.in_memory:
            if self.memory_storage[loc.idx] is None:
                self._ensure_usable(loc.path)

                if os.path.exists(loc.path):
                    logging.debug("Loading %s from disk to memory..." % loc.path)
                    # Loading from on-disk storage
                    with open(loc.path, 'rb') as f:
                        self.memory_storage[loc.idx] = SparseHashMap(f)
                    # TODO: can I make it look better?
                    recorded_usage = self.storage_usage[loc.path]
                    actual_usage = len(self.memory_storage[loc.idx])
                    if (
                            (self.config.get('fix_usage', False) is True) and
                            (recorded_usage != actual_usage)
                        ):
                        new_total = sum(value for key, value in self.storage_usage if key not in (usage_total_key, usage_memory_key))
                        new_memory = sum(value for key, value in self.storage_usage if key not in (usage_total_key, usage_memory_key) and self.config['in_memory'][self.storage_path.index(key)])

                        logging.warning("Fixing usage of %s. Own: %d->%d (%d). Total: %d->%d (%d). Memory: %d->%d (%d)." %
                            (
                                loc.path, recorded_usage, actual_usage, actual_usage - recorded_usage,
                                self.storage_usage[usage_total_key], new_total, new_total - self.storage_usage[usage_total_key],
                                self.storage_usage[usage_memory_key], new_memory, new_memory - self.storage_usage[usage_memory_key]
                            )
                        )

                        self.storage_usage[loc.path] = actual_usage
                        self.storage_usage[usage_total_key] = new_total
                        self.storage_usage[usage_memory_key] = new_memory
                    else:
                        assert recorded_usage == actual_usage
                else:
                    self.memory_storage[loc.idx] = SparseHashMap()
        else:

            if loc.path not in self.storage_handles:
                self._ensure_usable(loc.path)

                if not os.path.exists(loc.path):
                    self._reset_storage(loc)

                assert loc.path not in self.storage_handles

                fileobj = open(loc.path, 'r+b')
                mmapobj = mmap.mmap(fileobj.fileno(), loc.size * bytes_per_entry)

                self.storage_handles[loc.path] = Handle(fileobj, mmapobj)

            if loc.path != self.last_path and self._do_prefetch(loc):
                # os.posix_fadvise tells the kernel to prefetch the whole file in memory.
                # API reference: https://docs.python.org/3/library/os.html#os.posix_fadvise
                # According to http://linux.die.net/man/2/posix_fadvise , len=0 (3rd argument)
                # indicates the intention to access the whole file.
                if self.last_path is not None:
                    logging.debug("Telling the kernel that %s will be used instead of %s..." % (loc.path, self.last_path))
{'fraction_non_alphanumeric': 0.08886056884666628, 'fraction_numerical': 0.010890343509239413, 'mean_word_length': 2.3334620583124157, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 12, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3166801', 'n_tokens_mistral': 5425, 'n_tokens_neox': 5095, 'n_words': 1386}