天天看點

Golang:一個服務同時處理普通http請求和gRPC請求

核心部分

package main

import (
	"log"
	"net/http"
	"strings"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
	"google.golang.org/grpc"
)

var (
	EnableTLS bool
	CertFile  string
	KeyFile   string
	Address   string
)

type MyServer struct {
	grpcServer *grpc.Server
	httpServer *http.ServeMux
}

func (ms *MyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
		ms.grpcServer.ServeHTTP(w, r)
	} else {
		ms.httpServer.ServeHTTP(w, r)
	}
}

func main() {
	ms := &MyServer{
		grpcServer: grpc.NewServer(),
		httpServer: http.NewServeMux(),
	}

	server := &http.Server{
		Addr: Address,
	}
	var err error
	if EnableTLS {
		server.Handler = ms
		err = server.ListenAndServeTLS(CertFile, KeyFile)
	} else {
		server.Handler = h2c.NewHandler(ms, &http2.Server{})
		err = server.ListenAndServe()
	}
	if err != nil {
		log.Printf("serve failed")
	}
}