github.com/thiagoyeds/go-cloud@v0.26.0/server/health/sqlhealth/sqlhealth.go (about)

     1  // Copyright 2018 The Go Cloud Development Kit 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  //     https://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 sqlhealth provides a health check for a SQL database connection.
    16  package sqlhealth // import "gocloud.dev/server/health/sqlhealth"
    17  
    18  import (
    19  	"context"
    20  	"database/sql"
    21  	"errors"
    22  	"time"
    23  )
    24  
    25  // Checker checks the health of a SQL database.
    26  type Checker struct {
    27  	cancel context.CancelFunc
    28  
    29  	stopped <-chan struct{}
    30  	healthy bool
    31  }
    32  
    33  // New starts a new asynchronous ping of the SQL database. Pings will be sent
    34  // until one succeeds or Stop is called, whichever comes first.
    35  func New(db *sql.DB) *Checker {
    36  	// We create a context here because we are detaching.
    37  	ctx, cancel := context.WithCancel(context.Background())
    38  	stopped := make(chan struct{})
    39  	c := &Checker{
    40  		cancel:  cancel,
    41  		stopped: stopped,
    42  	}
    43  	go func() {
    44  		var timer *time.Timer
    45  		defer func() {
    46  			if timer != nil {
    47  				timer.Stop()
    48  			}
    49  			close(stopped)
    50  		}()
    51  
    52  		wait := 250 * time.Millisecond
    53  		const maxWait = 30 * time.Second
    54  		for {
    55  			if err := db.PingContext(ctx); err == nil {
    56  				c.healthy = true
    57  				return
    58  			}
    59  			if timer == nil {
    60  				timer = time.NewTimer(wait)
    61  			} else {
    62  				// Timer already fired, so resetting does not race.
    63  				timer.Reset(wait)
    64  			}
    65  			select {
    66  			case <-timer.C:
    67  				if wait < maxWait {
    68  					// Back off next ping.
    69  					wait *= 2
    70  					if wait > maxWait {
    71  						wait = maxWait
    72  					}
    73  				}
    74  			case <-ctx.Done():
    75  				return
    76  			}
    77  		}
    78  	}()
    79  	return c
    80  }
    81  
    82  // CheckHealth returns nil iff the ping started by New has returned
    83  // success.
    84  func (c *Checker) CheckHealth() error {
    85  	select {
    86  	case <-c.stopped:
    87  		if !c.healthy {
    88  			return errors.New("ping stopped before becoming healthy")
    89  		}
    90  		return nil
    91  	default:
    92  		return errors.New("still pinging database")
    93  	}
    94  }
    95  
    96  // Stop stops any ongoing ping of the database.
    97  func (c *Checker) Stop() {
    98  	c.cancel()
    99  	<-c.stopped
   100  }