hf_sync_journal.py

"""同目录HF待同步写前记录；调用方持业务文件锁，提交方另持公共提交锁。"""
import hashlib
import json
import os
from pathlib import Path
import stat
import tempfile
import uuid


class SyncJournal:
    def __init__(self,directory):
        self.directory=Path(directory)/'.hf-pending'

    @staticmethod
    def validate_name(name):
        if (not isinstance(name,str) or not name or name in {'.','..'}
                or os.path.normcase(name)==os.path.normcase('__hf_dataset_sync__')
                or any(char in name for char in '/\\:') or any(ord(char)<32 for char in name)):
            raise ValueError('invalid sync file name')

    def path(self,name):
        self.validate_name(name)
        key=hashlib.sha256(os.path.normcase(name).encode('utf-8')).hexdigest()
        return self.directory/(key+'.json')

    @staticmethod
    def check_regular(path,directory=False):
        info=os.lstat(path)
        if (stat.S_ISLNK(info.st_mode) or getattr(info,'st_file_attributes',0)&1024
                or not (stat.S_ISDIR(info.st_mode) if directory else stat.S_ISREG(info.st_mode))):
            raise ValueError('invalid sync journal path')

    def mark(self,name):
        target=self.path(name)
        self.directory.mkdir(exist_ok=True)
        self.check_regular(self.directory,directory=True)
        if target.exists() or target.is_symlink(): self.check_regular(target)
        generation=uuid.uuid4().hex
        data=json.dumps({'schema':1,'file':name,'generation':generation},ensure_ascii=False).encode('utf-8')
        if len(data)>4096: raise ValueError('sync journal record too large')
        descriptor,temporary=tempfile.mkstemp(prefix='.pending-',suffix='.tmp',dir=self.directory)
        try:
            with os.fdopen(descriptor,'wb') as stream:
                stream.write(data); stream.flush(); os.fsync(stream.fileno())
            os.replace(temporary,target)
        finally:
            try: os.unlink(temporary)
            except FileNotFoundError: pass
        return generation

    def read_record(self,path):
        self.check_regular(self.directory,directory=True)
        self.check_regular(path)
        descriptor=os.open(path,os.O_RDONLY|getattr(os,'O_NOFOLLOW',0))
        with os.fdopen(descriptor,'rb') as stream:
            if not stat.S_ISREG(os.fstat(stream.fileno()).st_mode): raise ValueError('invalid sync journal file')
            content=stream.read(4097)
        if len(content)>4096: raise ValueError('sync journal record too large')
        data=json.loads(content)
        if (not isinstance(data,dict) or set(data)!={'schema','file','generation'}
                or type(data['schema']) is not int or data['schema']!=1):
            raise ValueError('invalid sync journal schema')
        self.validate_name(data['file'])
        generation=data['generation']
        if (not isinstance(generation,str) or len(generation)!=32
                or any(char not in '0123456789abcdef' for char in generation)
                or self.path(data['file'])!=path):
            raise ValueError('invalid sync journal identity')
        return data

    def read(self,name):
        try: return self.read_record(self.path(name))['generation']
        except FileNotFoundError: return None

    def pending(self,on_invalid=None):
        """默认严格读取；同步方可隔离坏记录，目录本身故障仍整体拒绝。"""
        if not self.directory.exists():
            if self.directory.is_symlink(): raise ValueError('invalid sync journal directory')
            return set()
        self.check_regular(self.directory,directory=True)
        names=set()
        for path in self.directory.glob('*.json'):
            try: names.add(self.read_record(path)['file'])
            except FileNotFoundError: pass  # 其他提交方已确认该记录。
            except (OSError,ValueError) as error:
                if on_invalid is None: raise
                on_invalid(error)
        return names

    def acknowledge(self,name,generation):
        if (not isinstance(generation,str) or len(generation)!=32
                or any(char not in '0123456789abcdef' for char in generation)):
            return False
        if self.read(name)!=generation: return False
        os.unlink(self.path(name))
        return True