github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/utils/remotesrv/main.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"context"
    19  	"flag"
    20  	"fmt"
    21  	"log"
    22  	"net"
    23  	"net/http"
    24  	"os"
    25  	"os/signal"
    26  	"sync"
    27  
    28  	"google.golang.org/grpc"
    29  
    30  	remotesapi "github.com/dolthub/dolt/go/gen/proto/dolt/services/remotesapi/v1alpha1"
    31  	"github.com/dolthub/dolt/go/libraries/utils/filesys"
    32  )
    33  
    34  func main() {
    35  	dirParam := flag.String("dir", "", "root directory that this command will run in.")
    36  	grpcPortParam := flag.Int("grpc-port", -1, "root directory that this command will run in.")
    37  	httpPortParam := flag.Int("http-port", -1, "root directory that this command will run in.")
    38  	flag.Parse()
    39  
    40  	if dirParam != nil && len(*dirParam) > 0 {
    41  		err := os.Chdir(*dirParam)
    42  
    43  		if err != nil {
    44  			log.Fatalln("failed to chdir to:", *dirParam)
    45  			log.Fatalln("error:", err.Error())
    46  			os.Exit(1)
    47  		} else {
    48  			log.Println("cwd set to " + *dirParam)
    49  		}
    50  	} else {
    51  		log.Println("'dir' parameter not provided. Using the current working dir.")
    52  	}
    53  
    54  	httpHost := "localhost"
    55  
    56  	if *httpPortParam != -1 {
    57  		httpHost = fmt.Sprintf("%s:%d", httpHost, *httpPortParam)
    58  	} else {
    59  		*httpPortParam = 80
    60  		log.Println("'http-port' parameter not provided. Using default port 80")
    61  	}
    62  
    63  	if *grpcPortParam == -1 {
    64  		*grpcPortParam = 50051
    65  		log.Println("'grpc-port' parameter not provided. Using default port 50051")
    66  	}
    67  
    68  	stopChan, wg := startServer(httpHost, *httpPortParam, *grpcPortParam)
    69  	waitForSignal()
    70  
    71  	close(stopChan)
    72  	wg.Wait()
    73  }
    74  
    75  func waitForSignal() {
    76  	c := make(chan os.Signal)
    77  	signal.Notify(c, os.Interrupt)
    78  	signal.Notify(c, os.Kill)
    79  
    80  	<-c
    81  }
    82  
    83  func startServer(httpHost string, httpPort, grpcPort int) (chan interface{}, *sync.WaitGroup) {
    84  	wg := sync.WaitGroup{}
    85  	stopChan := make(chan interface{})
    86  
    87  	wg.Add(1)
    88  	go func() {
    89  		defer wg.Done()
    90  		httpServer(httpPort, stopChan)
    91  	}()
    92  
    93  	wg.Add(1)
    94  	go func() {
    95  		defer wg.Done()
    96  		grpcServer(httpHost, grpcPort, stopChan)
    97  	}()
    98  
    99  	return stopChan, &wg
   100  }
   101  
   102  func grpcServer(httpHost string, grpcPort int, stopChan chan interface{}) {
   103  	defer func() {
   104  		log.Println("exiting grpc Server go routine")
   105  	}()
   106  
   107  	dbCache := NewLocalCSCache(filesys.LocalFS)
   108  	chnkSt := NewHttpFSBackedChunkStore(httpHost, dbCache)
   109  
   110  	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", grpcPort))
   111  	if err != nil {
   112  		log.Fatalf("failed to listen: %v", err)
   113  	}
   114  
   115  	grpcServer := grpc.NewServer(grpc.MaxRecvMsgSize(128 * 1024 * 1024))
   116  	go func() {
   117  		remotesapi.RegisterChunkStoreServiceServer(grpcServer, chnkSt)
   118  
   119  		log.Println("Starting grpc server on port", grpcPort)
   120  		err := grpcServer.Serve(lis)
   121  		log.Println("grpc server exited. error:", err)
   122  	}()
   123  
   124  	<-stopChan
   125  	grpcServer.GracefulStop()
   126  }
   127  
   128  func httpServer(httpPort int, stopChan chan interface{}) {
   129  	defer func() {
   130  		log.Println("exiting http Server go routine")
   131  	}()
   132  
   133  	server := http.Server{
   134  		Addr:    fmt.Sprintf(":%d", httpPort),
   135  		Handler: http.HandlerFunc(ServeHTTP),
   136  	}
   137  
   138  	go func() {
   139  		log.Println("Starting http server on port ", httpPort)
   140  		err := server.ListenAndServe()
   141  		log.Println("http server exited. exit error:", err)
   142  	}()
   143  
   144  	<-stopChan
   145  	server.Shutdown(context.Background())
   146  }