github.com/cloudwego/hertz@v0.9.3/pkg/app/server/hertz.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     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 server
    18  
    19  import (
    20  	"context"
    21  	"os"
    22  	"os/signal"
    23  	"syscall"
    24  	"time"
    25  
    26  	"github.com/cloudwego/hertz/pkg/app/middlewares/server/recovery"
    27  	"github.com/cloudwego/hertz/pkg/common/config"
    28  	"github.com/cloudwego/hertz/pkg/common/errors"
    29  	"github.com/cloudwego/hertz/pkg/common/hlog"
    30  	"github.com/cloudwego/hertz/pkg/route"
    31  )
    32  
    33  // Hertz is the core struct of hertz.
    34  type Hertz struct {
    35  	*route.Engine
    36  	signalWaiter func(err chan error) error
    37  }
    38  
    39  // New creates a hertz instance without any default config.
    40  func New(opts ...config.Option) *Hertz {
    41  	options := config.NewOptions(opts)
    42  	h := &Hertz{
    43  		Engine: route.NewEngine(options),
    44  	}
    45  	return h
    46  }
    47  
    48  // Default creates a hertz instance with default middlewares.
    49  func Default(opts ...config.Option) *Hertz {
    50  	h := New(opts...)
    51  	h.Use(recovery.Recovery())
    52  
    53  	return h
    54  }
    55  
    56  // Spin runs the server until catching os.Signal or error returned by h.Run().
    57  func (h *Hertz) Spin() {
    58  	errCh := make(chan error)
    59  	h.initOnRunHooks(errCh)
    60  	go func() {
    61  		errCh <- h.Run()
    62  	}()
    63  
    64  	signalWaiter := waitSignal
    65  	if h.signalWaiter != nil {
    66  		signalWaiter = h.signalWaiter
    67  	}
    68  
    69  	if err := signalWaiter(errCh); err != nil {
    70  		hlog.SystemLogger().Errorf("Receive close signal: error=%v", err)
    71  		if err := h.Engine.Close(); err != nil {
    72  			hlog.SystemLogger().Errorf("Close error=%v", err)
    73  		}
    74  		return
    75  	}
    76  
    77  	hlog.SystemLogger().Infof("Begin graceful shutdown, wait at most num=%d seconds...", h.GetOptions().ExitWaitTimeout/time.Second)
    78  
    79  	ctx, cancel := context.WithTimeout(context.Background(), h.GetOptions().ExitWaitTimeout)
    80  	defer cancel()
    81  
    82  	if err := h.Shutdown(ctx); err != nil {
    83  		hlog.SystemLogger().Errorf("Shutdown error=%v", err)
    84  	}
    85  }
    86  
    87  // SetCustomSignalWaiter sets the signal waiter function.
    88  // If Default one is not met the requirement, set this function to customize.
    89  // Hertz will exit immediately if f returns an error, otherwise it will exit gracefully.
    90  func (h *Hertz) SetCustomSignalWaiter(f func(err chan error) error) {
    91  	h.signalWaiter = f
    92  }
    93  
    94  // Default implementation for signal waiter.
    95  // SIGTERM triggers immediately close.
    96  // SIGHUP|SIGINT triggers graceful shutdown.
    97  func waitSignal(errCh chan error) error {
    98  	signalToNotify := []os.Signal{syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM}
    99  	if signal.Ignored(syscall.SIGHUP) {
   100  		signalToNotify = []os.Signal{syscall.SIGINT, syscall.SIGTERM}
   101  	}
   102  
   103  	signals := make(chan os.Signal, 1)
   104  	signal.Notify(signals, signalToNotify...)
   105  
   106  	select {
   107  	case sig := <-signals:
   108  		switch sig {
   109  		case syscall.SIGTERM:
   110  			// force exit
   111  			return errors.NewPublic(sig.String()) // nolint
   112  		case syscall.SIGHUP, syscall.SIGINT:
   113  			hlog.SystemLogger().Infof("Received signal: %s\n", sig)
   114  			// graceful shutdown
   115  			return nil
   116  		}
   117  	case err := <-errCh:
   118  		// error occurs, exit immediately
   119  		return err
   120  	}
   121  
   122  	return nil
   123  }
   124  
   125  func (h *Hertz) initOnRunHooks(errChan chan error) {
   126  	// add register func to runHooks
   127  	opt := h.GetOptions()
   128  	h.OnRun = append(h.OnRun, func(ctx context.Context) error {
   129  		go func() {
   130  			// delay register 1s
   131  			time.Sleep(1 * time.Second)
   132  			if err := opt.Registry.Register(opt.RegistryInfo); err != nil {
   133  				hlog.SystemLogger().Errorf("Register error=%v", err)
   134  				// pass err to errChan
   135  				errChan <- err
   136  			}
   137  		}()
   138  		return nil
   139  	})
   140  }