from contextlib import contextmanager
from celery import shared_task
from celery.five import monotonic
from django.core.cache import cache
from mrbelvedereci.build.models import Build
from mrbelvedereci.salesforce.models import Org

LOCK_EXPIRE = 60 * 120  # Lock expires in 2 hours

@contextmanager
def lock_org(lock_id, oid):
    timeout_at = monotonic() + LOCK_EXPIRE - 3
    # cache.add fails if the key already exists
    status = cache.add(lock_id, oid, LOCK_EXPIRE)
    try:
        yield status
    finally:
        # memcache delete is very slow, but we have to use it to take
        # advantage of using add() for atomic locking
        if monotonic() < timeout_at:
            # don't release the lock if we exceeded the timeout
            # to lessen the chance of releasing an expired lock
            # owned by someone else.
            cache.delete(lock_id)

@shared_task
def run_build(build_id):
    build = Build.objects.get(id=build_id)
    build.run()
    return build.status

@shared_task(bind=True)
def check_queued_build(self, build_id):
    build = Build.objects.get(id = build_id)

    # Check for concurrency blocking
    try:
        org = Org.objects.get(name = build.trigger.org, repo = build.repo)
    except Org.DoesNotExist:
        return

    # If this is not a scratch org, ensure no builds are currently running against the org
    if not org.scratch:
        lock_id = 'mrbelvedereci-org-lock-{}'.format(org.id)
        with lock_org(lock_id, self.app.oid) as acquired:
            if acquired:
                # Run the build
                run_build.apply_async((build.id,), countdown=1)
            else:
                # Queue next check in 5 seconds
                check_queued_build.apply_async((build.id,), countdown=5)
    # For scratch orgs, don't worry about locking, just run the build
    else:
        run_build.apply_async((build.id,), countdown=1)

@@ -3,104 +3,32 @@
 om c
-ontextlib import contextmanager%0Afrom celery import shared_task%0Afrom celery.five import monotonic
+elery import shared_task
 %0Afro
@@ -158,515 +158,202 @@
 rg%0A%0A
-LOCK_EXPIRE = 60 * 120  # Lock expires in 2 hours%0A%0A@contextmanager%0Adef lock_org(lock_id, oid):%0A    timeout_at = monotonic() + LOCK_EXPIRE - 3%0A    # cache.add fails if the key already exists%0A    status = cache.add(lock_id, oid, LOCK_EXPIRE)%0A    try:%0A        yield status%0A    finally:%0A        # memcache delete is very slow, but we have to use it to take%0A        # advantage of using add() for atomic locking%0A        if monotonic() %3C timeout_at:%0A            # don't release the lock if we exceeded the timeout
+@shared_task%0Adef run_build(build_id, lock_id=None):%0A    try:%0A        build = Build.objects.get(id=build_id)%0A        build.run()%0A%0A    except:%0A        if lock_id:%0A            cache.delete(lock_id)
 %0A   
@@ -361,232 +361,70 @@
     
+%0A
     
-# to lessen the chance of releasing an expired lock%0A            # owned by someone else.%0A            cache.delete(lock_id)%0A%0A@shared_task%0Adef run_build(build_id):%0A    build = Build.objects.get(id=build_id)%0A    build.run()
+    raise%0A%0A    if lock_id:%0A        cache.delete(lock_id)%0A
 %0A   
@@ -729,26 +729,37 @@
     
-# If this is not a
+if org.scratch:%0A        # For
  scr
@@ -770,88 +770,203 @@
  org
+s
 , 
-ensure no builds are currently running against the org%0A    if not org.scratch:
+we don't need concurrency blocking logic, just run the build%0A        run_build.apply_async((build.id,), countdown=1)%0A    %0A    else:%0A        # For persistent orgs, use the cache to lock the org
 %0A   
@@ -1035,62 +1035,55 @@
     
-with lock_org(lock_id, self.app.oid) as acquired:%0A
+status = cache.add(lock_id, self.oid)%0A    
     
+%0A
     
@@ -1093,37 +1093,56 @@
  if 
-acquired:%0A                # R
+status == True:%0A            # Lock successful, r
 un t
@@ -1146,28 +1146,24 @@
 n the build%0A
-    
             
@@ -1186,32 +1186,39 @@
 async((build.id,
+lock_id
 ), countdown=1)%0A
@@ -1225,20 +1225,16 @@
         
-  
-  
 else:%0A  
@@ -1247,15 +1247,31 @@
     
-    # Q
+# Failed to get lock, q
 ueue
@@ -1307,20 +1307,16 @@
         
-   
- 
 check_qu
@@ -1368,140 +1368,4 @@
 =5)%0A
-    # For scratch orgs, don't worry about locking, just run the build%0A    else:%0A        run_build.apply_async((build.id,), countdown=1)%0A

Change lock manager to live inside method