github.com/snowflakedb/gosnowflake@v1.9.0/heartbeat.go (about)

     1  // Copyright (c) 2019-2022 Snowflake Computing Inc. All rights reserved.
     2  
     3  package gosnowflake
     4  
     5  import (
     6  	"context"
     7  	"encoding/json"
     8  	"fmt"
     9  	"io"
    10  	"net/http"
    11  	"net/url"
    12  	"time"
    13  )
    14  
    15  const (
    16  	// One hour interval should be good enough to renew tokens for four hours master token validity
    17  	heartBeatInterval = 3600 * time.Second
    18  )
    19  
    20  type heartbeat struct {
    21  	restful      *snowflakeRestful
    22  	shutdownChan chan bool
    23  }
    24  
    25  func (hc *heartbeat) run() {
    26  	hbTicker := time.NewTicker(heartBeatInterval)
    27  	defer hbTicker.Stop()
    28  	for {
    29  		select {
    30  		case <-hbTicker.C:
    31  			err := hc.heartbeatMain()
    32  			if err != nil {
    33  				logger.Error("failed to heartbeat")
    34  			}
    35  		case <-hc.shutdownChan:
    36  			logger.Info("stopping heartbeat")
    37  			return
    38  		}
    39  	}
    40  }
    41  
    42  func (hc *heartbeat) start() {
    43  	hc.shutdownChan = make(chan bool)
    44  	go hc.run()
    45  	logger.Info("heartbeat started")
    46  }
    47  
    48  func (hc *heartbeat) stop() {
    49  	hc.shutdownChan <- true
    50  	close(hc.shutdownChan)
    51  	logger.Info("heartbeat stopped")
    52  }
    53  
    54  func (hc *heartbeat) heartbeatMain() error {
    55  	logger.Info("Heartbeating!")
    56  	params := &url.Values{}
    57  	params.Add(requestIDKey, NewUUID().String())
    58  	params.Add(requestGUIDKey, NewUUID().String())
    59  	headers := getHeaders()
    60  	token, _, _ := hc.restful.TokenAccessor.GetTokens()
    61  	headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
    62  
    63  	fullURL := hc.restful.getFullURL(heartBeatPath, params)
    64  	timeout := hc.restful.RequestTimeout
    65  	resp, err := hc.restful.FuncPost(context.Background(), hc.restful, fullURL, headers, nil, timeout, defaultTimeProvider, nil)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	defer resp.Body.Close()
    70  	if resp.StatusCode == http.StatusOK {
    71  		logger.Infof("heartbeatMain: resp: %v", resp)
    72  		var respd execResponse
    73  		err = json.NewDecoder(resp.Body).Decode(&respd)
    74  		if err != nil {
    75  			logger.Infof("failed to decode JSON. err: %v", err)
    76  			return err
    77  		}
    78  		if respd.Code == sessionExpiredCode {
    79  			err = hc.restful.renewExpiredSessionToken(context.Background(), timeout, token)
    80  			if err != nil {
    81  				return err
    82  			}
    83  		}
    84  		return nil
    85  	}
    86  	b, err := io.ReadAll(resp.Body)
    87  	if err != nil {
    88  		logger.Errorf("failed to extract HTTP response body. err: %v", err)
    89  		return err
    90  	}
    91  	logger.Infof("HTTP: %v, URL: %v, Body: %v", resp.StatusCode, fullURL, b)
    92  	logger.Infof("Header: %v", resp.Header)
    93  	return &SnowflakeError{
    94  		Number:   ErrFailedToHeartbeat,
    95  		SQLState: SQLStateConnectionFailure,
    96  		Message:  "Failed to heartbeat.",
    97  	}
    98  }