full/all/internal/server/server.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())
}