github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/x/health.go (about)

     1  /*
     2   * Copyright 2017-2018 Dgraph Labs, Inc. and Contributors
     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 x
    18  
    19  import (
    20  	"sync/atomic"
    21  
    22  	"github.com/pkg/errors"
    23  )
    24  
    25  var (
    26  	// the drainingMode variable should be accessed through the atomic.Store and atomic.Load
    27  	// functions. The value 0 means the draining-mode is disabled, and the value 1 means the
    28  	// mode is enabled
    29  	drainingMode uint32
    30  
    31  	healthCheck     uint32
    32  	errHealth       = errors.New("Please retry again, server is not ready to accept requests")
    33  	errDrainingMode = errors.New("the server is in draining mode " +
    34  		"and client requests will only be allowed after exiting the mode " +
    35  		" by sending a POST request to /admin/draining?enable=false")
    36  )
    37  
    38  // UpdateHealthStatus updates the server's health status so it can start accepting requests.
    39  func UpdateHealthStatus(ok bool) {
    40  	setStatus(&healthCheck, ok)
    41  }
    42  
    43  // UpdateDrainingMode updates the server's draining mode
    44  func UpdateDrainingMode(enable bool) {
    45  	setStatus(&drainingMode, enable)
    46  }
    47  
    48  // HealthCheck returns whether the server is ready to accept requests or not
    49  // Load balancer would add the node to the endpoint once health check starts
    50  // returning true
    51  func HealthCheck() error {
    52  	if atomic.LoadUint32(&healthCheck) == 0 {
    53  		return errHealth
    54  	}
    55  	if atomic.LoadUint32(&drainingMode) == 1 {
    56  		return errDrainingMode
    57  	}
    58  	return nil
    59  }
    60  
    61  func setStatus(v *uint32, ok bool) {
    62  	if ok {
    63  		atomic.StoreUint32(v, 1)
    64  	} else {
    65  		atomic.StoreUint32(v, 0)
    66  	}
    67  }