diff --git a/lib/uidpool.py b/lib/uidpool.py index d216ec5cd0f9bb54dc0ceba3d1dde627c6466080..21c9113699f8f04be08bd43a1d3db15a7493f9f0 100644 --- a/lib/uidpool.py +++ b/lib/uidpool.py @@ -328,3 +328,31 @@ def ReleaseUid(uid): except OSError, err: raise errors.LockError("Failed to remove user-id lockfile" " for user-id %s: %s" % (uid, err)) + + +def ExecWithUnusedUid(fn, all_uids, *args, **kwargs): + """Execute a callable and provide an unused user-id in its kwargs. + + This wrapper function provides a simple way to handle the requesting, + unlocking and releasing a user-id. + "fn" is called by passing a "uid" keyword argument that + contains an unused user-id (as a string) selected from the set of user-ids + passed in all_uids. + If there is an error while executing "fn", the user-id is returned + to the pool. + + @param fn: a callable + @param all_uids: a set containing all user-ids in the user-id pool + + """ + uid = RequestUnusedUid(all_uids) + kwargs["uid"] = str(uid) + try: + return_value = fn(*args, **kwargs) + except: + # The failure of "callabe" means that starting a process with the uid + # failed, so let's put the uid back into the pool. + ReleaseUid(uid) + raise + uid.Unlock() + return return_value