Closing the producer goroutine from the consumer

The producer-and-consumer pattern is well used in Go concurrent programming. Whenthe consumer stops, we want to gracefully stop the producer as well.

Problem

When a gRPC server receives a streaming request, it usually calls afunction that returns a channel,reads the result from that channel and send the result to the client one by one.

Take the following code for instance: upon receiving a request, the main goroutineService calls launchJob. launchJob starts a separate goroutine as an anonymousfunction call and returns a channel. In the anonymous function, items will be sent tochannel. And Service on the otherside of the channel will reads from it.

  1. func Service(req *Request, stream *StreamResponse) error {
  2.   result := launchJob(req.Content)
  3.   for r := range result {
  4.     if e := stream.Send(result); e != nil {
  5.       // should we signal the running goroutine so it will stop sending?
  6.       return e
  7.     }
  8.   }
  9. }
  11. func launchJob(content string) chan Item {
  12.   c := make(chan Item)
  14.   go func() {
  15.     defer close(c)
  17.     acquireScarceResources()
  18.     defer releaseScarceResources()
  20.     ...
  21.     // if stream.Send(result) returns an error and the Service returns, this will be blocked
  22.     c <- Item{}
  23.     ...
  24.   }()
  26.   return c
  27. }

There is a major problem in this implementation. As pointed out by the comment,if the Send in Service returns an error, the Service function will return,leaving the anonymous function being blocked on c <- Item{} forever.

This problem is important because the leaking goroutine usually owns scarce systemresources such as network connection and memory.

Solution: pipeline explicit cancellation

Inspired by this blog post sectionExplicit cancellation, we can signal the cancellation via closing on a separatechannel. And we can follow the terminology as io.Pipe.

  1. package sql
  3. import (
  4.     "errors"
  5. )
  7. var ErrClosedPipe = errors.New("pipe: write on closed pipe")
  9. // pipe follows the design at https://blog.golang.org/pipelines
  10. // - wrCh: chan for piping data
  11. // - done: chan for signaling Close from Reader to Writer
  12. type pipe struct {
  13.     wrCh chan interface{}
  14.     done chan struct{}
  15. }
  17. // PipeReader reads real data
  18. type PipeReader struct {
  19.     p *pipe
  20. }
  22. // PipeWriter writes real data
  23. type PipeWriter struct {
  24.     p *pipe
  25. }
  27. // Pipe creates a synchronous in-memory pipe.
  28. //
  29. // It is safe to call Read and Write in parallel with each other or with Close.
  30. // Parallel calls to Read and parallel calls to Write are also safe:
  31. // the individual calls will be gated sequentially.
  32. func Pipe() (*PipeReader, *PipeWriter) {
  33.     p := &pipe{
  34.         wrCh: make(chan interface{}),
  35.         done: make(chan struct{})}
  36.     return &PipeReader{p}, &PipeWriter{p}
  37. }
  39. // Close closes the reader; subsequent writes to the
  40. func (r *PipeReader) Close() {
  41.     close(r.p.done)
  42. }
  44. // ReadAll returns the data chan. The caller should
  45. // use it as `for r := range pr.ReadAll()`
  46. func (r *PipeReader) ReadAll() chan interface{} {
  47.     return r.p.wrCh
  48. }
  50. // Close closes the writer; subsequent ReadAll from the
  51. // read half of the pipe will return a closed channel.
  52. func (w *PipeWriter) Close() {
  53.     close(w.p.wrCh)
  54. }
  56. // Write writes the item to the underlying data stream.
  57. // It returns ErrClosedPipe when the data stream is closed.
  58. func (w *PipeWriter) Write(item interface{}) error {
  59.     select {
  60.     case w.p.wrCh <- item:
  61.         return nil
  62.     case <-w.p.done:
  63.         return ErrClosedPipe
  64.     }
  65. }

And the consumer and producer be can implemented as

  1. func Service(req *Request, stream *StreamResponse) error {
  2.   pr := launchJob(req.Content)
  3.   defer pr.Close()
  4.   for r := range pr.ReadAll() {
  5.     if e := stream.Send(r); e != nil {
  6.       return e
  7.     }
  8.   }
  9. }
  11. func launchJob(content string) PipeReader {
  12.     pr, pw := Pipe()
  13.     go func() {
  14.         defer pw.Close()
  16.         if err := pw.Write(Item{}); err != nil {
  17.             return
  18.         }
  19.     }
  20.     return pr
  21. }

Further Reading