Transactional rename on Windows

client/sources/ok_test/models.py

@@ -3,6 +3,7 @@
 from client.sources.common import models
 from client.utils import format
 from client.utils import output
+from client.utils import storage
 import os
 
 ##########
@@ -179,15 +180,19 @@ def dump(self):
         with open(test_tmp, 'w', encoding='utf-8') as f:
             f.write('test = {}\n'.format(json))
 
-        # Try to use os.replace, but if on Windows manually remove then rename
-        # (ref issue #339)
-        if os.name == 'nt':
-        # TODO(colin) Add additional error handling in case process gets killed mid remove/rename
-            os.remove(self.file)
-            os.rename(test_tmp, self.file)
-        else:
-            # Use an atomic rename operation to prevent test corruption
-            os.replace(test_tmp, self.file)
+        try:
+            storage.replace_transactional(test_tmp, self.file)
+            success = True
+        except (NotImplementedError, OSError):
+            # Try to use os.replace, but if on Windows manually remove then rename
+            # (ref issue #339)
+            if os.name == 'nt':
+            # TODO(colin) Add additional error handling in case process gets killed mid remove/rename
+                os.remove(self.file)
+                os.rename(test_tmp, self.file)
+            else:
+                # Use an atomic rename operation to prevent test corruption
+                os.replace(test_tmp, self.file)
 
     @property
     def unique_id_prefix(self):

client/utils/storage.py

@@ -1,6 +1,23 @@
+import ctypes
 import shelve # persistance
 import hmac # security
 
+def set_foreign_function_type(func, restype, argtypes):
+    if func.argtypes is None:
+        func.argtypes = argtypes
+        func.restype = restype
+
+# Platform-specific imports
+windll = None
+try:
+    windll = ctypes.windll
+    from ctypes.wintypes import BOOL, BOOLEAN, BYTE, DWORD, HANDLE, LARGE_INTEGER, LPCWSTR, LPVOID, ULONG, WCHAR
+    set_foreign_function_type(windll.ktmw32.CreateTransaction, HANDLE, [LPVOID, LPVOID, DWORD, DWORD, DWORD]),
+    set_foreign_function_type(windll.ktmw32.CommitTransaction, BOOL, [HANDLE]),
+    set_foreign_function_type(windll.kernel32.MoveFileTransactedW, BOOL, [LPCWSTR, LPCWSTR, ctypes.WINFUNCTYPE(LARGE_INTEGER, LARGE_INTEGER, LARGE_INTEGER, LARGE_INTEGER, DWORD, DWORD, HANDLE, HANDLE, LPVOID), LPVOID, DWORD, HANDLE]),
+    set_foreign_function_type(windll.kernel32.CloseHandle, BOOL, [HANDLE]),
+except (AttributeError, ImportError, OSError): pass
+
 ##################
 # Secure Storage #
 ##################
@@ -33,3 +50,18 @@ def get(root, key, default=None):
         if not hmac.compare_digest(data['mac'], mac(data['value'])):
             raise ProtocolException('{} was tampered.  Reverse changes, or redownload assignment'.format(SHELVE_FILE))
     return data['value']
+
+def replace_transactional(source, destination):
+    # Like os.replace, but tries to be actually atomic when possible on Windows.
+    if windll:
+        error_code = 50  # ERROR_NOT_SUPPORTED
+        if windll.ktmw32:
+            tx = windll.ktmw32.CreateTransaction(None, None, 0, 0, 0)
+            if tx != HANDLE(-1).value:
+                try: error_code = 0 if windll.kernel32.MoveFileTransactedW(source, destination, windll.kernel32.MoveFileTransactedW.argtypes[2](), None, 0x1 | 0x2, tx) and windll.ktmw32.CommitTransaction(tx) else ctypes.GetLastError()
+                finally: windll.kernel32.CloseHandle(tx)
+            else: error_code = ctypes.GetLastError()
+        if error_code:
+            raise ctypes.WinError(error_code)
+    else:
+        raise NotImplementedError("transactional file systems not supported")