第38章 Capstone：把玩具代理改造成生产级骨架

package main

import (
	"context"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

// 缓冲池：复用转发缓冲，避免每连接分配（第31章）
var bufPool = sync.Pool{New: func() any { b := make([]byte, 32*1024); return &b }}

// 并发上限：满了快速拒绝（减载，第37章）
var sem = make(chan struct{}, 10000)

var active int64 // 可观测：活跃连接数（第10章）

// 到后端的连接池 + 超时分级（第03/31/35章）
var transport = &http.Transport{
	DialContext: (&net.Dialer{
		Timeout:   5 * time.Second,  // connect 超时（第35章 SYN_SENT）
		KeepAlive: 30 * time.Second,
	}).DialContext,
	MaxIdleConns:          1000,
	MaxIdleConnsPerHost:   100,              // 后端连接池大小
	IdleConnTimeout:       90 * time.Second,
	ResponseHeaderTimeout: 30 * time.Second, // read 超时（第35章 → 504）
}

var hopHeaders = []string{"Connection", "Proxy-Connection", "Keep-Alive",
	"Proxy-Authenticate", "Proxy-Authorization", "Te", "Trailer", "Transfer-Encoding", "Upgrade"}

func main() {
	srv := &http.Server{
		Addr:              ":8080",
		Handler:           http.HandlerFunc(handle),
		ReadHeaderTimeout: 10 * time.Second, // 防 Slowloris（第34/35章）
		IdleTimeout:       60 * time.Second, // idle 回收（第35章）
	}

	// 优雅关闭：停止 accept，等存量 drain，超时强制（第35章）
	go func() {
		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
		defer stop()
		<-ctx.Done()
		log.Println("关闭信号：停止 accept，drain 存量（≤30s）")
		sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		srv.Shutdown(sctx)
	}()

	log.Println("生产级骨架代理 :8080")
	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatal(err)
	}
}

func handle(w http.ResponseWriter, r *http.Request) {
	// 减载：并发满了直接 503，保住其余请求 SLA（第37章）
	select {
	case sem <- struct{}{}:
		defer func() { <-sem }()
	default:
		http.Error(w, "overloaded", http.StatusServiceUnavailable)
		return
	}
	atomic.AddInt64(&active, 1)
	defer atomic.AddInt64(&active, -1)

	// panic 隔离：单连接崩溃不波及整个代理
	defer func() {
		if v := recover(); v != nil {
			log.Printf("panic recovered: %v", v)
		}
	}()

	if r.Method == http.MethodConnect {
		handleConnect(w, r)
		return
	}

	// 明文 HTTP：经连接池转发（删 hop-by-hop 头，第04章）
	for _, h := range hopHeaders {
		r.Header.Del(h)
	}
	r.RequestURI = ""
	resp, err := transport.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway) // 502
		return
	}
	defer resp.Body.Close()
	for _, h := range hopHeaders {
		resp.Header.Del(h)
	}
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	buf := bufPool.Get().(*[]byte)
	defer bufPool.Put(buf)
	io.CopyBuffer(w, resp.Body, *buf) // 缓冲池复用 + io.Copy 自带背压（第34章）
}

func handleConnect(w http.ResponseWriter, r *http.Request) {
	dst, err := net.DialTimeout("tcp", r.Host, 5*time.Second) // connect 超时
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	hij, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "no hijack", http.StatusInternalServerError)
		dst.Close()
		return
	}
	src, _, err := hij.Hijack()
	if err != nil {
		dst.Close()
		return
	}
	src.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))

	// 双向转发：半关闭传播（第35章）+ 背压（io.Copy 自带，第34章）
	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); pipe(dst, src) }()
	go func() { defer wg.Done(); pipe(src, dst) }()
	wg.Wait()
	dst.Close()
	src.Close()
}

// 单向转发：EOF 时只半关闭【对端写方向】，不 close 整条（第35章）
func pipe(dst, src net.Conn) {
	buf := bufPool.Get().(*[]byte)
	defer bufPool.Put(buf)
	io.CopyBuffer(dst, src, *buf) // Write 阻塞自动停 Read = 背压（第34章）
	if c, ok := dst.(interface{ CloseWrite() error }); ok {
		c.CloseWrite() // 传播 FIN，不截断另一方向（第35章）
	}
}