diff --git a/chat_archiver/chat_archiver/main.py b/chat_archiver/chat_archiver/main.py index 7c7ab56..0dd2edb 100644 --- a/chat_archiver/chat_archiver/main.py +++ b/chat_archiver/chat_archiver/main.py @@ -13,12 +13,11 @@ from calendar import timegm from collections import defaultdict from datetime import datetime from itertools import count -from uuid import uuid4 import gevent.event import gevent.queue -from common import ensure_directory, listdir, rename +from common import atomic_write, listdir from common.chat import BATCH_INTERVAL, format_batch, get_batch_files, merge_messages from girc import Client @@ -275,14 +274,6 @@ class Archiver(object): self.client.stop() -def atomic_write(filepath, content): - temp_path = "{}.{}.temp".format(filepath, uuid4()) - ensure_directory(filepath) - with open(temp_path, 'wb') as f: - f.write(content) - rename(temp_path, filepath) - - _EMOTES_RUNNING = {} # map (base_dir, emote id, theme, scale) -> in-progress greenlet fetching that path def ensure_emotes(base_dir, emote_ids): """Tries to download given emote from twitch if it doesn't already exist. diff --git a/common/common/__init__.py b/common/common/__init__.py index 8f288c2..616d5b1 100644 --- a/common/common/__init__.py +++ b/common/common/__init__.py @@ -6,6 +6,7 @@ import logging import os import random from signal import SIGTERM +from uuid import uuid4 import gevent.event @@ -128,6 +129,21 @@ def writeall(write, value): value = value[n:] +def atomic_write(filepath, content): + """Writes content to filepath atomically, ie. replacing the file in one step + without potential for partial write. content may be str or bytes. + If the file already exists, it will silently do nothing as it is assumed a given + filename can only ever contain the same content. + """ + if isinstance(content, str): + content = content.encode("utf-8") + temp_path = "{}.{}.temp".format(filepath, uuid4()) + ensure_directory(filepath) + with open(temp_path, 'wb') as f: + writeall(f.write, content) + rename(temp_path, filepath) + + def serve_with_graceful_shutdown(server, stop_timeout=20): """Takes a gevent.WSGIServer and serves forever until SIGTERM is received, or the server errors. This is slightly tricky to do due to race conditions