context.go 658 B

12345678910111213141516171819202122232425262728293031
  1. package net
  2. import (
  3. "context"
  4. "net"
  5. )
  6. // SetupContextForConn is a helper function that starts connection I/O interrupter goroutine.
  7. func SetupContextForConn(ctx context.Context, conn net.Conn) (done func(*error)) {
  8. var (
  9. quit = make(chan struct{})
  10. interrupt = make(chan error, 1)
  11. )
  12. go func() {
  13. select {
  14. case <-quit:
  15. interrupt <- nil
  16. case <-ctx.Done():
  17. // Close the connection, discarding the error
  18. _ = conn.Close()
  19. interrupt <- ctx.Err()
  20. }
  21. }()
  22. return func(inputErr *error) {
  23. close(quit)
  24. if ctxErr := <-interrupt; ctxErr != nil && inputErr != nil {
  25. // Return context error to user.
  26. inputErr = &ctxErr
  27. }
  28. }
  29. }