source: fedd/federation/synch_store.py @ b745876

#!/usr/local/bin/python

from threading import Lock, Condition

class synch_store:
    """
    Serialized write-once dictionary.  Threads reading before the value for a
    key is assigned wait for the assignment, though polling is an option.  Keys
    may also be deleted, though deleting a value too closely after the set may
    deadlock threads waiting for the assignment.
    """

    # Thrown if a key is assigned more than once
    class CollisionError(Exception): pass

    # Thrown if a non-existant key is deleted
    class BadDeletionError(Exception): pass

    # Thrown if a key is specified with a white space
    class BadKeyError(Exception): pass

    # Thrown when a variable has been revoked.
    class RevokedKeyError(Exception): pass

    # A value in the dict.  Each has its own condition (and lock) for threads
    # to wait on.
    class synch_value:
        # The synch_value constructor. 
        def __init__(self, value=None):
            self.value = value
            self.condition = Condition()

        # Debugging representation
        def __str__(self):
            return "value %s cond %s" % (self.value, self.condition)

    # The synch_store constructor.  Note that the store has its own lock for
    # serializing additions to the dict.
    def __init__(self):
        self.main_lock = Lock()
        self.values = { }
        self.revoked = set([])

    # Set the value in the dict.  If threads are waiting, wake them.
    def set_value(self, key, value):
        if key.find(' ') != -1:
            raise self.BadKeyError("Whitespace is not allowed in keys")

        self.main_lock.acquire()
        if key in self.revoked:
            self.main_lock.release()
            raise synch_store.RevokedKeyError("Synch key %s revoked" % key)
        if self.values.has_key(key):
            v = self.values[key]
            v.condition.acquire()
            self.main_lock.release()
            if v.value is None:
                v.value = value
                v.condition.notifyAll()
                v.condition.release()
            else:
                v.condition.release()
                raise synch_store.CollisionError("%s already set" % key)
        else:
            self.values[key] = synch_store.synch_value(value)
            self.main_lock.release()

    # Ask for the value.  If wait is true, wait until the key is set, if it is
    # not already.  If wait is false, poll.
    def get_value(self, key, wait=True):
        self.main_lock.acquire()
        if key in self.revoked:
            self.main_lock.release()
            raise synch_store.RevokedKeyError("Synch key %s revoked" % key)
        if not self.values.has_key(key):
            self.values[key] = synch_store.synch_value()
        v = self.values[key]
        v.condition.acquire()
        self.main_lock.release()

        while v.value is None:
            if wait: 
                v.condition.wait()
                v.condition.release()
                self.main_lock.acquire()
                if key in self.revoked:
                    self.main_lock.release()
                    raise synch_store.RevokedKeyError("Synch key %s revoked" % key)
                else:
                    self.main_lock.release()
                v.condition.acquire()
            else:
                break
        rv = v.value
        v.condition.release()
        return rv

    # Remove the key from the shared store.  If the key is not present or has
    # never been set, fail.  Otherwise, wake any waiting threads (there
    # shouldn't be any) and remove the value and key from the store.  Deleteing
    # a revoked key removes its revocation.
    def del_value(self, key):
        self.main_lock.acquire()
        if key in self.revoked:
            self.revoked.remove(key)
            self.main_lock.release()
        else:
            if self.values.has_key(key):
                v = self.values[key]
                v.condition.acquire()
                if v.value is not None:
                    del self.values[key]
                    self.main_lock.release()
                    v.condition.notifyAll()
                    v.condition.release()
                else:
                    self.main_lock.release()
                    v.condition.release()
                    raise synch_store.BadDeletionError("%s has not been set" \
                            % key)
            else:
                self.main_lock.release()
                raise synch_store.BadDeletionError("%s does not exist yet" \
                        % key)

    # Invlaidate this synchronization point.  Those waiting for it will get
    # exceptions, and those trying to wait for it in the future will as well.
    def revoke_key(self, key):
        self.main_lock.acquire()
        self.revoked.add(key)
        if self.values.has_key(key):
            v = self.values[key]
            del self.values[key]
            self.main_lock.release()
            v.condition.acquire()
            v.condition.notifyAll()
            v.condition.release()
        else:
            self.main_lock.release()


    def all_keys(self):
        self.main_lock.acquire()
        rv = [ k for k in self.values.keys()]
        self.main_lock.release()
        return rv

    # Save the store contents to a file in a simple way
    def save(self, filename):
        self.main_lock.acquire()
        f = open(filename, 'w')
        for k, v in self.values.items():
            self.values[k].condition.acquire()
            if self.values[k].value is not None:
                print >> f, "%s %s" % ( k, str(self.values[k].value))
            else:
                print >> f, "%s" % k
            self.values[k].condition.release()
        self.main_lock.release()
        f.close()

    # Load the store contents from a file
    def load(self, filename):
        self.main_lock.acquire()
        self.values = { }
        f = open(filename, 'r')
        for l in f:
            i = l.find(' ')
            if i != -1: self.values[l[:i]] = synch_store.synch_value(l[i+1:])
            else: self.values[l] = synch_store.synch_value()
        f.close()
        self.main_lock.release()

    # A debugging representation
    def __str__(self):
        rv = ""
        self.main_lock.acquire()
        for k, v in self.values.items():
            self.values[k].condition.acquire()
            rv += "%s: %s\n" % ( k, str(self.values[k]))
            self.values[k].condition.release()
        self.main_lock.release()
        return rv

#debugging scaffolding, left here in case we have problems
if __name__ == '__main__':
    from threading import currentThread, Thread
    from time import sleep

    def get_val(k, ss, wait=True):
        print "%s: getting %s" % (currentThread(), k)
        try:
            v = ss.get_value(k, wait)
        except synch_store.RevokedKeyError, e:
            print "Error: %s" %e
            return
        print "%s: %s is %s" % (currentThread(), k, v)


    def set_val(k, v, ss):
        print "setting %s to %s" % (k, v)
        try:
            ss.set_value(k, v)
        except synch_store.CollisionError, e:
            print "Error: %s" %e
            return
        except synch_store.RevokedKeyError, e:
            print "Error: %s" %e
            return
        print "Done with set"

    def del_val(k, ss):
        print "deleteing %s" % k
        try:
            ss.del_value(k)
        except synch_store.BadDeletionError, e:
            print "Error: %s" %e
            return
        except synch_store.RevokedKeyError, e:
            print "Error: %s" %e
            return
        print "Done with del"

    def rev_val(k, ss):
        print "revoking %s" % k
        try:
            ss.revoke_key(k)
        except synch_store.BadDeletionError, e:
            print "Error: %s" %e
            return
        except synch_store.RevokedKeyError, e:
            print "Error: %s" %e
            return
        print "Done with revoking"


    ss = synch_store()
    ss.set_value('b', 'initial')
    print ss

    for i in range(0,10):
        t = Thread(target=get_val, args=('a', ss, i % 2 == 1))
        t.start()
    t = Thread(target=set_val, args=('a', '1', ss))
    t.start()
    for i in range(0,10):
        t = Thread(target=get_val, args=('a', ss))
        t.start()

    print ss
    t = Thread(target=del_val, args=('a', ss))
    t.start()
    print ss

    for i in range(0,10):
        t = Thread(target=get_val, args=('c', ss))
        t.start()

    
    t = Thread(target=rev_val, args=('c', ss))
    t.start()