自定义中间件路由

package main

import (
    "fmt"
    "log"
    "net/http"
    "time"
)

func product(w http.ResponseWriter, r *http.Request) {
    time.Sleep(300 * time.Millisecond)
    w.Write([]byte("hello world product"))
}

func goods(writer http.ResponseWriter, request *http.Request) {
    time.Sleep(200 * time.Millisecond)
    writer.Write([]byte("hello world goods"))
}

// 中间件
func timeMiddleWare(next http.Handler) http.Handler {
    return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
        begin := time.Now()
        requestPath := request.URL.Path
        //保护的路由
        next.ServeHTTP(writer, request)
        elapse := time.Since(begin)
        log.Printf("%s use time %d ms", requestPath, elapse.Milliseconds())
    })
}

var limitChan = make(chan struct{}, 100)

// 限制
func limitMiddleWare(next http.Handler) http.Handler {
    return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
        limitChan <- struct{}{}
        next.ServeHTTP(writer, request)
        <-limitChan
    })
}

// 自定义路由
type middleWare func(next http.Handler) http.Handler

// 定义一个结构体
type Router struct {
    middleWareChina []middleWare            //中间件切片
    mux             map[string]http.Handler //字符串对应的路由
}

// 构造路由
func NewRouter() *Router {
    return &Router{
        middleWareChina: make([]middleWare, 0),
        mux:             make(map[string]http.Handler),
    }
}

// 添加一个中间件
func (Router *Router) Use(m middleWare) {
    Router.middleWareChina = append(Router.middleWareChina, m)
}

// 添加一个路由
func (Router *Router) Add(path string, handler http.Handler) {
    var mergeHandler = handler //已经存在的中间件
    for i := len(Router.middleWareChina) - 1; i >= 0; i-- {
        //当前的中间件 router.middleWareChina[i]
        mv := Router.middleWareChina[i]
        //当前的中间件(已经存在的中间件)
        mergeHandler = mv(mergeHandler)
    }
    Router.mux[path] = mergeHandler
}

// Router实现ServeHTTP方法
func (Router *Router) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
    requestPath := request.URL.Path
    if handler, exists := Router.mux[requestPath]; exists {
        handler.ServeHTTP(writer, request)
    } else {
        http.NotFoundHandler().ServeHTTP(writer, request)
    }
}

func main() {
    //原生的路由
    //http.Handle("/goods", limitMiddleWare(timeMiddleWare(http.HandlerFunc(goods))))
    //http.Handle("/product", limitMiddleWare(timeMiddleWare(http.HandlerFunc(product))))

    //自定义路由
    Router := NewRouter()
    Router.Use(limitMiddleWare)
    Router.Use(timeMiddleWare)
    Router.Add("/goods", http.HandlerFunc(goods))
    Router.Add("/product", http.HandlerFunc(product))
    if err := http.ListenAndServe(":5656", Router); err != nil {
        fmt.Println(err)
    }
}

本作品采用《CC 协议》，转载必须注明作者和本文链接