dolthub--dolt/go/libraries/doltcore/remotesrv/server.go

// Copyright 2019 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package remotesrv

import (
	"context"
	"crypto/tls"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/sirupsen/logrus"
	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"
	"google.golang.org/grpc"

	remotesapi "github.com/dolthub/dolt/go/gen/proto/dolt/services/remotesapi/v1alpha1"
	"github.com/dolthub/dolt/go/libraries/utils/filesys"
)

// disabledFeaturesEnvVar is an undocumented internal test hook. If
// set, its value is parsed as a comma-separated list of
// remotesapi.Feature enum names; each resolved feature is appended to
// ServerArgs.DisabledFeatures before the server starts. Used by
// integration tests to force clients down legacy-RPC fallback paths
// against a server that actually supports the new RPC. Unknown names
// cause NewServer to log.Fatalln — loud failure on typos is the right
// default for a test hook.
const disabledFeaturesEnvVar = "DOLT_REMOTESAPI_DISABLED_FEATURES"

// MaxGRPCMessageSize is the maximum gRPC message size used for chunk store
// operations. This applies to both send and receive directions.
const MaxGRPCMessageSize = 128 * 1024 * 1024

type Server struct {
	wg       sync.WaitGroup
	stopChan chan struct{}

	grpcListenAddr string
	httpListenAddr string

	grpcSrv *grpc.Server
	httpSrv http.Server

	grpcHttpReqsWG sync.WaitGroup

	tlsConfig *tls.Config
}

func (s *Server) GracefulStop() {
	close(s.stopChan)
	s.wg.Wait()
}

type ServerArgs struct {
	Logger   *logrus.Entry
	HttpHost string

	HttpListenAddr string
	GrpcListenAddr string

	FS       filesys.Filesys
	DBCache  DBCache
	ReadOnly bool
	Options  []grpc.ServerOption

	ConcurrencyControl remotesapi.PushConcurrencyControl

	// Feature flags this server implements but will not advertise in
	// GetRepoMetadataResponse.features. Internal test hook — used by
	// integration tests to force clients down the older-RPC fallback
	// against a server that actually supports the new RPC. Go-level
	// tests set this directly; process-level tests set it indirectly
	// via the DOLT_REMOTESAPI_DISABLED_FEATURES env var, which
	// NewServer appends to this field at startup.
	DisabledFeatures []remotesapi.Feature

	HttpInterceptor func(http.Handler) http.Handler

	// If supplied, the listener(s) returned from Listeners() will be TLS
	// listeners. The scheme used in the URLs returned from the gRPC server
	// will be https.
	TLSConfig *tls.Config
}

func NewServer(args ServerArgs) (*Server, error) {
	if v := os.Getenv(disabledFeaturesEnvVar); v != "" {
		for _, name := range strings.Split(v, ",") {
			name = strings.TrimSpace(name)
			if name == "" {
				continue
			}
			id, ok := remotesapi.Feature_value[name]
			if !ok {
				log.Fatalln(disabledFeaturesEnvVar + ": unknown feature name: " + name)
			}
			args.DisabledFeatures = append(args.DisabledFeatures, remotesapi.Feature(id))
		}
	}

	if args.Logger == nil {
		args.Logger = logrus.NewEntry(logrus.StandardLogger())
	}

	s := new(Server)
	s.stopChan = make(chan struct{})

	sealer, err := NewSingleSymmetricKeySealer()
	if err != nil {
		return nil, err
	}

	scheme := "http"
	if args.TLSConfig != nil {
		scheme = "https"
	}
	s.tlsConfig = args.TLSConfig

	s.wg.Add(2)
	s.grpcListenAddr = args.GrpcListenAddr
	s.grpcSrv = grpc.NewServer(append([]grpc.ServerOption{grpc.MaxRecvMsgSize(MaxGRPCMessageSize)}, args.Options...)...)
	var chnkSt remotesapi.ChunkStoreServiceServer = NewHttpFSBackedChunkStore(args.Logger, args.HttpHost, args.DBCache, args.FS, scheme, args.ConcurrencyControl, sealer, args.DisabledFeatures)

	if args.ReadOnly {
		chnkSt = ReadOnlyChunkStore{chnkSt}
	}
	remotesapi.RegisterChunkStoreServiceServer(s.grpcSrv, chnkSt)

	var handler http.Handler = newFileHandler(args.Logger, args.DBCache, args.FS, args.ReadOnly, sealer)
	if args.HttpInterceptor != nil {
		handler = args.HttpInterceptor(handler)
	}
	if args.HttpListenAddr == args.GrpcListenAddr {
		handler = s.grpcMultiplexHandler(s.grpcSrv, handler)
	} else {
		s.wg.Add(2)
	}

	s.httpListenAddr = args.HttpListenAddr
	s.httpSrv = http.Server{
		Addr:    args.HttpListenAddr,
		Handler: handler,
	}

	return s, nil
}

func (s *Server) grpcMultiplexHandler(grpcSrv *grpc.Server, handler http.Handler) http.Handler {
	h2s := &http2.Server{}
	newHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") {
			s.grpcHttpReqsWG.Add(1)
			defer s.grpcHttpReqsWG.Done()
			grpcSrv.ServeHTTP(w, r)
		} else {
			handler.ServeHTTP(w, r)
		}
	})
	return h2c.NewHandler(newHandler, h2s)
}

type Listeners struct {
	http net.Listener
	grpc net.Listener
}

func (l Listeners) Close() error {
	if l.http != nil {
		err := l.http.Close()
		if err != nil {
			if l.grpc != nil {
				l.grpc.Close()
			}
			return err
		}
	}
	if l.grpc != nil {
		return l.grpc.Close()
	}
	return nil
}

func (s *Server) Listeners() (Listeners, error) {
	var httpListener net.Listener
	var grpcListener net.Listener
	var err error
	if s.tlsConfig != nil {
		httpListener, err = tls.Listen("tcp", s.httpListenAddr, s.tlsConfig)
	} else {
		httpListener, err = net.Listen("tcp", s.httpListenAddr)
	}
	if err != nil {
		return Listeners{}, err
	}
	if s.httpListenAddr == s.grpcListenAddr {
		return Listeners{http: httpListener}, nil
	}
	if s.tlsConfig != nil {
		grpcListener, err = tls.Listen("tcp", s.grpcListenAddr, s.tlsConfig)
	} else {
		grpcListener, err = net.Listen("tcp", s.grpcListenAddr)
	}
	if err != nil {
		httpListener.Close()
		return Listeners{}, err
	}
	return Listeners{http: httpListener, grpc: grpcListener}, nil
}

// Can be used to register more services on the server.
// Should only be accessed before `Serve` is called.
func (s *Server) GrpcServer() *grpc.Server {
	return s.grpcSrv
}

func (s *Server) Serve(listeners Listeners) {
	if listeners.grpc != nil {
		go func() {
			defer s.wg.Done()
			logrus.Println("Starting grpc server on", s.grpcListenAddr)
			err := s.grpcSrv.Serve(listeners.grpc)
			logrus.Println("grpc server exited. error:", err)
		}()
		go func() {
			defer s.wg.Done()
			<-s.stopChan
			logrus.Traceln("Calling grpcSrv.GracefulStop")
			s.grpcSrv.GracefulStop()
			logrus.Traceln("Finished calling grpcSrv.GracefulStop")
		}()
	}

	go func() {
		defer s.wg.Done()
		logrus.Println("Starting http server on", s.httpListenAddr)
		err := s.httpSrv.Serve(listeners.http)
		logrus.Println("http server exited. exit error:", err)
	}()
	go func() {
		defer s.wg.Done()
		<-s.stopChan
		logrus.Traceln("Calling httpSrv.Shutdown")
		s.httpSrv.Shutdown(context.Background())
		logrus.Traceln("Finished calling httpSrv.Shutdown")

		// If we are multiplexing HTTP and gRPC requests on the same
		// listener, we need to stop the gRPC server here as well. We
		// cannot stop it gracefully, but if we stop it forcefully
		// here, we guarantee all the handler threads are cleaned up
		// before we return.
		if listeners.grpc == nil {
			logrus.Traceln("Calling grpcSrv.Stop")
			s.grpcSrv.Stop()
			s.grpcHttpReqsWG.Wait()
			logrus.Traceln("Finished calling grpcSrv.Stop")
		}
	}()

	s.wg.Wait()
}