proxy/template/service/service.go

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()
}