Untitled

mail@pastecode.io avatar
unknown
plain_text
10 days ago
4.2 kB
3
Indexable
Never
from collections import deque
from in_memory_db import InMemoryDB

class InMemoryDBImpl(InMemoryDB):
    def __init__(self):
        self.db = {}
        self.modification_counts = {}  # To track the number of modifications
        self.locks = {}  # To track which user has locked a record
        self.lock_queues = {}  # To track waiting users for a lock

    def set_or_inc(self, key: str, field: str, value: int) -> int:
        if key in self.locks:
            return None  # Ignore if locked and no caller_id provided

        return self._set_or_inc_internal(key, field, value)

    def set_or_inc_by_caller(self, key: str, field: str, value: int, caller_id: str) -> int | None:
        if key in self.locks and self.locks[key] != caller_id:
            return self.get(key, field)  # Return current value if locked by another user
        
        return self._set_or_inc_internal(key, field, value)

    def _set_or_inc_internal(self, key: str, field: str, value: int) -> int:
        if key not in self.db:
            self.db[key] = {}
            self.modification_counts[key] = 0

        if field in self.db[key]:
            self.db[key][field] += value
        else:
            self.db[key][field] = value
        
        # Increment the modification count whenever a field is set or incremented
        self.modification_counts[key] += 1
        
        return self.db[key][field]

    def get(self, key: str, field: str) -> int | None:
        if key in self.db and field in self.db[key]:
            return self.db[key][field]
        return None

    def delete(self, key: str, field: str) -> bool:
        if key in self.locks:
            return False  # Ignore if locked and no caller_id provided

        return self._delete_internal(key, field)

    def delete_by_caller(self, key: str, field: str, caller_id: str) -> bool:
        if key in self.locks and self.locks[key] != caller_id:
            return False  # Ignore if locked by another user

        return self._delete_internal(key, field)

    def _delete_internal(self, key: str, field: str) -> bool:
        if key in self.db and field in self.db[key]:
            del self.db[key][field]
            
            # Increment the modification count whenever a field is deleted
            self.modification_counts[key] += 1
            
            if not self.db[key]:  # If the record has no fields left
                del self.db[key]
                del self.modification_counts[key]
                if key in self.locks:
                    del self.locks[key]
                if key in self.lock_queues:
                    del self.lock_queues[key]
            return True
        return False

    def top_n_keys(self, n: int) -> list[str]:
        sorted_keys = sorted(
            self.modification_counts.items(), 
            key=lambda item: (-item[1], item[0])  # Sort by modification count desc, then by key name asc
        )
        return [f"{key}({count})" for key, count in sorted_keys[:n]]

    def lock(self, caller_id: str, key: str) -> str | None:
        if key not in self.db:
            return "invalid_request"
        
        if key in self.locks:
            if self.locks[key] == caller_id:
                return None  # Already locked by the same user
            if caller_id in self.lock_queues.get(key, deque()):
                return None  # Already in the queue

            if key not in self.lock_queues:
                self.lock_queues[key] = deque()
            self.lock_queues[key].append(caller_id)
            return "wait"
        
        # Lock the key for the caller
        self.locks[key] = caller_id
        return "acquired"

    def unlock(self, key: str) -> str | None:
        if key not in self.locks:
            return None
        
        released_user = self.locks.pop(key)
        if key not in self.lock_queues or not self.lock_queues[key]:
            # No one is waiting in the queue
            return "released"
        
        # Assign the lock to the next user in the queue
        next_user = self.lock_queues[key].popleft()
        self.locks[key] = next_user
        return "released"

Leave a Comment