package service import ( "context" "errors" "io" "log" "net" "strconv" "sync" "time" ) type Config struct { Host string Port uint16 CloseWait time.Duration } type Server struct { config *Config ctx context.Context cancel context.CancelFunc wg sync.WaitGroup } func New(conf *Config) (*Server, error) { if conf.Host == "" { conf.Host = "localhost" } if conf.Port == 0 { return nil, errors.New("port is required") } ctx, cancel := context.WithCancel(context.Background()) return &Server{ conf, ctx, cancel, sync.WaitGroup{}, }, nil } func (s *Server) Run() error { // start listen addr := net.JoinHostPort(s.config.Host, strconv.Itoa(int(s.config.Port))) ls, err := net.Listen("tcp", addr) if err != nil { return errors.New("failed to listen") } defer closeRes(ls) // wait accept connCh := make(chan net.Conn) defer close(connCh) go func() { for { conn, err := ls.Accept() if err != nil { if !errors.Is(err, net.ErrClosed) { log.Println("accept failed", err) } // retry on temporary error var ne net.Error if errors.As(err, &ne) && ne.Temporary() { continue } return } select { case <-s.ctx.Done(): closeRes(conn) return case connCh <- conn: } } }() // handle accept func() { for { select { case <-s.ctx.Done(): return case conn, ok := <-connCh: if !ok { return } s.wg.Add(1) go func() { defer s.wg.Done() _ = s.handle(conn) }() } } }() // close timeout, cancel := context.WithTimeout(context.Background(), s.config.CloseWait) defer cancel() waitCh := make(chan struct{}) defer close(waitCh) go func() { s.wg.Wait() select { case <-timeout.Done(): // waitCh may be closed case waitCh <- struct{}{}: } }() err = nil select { case <-timeout.Done(): err = errors.New("close timeout") case <-waitCh: } return err } func (s *Server) Close() { s.cancel() } func (s *Server) handle(conn net.Conn) error { defer closeRes(conn) return nil } func closeRes[T io.Closer](res T) { _ = res.Close() }