go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/warmup/warmup.go (about)

     1  // Copyright 2017 The LUCI Authors.
     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 warmup allows to register hooks executed during the server warmup.
    16  //
    17  // All registered hooks should be optional. The warmup stage can be skipped.
    18  package warmup
    19  
    20  import (
    21  	"context"
    22  	"net/http"
    23  	"sync"
    24  
    25  	"go.chromium.org/luci/common/errors"
    26  	"go.chromium.org/luci/common/logging"
    27  
    28  	"go.chromium.org/luci/server/router"
    29  )
    30  
    31  // Callback will be called during warmup.
    32  type Callback func(c context.Context) error
    33  
    34  var state struct {
    35  	sync.Mutex
    36  	callbacks []callbackWithName
    37  }
    38  
    39  type callbackWithName struct {
    40  	Callback
    41  	name string
    42  }
    43  
    44  // Register adds a callback called during warmup.
    45  func Register(name string, cb Callback) {
    46  	if name == "" {
    47  		panic("warmup callback name is required")
    48  	}
    49  	state.Lock()
    50  	defer state.Unlock()
    51  	state.callbacks = append(state.callbacks, callbackWithName{cb, name})
    52  }
    53  
    54  // Warmup executes all registered warmup callbacks, sequentially.
    55  //
    56  // Doesn't abort on individual callback errors, just collects and returns them
    57  // all.
    58  func Warmup(c context.Context) error {
    59  	state.Lock()
    60  	defer state.Unlock()
    61  
    62  	var merr errors.MultiError
    63  	for _, cb := range state.callbacks {
    64  		logging.Infof(c, "Warming up %q", cb.name)
    65  		if err := cb.Callback(c); err != nil {
    66  			logging.Errorf(c, "Error when warming up %q: %s", cb.name, err)
    67  			merr = append(merr, err)
    68  		}
    69  	}
    70  
    71  	logging.Infof(c, "Finished warming up")
    72  	if len(merr) == 0 {
    73  		return nil
    74  	}
    75  	return merr
    76  }
    77  
    78  // InstallHandlersDeprecated installs HTTP handlers for warmup /_ah/* routes.
    79  //
    80  // It is deprecated. Do not use. On GAEv1 it is called by the framework code.
    81  // On GAEv2 or on GKE use NewModuleFromFlags() and register the warmup module
    82  // when starting the server.
    83  func InstallHandlersDeprecated(r *router.Router, base router.MiddlewareChain) {
    84  	r.GET("/_ah/warmup", base, httpHandler)
    85  	r.GET("/_ah/start", base, httpHandler)
    86  }
    87  
    88  func httpHandler(c *router.Context) {
    89  	status := http.StatusOK
    90  	if err := Warmup(c.Request.Context()); err != nil {
    91  		status = http.StatusInternalServerError
    92  	}
    93  	c.Writer.WriteHeader(status)
    94  }