Repository: bokeh/bokeh

Issue:
[BUG] Unable to use async callbacks in `add_next_tick_callback` and the like
### ALL software version info (bokeh, python, notebook, OS, browser, any other relevant packages)
Bokeh 2.0.0dev6

#### Description of expected behavior and the observed behavior
It should be possible to use async callbacks in `add_next_tick_callback`. However, such callbacks are never executed and result in `RuntimeWarning: coroutine 'f' was never awaited`.

#### Complete, minimal, self-contained example code that reproduces the issue
Should be used in an app for `bokeh.server.server.Server`:
```python
async def f():
    print('hi')

doc.add_next_tick_callback(f)
```

#### Stack traceback and/or browser JavaScript console output
```
/.../python3.7/asyncio/base_events.py:1760: RuntimeWarning: coroutine 'setup_theta_analysis.<locals>.f' was never awaited
Coroutine created at (most recent call last)
  File "/.../python3.7/site-packages/tornado/platform/asyncio.py", line 148, in start
    self.asyncio_loop.run_forever()
  File "/.../python3.7/asyncio/base_events.py", line 539, in run_forever
    self._run_once()
  File "/.../python3.7/asyncio/base_events.py", line 1767, in _run_once
    handle._run()
  File "/.../python3.7/asyncio/events.py", line 88, in _run
    self._context.run(self._callback, *self._args)
  File "/.../python3.7/site-packages/bokeh/server/session.py", line 65, in _needs_document_lock_wrapper
    result = func(self, *args, **kwargs)
  File "/.../python3.7/site-packages/bokeh/server/session.py", line 181, in with_document_locked
    return func(*args, **kwargs)
  File "/.../python3.7/site-packages/bokeh/document/document.py", line 1160, in wrapper
    return doc._with_self_as_curdoc(invoke)
  File "/.../python3.7/site-packages/bokeh/document/document.py", line 1146, in _with_self_as_curdoc
    return f()
  File "/.../python3.7/site-packages/bokeh/document/document.py", line 1159, in invoke
    return f(*args, **kwargs)
  File "/.../python3.7/site-packages/bokeh/document/document.py", line 949, in remove_then_invoke
    return callback(*args, **kwargs)
  handle = self._ready.popleft()
```


Discussion:

Resolving patch:
diff --git a/bokeh/server/contexts.py b/bokeh/server/contexts.py
--- a/bokeh/server/contexts.py
+++ b/bokeh/server/contexts.py
@@ -198,12 +198,12 @@ async def create_session_if_needed(self, session_id, request=None, token=None):
                                                   self.server_context,
                                                   doc,
                                                   logout_url=self._logout_url)
-
-            # using private attr so users only have access to a read-only property
-            payload = get_token_payload(token) if token else {}
-            session_context._request = _RequestProxy(request,
-                                                     cookies=payload.get('cookies'),
-                                                     headers=payload.get('headers'))
+            if request is not None:
+                payload = get_token_payload(token) if token else {}
+                # using private attr so users only have access to a read-only property
+                session_context._request = _RequestProxy(request,
+                                                         cookies=payload.get('cookies'),
+                                                         headers=payload.get('headers'))
             session_context._token = token
 
             # expose the session context to the document
diff --git a/bokeh/server/session.py b/bokeh/server/session.py
--- a/bokeh/server/session.py
+++ b/bokeh/server/session.py
@@ -19,6 +19,7 @@
 #-----------------------------------------------------------------------------
 
 # Standard library imports
+import inspect
 import time
 
 # External imports
@@ -64,6 +65,10 @@ async def _needs_document_lock_wrapper(self, *args, **kwargs):
                 self._pending_writes = []
                 try:
                     result = func(self, *args, **kwargs)
+                    if inspect.isawaitable(result):
+                        # Note that this must not be outside of the critical section.
+                        # Otherwise, the async callback will be ran without document locking.
+                        result = await result
                 finally:
                     # we want to be very sure we reset this or we'll
                     # keep hitting the RuntimeError above as soon as


Test patch:
diff --git a/tests/unit/bokeh/server/test_contexts.py b/tests/unit/bokeh/server/test_contexts.py
--- a/tests/unit/bokeh/server/test_contexts.py
+++ b/tests/unit/bokeh/server/test_contexts.py
@@ -14,6 +14,12 @@
 # Imports
 #-----------------------------------------------------------------------------
 
+# Standard library imports
+import asyncio
+
+# External imports
+from tornado.ioloop import IOLoop
+
 # Bokeh imports
 from bokeh.application import Application
 
@@ -106,29 +112,20 @@ def test_get_session_failure(self) -> None:
     async def test_create_session_if_needed_new(self) -> None:
         app = Application()
         c = bsc.ApplicationContext(app, io_loop="ioloop")
-        class FakeRequest(object):
-            arguments = dict(foo=10)
-        req = FakeRequest()
-        s = await c.create_session_if_needed("foo", request=req)
+        s = await c.create_session_if_needed("foo")
         assert c.get_session("foo") == s
 
     async def test_create_session_if_needed_exists(self) -> None:
         app = Application()
         c = bsc.ApplicationContext(app, io_loop="ioloop")
-        class FakeRequest(object):
-            arguments = dict(foo=10)
-        req = FakeRequest()
-        s1 = await c.create_session_if_needed("foo", request=req)
-        s2 = await c.create_session_if_needed("foo", request=req)
+        s1 = await c.create_session_if_needed("foo")
+        s2 = await c.create_session_if_needed("foo")
         assert s1 == s2
 
     async def test_create_session_if_needed_bad_sessionid(self) -> None:
         app = Application()
         c = bsc.ApplicationContext(app, io_loop="ioloop")
-        class FakeRequest(object):
-            arguments = dict(foo=10)
-        req = FakeRequest()
-        r = c.create_session_if_needed("", request=req)
+        r = c.create_session_if_needed("")
         with pytest.raises(bsc.ProtocolError) as e:
             await r
         assert str(e.value).endswith("Session ID must not be empty")
@@ -136,14 +133,31 @@ class FakeRequest(object):
     async def test_create_session_if_needed_logout_url(self) -> None:
         app = Application()
         c = bsc.ApplicationContext(app, io_loop="ioloop", logout_url="/logout")
-        class FakeRequest(object):
-            arguments = dict(foo=10)
-        req = FakeRequest()
-        s = await c.create_session_if_needed("foo", request=req)
+        s = await c.create_session_if_needed("foo")
         session = c.get_session("foo")
         assert session == s
         assert c._session_contexts[session.id].logout_url == "/logout"
 
+    async def test_async_next_tick_callback_is_called(self) -> None:
+        app = Application()
+        c = bsc.ApplicationContext(app, io_loop=IOLoop.current())
+
+        s = await c.create_session_if_needed("foo")
+
+        latch_f = asyncio.Future()
+        result_f = asyncio.Future()
+
+        async def cb():
+            m = await latch_f
+            result_f.set_result(m)
+
+        s.document.add_next_tick_callback(cb)
+
+        message = 'Done'
+        latch_f.set_result(message)
+        result = await asyncio.wait_for(result_f, 1)
+        assert result == message
+
 #-----------------------------------------------------------------------------
 # Dev API
 #-----------------------------------------------------------------------------