Repository: tornadoweb/tornado

Issue:
gen.WaitIterator mess up current_index
Hi, if we have some coroutines to return very quickly (using unittest mock, or quick sql query),
the result of current_index may be different in the `async for` loop, which will have less result return

currently I need to use `gen.Multi` to avoid this
```python
quick_return_coros = {'name': coro, ...} # some coroutines return very quick
iterator = gen.WaitIterator(
    **{k: ensure_future(x()) for x in quick_return_coros}
)
res = {}
async for result in iterator:
  res[iterator.current_index] = result
# res may be less than quick_return_coros   
```


Discussion:
Can you post a complete example? I don't see how "returning quickly" could be enough for this to happen. Are threads involved? Are you using the same Future object twice anywhere?
@bdarnell here's the sample code
```python
import asyncio
from asyncio import ensure_future

import tornado.ioloop
import tornado.web
from tornado import gen
from tornado.platform.asyncio import AsyncIOMainLoop


class MainHandler(tornado.web.RequestHandler):
  async def get_result(self, text):
    return text

  async def get(self):
    k = [
      'a',
      'b',
      'c'
    ]
    res = {}
    iterator = gen.WaitIterator(
      **{x: ensure_future(self.get_result(x)) for x in k})
    async for i in iterator:
      res[iterator.current_index] = i
    self.finish(res)

  async def post(self):
    k = [
      'a',
      'b',
      'c'
    ]
    res = await gen.multi({x: ensure_future(self.get_result(x)) for x in k})
    self.finish(res)


def make_app():
  return tornado.web.Application([
    (r"/", MainHandler),
  ])


if __name__ == "__main__":
  AsyncIOMainLoop().install()
  loop = asyncio.get_event_loop()
  app = make_app()
  app.listen(8888)
  tornado.ioloop.IOLoop.current().start()
```
a GET will has one return only, while the POST works correctly
```json
# get
{
  "c": "a"
}
# post
{
  "c": "c",
  "a": "a",
  "b": "b"
}
```
Ah, got it. This has to do with `asyncio.Future` vs Tornado `Future`: In Tornado the `set_result` method is synchronous, but in `asyncio` it schedules work on the next IOLoop iteration. This is another thing to work out in the 5.0 refactoring of coroutines to do less mixing-and-matching of parts from the two frameworks. In the meantime, you can work around this by using `tornado.gen.convert_yielded` instead of `asyncio.ensure_future`. 

Resolving patch:
diff --git a/tornado/gen.py b/tornado/gen.py
--- a/tornado/gen.py
+++ b/tornado/gen.py
@@ -400,7 +400,7 @@ def next(self) -> Future:
         self._running_future = Future()
 
         if self._finished:
-            self._return_result(self._finished.popleft())
+            return self._return_result(self._finished.popleft())
 
         return self._running_future
 
@@ -410,7 +410,7 @@ def _done_callback(self, done: Future) -> None:
         else:
             self._finished.append(done)
 
-    def _return_result(self, done: Future) -> None:
+    def _return_result(self, done: Future) -> Future:
         """Called set the returned future's state that of the future
         we yielded, and set the current future for the iterator.
         """
@@ -418,9 +418,13 @@ def _return_result(self, done: Future) -> None:
             raise Exception("no future is running")
         chain_future(done, self._running_future)
 
+        res = self._running_future
+        self._running_future = None
         self.current_future = done
         self.current_index = self._unfinished.pop(done)
 
+        return res
+
     def __aiter__(self) -> typing.AsyncIterator:
         return self
 


Test patch: