acceptor.go raw

   1  package io
   2  
   3  import (
   4  	"errors"
   5  	"fmt"
   6  	"net"
   7  	"strconv"
   8  	"time"
   9  
  10  	"git.smesh.lol/torly/pkg/unix"
  11  
  12  	"git.smesh.lol/torly/pkg/wire"
  13  )
  14  
  15  // Acceptor owns the listening socket. It polls the socket with a bounded
  16  // timeout so the shut-down breaker is observed promptly, and hands accepted
  17  // sockets to reader workers in round-robin order.
  18  func Acceptor(data AcceptorData, ch AcceptorChans) {
  19  	fd, bound, err := listenTCP(data.Addr)
  20  	if err != nil {
  21  		ch.Ready <- wire.Ready{Err: err.Error()}
  22  		return
  23  	}
  24  	ch.Ready <- wire.Ready{OK: true, Addr: bound}
  25  	defer unix.Close(fd)
  26  
  27  	if len(ch.Workers) == 0 {
  28  		return
  29  	}
  30  	next := 0
  31  	var id wire.ConnID
  32  	for {
  33  		if ch.Stop.Closed() {
  34  			return
  35  		}
  36  		pfd := []unix.PollFd{{Fd: int32(fd), Events: int16(unix.POLLIN)}}
  37  		n, err := unix.Poll(pfd, 200)
  38  		if err != nil {
  39  			if errors.Is(err, unix.EINTR) {
  40  				continue
  41  			}
  42  			// A transient poll error must not kill the only acceptor.
  43  			time.Sleep(10 * time.Millisecond)
  44  			continue
  45  		}
  46  		if n == 0 {
  47  			continue
  48  		}
  49  		nfd, sa, err := unix.Accept4(fd, unix.SOCK_NONBLOCK|unix.SOCK_CLOEXEC)
  50  		if err != nil {
  51  			if errors.Is(err, unix.EAGAIN) || errors.Is(err, unix.EINTR) {
  52  				continue
  53  			}
  54  			// ECONNABORTED, EMFILE and similar errors are transient under
  55  			// connection churn. Exiting here left the backlog permanently full
  56  			// and made the relay look hung.
  57  			time.Sleep(10 * time.Millisecond)
  58  			continue
  59  		}
  60  		id++
  61  		msg := wire.Assign{ID: id, Socket: &wire.Socket{Conn: newFDConn(nfd, sockaddrString(sa))}}
  62  		sent := false
  63  		for i := 0; i < len(ch.Workers); i++ {
  64  			idx := (next + i) % len(ch.Workers)
  65  			if ch.Workers[idx].TrySend(msg) {
  66  				next = idx + 1
  67  				sent = true
  68  				break
  69  			}
  70  		}
  71  		if !sent {
  72  			idx := next % len(ch.Workers)
  73  			select {
  74  			case ch.Workers[idx].Tx <- msg:
  75  				next = idx + 1
  76  			case <-ch.Stop.Done():
  77  				msg.Socket.Conn.Close()
  78  				return
  79  			}
  80  		}
  81  	}
  82  }
  83  
  84  func listenTCP(addr string) (int, string, error) {
  85  	host, portStr, err := net.SplitHostPort(addr)
  86  	if err != nil {
  87  		return -1, "", err
  88  	}
  89  	port, err := strconv.Atoi(portStr)
  90  	if err != nil || port < 0 || port > 65535 {
  91  		return -1, "", fmt.Errorf("bad listen port %q", portStr)
  92  	}
  93  	var ip [4]byte
  94  	if host != "" && host != "0.0.0.0" {
  95  		parsed := net.ParseIP(host)
  96  		if parsed == nil || parsed.To4() == nil {
  97  			return -1, "", fmt.Errorf("unsupported listen address %q", host)
  98  		}
  99  		copy(ip[:], parsed.To4())
 100  	}
 101  	fd, err := unix.Socket(unix.AF_INET, unix.SOCK_STREAM|unix.SOCK_NONBLOCK, 0)
 102  	if err != nil {
 103  		return -1, "", err
 104  	}
 105  	if err := unix.SetsockoptInt(fd, unix.SOL_SOCKET, unix.SO_REUSEADDR, 1); err != nil {
 106  		unix.Close(fd)
 107  		return -1, "", err
 108  	}
 109  	sa := &unix.SockaddrInet4{Port: port, Addr: ip}
 110  	if err := unix.Bind(fd, sa); err != nil {
 111  		unix.Close(fd)
 112  		return -1, "", err
 113  	}
 114  	if err := unix.Listen(fd, 1024); err != nil {
 115  		unix.Close(fd)
 116  		return -1, "", err
 117  	}
 118  	name, err := unix.Getsockname(fd)
 119  	if err != nil {
 120  		unix.Close(fd)
 121  		return -1, "", err
 122  	}
 123  	boundPort := port
 124  	if got, ok := name.(*unix.SockaddrInet4); ok {
 125  		boundPort = got.Port
 126  	}
 127  	return fd, fmt.Sprintf("%s:%d", hostOrAny(host), boundPort), nil
 128  }
 129  
 130  func hostOrAny(host string) string {
 131  	if host == "" || host == "0.0.0.0" {
 132  		return "0.0.0.0"
 133  	}
 134  	return host
 135  }
 136  
 137  func sockaddrString(sa unix.Sockaddr) string {
 138  	switch a := sa.(type) {
 139  	case *unix.SockaddrInet4:
 140  		ip := net.IP(a.Addr[:]).String()
 141  		return ip + ":" + strconv.Itoa(a.Port)
 142  	default:
 143  		return "unknown"
 144  	}
 145  }
 146