Question about tokio and tasks

Hey there!

I’m currently working on a proxy using Tokio that multiplexes M requests on N connections (where M > N). The goal of this proxy is to limit the number of concurrent connections to a given server.

Good news is I got it working. Bad news is I have no clue what I’m doing.

The proxy creates a TCP listener and iterates on the stream of incoming connections. When a client connects to it, I ask a server connection pool to return me an available connection. This pool returns a future (FutureConnection). When polled, that future will return Ok(Async::Ready(connection)) if it’s available or Async::NotReady if it’s not.

When I first implemented this future, I forgot to save the current task somewhere. If server connections were all busy, the future would simply be dropped and so would the client connection.

Then, after skimming through the internets, I found out about task::current().

So now, from the poll method, if the future is not ready, I simply push this task to a vector on the server connection pool, and when a connection becomes available, I do something like self.tasks.drain(..).for_each(|task| task.notify()).

It seems to work, although, by the look of it when I squint at my screen, I have a hunch this does not follow best practices.

Any idea how to achieve this, the proper way?

I changed my solution to something I feel is cleaner, but that’s up to debate.

My connection pool is now initialized with a FIFO for tasks using a VecDeque.

The first time the future is polled, if there is no connection available, it pushes its task (with task::current()) to the pool queue and marks itself as queued, literally using a field on the future struct. On subsequent calls of poll, since it was marked as queued, its task won’t be queued again.

Then on the pool side, once a connection becomes available, it pops the dequeues from the FIFO and notifies that task.

Let me know what you think.

As you discovered initially, you shouldn’t return NotReady without ensuring your task will be woken up later to be polled again. Since you’re mapping M clients to (smaller) N network event sources for wakeups, you’re going to have to do something like this to fan out the wakeup events to all the other tasks.

This sounds workable, given that you’re returning back connection handles. Are these connections cloned, persist with the client after the first request from the pool, and handle the multiplexing internally? Or are they handed out for single use, and the connection is “busy” until completed and the client comes back for a connection again next time?

Since you talk about “all connections busy”, I assume it’s the latter. You possibly only want to pull one task off the FIFO per released connection, otherwise the others will just end up re-queueing themselves immediately.

But since you also talk about handing out backend connections on initial client connection, it could also be the former design. In this case another alternative would be to have the pool be more protocol-aware, and passed requests, return a future for a response, and handle all the load balancing and back-end connection management. That’s a different design, but allows for recovery from back end connection-failure, perhaps some more flexible scale up/down of N, and maybe some other smarter load-balancing options.

1 Like

Currently I’m toying with a little prototype and I am indeed implementing the latter.

You possibly only want to pull one task off the FIFO per released connection, otherwise the others will just end up re-queueing themselves immediately.

This is what I ended up doing as described in my second message :slight_smile:

Also, I might eventually go for the “protocol-aware pool”, it’s indeed more powerful. Thanks for your answer!

1 Like