feat: add bounded native http server

2026-09-21 01:58:50 +08:00
parent 3bbf4627a1
commit 05b4088a25
14 changed files with 1058 additions and 19 deletions
--- a/transport/native/server.go
+++ b/transport/native/server.go
@@ -0,0 +1,655 @@
+// Package native provides GoOmni's first project-owned HTTP/1.1 server.
+//
+// It deliberately keeps the public handler boundary as net/http so existing
+// GoOmni business code remains unchanged while connection management and
+// protocol I/O move out of net/http.Server. The parser currently supports
+// Content-Length request bodies and rejects chunked requests.
+package native
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+const (
+	defaultReadHeaderTimeout = 5 * time.Second
+	defaultReadTimeout       = 30 * time.Second
+	defaultWriteTimeout      = 30 * time.Second
+	defaultIdleTimeout       = 60 * time.Second
+	defaultMaxHeaderBytes    = 1 << 20
+	defaultMaxBodyBytes      = 8 << 20
+	defaultReadBufferSize    = 4 << 10
+	maxPooledResponseBytes   = 64 << 10
+)
+
+// Server is a small HTTP/1.1 server with project-owned connection lifecycle.
+// It supports keep-alive, graceful shutdown, bounded concurrency and standard
+// http.Handler values. Zero MaxConnections and MaxConcurrentRequests mean
+// unlimited; pooled protocol objects are reused automatically.
+type Server struct {
+	ReadHeaderTimeout     time.Duration
+	ReadTimeout           time.Duration
+	WriteTimeout          time.Duration
+	IdleTimeout           time.Duration
+	ReadBufferSize        int
+	MaxHeaderBytes        int
+	MaxBodyBytes          int64
+	MaxConnections        int
+	MaxConcurrentRequests int
+
+	mu             sync.Mutex
+	handler        http.Handler
+	listener       net.Listener
+	conns          map[*connection]struct{}
+	closing        bool
+	cancel         context.CancelFunc
+	wg             sync.WaitGroup
+	connSlots      chan struct{}
+	requestSlots   chan struct{}
+	readerPool     sync.Pool
+	writerPool     sync.Pool
+	connectionPool sync.Pool
+}
+
+type connection struct {
+	server        *Server
+	conn          net.Conn
+	active        atomic.Bool
+	closeWhenIdle atomic.Bool
+	closed        atomic.Bool
+	slotHeld      bool
+	closeDone     chan struct{}
+}
+
+// New creates a native server. The handler may also be supplied later through
+// App.RunWith, which is useful when the App owns the lifecycle.
+func New(handler http.Handler) *Server {
+	server := &Server{
+		handler:           handler,
+		ReadHeaderTimeout: defaultReadHeaderTimeout,
+		ReadTimeout:       defaultReadTimeout,
+		WriteTimeout:      defaultWriteTimeout,
+		IdleTimeout:       defaultIdleTimeout,
+		ReadBufferSize:    defaultReadBufferSize,
+		MaxHeaderBytes:    defaultMaxHeaderBytes,
+		MaxBodyBytes:      defaultMaxBodyBytes,
+		conns:             make(map[*connection]struct{}),
+	}
+	server.readerPool.New = func() any { return bufio.NewReaderSize(nil, server.readBufferSize()) }
+	server.writerPool.New = func() any { return newResponseWriter("") }
+	server.connectionPool.New = func() any { return &connection{server: server} }
+	return server
+}
+
+func (s *Server) readBufferSize() int {
+	if s.ReadBufferSize > 0 {
+		return s.ReadBufferSize
+	}
+	return defaultReadBufferSize
+}
+
+// Run implements app.Runner without importing the app package.
+func (s *Server) Run(ctx context.Context, handler http.Handler, address string) error {
+	if handler != nil {
+		s.handler = handler
+	}
+	if s.handler == nil {
+		return errors.New("go-omni/native: nil handler")
+	}
+	if address == "" {
+		address = ":8080"
+	}
+	listener, err := net.Listen("tcp", address)
+	if err != nil {
+		return err
+	}
+	return s.Serve(ctx, listener)
+}
+
+// Serve accepts connections from listener until the context is canceled or an
+// unrecoverable accept error occurs. The listener is owned by the Server.
+func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
+	if listener == nil {
+		return errors.New("go-omni/native: nil listener")
+	}
+	if s.handler == nil {
+		return errors.New("go-omni/native: nil handler")
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	s.mu.Lock()
+	if s.listener != nil {
+		s.mu.Unlock()
+		return errors.New("go-omni/native: server is already running")
+	}
+	s.listener = listener
+	s.closing = false
+	if s.MaxConnections > 0 {
+		s.connSlots = make(chan struct{}, s.MaxConnections)
+	} else {
+		s.connSlots = nil
+	}
+	if s.MaxConcurrentRequests > 0 {
+		s.requestSlots = make(chan struct{}, s.MaxConcurrentRequests)
+	} else {
+		s.requestSlots = nil
+	}
+	baseCtx, cancel := context.WithCancel(context.Background())
+	s.cancel = cancel
+	s.mu.Unlock()
+	defer cancel()
+
+	shutdownDone := make(chan struct{})
+	go func() {
+		select {
+		case <-ctx.Done():
+			timeout := s.shutdownTimeout()
+			shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+			_ = s.Shutdown(shutdownCtx)
+			cancel()
+		case <-shutdownDone:
+		}
+	}()
+
+	var serveErr error
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			if ctx.Err() != nil || s.isClosing() {
+				break
+			}
+			serveErr = err
+			break
+		}
+		if slots := s.connSlots; slots != nil {
+			select {
+			case slots <- struct{}{}:
+			default:
+				_ = conn.Close()
+				continue
+			}
+		}
+		state := s.connectionPool.Get().(*connection)
+		state.server = s
+		state.conn = conn
+		state.active.Store(false)
+		state.closeWhenIdle.Store(false)
+		state.closed.Store(false)
+		state.slotHeld = s.connSlots != nil
+		state.closeDone = make(chan struct{})
+		s.mu.Lock()
+		if s.closing {
+			s.mu.Unlock()
+			state.close()
+			s.recycleConnection(state)
+			continue
+		}
+		s.conns[state] = struct{}{}
+		s.wg.Add(1)
+		s.mu.Unlock()
+		go s.serveConnection(baseCtx, state)
+	}
+
+	close(shutdownDone)
+	s.mu.Lock()
+	if s.listener == listener {
+		s.listener = nil
+	}
+	s.mu.Unlock()
+	if ctx.Err() != nil || s.isClosing() {
+		timeout := s.shutdownTimeout()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+		shutdownErr := s.Shutdown(shutdownCtx)
+		cancel()
+		return errors.Join(shutdownErr, serveErr)
+	}
+	_ = listener.Close()
+	timeout := s.shutdownTimeout()
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
+	shutdownErr := s.Shutdown(shutdownCtx)
+	cancel()
+	return errors.Join(shutdownErr, serveErr)
+}
+
+// Shutdown stops accepting requests, cancels active request contexts and waits
+// for connections to finish. On timeout, active sockets are closed.
+func (s *Server) Shutdown(ctx context.Context) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	s.mu.Lock()
+	s.closing = true
+	if s.cancel != nil {
+		s.cancel()
+	}
+	listener := s.listener
+	states := make([]*connection, 0, len(s.conns))
+	for state := range s.conns {
+		state.closeWhenIdle.Store(true)
+		states = append(states, state)
+	}
+	s.listener = nil
+	s.mu.Unlock()
+	if listener != nil {
+		_ = listener.Close()
+	}
+	for _, state := range states {
+		if !state.active.Load() {
+			state.close()
+		}
+	}
+
+	done := make(chan struct{})
+	go func() {
+		s.wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		for _, state := range states {
+			state.close()
+		}
+		<-done
+		return ctx.Err()
+	}
+}
+
+func (s *Server) shutdownTimeout() time.Duration {
+	if value := s.WriteTimeout; value > 0 {
+		return value
+	}
+	return defaultWriteTimeout
+}
+
+func (s *Server) isClosing() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.closing
+}
+
+func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
+	defer s.wg.Done()
+	defer func() {
+		state.close()
+		s.recycleConnection(state)
+	}()
+	reader := s.readerPool.Get().(*bufio.Reader)
+	reader.Reset(state.conn)
+	defer func() {
+		reader.Reset(nil)
+		s.readerPool.Put(reader)
+	}()
+	firstRequest := true
+	for {
+		readTimeout := s.ReadHeaderTimeout
+		if readTimeout <= 0 {
+			readTimeout = defaultReadHeaderTimeout
+		}
+		if !firstRequest && s.IdleTimeout > 0 {
+			readTimeout = s.IdleTimeout
+		} else if s.IdleTimeout > 0 && s.IdleTimeout < readTimeout {
+			readTimeout = s.IdleTimeout
+		}
+		_ = state.conn.SetReadDeadline(time.Now().Add(readTimeout))
+		request, closeAfter, err := readRequest(reader, state.conn, s.MaxHeaderBytes, s.MaxBodyBytes, func() {
+			if s.ReadTimeout > 0 {
+				_ = state.conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
+			} else {
+				_ = state.conn.SetReadDeadline(time.Time{})
+			}
+		})
+		if err != nil {
+			if !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) && !isTimeout(err) {
+				status := http.StatusBadRequest
+				message := "bad request\n"
+				var protocolErr *protocolError
+				if errors.As(err, &protocolErr) {
+					status = protocolErr.status
+					message = protocolErr.message + "\n"
+				}
+				_ = writeSimple(state.conn, status, message, true, s.WriteTimeout)
+			}
+			return
+		}
+		_ = state.conn.SetReadDeadline(time.Time{})
+		firstRequest = false
+		requestContext, cancel := context.WithCancel(baseCtx)
+		request = request.WithContext(requestContext)
+		state.active.Store(true)
+		writer := s.writerPool.Get().(*responseWriter)
+		writer.reset(request.Method)
+		requestSlot := false
+		if slots := s.requestSlots; slots != nil {
+			select {
+			case slots <- struct{}{}:
+				requestSlot = true
+			case <-baseCtx.Done():
+				cancel()
+				writer.reset("")
+				s.writerPool.Put(writer)
+				state.active.Store(false)
+				return
+			default:
+				writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
+				writer.WriteHeader(http.StatusServiceUnavailable)
+				writer.body.WriteString("server busy\n")
+				closing := s.isClosing() || state.closeWhenIdle.Load()
+				_ = writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout)
+				cancel()
+				writer.reset("")
+				s.writerPool.Put(writer)
+				state.active.Store(false)
+				if closeAfter || closing {
+					return
+				}
+				continue
+			}
+		}
+		handler := s.handler
+		handler.ServeHTTP(writer, request)
+		cancel()
+
+		closing := s.isClosing() || state.closeWhenIdle.Load()
+		if err := writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout); err != nil {
+			if requestSlot {
+				<-s.requestSlots
+			}
+			writer.reset("")
+			s.writerPool.Put(writer)
+			state.active.Store(false)
+			return
+		}
+		if requestSlot {
+			<-s.requestSlots
+		}
+		writer.reset("")
+		s.writerPool.Put(writer)
+		state.active.Store(false)
+		if closeAfter || closing {
+			return
+		}
+	}
+}
+
+func (c *connection) close() {
+	if c.closed.CompareAndSwap(false, true) {
+		defer close(c.closeDone)
+		_ = c.conn.Close()
+		c.server.mu.Lock()
+		delete(c.server.conns, c)
+		c.server.mu.Unlock()
+	}
+}
+
+func (s *Server) recycleConnection(c *connection) {
+	<-c.closeDone
+	if c.slotHeld {
+		if slots := s.connSlots; slots != nil {
+			<-slots
+		}
+		c.slotHeld = false
+	}
+	s.connectionPool.Put(c)
+}
+
+type protocolError struct {
+	status  int
+	message string
+}
+
+func (e *protocolError) Error() string { return e.message }
+
+func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBodyBytes int64, setBodyDeadline func()) (*http.Request, bool, error) {
+	if maxHeaderBytes <= 0 {
+		maxHeaderBytes = defaultMaxHeaderBytes
+	}
+	if maxBodyBytes <= 0 {
+		maxBodyBytes = defaultMaxBodyBytes
+	}
+	used := 0
+	readLine := func() (string, error) {
+		line, err := reader.ReadString('\n')
+		used += len(line)
+		if used > maxHeaderBytes {
+			return "", &protocolError{status: http.StatusRequestHeaderFieldsTooLarge, message: "request headers too large"}
+		}
+		if err != nil {
+			return "", err
+		}
+		if !strings.HasSuffix(line, "\n") {
+			return "", &protocolError{status: http.StatusBadRequest, message: "malformed request line"}
+		}
+		return strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r"), nil
+	}
+
+	line, err := readLine()
+	if err != nil {
+		return nil, false, err
+	}
+	parts := strings.Split(line, " ")
+	if len(parts) != 3 || !validToken(parts[0]) {
+		return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed request line"}
+	}
+	protoMajor, protoMinor, ok := parseProtocol(parts[2])
+	if !ok {
+		return nil, true, &protocolError{status: http.StatusHTTPVersionNotSupported, message: "unsupported HTTP version"}
+	}
+	target := parts[1]
+	var requestURL *url.URL
+	if target == "*" {
+		requestURL = &url.URL{Path: "*"}
+	} else {
+		requestURL, err = url.ParseRequestURI(target)
+		if err != nil || requestURL.Path == "" || !strings.HasPrefix(requestURL.Path, "/") {
+			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed request target"}
+		}
+	}
+
+	header := make(http.Header)
+	for {
+		line, err = readLine()
+		if err != nil {
+			return nil, true, err
+		}
+		if line == "" {
+			break
+		}
+		if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
+			return nil, true, &protocolError{status: http.StatusBadRequest, message: "folded headers are not supported"}
+		}
+		index := strings.IndexByte(line, ':')
+		if index <= 0 {
+			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed header"}
+		}
+		name := strings.TrimSpace(line[:index])
+		value := strings.TrimSpace(line[index+1:])
+		if !validToken(name) || strings.ContainsAny(value, "\r\n") {
+			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed header"}
+		}
+		header.Add(name, value)
+	}
+
+	if value := header.Get("Expect"); value != "" {
+		return nil, true, &protocolError{status: http.StatusExpectationFailed, message: "unsupported expectation"}
+	}
+	if value := header.Get("Transfer-Encoding"); value != "" && !strings.EqualFold(value, "identity") {
+		return nil, true, &protocolError{status: http.StatusNotImplemented, message: "chunked requests are not supported"}
+	}
+	contentLength := int64(0)
+	if value := header.Get("Content-Length"); value != "" {
+		contentLength, err = strconv.ParseInt(strings.TrimSpace(value), 10, 64)
+		if err != nil || contentLength < 0 {
+			return nil, true, &protocolError{status: http.StatusBadRequest, message: "invalid content length"}
+		}
+	}
+	if contentLength > maxBodyBytes {
+		return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
+	}
+	body := make([]byte, contentLength)
+	if contentLength > 0 {
+		if setBodyDeadline != nil {
+			setBodyDeadline()
+		}
+		if _, err := io.ReadFull(reader, body); err != nil {
+			return nil, true, err
+		}
+	}
+	request := &http.Request{
+		Method:        parts[0],
+		URL:           requestURL,
+		Proto:         parts[2],
+		ProtoMajor:    protoMajor,
+		ProtoMinor:    protoMinor,
+		Header:        header,
+		Body:          io.NopCloser(bytes.NewReader(body)),
+		ContentLength: contentLength,
+		Host:          header.Get("Host"),
+		RequestURI:    target,
+		RemoteAddr:    conn.RemoteAddr().String(),
+	}
+	closeAfter := strings.EqualFold(header.Get("Connection"), "close") ||
+		(protoMajor == 1 && protoMinor == 0 && !strings.Contains(strings.ToLower(header.Get("Connection")), "keep-alive"))
+	request.Close = closeAfter
+	return request, closeAfter, nil
+}
+
+func parseProtocol(value string) (int, int, bool) {
+	if value != "HTTP/1.1" && value != "HTTP/1.0" {
+		return 0, 0, false
+	}
+	return 1, int(value[len(value)-1] - '0'), true
+}
+
+func validToken(value string) bool {
+	if value == "" {
+		return false
+	}
+	for _, char := range value {
+		if char <= ' ' || strings.ContainsRune("()<>@,;:\\\"/[]?={} \t", char) {
+			return false
+		}
+	}
+	return true
+}
+
+func isTimeout(err error) bool {
+	var netErr net.Error
+	return errors.As(err, &netErr) && netErr.Timeout()
+}
+
+func writeSimple(conn net.Conn, status int, body string, closeAfter bool, timeout time.Duration) error {
+	writer := newResponseWriter("")
+	writer.status = status
+	writer.wroteHeader = true
+	writer.body.WriteString(body)
+	writer.header.Set("Content-Type", "text/plain; charset=utf-8")
+	return writer.writeTo(conn, closeAfter, timeout)
+}
+
+type responseWriter struct {
+	header        http.Header
+	status        int
+	wroteHeader   bool
+	body          bytes.Buffer
+	requestMethod string
+}
+
+func newResponseWriter(method string) *responseWriter {
+	return &responseWriter{header: make(http.Header), requestMethod: method}
+}
+
+func (w *responseWriter) reset(method string) {
+	for key := range w.header {
+		delete(w.header, key)
+	}
+	w.status = 0
+	w.wroteHeader = false
+	if w.body.Cap() > maxPooledResponseBytes {
+		w.body = bytes.Buffer{}
+	} else {
+		w.body.Reset()
+	}
+	w.requestMethod = method
+}
+
+func (w *responseWriter) Header() http.Header { return w.header }
+
+func (w *responseWriter) WriteHeader(status int) {
+	if w.wroteHeader {
+		return
+	}
+	w.status = status
+	w.wroteHeader = true
+}
+
+func (w *responseWriter) Write(data []byte) (int, error) {
+	if !w.wroteHeader {
+		w.WriteHeader(http.StatusOK)
+	}
+	return w.body.Write(data)
+}
+
+func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Duration) error {
+	if timeout > 0 {
+		_ = conn.SetWriteDeadline(time.Now().Add(timeout))
+	}
+	if !w.wroteHeader {
+		w.status = http.StatusOK
+		w.wroteHeader = true
+	}
+	body := w.body.Bytes()
+	noBody := w.requestMethod == http.MethodHead || w.status == http.StatusNoContent ||
+		w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200)
+	if noBody {
+		body = nil
+	}
+	if w.header.Get("Content-Length") == "" && w.status != http.StatusNoContent && w.status != http.StatusNotModified {
+		w.header.Set("Content-Length", strconv.Itoa(len(w.body.Bytes())))
+	}
+	if closeAfter {
+		w.header.Set("Connection", "close")
+	}
+	keys := make([]string, 0, len(w.header))
+	for key := range w.header {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	var output bytes.Buffer
+	output.WriteString("HTTP/1.1 ")
+	output.WriteString(strconv.Itoa(w.status))
+	output.WriteByte(' ')
+	output.WriteString(http.StatusText(w.status))
+	output.WriteString("\r\n")
+	for _, key := range keys {
+		for _, value := range w.header.Values(key) {
+			if strings.ContainsAny(value, "\r\n") {
+				return errors.New("invalid response header")
+			}
+			output.WriteString(key)
+			output.WriteString(": ")
+			output.WriteString(value)
+			output.WriteString("\r\n")
+		}
+	}
+	output.WriteString("\r\n")
+	output.Write(body)
+	_, err := conn.Write(output.Bytes())
+	_ = conn.SetWriteDeadline(time.Time{})
+	return err
+}