workercmd.go

Documentation: github.com/sassoftware/relic/cmdline/workercmd

     1  //
     2  // Copyright (c) SAS Institute Inc.
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  //
    16  
    17  package workercmd
    18  
    19  import (
    20  	"context"
    21  	"io/ioutil"
    22  	"log"
    23  	"net/http"
    24  	"os"
    25  	"os/signal"
    26  	"sync"
    27  	"syscall"
    28  
    29  	"github.com/sassoftware/relic/cmdline/shared"
    30  	"github.com/sassoftware/relic/internal/activation"
    31  	"github.com/sassoftware/relic/internal/activation/activatecmd"
    32  	"github.com/sassoftware/relic/token/open"
    33  	"github.com/spf13/cobra"
    34  )
    35  
    36  func init() {
    37  	AddWorkerCommand(shared.RootCmd)
    38  }
    39  
    40  func AddWorkerCommand(parent *cobra.Command) {
    41  	wc := &cobra.Command{
    42  		Use:    "worker",
    43  		Run:    worker,
    44  		Hidden: true,
    45  	}
    46  	parent.AddCommand(wc)
    47  }
    48  
    49  func worker(cmd *cobra.Command, args []string) {
    50  	log.SetFlags(0)
    51  	if len(args) != 2 {
    52  		log.Fatalln("invalid argument")
    53  	}
    54  	shared.ArgConfig = args[0]
    55  	tokenName := args[1]
    56  	if err := runWorker(tokenName); err != nil {
    57  		log.Fatalf("error: worker for token \"%s\": %+v", tokenName, err)
    58  	}
    59  }
    60  
    61  func runWorker(tokenName string) error {
    62  	if err := shared.InitConfig(); err != nil {
    63  		return err
    64  	}
    65  	cookie, err := ioutil.ReadAll(os.Stdin)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	lis, err := activation.GetListener(0, "tcp", "")
    70  	if err != nil {
    71  		return err
    72  	}
    73  	tok, err := open.Token(shared.CurrentConfig, tokenName, nil)
    74  	if err != nil {
    75  		return err
    76  	}
    77  	handler := &handler{
    78  		token:  tok,
    79  		cookie: cookie,
    80  	}
    81  	srv := &http.Server{Handler: handler}
    82  	wg := new(sync.WaitGroup)
    83  	handler.shutdown = func() {
    84  		// keep the main goroutine from exiting until all requests are served
    85  		wg.Add(1)
    86  		// notify parent that this process is doomed and to start another one
    87  		activatecmd.DaemonStopping()
    88  		// stop accepting requests and wait for ongoing ones to finish
    89  		srv.Shutdown(context.Background())
    90  		wg.Done()
    91  	}
    92  	go handler.watchSignals()
    93  	go handler.healthCheck()
    94  	activation.DaemonReady()
    95  	err = srv.Serve(lis)
    96  	if err == http.ErrServerClosed {
    97  		err = nil
    98  	}
    99  	wg.Wait() // wait for shutdown to finish
   100  	return err
   101  }
   102  
   103  func (h *handler) watchSignals() {
   104  	ch := make(chan os.Signal, 1)
   105  	signal.Notify(ch,
   106  		syscall.SIGINT,
   107  		syscall.SIGTERM,
   108  		syscall.SIGQUIT,
   109  	)
   110  	<-ch
   111  	signal.Stop(ch)
   112  	h.shutdown()
   113  }
   114
View as plain text