107 lines
2.4 KiB
Go
107 lines
2.4 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
gwRuntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
|
|
"golang.org/x/net/http2"
|
|
"golang.org/x/net/http2/h2c"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/reflection"
|
|
)
|
|
|
|
type Server struct {
|
|
GRPC *grpc.Server
|
|
Gateway *gwRuntime.ServeMux
|
|
HTTP *gin.Engine
|
|
server *http.Server
|
|
}
|
|
|
|
func New() *Server {
|
|
grpcServer := grpc.NewServer()
|
|
reflection.Register(grpcServer)
|
|
engine := gin.New()
|
|
engine.Use(gin.Logger(), gin.Recovery())
|
|
return &Server{
|
|
GRPC: grpcServer,
|
|
Gateway: gwRuntime.NewServeMux(),
|
|
HTTP: engine,
|
|
}
|
|
}
|
|
|
|
func (s *Server) Start(addr string) error {
|
|
handler := h2c.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
|
|
s.GRPC.ServeHTTP(w, r)
|
|
return
|
|
}
|
|
recorder := newBufferedResponse()
|
|
s.Gateway.ServeHTTP(recorder, r)
|
|
if recorder.status != http.StatusNotFound {
|
|
recorder.flush(w)
|
|
return
|
|
}
|
|
s.HTTP.ServeHTTP(w, r)
|
|
}), &http2.Server{})
|
|
s.server = &http.Server{
|
|
Addr: addr,
|
|
Handler: handler,
|
|
ReadHeaderTimeout: 10 * time.Second,
|
|
IdleTimeout: 120 * time.Second,
|
|
MaxHeaderBytes: 1 << 20,
|
|
}
|
|
listener, err := net.Listen("tcp", addr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
fmt.Printf("all services listening on %s (gRPC + HTTP)\n", addr)
|
|
return s.server.Serve(listener)
|
|
}
|
|
|
|
func (s *Server) Stop(ctx context.Context) error {
|
|
stopped := make(chan struct{})
|
|
go func() {
|
|
s.GRPC.GracefulStop()
|
|
close(stopped)
|
|
}()
|
|
select {
|
|
case <-stopped:
|
|
case <-ctx.Done():
|
|
s.GRPC.Stop()
|
|
}
|
|
if s.server == nil {
|
|
return ctx.Err()
|
|
}
|
|
return s.server.Shutdown(ctx)
|
|
}
|
|
|
|
type bufferedResponse struct {
|
|
header http.Header
|
|
body bytes.Buffer
|
|
status int
|
|
}
|
|
|
|
func newBufferedResponse() *bufferedResponse {
|
|
return &bufferedResponse{header: make(http.Header), status: http.StatusOK}
|
|
}
|
|
|
|
func (r *bufferedResponse) Header() http.Header { return r.header }
|
|
func (r *bufferedResponse) WriteHeader(status int) { r.status = status }
|
|
func (r *bufferedResponse) Write(data []byte) (int, error) { return r.body.Write(data) }
|
|
func (r *bufferedResponse) flush(w http.ResponseWriter) {
|
|
for key, values := range r.header {
|
|
for _, value := range values {
|
|
w.Header().Add(key, value)
|
|
}
|
|
}
|
|
w.WriteHeader(r.status)
|
|
_, _ = w.Write(r.body.Bytes())
|
|
}
|