# Copyright 2018 VMware, Inc.
# All rights reserved. -- VMware Confidential

"""Socket locks

"Abstract sockets automatically disappear when all open references to the
socket are closed" (man 7 unix). This works well with Python's contextmanager.

On ESXi this is the default behaviour for unix sockets w/o the need for a '\0'.

Be aware: Locking with abstract sockets needs to be used by all users in order
to synchronize access due to it is not a filesystem-supported locking mechanism.
However, the nature of abstract sockets makes sure that we cannot have a stale
lock if the process dies or gets killed.
"""

from contextlib import contextmanager
import errno
from hashlib import sha256
import os
import socket
from time import sleep

if os.name == 'posix' and os.uname()[0] == 'VMkernel':
   # ESXi
   ABSTRACT_SOCKET_PREFIX = '/com/vmware/'
else:
   # Linux
   ABSTRACT_SOCKET_PREFIX = '\0/com/vmware/'

class LockAcquisitionError(OSError):
   pass

@contextmanager
def acquire(socketName, blocking=False):
   """Acquire an abstract socket within a context manager.
   """
   with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
      sock.setblocking(blocking)
      # sha256sum in order to control socket name length
      lockId = os.path.join(ABSTRACT_SOCKET_PREFIX,
                            sha256(socketName.encode()).hexdigest())
      try:
         sock.bind(lockId)
      except OSError as e:
         if e.errno == errno.EADDRINUSE or e.errno == errno.EEXIST:
            raise LockAcquisitionError(socketName + " => " + str(lockId))
         raise
      yield sock

@contextmanager
def openWithLock(fileName, mode='r', blocking=False):
   """Open file with an abstract socket lock.
   """
   with acquire(os.path.realpath(fileName)) as lock:
      with open(fileName, mode) as f:
         yield f

@contextmanager
def openWithLockAndRetry(fileName, mode='r', retry=10, retrySleep=0.1):
   """Open a file with an abstract socket lock and retry on failure.
   """
   while retry > 0:
      try:
         with openWithLock(fileName, mode) as f:
            yield f
      except LockAcquisitionError:
         retry -= 1
         if retry <= 0:
            raise
         sleep(retrySleep)
      else:
         break

