From c2b9d1426869437a4843d9d4af37444d0d8fa342 Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Sun, 13 Apr 2014 18:15:36 -0700
Subject: [PATCH] Fix up cancel and timer a bit

---
 capnp/helpers/helpers.pxd |  2 +-
 capnp/lib/capnp.pyx       | 29 +++++++++++++++++++++++++++--
 test/test_capability.py   | 22 ++++++++++++++++++++++
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/capnp/helpers/helpers.pxd b/capnp/helpers/helpers.pxd
index fe5b7c6..1e842c2 100644
--- a/capnp/helpers/helpers.pxd
+++ b/capnp/helpers/helpers.pxd
@@ -38,4 +38,4 @@ cdef extern from "../helpers/serialize.h":
 
 cdef extern from "../helpers/asyncHelper.h":
     void waitNeverDone(WaitScope&)
-    Timer * getTimer(AsyncIoContext *)
+    Timer * getTimer(AsyncIoContext *) except +reraise_kj_exception
diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index 0a166ca..b6ad05a 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -1357,7 +1357,7 @@ cdef class Timer:
         self.thisptr = timer
         return self
 
-    cpdef after_delay(self, time):
+    cpdef after_delay(self, time) except +reraise_kj_exception:
         return _VoidPromise()._init(self.thisptr.afterDelay(capnp.Duration(time)))
 
 def getTimer():
@@ -1434,7 +1434,7 @@ cdef class Promise:
         else:
             self.is_consumed = False
             self._obj = obj
-            Py_INCREF(obj)
+            Py_INCREF(obj) # TODO: MEM: fix leak
             self.thisptr = new PyPromise(<PyObject *>obj)
 
         self._event_loop = C_DEFAULT_EVENT_LOOP_GETTER()
@@ -1476,6 +1476,15 @@ cdef class Promise:
 
         return ret
 
+    cpdef cancel(self) except +reraise_kj_exception:
+        if self.is_consumed:
+            raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
+            
+        self.is_consumed = True
+        del self.thisptr
+        self.thisptr = NULL
+
+
 cdef class _VoidPromise:
     cdef VoidPromise * thisptr
     cdef public bint is_consumed
@@ -1522,6 +1531,14 @@ cdef class _VoidPromise:
 
         return ret
 
+    cpdef cancel(self) except +reraise_kj_exception:
+        if self.is_consumed:
+            raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
+            
+        self.is_consumed = True
+        del self.thisptr
+        self.thisptr = NULL
+
 cdef class _RemotePromise:
     cdef RemotePromise * thisptr
     cdef public bint is_consumed
@@ -1589,6 +1606,14 @@ cdef class _RemotePromise:
     def to_dict(self, verbose=False):
         return _to_dict(self, verbose)
 
+    cpdef cancel(self) except +reraise_kj_exception:
+        if self.is_consumed:
+            raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
+            
+        self.is_consumed = True
+        del self.thisptr
+        self.thisptr = NULL
+
     # def attach(self, *args):
     #     if self.is_consumed:
     #         raise ValueError('Promise was already used in a consuming operation. You can no longer use this Promise object')
diff --git a/test/test_capability.py b/test/test_capability.py
index b5de8a6..6f1e79d 100644
--- a/test/test_capability.py
+++ b/test/test_capability.py
@@ -1,6 +1,7 @@
 import pytest
 import capnp
 import os
+import time
 
 import test_capability_capnp as capability
 
@@ -257,6 +258,19 @@ def test_tail_call():
     assert caller_server.count == 1
 
 
+def test_cancel():
+    client = capability.TestInterface._new_client(Server())
+
+    req = client._request('foo')
+    req.i = 5
+
+    remote = req.send()
+    remote.cancel()
+
+    with pytest.raises(ValueError):
+        remote.wait()
+
+
 def test_timer():
     global test_timer_var
     test_timer_var = False
@@ -267,3 +281,11 @@ def test_timer():
     capnp.getTimer().after_delay(1).then(set_timer_var).wait()
 
     assert test_timer_var is True
+
+    promise = capnp.Promise(0).then(lambda x: time.sleep(.1)).then(lambda x: time.sleep(.1))
+
+    canceller = capnp.getTimer().after_delay(1000).then(lambda: promise.cancel())
+
+    joined = capnp.join_promises([promise, canceller])
+    with pytest.raises(Exception):
+        joined.wait()