<commit_before>package config

import (
	"fmt"
	"time"

	"github.com/hashicorp/vault/api"
)

const (
	// DefaultVaultGrace is the default grace period before which to read a new
	// secret from Vault. If a lease is due to expire in 15 seconds, Consul
	// Template will read a new secret at that time minus this value.
	DefaultVaultGrace = 15 * time.Second

	// DefaultVaultRenewToken is the default value for if the Vault token should
	// be renewed.
	DefaultVaultRenewToken = true

	// DefaultVaultUnwrapToken is the default value for if the Vault token should
	// be unwrapped.
	DefaultVaultUnwrapToken = false

	// DefaultVaultRetryBase is the default value for the base time to use for
	// exponential backoff.
	DefaultVaultRetryBase = 250 * time.Millisecond

	// DefaultVaultRetryMaxAttempts is the default maximum number of attempts to
	// retry before quitting.
	DefaultVaultRetryMaxAttempts = 5
)

// VaultConfig is the configuration for connecting to a vault server.
type VaultConfig struct {
	// Address is the URI to the Vault server.
	Address *string `mapstructure:"address"`

	// Enabled controls whether the Vault integration is active.
	Enabled *bool `mapstructure:"enabled"`

	// Grace is the amount of time before a lease is about to expire to force a
	// new secret to be read.
	Grace *time.Duration `mapstructure:"grace"`

	// Namespace is the Vault namespace to use for reading/writing secrets. This can
	// also be set via the VAULT_NAMESPACE environment variable.
	Namespace *string `mapstructure:"namespace"`

	// RenewToken renews the Vault token.
	RenewToken *bool `mapstructure:"renew_token"`

	// Retry is the configuration for specifying how to behave on failure.
	Retry *RetryConfig `mapstructure:"retry"`

	// SSL indicates we should use a secure connection while talking to Vault.
	SSL *SSLConfig `mapstructure:"ssl"`

	// Token is the Vault token to communicate with for requests. It may be
	// a wrapped token or a real token. This can also be set via the VAULT_TOKEN
	// environment variable, or via the VaultAgentTokenFile.
	Token *string `mapstructure:"token" json:"-"`

	// VaultAgentTokenFile is the path of file that contains a Vault Agent token.
	// If vault_agent_token_file is specified:
	//   - Consul Template will not try to renew the Vault token.
	//   - Consul Template will periodically stat the file and update the token if it has
	// changed.
	VaultAgentTokenFile *string `mapstructure:"vault_agent_token_file" json:"-"`

	// Transport configures the low-level network connection details.
	Transport *TransportConfig `mapstructure:"transport"`

	// UnwrapToken unwraps the provided Vault token as a wrapped token.
	UnwrapToken *bool `mapstructure:"unwrap_token"`
}

// DefaultVaultConfig returns a configuration that is populated with the
// default values.
func DefaultVaultConfig() *VaultConfig {
	v := &VaultConfig{
		Retry:     DefaultRetryConfig(),
		SSL:       DefaultSSLConfig(),
		Transport: DefaultTransportConfig(),
	}

	// Force SSL when communicating with Vault.
	v.SSL.Enabled = Bool(true)

	return v
}

// Copy returns a deep copy of this configuration.
func (c *VaultConfig) Copy() *VaultConfig {
	if c == nil {
		return nil
	}

	var o VaultConfig
	o.Address = c.Address

	o.Enabled = c.Enabled

	o.Grace = c.Grace

	o.Namespace = c.Namespace

	o.RenewToken = c.RenewToken

	if c.Retry != nil {
		o.Retry = c.Retry.Copy()
	}

	if c.SSL != nil {
		o.SSL = c.SSL.Copy()
	}

	o.Token = c.Token

	o.VaultAgentTokenFile = c.VaultAgentTokenFile

	if c.Transport != nil {
		o.Transport = c.Transport.Copy()
	}

	o.UnwrapToken = c.UnwrapToken

	return &o
}

// Merge combines all values in this configuration with the values in the other
// configuration, with values in the other configuration taking precedence.
// Maps and slices are merged, most other values are overwritten. Complex
// structs define their own merge functionality.
func (c *VaultConfig) Merge(o *VaultConfig) *VaultConfig {
	if c == nil {
		if o == nil {
			return nil
		}
		return o.Copy()
	}

	if o == nil {
		return c.Copy()
	}

	r := c.Copy()

	if o.Address != nil {
		r.Address = o.Address
	}

	if o.Enabled != nil {
		r.Enabled = o.Enabled
	}

	if o.Grace != nil {
		r.Grace = o.Grace
	}

	if o.Namespace != nil {
		r.Namespace = o.Namespace
	}

	if o.RenewToken != nil {
		r.RenewToken = o.RenewToken
	}

	if o.Retry != nil {
		r.Retry = r.Retry.Merge(o.Retry)
	}

	if o.SSL != nil {
		r.SSL = r.SSL.Merge(o.SSL)
	}

	if o.Token != nil {
		r.Token = o.Token
	}

	if o.VaultAgentTokenFile != nil {
		r.VaultAgentTokenFile = o.VaultAgentTokenFile
	}

	if o.Transport != nil {
		r.Transport = r.Transport.Merge(o.Transport)
	}

	if o.UnwrapToken != nil {
		r.UnwrapToken = o.UnwrapToken
	}

	return r
}

// Finalize ensures there no nil pointers.
func (c *VaultConfig) Finalize() {
	if c.Address == nil {
		c.Address = stringFromEnv([]string{
			api.EnvVaultAddress,
		}, "")
	}

	if c.Grace == nil {
		c.Grace = TimeDuration(DefaultVaultGrace)
	}

	if c.Namespace == nil {
		c.Namespace = stringFromEnv([]string{"VAULT_NAMESPACE"}, "")
	}

	if c.RenewToken == nil {
		c.RenewToken = boolFromEnv([]string{
			"VAULT_RENEW_TOKEN",
		}, DefaultVaultRenewToken)
	}

	if c.Retry == nil {
		c.Retry = DefaultRetryConfig()
	}
	c.Retry.Finalize()

	// Vault has custom SSL settings
	if c.SSL == nil {
		c.SSL = DefaultSSLConfig()
	}
	if c.SSL.Enabled == nil {
		c.SSL.Enabled = Bool(true)
	}
	if c.SSL.CaCert == nil {
		c.SSL.CaCert = stringFromEnv([]string{api.EnvVaultCACert}, "")
	}
	if c.SSL.CaPath == nil {
		c.SSL.CaPath = stringFromEnv([]string{api.EnvVaultCAPath}, "")
	}
	if c.SSL.Cert == nil {
		c.SSL.Cert = stringFromEnv([]string{api.EnvVaultClientCert}, "")
	}
	if c.SSL.Key == nil {
		c.SSL.Key = stringFromEnv([]string{api.EnvVaultClientKey}, "")
	}
	if c.SSL.ServerName == nil {
		c.SSL.ServerName = stringFromEnv([]string{api.EnvVaultTLSServerName}, "")
	}
	if c.SSL.Verify == nil {
		c.SSL.Verify = antiboolFromEnv([]string{api.EnvVaultInsecure}, true)
	}
	c.SSL.Finalize()

	// Order of precedence
	// 1. `vault_agent_token_file` configuration value
	// 2. `token` configuration value`
	// 3. `VAULT_TOKEN` environment variable
	if c.Token == nil {
		c.Token = stringFromEnv([]string{
			"VAULT_TOKEN",
		}, "")
	}

	if c.VaultAgentTokenFile == nil {
		if StringVal(c.Token) == "" {
			if homePath != "" {
				c.Token = stringFromFile([]string{
					homePath + "/.vault-token",
				}, "")
			}
		}
	} else {
		c.Token = stringFromFile([]string{*c.VaultAgentTokenFile}, "")
		c.RenewToken = Bool(false)
	}

	if c.Transport == nil {
		c.Transport = DefaultTransportConfig()
	}
	c.Transport.Finalize()

	if c.UnwrapToken == nil {
		c.UnwrapToken = boolFromEnv([]string{
			"VAULT_UNWRAP_TOKEN",
		}, DefaultVaultUnwrapToken)
	}

	if c.Enabled == nil {
		c.Enabled = Bool(StringPresent(c.Address))
	}
}

// GoString defines the printable version of this struct.
func (c *VaultConfig) GoString() string {
	if c == nil {
		return "(*VaultConfig)(nil)"
	}

	return fmt.Sprintf("&VaultConfig{"+
		"Address:%s, "+
		"Enabled:%s, "+
		"Grace:%s, "+
		"Namespace:%s,"+
		"RenewToken:%s, "+
		"Retry:%#v, "+
		"SSL:%#v, "+
		"Token:%t, "+
		"VaultAgentTokenFile:%t, "+
		"Transport:%#v, "+
		"UnwrapToken:%s"+
		"}",
		StringGoString(c.Address),
		BoolGoString(c.Enabled),
		TimeDurationGoString(c.Grace),
		StringGoString(c.Namespace),
		BoolGoString(c.RenewToken),
		c.Retry,
		c.SSL,
		StringPresent(c.Token),
		StringPresent(c.VaultAgentTokenFile),
		c.Transport,
		BoolGoString(c.UnwrapToken),
	)
}
<commit_msg>remove the forcable setting of the renew token flag<commit_after>package config

import (
	"fmt"
	"time"

	"github.com/hashicorp/vault/api"
)

const (
	// DefaultVaultGrace is the default grace period before which to read a new
	// secret from Vault. If a lease is due to expire in 15 seconds, Consul
	// Template will read a new secret at that time minus this value.
	DefaultVaultGrace = 15 * time.Second

	// DefaultVaultRenewToken is the default value for if the Vault token should
	// be renewed.
	DefaultVaultRenewToken = true

	// DefaultVaultUnwrapToken is the default value for if the Vault token should
	// be unwrapped.
	DefaultVaultUnwrapToken = false

	// DefaultVaultRetryBase is the default value for the base time to use for
	// exponential backoff.
	DefaultVaultRetryBase = 250 * time.Millisecond

	// DefaultVaultRetryMaxAttempts is the default maximum number of attempts to
	// retry before quitting.
	DefaultVaultRetryMaxAttempts = 5
)

// VaultConfig is the configuration for connecting to a vault server.
type VaultConfig struct {
	// Address is the URI to the Vault server.
	Address *string `mapstructure:"address"`

	// Enabled controls whether the Vault integration is active.
	Enabled *bool `mapstructure:"enabled"`

	// Grace is the amount of time before a lease is about to expire to force a
	// new secret to be read.
	Grace *time.Duration `mapstructure:"grace"`

	// Namespace is the Vault namespace to use for reading/writing secrets. This can
	// also be set via the VAULT_NAMESPACE environment variable.
	Namespace *string `mapstructure:"namespace"`

	// RenewToken renews the Vault token.
	RenewToken *bool `mapstructure:"renew_token"`

	// Retry is the configuration for specifying how to behave on failure.
	Retry *RetryConfig `mapstructure:"retry"`

	// SSL indicates we should use a secure connection while talking to Vault.
	SSL *SSLConfig `mapstructure:"ssl"`

	// Token is the Vault token to communicate with for requests. It may be
	// a wrapped token or a real token. This can also be set via the VAULT_TOKEN
	// environment variable, or via the VaultAgentTokenFile.
	Token *string `mapstructure:"token" json:"-"`

	// VaultAgentTokenFile is the path of file that contains a Vault Agent token.
	// If vault_agent_token_file is specified:
	//   - Consul Template will not try to renew the Vault token.
	//   - Consul Template will periodically stat the file and update the token if it has
	// changed.
	VaultAgentTokenFile *string `mapstructure:"vault_agent_token_file" json:"-"`

	// Transport configures the low-level network connection details.
	Transport *TransportConfig `mapstructure:"transport"`

	// UnwrapToken unwraps the provided Vault token as a wrapped token.
	UnwrapToken *bool `mapstructure:"unwrap_token"`
}

// DefaultVaultConfig returns a configuration that is populated with the
// default values.
func DefaultVaultConfig() *VaultConfig {
	v := &VaultConfig{
		Retry:     DefaultRetryConfig(),
		SSL:       DefaultSSLConfig(),
		Transport: DefaultTransportConfig(),
	}

	// Force SSL when communicating with Vault.
	v.SSL.Enabled = Bool(true)

	return v
}

// Copy returns a deep copy of this configuration.
func (c *VaultConfig) Copy() *VaultConfig {
	if c == nil {
		return nil
	}

	var o VaultConfig
	o.Address = c.Address

	o.Enabled = c.Enabled

	o.Grace = c.Grace

	o.Namespace = c.Namespace

	o.RenewToken = c.RenewToken

	if c.Retry != nil {
		o.Retry = c.Retry.Copy()
	}

	if c.SSL != nil {
		o.SSL = c.SSL.Copy()
	}

	o.Token = c.Token

	o.VaultAgentTokenFile = c.VaultAgentTokenFile

	if c.Transport != nil {
		o.Transport = c.Transport.Copy()
	}

	o.UnwrapToken = c.UnwrapToken

	return &o
}

// Merge combines all values in this configuration with the values in the other
// configuration, with values in the other configuration taking precedence.
// Maps and slices are merged, most other values are overwritten. Complex
// structs define their own merge functionality.
func (c *VaultConfig) Merge(o *VaultConfig) *VaultConfig {
	if c == nil {
		if o == nil {
			return nil
		}
		return o.Copy()
	}

	if o == nil {
		return c.Copy()
	}

	r := c.Copy()

	if o.Address != nil {
		r.Address = o.Address
	}

	if o.Enabled != nil {
		r.Enabled = o.Enabled
	}

	if o.Grace != nil {
		r.Grace = o.Grace
	}

	if o.Namespace != nil {
		r.Namespace = o.Namespace
	}

	if o.RenewToken != nil {
		r.RenewToken = o.RenewToken
	}

	if o.Retry != nil {
		r.Retry = r.Retry.Merge(o.Retry)
	}

	if o.SSL != nil {
		r.SSL = r.SSL.Merge(o.SSL)
	}

	if o.Token != nil {
		r.Token = o.Token
	}

	if o.VaultAgentTokenFile != nil {
		r.VaultAgentTokenFile = o.VaultAgentTokenFile
	}

	if o.Transport != nil {
		r.Transport = r.Transport.Merge(o.Transport)
	}

	if o.UnwrapToken != nil {
		r.UnwrapToken = o.UnwrapToken
	}

	return r
}

// Finalize ensures there no nil pointers.
func (c *VaultConfig) Finalize() {
	if c.Address == nil {
		c.Address = stringFromEnv([]string{
			api.EnvVaultAddress,
		}, "")
	}

	if c.Grace == nil {
		c.Grace = TimeDuration(DefaultVaultGrace)
	}

	if c.Namespace == nil {
		c.Namespace = stringFromEnv([]string{"VAULT_NAMESPACE"}, "")
	}

	if c.RenewToken == nil {
		c.RenewToken = boolFromEnv([]string{
			"VAULT_RENEW_TOKEN",
		}, DefaultVaultRenewToken)
	}

	if c.Retry == nil {
		c.Retry = DefaultRetryConfig()
	}
	c.Retry.Finalize()

	// Vault has custom SSL settings
	if c.SSL == nil {
		c.SSL = DefaultSSLConfig()
	}
	if c.SSL.Enabled == nil {
		c.SSL.Enabled = Bool(true)
	}
	if c.SSL.CaCert == nil {
		c.SSL.CaCert = stringFromEnv([]string{api.EnvVaultCACert}, "")
	}
	if c.SSL.CaPath == nil {
		c.SSL.CaPath = stringFromEnv([]string{api.EnvVaultCAPath}, "")
	}
	if c.SSL.Cert == nil {
		c.SSL.Cert = stringFromEnv([]string{api.EnvVaultClientCert}, "")
	}
	if c.SSL.Key == nil {
		c.SSL.Key = stringFromEnv([]string{api.EnvVaultClientKey}, "")
	}
	if c.SSL.ServerName == nil {
		c.SSL.ServerName = stringFromEnv([]string{api.EnvVaultTLSServerName}, "")
	}
	if c.SSL.Verify == nil {
		c.SSL.Verify = antiboolFromEnv([]string{api.EnvVaultInsecure}, true)
	}
	c.SSL.Finalize()

	// Order of precedence
	// 1. `vault_agent_token_file` configuration value
	// 2. `token` configuration value`
	// 3. `VAULT_TOKEN` environment variable
	if c.Token == nil {
		c.Token = stringFromEnv([]string{
			"VAULT_TOKEN",
		}, "")
	}

	if c.VaultAgentTokenFile == nil {
		if StringVal(c.Token) == "" {
			if homePath != "" {
				c.Token = stringFromFile([]string{
					homePath + "/.vault-token",
				}, "")
			}
		}
	} else {
		c.Token = stringFromFile([]string{*c.VaultAgentTokenFile}, "")
	}

	if c.Transport == nil {
		c.Transport = DefaultTransportConfig()
	}
	c.Transport.Finalize()

	if c.UnwrapToken == nil {
		c.UnwrapToken = boolFromEnv([]string{
			"VAULT_UNWRAP_TOKEN",
		}, DefaultVaultUnwrapToken)
	}

	if c.Enabled == nil {
		c.Enabled = Bool(StringPresent(c.Address))
	}
}

// GoString defines the printable version of this struct.
func (c *VaultConfig) GoString() string {
	if c == nil {
		return "(*VaultConfig)(nil)"
	}

	return fmt.Sprintf("&VaultConfig{"+
		"Address:%s, "+
		"Enabled:%s, "+
		"Grace:%s, "+
		"Namespace:%s,"+
		"RenewToken:%s, "+
		"Retry:%#v, "+
		"SSL:%#v, "+
		"Token:%t, "+
		"VaultAgentTokenFile:%t, "+
		"Transport:%#v, "+
		"UnwrapToken:%s"+
		"}",
		StringGoString(c.Address),
		BoolGoString(c.Enabled),
		TimeDurationGoString(c.Grace),
		StringGoString(c.Namespace),
		BoolGoString(c.RenewToken),
		c.Retry,
		c.SSL,
		StringPresent(c.Token),
		StringPresent(c.VaultAgentTokenFile),
		c.Transport,
		BoolGoString(c.UnwrapToken),
	)
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2014 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package schnorr

import (
	"fmt"
	"math/big"
)

// Signature is a type representing a Schnorr signature.
type Signature struct {
	R *big.Int
	S *big.Int
}

// SignatureSize is the size of an encoded Schnorr signature.
const SignatureSize = 64

// NewSignature instantiates a new signature given some R,S values.
func NewSignature(r, s *big.Int) *Signature {
	return &Signature{r, s}
}

// Serialize returns the Schnorr signature in the more strict format.
//
// The signatures are encoded as
//   sig[0:32]  R, a point encoded as big endian
//   sig[32:64] S, scalar multiplication/addition results = (ab+c) mod l
//     encoded also as big endian
func (sig Signature) Serialize() []byte {
	rBytes := bigIntToEncodedBytes(sig.R)
	sBytes := bigIntToEncodedBytes(sig.S)

	all := append(rBytes[:], sBytes[:]...)

	return all
}

func parseSig(sigStr []byte) (*Signature, error) {
	if len(sigStr) != SignatureSize {
		return nil, fmt.Errorf("bad signature size; have %v, want %v",
			len(sigStr), SignatureSize)
	}

	rBytes := copyBytes(sigStr[0:32])
	r := encodedBytesToBigInt(rBytes)
	sBytes := copyBytes(sigStr[32:64])
	s := encodedBytesToBigInt(sBytes)

	return &Signature{r, s}, nil
}

// ParseSignature parses a signature in BER format for the curve type `curve'
// into a Signature type, perfoming some basic sanity checks.
func ParseSignature(sigStr []byte) (*Signature, error) {
	return parseSig(sigStr)
}

// GetR satisfies the chainec PublicKey interface.
func (sig Signature) GetR() *big.Int {
	return sig.R
}

// GetS satisfies the chainec PublicKey interface.
func (sig Signature) GetS() *big.Int {
	return sig.S
}

// GetType satisfies the chainec Signature interface.
func (sig Signature) GetType() int {
	return ecTypeSecSchnorr
}
<commit_msg>schnorr: add signature IsEqual and Verify methods<commit_after>// Copyright (c) 2013-2014 The btcsuite developers
// Copyright (c) 2015-2016 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package schnorr

import (
	"fmt"
	"math/big"

	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/dcrec/secp256k1"
)

// Signature is a type representing a Schnorr signature.
type Signature struct {
	R *big.Int
	S *big.Int
}

// SignatureSize is the size of an encoded Schnorr signature.
const SignatureSize = 64

// NewSignature instantiates a new signature given some R,S values.
func NewSignature(r, s *big.Int) *Signature {
	return &Signature{r, s}
}

// Serialize returns the Schnorr signature in the more strict format.
//
// The signatures are encoded as
//   sig[0:32]  R, a point encoded as big endian
//   sig[32:64] S, scalar multiplication/addition results = (ab+c) mod l
//     encoded also as big endian
func (sig Signature) Serialize() []byte {
	rBytes := bigIntToEncodedBytes(sig.R)
	sBytes := bigIntToEncodedBytes(sig.S)

	all := append(rBytes[:], sBytes[:]...)

	return all
}

func parseSig(sigStr []byte) (*Signature, error) {
	if len(sigStr) != SignatureSize {
		return nil, fmt.Errorf("bad signature size; have %v, want %v",
			len(sigStr), SignatureSize)
	}

	rBytes := copyBytes(sigStr[0:32])
	r := encodedBytesToBigInt(rBytes)
	sBytes := copyBytes(sigStr[32:64])
	s := encodedBytesToBigInt(sBytes)

	return &Signature{r, s}, nil
}

// ParseSignature parses a signature in BER format for the curve type `curve'
// into a Signature type, perfoming some basic sanity checks.
func ParseSignature(sigStr []byte) (*Signature, error) {
	return parseSig(sigStr)
}

// GetR satisfies the chainec PublicKey interface.
func (sig Signature) GetR() *big.Int {
	return sig.R
}

// GetS satisfies the chainec PublicKey interface.
func (sig Signature) GetS() *big.Int {
	return sig.S
}

// GetType satisfies the chainec Signature interface.
func (sig Signature) GetType() int {
	return ecTypeSecSchnorr
}

// IsEqual compares this Signature instance to the one passed, returning true
// if both Signatures are equivalent. A signature is equivalent to another, if
// they both have the same scalar value for R and S.
func (sig Signature) IsEqual(otherSig *Signature) bool {
	return sig.R.Cmp(otherSig.R) == 0 &&
		sig.S.Cmp(otherSig.S) == 0
}

// Verify is the generalized and exported function for the verification of a
// secp256k1 Schnorr signature. BLAKE256 is used as the hashing function.
func (sig Signature) Verify(msg []byte, pubkey *secp256k1.PublicKey) bool {
	ok, _ := schnorrVerify(sig.Serialize(), pubkey, msg,
		chainhash.HashB)

	return ok
}
<|endoftext|>


<commit_before><commit_msg>Added test for import path parsing and ast file finding<commit_after><|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type measurement struct {
	Id                string  `json:"id"`
	CheckId           string  `json:"check_id"`
	Location          string  `json:"location"`
	Url               string  `json:"url"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func location() string {
	l := os.Getenv("LOCATION")
	if len(l) == 0 {
		fmt.Fprintf(os.Stderr, "LOCATION not defined in ENV\n")
		os.Exit(1)
	}

	return l
}

func checks_url() string {
	l := os.Getenv("CHECKS_URL")
	if len(l) == 0 {
		fmt.Fprintf(os.Stderr, "CHECKS_URL not defined in ENV\n")
		os.Exit(1)
	}

	return l
}

func measure(c check) measurement {
	var m measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.CheckId = c.Id
	m.Location = location()

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	m.Url = c.Url

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func measurer(checks chan check, measurements chan measurement) {
	for {
		c := <-checks
		m := measure(c)

		measurements <- m
	}
}

func recorder(measurements chan measurement) {
	for {
		m := <-measurements

		s, err := json.Marshal(m)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", "http://localhost:5000/measurements", body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		resp.Body.Close()
		if err != nil {
			panic(err)
		}

		fmt.Println(resp)
	}
}

func get_checks() []check {
	url := checks_url()

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	check_list := get_checks()

	checks := make(chan check)
	measurements := make(chan measurement)

	go measurer(checks, measurements)
	go recorder(measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<commit_msg>send an array of measurements, starting with just 1<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type measurement struct {
	Id                string  `json:"id"`
	CheckId           string  `json:"check_id"`
	Location          string  `json:"location"`
	Url               string  `json:"url"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func location() string {
	l := os.Getenv("LOCATION")
	if len(l) == 0 {
		fmt.Fprintf(os.Stderr, "LOCATION not defined in ENV\n")
		os.Exit(1)
	}

	return l
}

func checks_url() string {
	l := os.Getenv("CHECKS_URL")
	if len(l) == 0 {
		fmt.Fprintf(os.Stderr, "CHECKS_URL not defined in ENV\n")
		os.Exit(1)
	}

	return l
}

func measure(c check) measurement {
	var m measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.CheckId = c.Id
	m.Location = location()

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	m.Url = c.Url

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func measurer(checks chan check, measurements chan measurement) {
	for {
		c := <-checks
		m := measure(c)

		measurements <- m
	}
}

func recorder(measurements chan measurement) {
	var payload []measurement
	for {
		m := <-measurements
		payload = payload.append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", "http://localhost:5000/measurements", body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		resp.Body.Close()
		if err != nil {
			panic(err)
		}

		fmt.Println(resp)
	}
}

func get_checks() []check {
	url := checks_url()

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	check_list := get_checks()

	checks := make(chan check)
	measurements := make(chan measurement)

	go measurer(checks, measurements)
	go recorder(measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/stevencorona/elastic-haproxy/elb"
	"github.com/stevencorona/elastic-haproxy/haproxy"
	"github.com/stevencorona/elastic-haproxy/statsd"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var defaultConfigFile = "config/elastic.toml"
var flagConfigFile string

func main() {

	// Read in the configuration file settings. I think we should instead
	// serialize to JSON and treat this like a view.
	haproxy.Transform()

	flag.StringVar(&flagConfigFile, "configFile", defaultConfigFile, "Path to toml file")
	flag.Parse()

	conf := LoadConfig(flagConfigFile)

	server := new(haproxy.Server)

	// We use two channels— one to send actions to the server and one to recieve
	// notifications from it. Create them right now.
	actionChan := make(chan haproxy.Action)
	notificationChan := make(chan haproxy.Event)

	// Handle signals gracefully in another goroutine
	go gracefulSignals(server, actionChan, notificationChan)

	// Start up the HAProxy Server
	go server.Start(notificationChan, actionChan)

	// Setup the ELB HTTP Handlers
	go elb.InitApiHandlers()

	// Fire up statsd goroutine if statsd is enabled. This might be better off in
	// a seperate binary to monitor HAProxy.
	if conf.Statsd.Enabled {
		go statsd.SendMetrics(server)
	}

	// Event loop for handling events from the HAProxy server
	// (right now, it only sends start/stop notifications)
	for _ = range notificationChan {
		log.Println("Received a notification")
		time.Sleep(2 * time.Second)

		server.Socket = conf.Haproxy.Socket
		serverInfo := server.GetInfo()
		log.Println(serverInfo)
	}

}

// Capture, capture, capture those signals. We gracefully shutdown on a
// SIGQUIT (this is what docker sends w/ docker stop) and everything else
// causes a graceful (zero-downtime) restart of HAProxy. Remember, HAProxy
// does not support reloads or graceful restarts without some... trickery.
func gracefulSignals(server *haproxy.Server, actionChan chan haproxy.Action, notificationChan chan haproxy.Event) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)

	for s := range signals {
		log.Println("Received a signal", s)

		switch s {
		case syscall.SIGQUIT:
			log.Println("Caught SIGQUIT, Stopping HAProxy")
			actionChan <- haproxy.WantsStop
			<-notificationChan
			os.Exit(1)
		default:
			actionChan <- haproxy.WantsReload
		}
	}
}
<commit_msg>dont sleep when recieving a notification<commit_after>package main

import (
	"flag"
	"github.com/stevencorona/elastic-haproxy/elb"
	"github.com/stevencorona/elastic-haproxy/haproxy"
	"github.com/stevencorona/elastic-haproxy/statsd"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var defaultConfigFile = "config/elastic.toml"
var flagConfigFile string

func main() {

	// Read in the configuration file settings. I think we should instead
	// serialize to JSON and treat this like a view.
	haproxy.Transform()

	flag.StringVar(&flagConfigFile, "configFile", defaultConfigFile, "Path to toml file")
	flag.Parse()

	conf := LoadConfig(flagConfigFile)

	server := new(haproxy.Server)

	// We use two channels— one to send actions to the server and one to recieve
	// notifications from it. Create them right now.
	actionChan := make(chan haproxy.Action)
	notificationChan := make(chan haproxy.Event)

	// Handle signals gracefully in another goroutine
	go gracefulSignals(server, actionChan, notificationChan)

	// Start up the HAProxy Server
	go server.Start(notificationChan, actionChan)

	// Setup the ELB HTTP Handlers
	go elb.InitApiHandlers()

	// Fire up statsd goroutine if statsd is enabled. This might be better off in
	// a seperate binary to monitor HAProxy.
	if conf.Statsd.Enabled {
		go statsd.SendMetrics(server)
	}

	// Event loop for handling events from the HAProxy server
	// (right now, it only sends start/stop notifications)
	for _ = range notificationChan {
		log.Println("Received a notification")

		server.Socket = conf.Haproxy.Socket
		serverInfo := server.GetInfo()
		log.Println(serverInfo)
	}

}

// Capture, capture, capture those signals. We gracefully shutdown on a
// SIGQUIT (this is what docker sends w/ docker stop) and everything else
// causes a graceful (zero-downtime) restart of HAProxy. Remember, HAProxy
// does not support reloads or graceful restarts without some... trickery.
func gracefulSignals(server *haproxy.Server, actionChan chan haproxy.Action, notificationChan chan haproxy.Event) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)

	for s := range signals {
		log.Println("Received a signal", s)

		switch s {
		case syscall.SIGQUIT:
			log.Println("Caught SIGQUIT, Stopping HAProxy")
			actionChan <- haproxy.WantsStop
			<-notificationChan
			os.Exit(1)
		default:
			actionChan <- haproxy.WantsReload
		}
	}
}
<|endoftext|>


<commit_before>package rendezvous

import (
	"reflect"
	"testing"
)

type getTestcase struct {
	key          string
	expectedNode string
}

func TestHashGet(t *testing.T) {
	hash := New()

	gotNode := hash.Get("foo")
	if len(gotNode) != 0 {
		t.Errorf("got: %#v, expected: %#v", gotNode, "")
	}

	hash.Add("a", "b", "c", "d", "e")

	testcases := []getTestcase{
		{"", "d"},
		{"foo", "e"},
		{"bar", "c"},
	}

	for _, testcase := range testcases {
		gotNode := hash.Get(testcase.key)
		if gotNode != testcase.expectedNode {
			t.Errorf("got: %#v, expected: %#v", gotNode, testcase.expectedNode)
		}
	}
}

type getNTestcase struct {
	count         int
	key           string
	expectedNodes []string
}

func Test_Hash_GetN(t *testing.T) {
	hash := New()

	gotNodes := hash.GetN(2, "foo")
	if len(gotNodes) != 0 {
		t.Errorf("got: %#v, expected: %#v", gotNodes, []string{})
	}

	hash.Add("a", "b", "c", "d", "e")

	testcases := []getNTestcase{
		{1, "foo", []string{"e"}},
		{2, "bar", []string{"c", "e"}},
		{3, "baz", []string{"d", "a", "b"}},
		{2, "biz", []string{"b", "a"}},
		{0, "boz", []string{}},
		{100, "floo", []string{"d", "a", "b", "c", "e"}},
	}

	for _, testcase := range testcases {
		gotNodes := hash.GetN(testcase.count, testcase.key)
		if !reflect.DeepEqual(gotNodes, testcase.expectedNodes) {
			t.Errorf("got: %#v, expected: %#v", gotNodes, testcase.expectedNodes)
		}
	}
}
<commit_msg>Add basic benchmarks.<commit_after>package rendezvous

import (
	"reflect"
	"testing"
)

var (
	sampleKeys = []string{
		"352DAB08-C1FD-4462-B573-7640B730B721",
		"382080D3-B847-4BB5-AEA8-644C3E56F4E1",
		"2B340C12-7958-4DBE-952C-67496E15D0C8",
		"BE05F82B-902E-4868-8CC9-EE50A6C64636",
		"C7ECC571-E924-4523-A313-951DFD5D8073",
	}
)

type getTestcase struct {
	key          string
	expectedNode string
}

func TestHashGet(t *testing.T) {
	hash := New()

	gotNode := hash.Get("foo")
	if len(gotNode) != 0 {
		t.Errorf("got: %#v, expected: %#v", gotNode, "")
	}

	hash.Add("a", "b", "c", "d", "e")

	testcases := []getTestcase{
		{"", "d"},
		{"foo", "e"},
		{"bar", "c"},
	}

	for _, testcase := range testcases {
		gotNode := hash.Get(testcase.key)
		if gotNode != testcase.expectedNode {
			t.Errorf("got: %#v, expected: %#v", gotNode, testcase.expectedNode)
		}
	}
}

func BenchmarkHashGet(b *testing.B) {
	hash := New("a", "b", "c", "d", "e")
	for i := 0; i < b.N; i++ {
		hash.Get(sampleKeys[i%len(sampleKeys)])
	}
}

type getNTestcase struct {
	count         int
	key           string
	expectedNodes []string
}

func Test_Hash_GetN(t *testing.T) {
	hash := New()

	gotNodes := hash.GetN(2, "foo")
	if len(gotNodes) != 0 {
		t.Errorf("got: %#v, expected: %#v", gotNodes, []string{})
	}

	hash.Add("a", "b", "c", "d", "e")

	testcases := []getNTestcase{
		{1, "foo", []string{"e"}},
		{2, "bar", []string{"c", "e"}},
		{3, "baz", []string{"d", "a", "b"}},
		{2, "biz", []string{"b", "a"}},
		{0, "boz", []string{}},
		{100, "floo", []string{"d", "a", "b", "c", "e"}},
	}

	for _, testcase := range testcases {
		gotNodes := hash.GetN(testcase.count, testcase.key)
		if !reflect.DeepEqual(gotNodes, testcase.expectedNodes) {
			t.Errorf("got: %#v, expected: %#v", gotNodes, testcase.expectedNodes)
		}
	}
}

func BenchmarkHashGetN(b *testing.B) {
	hash := New("a", "b", "c", "d", "e")
	for i := 0; i < b.N; i++ {
		hash.GetN(3, sampleKeys[i%len(sampleKeys)])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/ninjasphere/app-scheduler/controller"
	"github.com/ninjasphere/app-scheduler/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/logger"
	nmodel "github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc"
	"time"
)

//SchedulerService is a facade used to prevent the app event handler being replaced
//by the service event handler.
type SchedulerService struct {
	scheduler   *controller.Scheduler
	log         *logger.Logger
	conn        *ninja.Connection
	model       *model.Schedule
	service     *rpc.ExportedService
	configStore func(m *model.Schedule)
}

func (s *SchedulerService) init(moduleID string) error {
	s.scheduler = &controller.Scheduler{}
	s.scheduler.SetLogger(s.log)
	s.scheduler.SetConnection(s.conn, time.Millisecond*time.Duration(config.Int(10000, "scheduler", "timeout")))

	var err error
	topic := fmt.Sprintf("$node/%s/app/%s/service/%s", config.Serial(), moduleID, "scheduler")
	announcement := &nmodel.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/scheduler",
	}
	if s.service, err = s.conn.ExportService(s, topic, announcement); err != nil {
		return err
	}
	if err := s.scheduler.Start(s.model); err != nil {
		return err
	}

	s.configStore(s.model)

	return nil
}

// Schedule a new task or re-schedules an existing task.
func (s *SchedulerService) Schedule(task *model.Task) (*string, error) {
	if s.scheduler != nil {
		err := s.Cancel(task.ID)
		if err != nil {
			s.log.Warningf("cancel failed %s", err)
		}
		err = s.scheduler.Schedule(task)
		if err == nil {
			s.model.Tasks = append(s.model.Tasks, task)
			s.configStore(s.model)
		}
		copy := task.ID
		return &copy, err
	}
	return nil, fmt.Errorf("cannot schedule a task while the scheduler is stopped")
}

// Cancel an existing task.
func (s *SchedulerService) Cancel(taskID string) error {
	if s.scheduler != nil {
		var err error
		found := -1
		for i, t := range s.model.Tasks {
			if t.ID == taskID {
				found = i
				break
			}
		}
		if found > -1 {
			err = s.scheduler.Cancel(taskID)
			if err == nil {
				s.model.Tasks = append(s.model.Tasks[0:found], s.model.Tasks[found+1:]...)
				s.configStore(s.model)
			}
		} else {
			err = nil
		}
		return err
	}
	return fmt.Errorf("cannot cancel a task while the scheduler is stopped")
}
<commit_msg>Add support for SetTimeZone and SetCoordinates methods.<commit_after>package main

import (
	"fmt"
	"github.com/ninjasphere/app-scheduler/controller"
	"github.com/ninjasphere/app-scheduler/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/config"
	"github.com/ninjasphere/go-ninja/logger"
	nmodel "github.com/ninjasphere/go-ninja/model"
	"github.com/ninjasphere/go-ninja/rpc"
	"time"
)

//SchedulerService is a facade used to prevent the app event handler being replaced
//by the service event handler.
type SchedulerService struct {
	scheduler   *controller.Scheduler
	log         *logger.Logger
	conn        *ninja.Connection
	model       *model.Schedule
	service     *rpc.ExportedService
	configStore func(m *model.Schedule)
}

func (s *SchedulerService) init(moduleID string) error {
	s.scheduler = &controller.Scheduler{}
	s.scheduler.SetLogger(s.log)
	s.scheduler.SetConnection(s.conn, time.Millisecond*time.Duration(config.Int(10000, "scheduler", "timeout")))

	var err error
	topic := fmt.Sprintf("$node/%s/app/%s/service/%s", config.Serial(), moduleID, "scheduler")
	announcement := &nmodel.ServiceAnnouncement{
		Schema: "http://schema.ninjablocks.com/service/scheduler",
	}
	if s.service, err = s.conn.ExportService(s, topic, announcement); err != nil {
		return err
	}
	if err := s.scheduler.Start(s.model); err != nil {
		return err
	}

	s.configStore(s.model)

	return nil
}

// Schedule a new task or re-schedules an existing task.
func (s *SchedulerService) Schedule(task *model.Task) (*string, error) {
	if s.scheduler != nil {
		err := s.Cancel(task.ID)
		if err != nil {
			s.log.Warningf("cancel failed %s", err)
		}
		err = s.scheduler.Schedule(task)
		if err == nil {
			s.model.Tasks = append(s.model.Tasks, task)
			s.configStore(s.model)
		}
		copy := task.ID
		return &copy, err
	}
	return nil, fmt.Errorf("cannot schedule a task while the scheduler is stopped")
}

// Cancel an existing task.
func (s *SchedulerService) Cancel(taskID string) error {
	if s.scheduler != nil {
		var err error
		found := -1
		for i, t := range s.model.Tasks {
			if t.ID == taskID {
				found = i
				break
			}
		}
		if found > -1 {
			err = s.scheduler.Cancel(taskID)
			if err == nil {
				s.model.Tasks = append(s.model.Tasks[0:found], s.model.Tasks[found+1:]...)
				s.configStore(s.model)
			}
		} else {
			err = nil
		}
		return err
	}
	return fmt.Errorf("cannot cancel a task while the scheduler is stopped")
}

// SetTimeZone sets the scheduler timezone. The scheduler will be restarted.
func (s *SchedulerService) SetTimeZone(timezone string) error {
	if s.scheduler != nil {
		if err := s.scheduler.Stop(); err != nil {
			return err
		}
		s.model.TimeZone = timezone
		if err := s.scheduler.Start(s.model); err != nil {
			return err
		}
		s.configStore(s.model)
	}
	return nil
}

// SetCoordinates set the location coordinates of the schedule. The scheduler will be restarted.
func (s *SchedulerService) SetCoordinates(coordinates *model.Location) error {
	if s.scheduler != nil {
		if err := s.scheduler.Stop(); err != nil {
			return err
		}
		s.model.Location = coordinates
		if err := s.scheduler.Start(s.model); err != nil {
			return err
		}
		s.configStore(s.model)
	}
	return nil
}
<|endoftext|>


<commit_before>package pluto

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/aukbit/fibonacci"
	"github.com/aukbit/pluto/client"
	"github.com/aukbit/pluto/common"
	"github.com/aukbit/pluto/datastore"
	"github.com/aukbit/pluto/server"
	"github.com/aukbit/pluto/server/router"
	"github.com/rs/zerolog"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

const (
	defaultName       = "pluto"
	defaultHealthAddr = ":9090"
)

var (
	ErrDatastoreNotInitialized = errors.New("datastore not initialized")
)

// Service representacion of a pluto service
type Service struct {
	close  chan bool
	cfg    Config
	wg     *sync.WaitGroup
	health *health.Server
	logger zerolog.Logger
}

func init() {
	zerolog.TimestampFieldName = "timestamp"
	zerolog.LevelFieldName = "severity"
	zerolog.MessageFieldName = "message"
	switch os.Getenv("LOGGER_LEVEL") {
	case "debug":
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	default:
		zerolog.SetGlobalLevel(zerolog.InfoLevel)
	}

}

// New returns a new pluto service with Options passed in
func New(opts ...Option) *Service {
	return newService(opts...)
}

func newService(opts ...Option) *Service {
	s := &Service{
		cfg:    newConfig(),
		close:  make(chan bool),
		wg:     &sync.WaitGroup{},
		health: health.NewServer(),
	}
	s.logger = zerolog.New(os.Stderr).With().Timestamp().Logger()
	if len(opts) > 0 {
		s = s.WithOptions(opts...)
	}
	return s
}

// WithOptions clones the current Service, applies the supplied Options, and
// returns the resulting Service. It's safe to use concurrently.
func (s *Service) WithOptions(opts ...Option) *Service {
	for _, opt := range opts {
		opt.apply(s)
	}
	return s
}

// Run starts service
func (s *Service) Run() error {
	s.logger = s.logger.With().Str("id", s.cfg.ID).Str("name", s.cfg.Name).Logger()
	// set health server
	s.setHealthServer()
	// start service
	if err := s.start(); err != nil {
		return err
	}
	// hook run after start
	if err := s.hookAfterStart(); err != nil {
		return err
	}
	// wait for all go routines to finish
	s.wg.Wait()
	s.logger.Warn().Msg(fmt.Sprintf("%s finished", s.Name()))
	return nil
}

// Stop stops service
func (s *Service) Stop() {
	s.logger.Info().Msg(fmt.Sprintf("%s stopping", s.Name()))
	s.close <- true
}

// Push allows to start additional options while service is running
func (s *Service) Push(opts ...Option) {
	for _, opt := range opts {
		opt.apply(s)
	}
}

// Server returns a server instance by name if initialized in service
func (s *Service) Server(name string) (srv *server.Server, ok bool) {
	name = common.SafeName(name, server.DefaultName)
	if srv, ok = s.cfg.Servers[name]; !ok {
		return
	}
	return srv, true
}

// Client returns a client instance by name if initialized in service
func (s *Service) Client(name string) (clt *client.Client, ok bool) {
	name = common.SafeName(name, client.DefaultName)
	if clt, ok = s.cfg.Clients[name]; !ok {
		return
	}
	return clt, true
}

// Datastore returns the datastore instance in initialize in service
func (s *Service) Datastore() (*datastore.Datastore, error) {
	if s.cfg.Datastore != nil {
		return s.cfg.Datastore, nil
	}
	return nil, ErrDatastoreNotInitialized
}

// Health ...
func (s *Service) Health() *healthpb.HealthCheckResponse {
	hcr, err := s.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{Service: s.cfg.ID})
	if err != nil {
		s.logger.Error().Msg(fmt.Sprintf("%s Health() %v", s.Name(), err.Error()))
	}
	return hcr
}

// Name returns service name
func (s *Service) Name() string {
	return s.cfg.Name
}

func (s *Service) setHealthServer() {
	s.health.SetServingStatus(s.cfg.ID, 1)
	// Define Router
	mux := router.New()
	mux.GET("/_health/:module/:name", healthHandler)
	// Define server
	srv := server.New(
		server.Name(s.cfg.Name+"_health"),
		server.Addr(s.cfg.HealthAddr),
		server.Mux(mux),
		// server.Logger(s.logger),
	)
	s.cfg.Servers[srv.Name()] = srv
}

func (s *Service) start() error {
	s.logger.Info().Str("ip4", common.IPaddress()).Int("servers", len(s.cfg.Servers)).Int("clients", len(s.cfg.Clients)).Msg(fmt.Sprintf("%s starting", s.Name()))
	err := s.initDatastore()
	if err != nil {
		return err
	}
	// run servers
	s.startServers()
	// dial clients
	s.startClients()
	// add go routine to WaitGroup
	s.wg.Add(1)
	go s.waitUntilStopOrSig()
	s.logger.Info().Msg(fmt.Sprintf("%s started", s.Name()))
	return nil
}

func (s *Service) hookAfterStart() error {
	hooks, ok := s.cfg.Hooks["after_start"]
	if !ok {
		return nil
	}
	ctx := context.Background()
	ctx = s.WithContext(ctx)
	// ctx = context.WithValue(ctx, Key("logger"), s.logger)
	for _, h := range hooks {
		if err := h(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (s *Service) initDatastore() error {
	db, err := s.Datastore()
	if err == ErrDatastoreNotInitialized {
		return nil
	}
	err = db.Init(
		datastore.Logger(s.logger),
	)
	if err != nil {
		return err
	}
	return nil
}

func (s *Service) startServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(srv *server.Server) {
			defer s.wg.Done()
			f := fibonacci.F()
			for {
				err := srv.Run(
					server.Middlewares(
						datastoreContextMiddleware(s),
						serviceContextMiddleware(s),
					),
					server.UnaryServerInterceptors(
						datastoreContextUnaryServerInterceptor(s),
						serviceContextUnaryServerInterceptor(s),
					),
					server.StreamServerInterceptors(
						datastoreContextStreamServerInterceptor(s),
						serviceContextStreamServerInterceptore(s),
					),
					server.Logger(s.logger),
				)
				if err == nil {
					return
				}
				srv.Logger().Error().Msg(fmt.Sprintf("run failed on server: %v - error: %v", srv.Name(), err.Error()))
				time.Sleep(time.Duration(f()) * time.Second)
			}
		}(srv)
	}
}

// startClients listen to the clientsCh
func (s *Service) startClients() {
	go func() {
		for {
			select {
			case clt, ok := <-s.cfg.clientsCh:
				if !ok {
					break
				}
				s.startClient(clt)
			default:
				time.Sleep(500 * time.Millisecond)
				continue
			}
		}
	}()
}

func (s *Service) startClient(clt *client.Client) {
	go func(clt *client.Client) {
		clt.Init(client.Logger(s.logger))
	}(clt)
}

// waitUntilStopOrSig waits for close channel or syscall Signal
func (s *Service) waitUntilStopOrSig() {
	defer s.wg.Done()
	//  Stop also in case of any host signal
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGTERM, syscall.SIGINT)

outer:
	for {
		select {
		case <-s.close:
			// Waits for call to stop
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		case sig := <-sigch:
			// Waits for signal to stop
			s.logger.Info().Msg(fmt.Sprintf("%s signal %v received", s.Name(), sig))
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		default:
			// keep on looping, non-blocking channel operations
			time.Sleep(50 * time.Millisecond)
			continue
		}
	}
}

func (s *Service) closeClients() {
	close(s.cfg.clientsCh)
	// for _, clt := range s.cfg.Clients {
	// 	// add go routine to WaitGroup
	// 	s.wg.Add(1)
	// 	go func(clt *client.Client) {
	// 		defer s.wg.Done()
	// 		clt.Close()
	// 	}(clt)
	// }
}

func (s *Service) stopServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(srv *server.Server) {
			defer s.wg.Done()
			srv.Stop()
		}(srv)
	}
}
<commit_msg>remove logger comment<commit_after>package pluto

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/aukbit/fibonacci"
	"github.com/aukbit/pluto/client"
	"github.com/aukbit/pluto/common"
	"github.com/aukbit/pluto/datastore"
	"github.com/aukbit/pluto/server"
	"github.com/aukbit/pluto/server/router"
	"github.com/rs/zerolog"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

const (
	defaultName       = "pluto"
	defaultHealthAddr = ":9090"
)

var (
	ErrDatastoreNotInitialized = errors.New("datastore not initialized")
)

// Service representacion of a pluto service
type Service struct {
	close  chan bool
	cfg    Config
	wg     *sync.WaitGroup
	health *health.Server
	logger zerolog.Logger
}

func init() {
	zerolog.TimestampFieldName = "timestamp"
	zerolog.LevelFieldName = "severity"
	zerolog.MessageFieldName = "message"
	switch os.Getenv("LOGGER_LEVEL") {
	case "debug":
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	default:
		zerolog.SetGlobalLevel(zerolog.InfoLevel)
	}

}

// New returns a new pluto service with Options passed in
func New(opts ...Option) *Service {
	return newService(opts...)
}

func newService(opts ...Option) *Service {
	s := &Service{
		cfg:    newConfig(),
		close:  make(chan bool),
		wg:     &sync.WaitGroup{},
		health: health.NewServer(),
	}
	s.logger = zerolog.New(os.Stderr).With().Timestamp().Logger()
	if len(opts) > 0 {
		s = s.WithOptions(opts...)
	}
	return s
}

// WithOptions clones the current Service, applies the supplied Options, and
// returns the resulting Service. It's safe to use concurrently.
func (s *Service) WithOptions(opts ...Option) *Service {
	for _, opt := range opts {
		opt.apply(s)
	}
	return s
}

// Run starts service
func (s *Service) Run() error {
	s.logger = s.logger.With().Str("id", s.cfg.ID).Str("name", s.cfg.Name).Logger()
	// set health server
	s.setHealthServer()
	// start service
	if err := s.start(); err != nil {
		return err
	}
	// hook run after start
	if err := s.hookAfterStart(); err != nil {
		return err
	}
	// wait for all go routines to finish
	s.wg.Wait()
	s.logger.Warn().Msg(fmt.Sprintf("%s finished", s.Name()))
	return nil
}

// Stop stops service
func (s *Service) Stop() {
	s.logger.Info().Msg(fmt.Sprintf("%s stopping", s.Name()))
	s.close <- true
}

// Push allows to start additional options while service is running
func (s *Service) Push(opts ...Option) {
	for _, opt := range opts {
		opt.apply(s)
	}
}

// Server returns a server instance by name if initialized in service
func (s *Service) Server(name string) (srv *server.Server, ok bool) {
	name = common.SafeName(name, server.DefaultName)
	if srv, ok = s.cfg.Servers[name]; !ok {
		return
	}
	return srv, true
}

// Client returns a client instance by name if initialized in service
func (s *Service) Client(name string) (clt *client.Client, ok bool) {
	name = common.SafeName(name, client.DefaultName)
	if clt, ok = s.cfg.Clients[name]; !ok {
		return
	}
	return clt, true
}

// Datastore returns the datastore instance in initialize in service
func (s *Service) Datastore() (*datastore.Datastore, error) {
	if s.cfg.Datastore != nil {
		return s.cfg.Datastore, nil
	}
	return nil, ErrDatastoreNotInitialized
}

// Health ...
func (s *Service) Health() *healthpb.HealthCheckResponse {
	hcr, err := s.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{Service: s.cfg.ID})
	if err != nil {
		s.logger.Error().Msg(fmt.Sprintf("%s Health() %v", s.Name(), err.Error()))
	}
	return hcr
}

// Name returns service name
func (s *Service) Name() string {
	return s.cfg.Name
}

func (s *Service) setHealthServer() {
	s.health.SetServingStatus(s.cfg.ID, 1)
	// Define Router
	mux := router.New()
	mux.GET("/_health/:module/:name", healthHandler)
	// Define server
	srv := server.New(
		server.Name(s.cfg.Name+"_health"),
		server.Addr(s.cfg.HealthAddr),
		server.Mux(mux),
		// server.Logger(s.logger),
	)
	s.cfg.Servers[srv.Name()] = srv
}

func (s *Service) start() error {
	s.logger.Info().Str("ip4", common.IPaddress()).Int("servers", len(s.cfg.Servers)).Int("clients", len(s.cfg.Clients)).Msg(fmt.Sprintf("%s starting", s.Name()))
	err := s.initDatastore()
	if err != nil {
		return err
	}
	// run servers
	s.startServers()
	// dial clients
	s.startClients()
	// add go routine to WaitGroup
	s.wg.Add(1)
	go s.waitUntilStopOrSig()
	s.logger.Info().Msg(fmt.Sprintf("%s started", s.Name()))
	return nil
}

func (s *Service) hookAfterStart() error {
	hooks, ok := s.cfg.Hooks["after_start"]
	if !ok {
		return nil
	}
	ctx := context.Background()
	ctx = s.WithContext(ctx)
	for _, h := range hooks {
		if err := h(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (s *Service) initDatastore() error {
	db, err := s.Datastore()
	if err == ErrDatastoreNotInitialized {
		return nil
	}
	err = db.Init(
		datastore.Logger(s.logger),
	)
	if err != nil {
		return err
	}
	return nil
}

func (s *Service) startServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(srv *server.Server) {
			defer s.wg.Done()
			f := fibonacci.F()
			for {
				err := srv.Run(
					server.Middlewares(
						datastoreContextMiddleware(s),
						serviceContextMiddleware(s),
					),
					server.UnaryServerInterceptors(
						datastoreContextUnaryServerInterceptor(s),
						serviceContextUnaryServerInterceptor(s),
					),
					server.StreamServerInterceptors(
						datastoreContextStreamServerInterceptor(s),
						serviceContextStreamServerInterceptore(s),
					),
					server.Logger(s.logger),
				)
				if err == nil {
					return
				}
				srv.Logger().Error().Msg(fmt.Sprintf("run failed on server: %v - error: %v", srv.Name(), err.Error()))
				time.Sleep(time.Duration(f()) * time.Second)
			}
		}(srv)
	}
}

// startClients listen to the clientsCh
func (s *Service) startClients() {
	go func() {
		for {
			select {
			case clt, ok := <-s.cfg.clientsCh:
				if !ok {
					break
				}
				s.startClient(clt)
			default:
				time.Sleep(500 * time.Millisecond)
				continue
			}
		}
	}()
}

func (s *Service) startClient(clt *client.Client) {
	go func(clt *client.Client) {
		clt.Init(client.Logger(s.logger))
	}(clt)
}

// waitUntilStopOrSig waits for close channel or syscall Signal
func (s *Service) waitUntilStopOrSig() {
	defer s.wg.Done()
	//  Stop also in case of any host signal
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGTERM, syscall.SIGINT)

outer:
	for {
		select {
		case <-s.close:
			// Waits for call to stop
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		case sig := <-sigch:
			// Waits for signal to stop
			s.logger.Info().Msg(fmt.Sprintf("%s signal %v received", s.Name(), sig))
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		default:
			// keep on looping, non-blocking channel operations
			time.Sleep(50 * time.Millisecond)
			continue
		}
	}
}

func (s *Service) closeClients() {
	close(s.cfg.clientsCh)
	// for _, clt := range s.cfg.Clients {
	// 	// add go routine to WaitGroup
	// 	s.wg.Add(1)
	// 	go func(clt *client.Client) {
	// 		defer s.wg.Done()
	// 		clt.Close()
	// 	}(clt)
	// }
}

func (s *Service) stopServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(srv *server.Server) {
			defer s.wg.Done()
			srv.Stop()
		}(srv)
	}
}
<|endoftext|>


<commit_before>package pluto

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"syscall"
	"time"

	"github.com/aukbit/fibonacci"
	"github.com/aukbit/pluto/v6/auth/jwt"
	"github.com/aukbit/pluto/v6/client"
	"github.com/aukbit/pluto/v6/common"
	"github.com/aukbit/pluto/v6/server"
	"github.com/aukbit/pluto/v6/server/router"
	"github.com/rs/zerolog"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

const (
	defaultName       = "pluto"
	defaultHealthAddr = ":9090"
)

var (
	ErrDatastoreNotInitialized = errors.New("datastore not initialized")
)

// Service representacion of a pluto service
type Service struct {
	close  chan bool
	cfg    Config
	wg     *sync.WaitGroup
	health *health.Server
	logger zerolog.Logger
}

func init() {
	zerolog.TimestampFieldName = "timestamp"
	zerolog.TimeFieldFormat = time.RFC3339Nano
	zerolog.LevelFieldName = "severity"
	zerolog.MessageFieldName = "message"
	switch os.Getenv("LOGGER_LEVEL") {
	case "debug":
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	default:
		zerolog.SetGlobalLevel(zerolog.InfoLevel)
	}

}

// New returns a new pluto service with Options passed in
func New(opts ...Option) *Service {
	return newService(opts...)
}

func newService(opts ...Option) *Service {
	s := &Service{
		cfg:    newConfig(),
		close:  make(chan bool),
		wg:     &sync.WaitGroup{},
		health: health.NewServer(),
	}
	s.logger = zerolog.New(os.Stderr).With().Timestamp().Logger()
	if len(opts) > 0 {
		s = s.WithOptions(opts...)
	}
	return s
}

// WithOptions clones the current Service, applies the supplied Options, and
// returns the resulting Service. It's safe to use concurrently.
func (s *Service) WithOptions(opts ...Option) *Service {
	for _, opt := range opts {
		opt.apply(s)
	}
	return s
}

// Run starts service
func (s *Service) Run() error {
	s.logger = s.logger.With().Str("id", s.cfg.ID).Str("name", s.cfg.Name).Logger()
	// // set health server
	// s.setHealthServer()
	// start service
	if err := s.start(); err != nil {
		return err
	}
	// hook run after start
	if err := s.hookAfterStart(); err != nil {
		return err
	}
	// wait for all go routines to finish
	s.wg.Wait()
	s.logger.Warn().Msg(fmt.Sprintf("%s has just exited", s.Name()))
	return nil
}

// Stop stops service
func (s *Service) Stop() {
	s.logger.Info().Msg(fmt.Sprintf("shutting down %s", s.Name()))
	s.close <- true
}

// Push allows to start additional options while service is running
func (s *Service) Push(opts ...Option) {
	for _, opt := range opts {
		opt.apply(s)
	}
}

// Server returns a server instance by name if initialized in service
func (s *Service) Server(name string) (srv *server.Server, ok bool) {
	name = common.SafeName(name, server.DefaultName)
	if srv, ok = s.cfg.Servers[name]; !ok {
		return
	}
	return srv, true
}

// Client returns a client instance by name if initialized in service
func (s *Service) Client(name string) (clt *client.Client, ok bool) {
	name = common.SafeName(name, client.DefaultName)
	if clt, ok = s.cfg.Clients[name]; !ok {
		return
	}
	return clt, true
}

// Health ...
func (s *Service) Health() *healthpb.HealthCheckResponse {
	hcr, err := s.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{Service: s.cfg.ID})
	if err != nil {
		s.logger.Error().Msg(err.Error())
	}
	return hcr
}

// ID returns service id
func (s *Service) ID() string {
	return s.cfg.ID
}

// Name returns service name
func (s *Service) Name() string {
	return s.cfg.Name
}

// Logger returns service logger
func (s *Service) Logger() zerolog.Logger {
	return s.logger
}

// WaitUntilFinish returns service name
func (s *Service) WaitUntilFinish() {
	<-s.close
}

func (s *Service) setHealthServer() {
	s.health.SetServingStatus(s.cfg.ID, 1)
	// Define Router
	mux := router.New()
	mux.GET("/healthz/ready", readyHealthHandler)
	mux.GET("/healthz/live", liveHealthHandler)
	mux.GET("/_health/:module/:name", healthHandler)
	// Define server
	srv := server.New(
		server.Name(s.cfg.Name+"_health"),
		server.Addr(s.cfg.HealthAddr),
		server.Mux(mux),
		// server.Logger(s.logger),
	)
	s.cfg.Servers[srv.Name()] = srv
}

func (s *Service) start() error {
	s.logger.Info().Str("ip4", common.IPaddress()).
		Int("servers", len(s.cfg.Servers)).
		Int("clients", len(s.cfg.Clients)).
		Msg(fmt.Sprintf("starting %s, servers: %d clients: %d", s.Name(), len(s.cfg.Servers), len(s.cfg.Clients)))
	// run servers
	s.startServers()
	// dial clients
	s.startClients()
	// add go routine to WaitGroup
	s.wg.Add(1)
	go s.waitUntilStopOrSig()
	return nil
}

func (s *Service) hookAfterStart() error {
	hooks, ok := s.cfg.Hooks["after_start"]
	if !ok {
		return nil
	}
	ctx := context.Background()
	// make service available in hooks context
	ctx = s.WithContext(ctx)
	// make logger available in hooks context
	sublogger := s.logger.With().Str("hook", "after_start").Logger()
	ctx = sublogger.WithContext(ctx)
	for _, h := range hooks {
		if err := h(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (s *Service) startServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(s *Service, srv *server.Server) {
			defer s.wg.Done()
			f := fibonacci.F()
			for {
				err := srv.Run(
					server.Middlewares(
						serviceContextMiddleware(s),
					),
					server.UnaryServerInterceptors(
						serviceContextUnaryServerInterceptor(s),
					),
					server.StreamServerInterceptors(
						serviceContextStreamServerInterceptore(s),
					),
					server.Logger(s.logger),
				)
				if err == nil {
					return
				}
				l := srv.Logger()
				l.Error().Msg(fmt.Sprintf("%v failed to start, error: %v", srv.Name(), err.Error()))
				time.Sleep(time.Duration(f()) * time.Second)
			}
		}(s, srv)
	}
}

// startClients listen to the clientsCh
func (s *Service) startClients() {
	go func() {
		for {
			select {
			case clt, ok := <-s.cfg.clientsCh:
				if !ok {
					break
				}
				s.startClient(clt)
			default:
				time.Sleep(500 * time.Millisecond)
				continue
			}
		}
	}()
}

func (s *Service) startClient(clt *client.Client) {
	go func(clt *client.Client) {
		clt.Init(client.Logger(s.logger))
	}(clt)
}

// waitUntilStopOrSig waits for close channel or syscall Signal
func (s *Service) waitUntilStopOrSig() {
	defer s.wg.Done()
	//  Stop also in case of any host signal
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGTERM, syscall.SIGINT)

outer:
	for {
		select {
		case <-s.close:
			// Waits for call to stop
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		case sig := <-sigch:
			// Waits for signal to stop
			s.logger.Info().Msg(fmt.Sprintf("shutting down, got signal: %v", sig))
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		default:
			// keep on looping, non-blocking channel operations
			time.Sleep(50 * time.Millisecond)
			continue
		}
	}
}

func (s *Service) closeClients() {
	close(s.cfg.clientsCh)
	// for _, clt := range s.cfg.Clients {
	// 	// add go routine to WaitGroup
	// 	s.wg.Add(1)
	// 	go func(clt *client.Client) {
	// 		defer s.wg.Done()
	// 		clt.Close()
	// 	}(clt)
	// }
}

func (s *Service) stopServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(s *Service, srv *server.Server) {
			defer s.wg.Done()
			srv.Stop()
		}(s, srv)
	}
}

// Call invoque's the methodName in the specific clientName initialize
func Call(ctx context.Context, clientName, methodName string, args ...interface{}) (interface{}, error) {

	c, ok := FromContext(ctx).Client(clientName)
	if !ok {
		return nil, fmt.Errorf("grpc client: %v not available", clientName)
	}
	conn, err := c.Dial(client.Timeout(2 * time.Second))
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	inputs := make([]reflect.Value, len(args)+1)
	// add context to the first position
	inputs[0] = reflect.ValueOf(ctx)
	for i := range args {
		inputs[i+1] = reflect.ValueOf(args[i])
	}

	resp := reflect.ValueOf(c.Stub(conn)).MethodByName(methodName).Call(inputs)

	if resp[1].Interface() != nil {
		return nil, resp[1].Interface().(error)
	}

	return resp[0].Interface(), nil
}

// CallWithCredentials invoque's the methodName in the specific clientName initialize
func CallWithCredentials(ctx context.Context, clientName, methodName string, args ...interface{}) (interface{}, error) {

	// Get token from context
	t, ok := jwt.TokenFromContext(ctx)
	if !ok {
		return "", fmt.Errorf("token not available in context")
	}

	c, ok := FromContext(ctx).Client(clientName)
	if !ok {
		return nil, fmt.Errorf("grpc client: %v not available", clientName)
	}

	conn, err := c.DialWithCredentials(t, client.Timeout(2*time.Second))
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	inputs := make([]reflect.Value, len(args)+1)
	// add context to the first position
	inputs[0] = reflect.ValueOf(ctx)
	for i := range args {
		inputs[i+1] = reflect.ValueOf(args[i])
	}

	resp := reflect.ValueOf(c.Stub(conn)).MethodByName(methodName).Call(inputs)

	if resp[1].Interface() != nil {
		return nil, resp[1].Interface().(error)
	}

	return resp[0].Interface(), nil
}
<commit_msg>fix last commit<commit_after>package pluto

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"syscall"
	"time"

	"github.com/aukbit/fibonacci"
	"github.com/aukbit/pluto/v6/auth/jwt"
	"github.com/aukbit/pluto/v6/client"
	"github.com/aukbit/pluto/v6/common"
	"github.com/aukbit/pluto/v6/server"
	"github.com/aukbit/pluto/v6/server/router"
	"github.com/rs/zerolog"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

const (
	defaultName       = "pluto"
	defaultHealthAddr = ":9090"
)

var (
	ErrDatastoreNotInitialized = errors.New("datastore not initialized")
)

// Service representacion of a pluto service
type Service struct {
	close  chan bool
	cfg    Config
	wg     *sync.WaitGroup
	health *health.Server
	logger zerolog.Logger
}

func init() {
	zerolog.TimestampFieldName = "timestamp"
	zerolog.TimeFieldFormat = time.RFC3339Nano
	zerolog.LevelFieldName = "severity"
	zerolog.MessageFieldName = "message"
	switch os.Getenv("LOGGER_LEVEL") {
	case "debug":
		zerolog.SetGlobalLevel(zerolog.DebugLevel)
	default:
		zerolog.SetGlobalLevel(zerolog.InfoLevel)
	}

}

// New returns a new pluto service with Options passed in
func New(opts ...Option) *Service {
	return newService(opts...)
}

func newService(opts ...Option) *Service {
	s := &Service{
		cfg:    newConfig(),
		close:  make(chan bool),
		wg:     &sync.WaitGroup{},
		health: health.NewServer(),
	}
	s.logger = zerolog.New(os.Stderr).With().Timestamp().Logger()
	if len(opts) > 0 {
		s = s.WithOptions(opts...)
	}
	return s
}

// WithOptions clones the current Service, applies the supplied Options, and
// returns the resulting Service. It's safe to use concurrently.
func (s *Service) WithOptions(opts ...Option) *Service {
	for _, opt := range opts {
		opt.apply(s)
	}
	return s
}

// Run starts service
func (s *Service) Run() error {
	s.logger = s.logger.With().Str("id", s.cfg.ID).Str("name", s.cfg.Name).Logger()
	// set health server
	s.setHealthServer()
	// start service
	if err := s.start(); err != nil {
		return err
	}
	// hook run after start
	if err := s.hookAfterStart(); err != nil {
		return err
	}
	// wait for all go routines to finish
	s.wg.Wait()
	s.logger.Warn().Msg(fmt.Sprintf("%s has just exited", s.Name()))
	return nil
}

// Stop stops service
func (s *Service) Stop() {
	s.logger.Info().Msg(fmt.Sprintf("shutting down %s", s.Name()))
	s.close <- true
}

// Push allows to start additional options while service is running
func (s *Service) Push(opts ...Option) {
	for _, opt := range opts {
		opt.apply(s)
	}
}

// Server returns a server instance by name if initialized in service
func (s *Service) Server(name string) (srv *server.Server, ok bool) {
	name = common.SafeName(name, server.DefaultName)
	if srv, ok = s.cfg.Servers[name]; !ok {
		return
	}
	return srv, true
}

// Client returns a client instance by name if initialized in service
func (s *Service) Client(name string) (clt *client.Client, ok bool) {
	name = common.SafeName(name, client.DefaultName)
	if clt, ok = s.cfg.Clients[name]; !ok {
		return
	}
	return clt, true
}

// Health ...
func (s *Service) Health() *healthpb.HealthCheckResponse {
	hcr, err := s.health.Check(
		context.Background(), &healthpb.HealthCheckRequest{Service: s.cfg.ID})
	if err != nil {
		s.logger.Error().Msg(err.Error())
	}
	return hcr
}

// ID returns service id
func (s *Service) ID() string {
	return s.cfg.ID
}

// Name returns service name
func (s *Service) Name() string {
	return s.cfg.Name
}

// Logger returns service logger
func (s *Service) Logger() zerolog.Logger {
	return s.logger
}

// WaitUntilFinish returns service name
func (s *Service) WaitUntilFinish() {
	<-s.close
}

func (s *Service) setHealthServer() {
	s.health.SetServingStatus(s.cfg.ID, 1)
	// Define Router
	mux := router.New()
	mux.GET("/healthz/ready", readyHealthHandler)
	mux.GET("/healthz/live", liveHealthHandler)
	mux.GET("/_health/:module/:name", healthHandler)
	// Define server
	srv := server.New(
		server.Name(s.cfg.Name+"_health"),
		server.Addr(s.cfg.HealthAddr),
		server.Mux(mux),
		// server.Logger(s.logger),
	)
	s.cfg.Servers[srv.Name()] = srv
}

func (s *Service) start() error {
	s.logger.Info().Str("ip4", common.IPaddress()).
		Int("servers", len(s.cfg.Servers)).
		Int("clients", len(s.cfg.Clients)).
		Msg(fmt.Sprintf("starting %s, servers: %d clients: %d", s.Name(), len(s.cfg.Servers), len(s.cfg.Clients)))
	// run servers
	s.startServers()
	// dial clients
	s.startClients()
	// add go routine to WaitGroup
	s.wg.Add(1)
	go s.waitUntilStopOrSig()
	return nil
}

func (s *Service) hookAfterStart() error {
	hooks, ok := s.cfg.Hooks["after_start"]
	if !ok {
		return nil
	}
	ctx := context.Background()
	// make service available in hooks context
	ctx = s.WithContext(ctx)
	// make logger available in hooks context
	sublogger := s.logger.With().Str("hook", "after_start").Logger()
	ctx = sublogger.WithContext(ctx)
	for _, h := range hooks {
		if err := h(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (s *Service) startServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(s *Service, srv *server.Server) {
			defer s.wg.Done()
			f := fibonacci.F()
			for {
				err := srv.Run(
					server.Middlewares(
						serviceContextMiddleware(s),
					),
					server.UnaryServerInterceptors(
						serviceContextUnaryServerInterceptor(s),
					),
					server.StreamServerInterceptors(
						serviceContextStreamServerInterceptore(s),
					),
					server.Logger(s.logger),
				)
				if err == nil {
					return
				}
				l := srv.Logger()
				l.Error().Msg(fmt.Sprintf("%v failed to start, error: %v", srv.Name(), err.Error()))
				time.Sleep(time.Duration(f()) * time.Second)
			}
		}(s, srv)
	}
}

// startClients listen to the clientsCh
func (s *Service) startClients() {
	go func() {
		for {
			select {
			case clt, ok := <-s.cfg.clientsCh:
				if !ok {
					break
				}
				s.startClient(clt)
			default:
				time.Sleep(500 * time.Millisecond)
				continue
			}
		}
	}()
}

func (s *Service) startClient(clt *client.Client) {
	go func(clt *client.Client) {
		clt.Init(client.Logger(s.logger))
	}(clt)
}

// waitUntilStopOrSig waits for close channel or syscall Signal
func (s *Service) waitUntilStopOrSig() {
	defer s.wg.Done()
	//  Stop also in case of any host signal
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, syscall.SIGTERM, syscall.SIGINT)

outer:
	for {
		select {
		case <-s.close:
			// Waits for call to stop
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		case sig := <-sigch:
			// Waits for signal to stop
			s.logger.Info().Msg(fmt.Sprintf("shutting down, got signal: %v", sig))
			s.health.SetServingStatus(s.cfg.ID, 2)
			s.unregister()
			s.closeClients()
			s.stopServers()
			break outer
		default:
			// keep on looping, non-blocking channel operations
			time.Sleep(50 * time.Millisecond)
			continue
		}
	}
}

func (s *Service) closeClients() {
	close(s.cfg.clientsCh)
	// for _, clt := range s.cfg.Clients {
	// 	// add go routine to WaitGroup
	// 	s.wg.Add(1)
	// 	go func(clt *client.Client) {
	// 		defer s.wg.Done()
	// 		clt.Close()
	// 	}(clt)
	// }
}

func (s *Service) stopServers() {
	for _, srv := range s.cfg.Servers {
		// add go routine to WaitGroup
		s.wg.Add(1)
		go func(s *Service, srv *server.Server) {
			defer s.wg.Done()
			srv.Stop()
		}(s, srv)
	}
}

// Call invoque's the methodName in the specific clientName initialize
func Call(ctx context.Context, clientName, methodName string, args ...interface{}) (interface{}, error) {

	c, ok := FromContext(ctx).Client(clientName)
	if !ok {
		return nil, fmt.Errorf("grpc client: %v not available", clientName)
	}
	conn, err := c.Dial(client.Timeout(2 * time.Second))
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	inputs := make([]reflect.Value, len(args)+1)
	// add context to the first position
	inputs[0] = reflect.ValueOf(ctx)
	for i := range args {
		inputs[i+1] = reflect.ValueOf(args[i])
	}

	resp := reflect.ValueOf(c.Stub(conn)).MethodByName(methodName).Call(inputs)

	if resp[1].Interface() != nil {
		return nil, resp[1].Interface().(error)
	}

	return resp[0].Interface(), nil
}

// CallWithCredentials invoque's the methodName in the specific clientName initialize
func CallWithCredentials(ctx context.Context, clientName, methodName string, args ...interface{}) (interface{}, error) {

	// Get token from context
	t, ok := jwt.TokenFromContext(ctx)
	if !ok {
		return "", fmt.Errorf("token not available in context")
	}

	c, ok := FromContext(ctx).Client(clientName)
	if !ok {
		return nil, fmt.Errorf("grpc client: %v not available", clientName)
	}

	conn, err := c.DialWithCredentials(t, client.Timeout(2*time.Second))
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	inputs := make([]reflect.Value, len(args)+1)
	// add context to the first position
	inputs[0] = reflect.ValueOf(ctx)
	for i := range args {
		inputs[i+1] = reflect.ValueOf(args[i])
	}

	resp := reflect.ValueOf(c.Stub(conn)).MethodByName(methodName).Call(inputs)

	if resp[1].Interface() != nil {
		return nil, resp[1].Interface().(error)
	}

	return resp[0].Interface(), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
)

// Session is the interface used by users to interact with the database.
//
// It extends the Node interface by adding a convinient query builder and
// automatically sets a default consinstency level on all operations
// that do not have a consistency level set.
type Session struct {
	Node     Node
	cons     Consistency
	pageSize int
	prefetch float64
	trace    Tracer
	mu       sync.RWMutex
}

// NewSession wraps an existing Node.
func NewSession(node Node) *Session {
	return &Session{Node: node, cons: Quorum, prefetch: 0.25}
}

// SetConsistency sets the default consistency level for this session. This
// setting can also be changed on a per-query basis and the default value
// is Quorum.
func (s *Session) SetConsistency(cons Consistency) {
	s.mu.Lock()
	s.cons = cons
	s.mu.Unlock()
}

// SetPageSize sets the default page size for this session. A value <= 0 will
// disable paging. This setting can also be changed on a per-query basis.
func (s *Session) SetPageSize(n int) {
	s.mu.Lock()
	s.pageSize = n
	s.mu.Unlock()
}

// SetPrefetch sets the default threshold for pre-fetching new pages. If
// there are only p*pageSize rows remaining, the next page will be requested
// automatically. This value can also be changed on a per-query basis and
// the default value is 0.25.
func (s *Session) SetPrefetch(p float64) {
	s.mu.Lock()
	s.prefetch = p
	s.mu.Unlock()
}

// SetTrace sets the default tracer for this session. This setting can also
// be changed on a per-query basis.
func (s *Session) SetTrace(trace Tracer) {
	s.mu.Lock()
	s.trace = trace
	s.mu.Unlock()
}

// Query generates a new query object for interacting with the database.
// Further details of the query may be tweaked using the resulting query
// value before the query is executed.
func (s *Session) Query(stmt string, values ...interface{}) *Query {
	s.mu.RLock()
	qry := &Query{stmt: stmt, values: values, cons: s.cons,
		session: s, pageSize: s.pageSize, trace: s.trace,
		prefetch: s.prefetch}
	s.mu.RUnlock()
	return qry
}

// Close closes all connections. The session is unusable after this
// operation.
func (s *Session) Close() {
	s.Node.Close()
}

func (s *Session) executeQuery(qry *Query) *Iter {
	conn := s.Node.Pick(nil)
	if conn == nil {
		return &Iter{err: ErrUnavailable}
	}
	return conn.executeQuery(qry)
}

func (s *Session) ExecuteBatch(batch *Batch) error {
	conn := s.Node.Pick(nil)
	if conn == nil {
		return ErrUnavailable
	}
	return conn.executeBatch(batch)
}

// Query represents a CQL statement that can be executed.
type Query struct {
	stmt      string
	values    []interface{}
	cons      Consistency
	pageSize  int
	pageState []byte
	prefetch  float64
	trace     Tracer
	session   *Session
}

// Consistency sets the consistency level for this query. If no consistency
// level have been set, the default consistency level of the cluster
// is used.
func (q *Query) Consistency(c Consistency) *Query {
	q.cons = c
	return q
}

// Trace enables tracing of this query. Look at the documentation of the
// Tracer interface to learn more about tracing.
func (q *Query) Trace(trace Tracer) *Query {
	q.trace = trace
	return q
}

// PageSize will tell the iterator to fetch the result in pages of size n.
// This is useful for iterating over large result sets, but setting the
// page size to low might decrease the performance. This feature is only
// available in Cassandra 2 and onwards.
func (q *Query) PageSize(n int) *Query {
	q.pageSize = n
	return q
}

// SetPrefetch sets the default threshold for pre-fetching new pages. If
// there are only p*pageSize rows remaining, the next page will be requested
// automatically.
func (q *Query) Prefetch(p float64) *Query {
	q.prefetch = p
	return q
}

// Exec executes the query without returning any rows.
func (q *Query) Exec() error {
	iter := q.session.executeQuery(q)
	return iter.err
}

// Iter executes the query and returns an iterator capable of iterating
// over all results.
func (q *Query) Iter() *Iter {
	return q.session.executeQuery(q)
}

// Scan executes the query, copies the columns of the first selected
// row into the values pointed at by dest and discards the rest. If no rows
// were selected, ErrNotFound is returned.
func (q *Query) Scan(dest ...interface{}) error {
	iter := q.Iter()
	iter.Scan(dest...)
	return iter.Close()
}

// Iter represents an iterator that can be used to iterate over all rows that
// were returned by a query. The iterator might send additional queries to the
// database during the iteration if paging was enabled.
type Iter struct {
	err     error
	pos     int
	rows    [][][]byte
	columns []ColumnInfo
	next    *nextIter
}

// Columns returns the name and type of the selected columns.
func (iter *Iter) Columns() []ColumnInfo {
	return iter.columns
}

// Scan consumes the next row of the iterator and copies the columns of the
// current row into the values pointed at by dest. Scan might send additional
// queries to the database to retrieve the next set of rows if paging was
// enabled.
//
// Scan returns true if the row was successfully unmarshaled or false if the
// end of the result set was reached or if an error occurred. Close should
// be called afterwards to retrieve any potential errors.
func (iter *Iter) Scan(dest ...interface{}) bool {
	if iter.err != nil {
		return false
	}
	if iter.pos >= len(iter.rows) {
		if iter.next != nil {
			*iter = *iter.next.fetch()
			return iter.Scan(dest...)
		}
		return false
	}
	if iter.next != nil && iter.pos == iter.next.pos {
		go iter.next.fetch()
	}
	if len(dest) != len(iter.columns) {
		iter.err = errors.New("count mismatch")
		return false
	}
	for i := 0; i < len(iter.columns); i++ {
		err := Unmarshal(iter.columns[i].TypeInfo, iter.rows[iter.pos][i], dest[i])
		if err != nil {
			iter.err = err
			return false
		}
	}
	iter.pos++
	return true
}

// Close closes the iterator and returns any errors that happened during
// the query or the iteration.
func (iter *Iter) Close() error {
	return iter.err
}

type nextIter struct {
	qry  Query
	pos  int
	once sync.Once
	next *Iter
}

func (n *nextIter) fetch() *Iter {
	n.once.Do(func() {
		n.next = n.qry.session.executeQuery(&n.qry)
	})
	return n.next
}

type Batch struct {
	Type    BatchType
	Entries []BatchEntry
	Cons    Consistency
}

func NewBatch(typ BatchType) *Batch {
	return &Batch{Type: typ}
}

func (b *Batch) Query(stmt string, args ...interface{}) {
	b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, Args: args})
}

type BatchType int

const (
	LoggedBatch   BatchType = 0
	UnloggedBatch BatchType = 1
	CounterBatch  BatchType = 2
)

type BatchEntry struct {
	Stmt string
	Args []interface{}
}

type Consistency int

const (
	Any Consistency = 1 + iota
	One
	Two
	Three
	Quorum
	All
	LocalQuorum
	EachQuorum
	Serial
	LocalSerial
)

var consinstencyNames = []string{
	0:           "default",
	Any:         "any",
	One:         "one",
	Two:         "two",
	Three:       "three",
	Quorum:      "quorum",
	All:         "all",
	LocalQuorum: "localquorum",
	EachQuorum:  "eachquorum",
	Serial:      "serial",
	LocalSerial: "localserial",
}

func (c Consistency) String() string {
	return consinstencyNames[c]
}

type ColumnInfo struct {
	Keyspace string
	Table    string
	Name     string
	TypeInfo *TypeInfo
}

// Tracer is the interface implemented by query tracers. Tracers have the
// ability to obtain a detailed event log of all events that happened during
// the execution of a query from Cassandra. Gathering this information might
// be essential for debugging and optimizing queries, but this feature should
// not be used on production systems with very high load.
type Tracer interface {
	Trace(traceId []byte)
}

type traceWriter struct {
	session *Session
	w       io.Writer
	mu      sync.Mutex
}

// NewTraceWriter returns a simple Tracer implementation that outputs
// the event log in a textual format.
func NewTraceWriter(session *Session, w io.Writer) Tracer {
	return traceWriter{session: session, w: w}
}

func (t traceWriter) Trace(traceId []byte) {
	var (
		coordinator string
		duration    int
	)
	t.session.Query(`SELECT coordinator, duration
			FROM system_traces.sessions
			WHERE session_id = ?`, traceId).
		Consistency(One).Scan(&coordinator, &duration)

	iter := t.session.Query(`SELECT event_id, activity, source, source_elapsed
			FROM system_traces.events
			WHERE session_id = ?`, traceId).
		Consistency(One).Iter()
	var (
		timestamp time.Time
		activity  string
		source    string
		elapsed   int
	)
	t.mu.Lock()
	defer t.mu.Unlock()
	fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
		traceId, coordinator, time.Duration(duration)*time.Microsecond)
	for iter.Scan(&timestamp, &activity, &source, &elapsed) {
		fmt.Fprintf(t.w, "%s: %s (source: %s, elapsed: %d)\n",
			timestamp.Format("2006/01/02 15:04:05.999999"), activity, source, elapsed)
	}
	if err := iter.Close(); err != nil {
		fmt.Fprintln(t.w, "Error:", err)
	}
}

type Error struct {
	Code    int
	Message string
}

func (e Error) Error() string {
	return e.Message
}

var (
	ErrNotFound    = errors.New("not found")
	ErrUnavailable = errors.New("unavailable")
	ErrProtocol    = errors.New("protocol error")
	ErrUnsupported = errors.New("feature not supported")
)
<commit_msg>Return ErrNotFound when no row matches the query<commit_after>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
)

// Session is the interface used by users to interact with the database.
//
// It extends the Node interface by adding a convinient query builder and
// automatically sets a default consinstency level on all operations
// that do not have a consistency level set.
type Session struct {
	Node     Node
	cons     Consistency
	pageSize int
	prefetch float64
	trace    Tracer
	mu       sync.RWMutex
}

// NewSession wraps an existing Node.
func NewSession(node Node) *Session {
	return &Session{Node: node, cons: Quorum, prefetch: 0.25}
}

// SetConsistency sets the default consistency level for this session. This
// setting can also be changed on a per-query basis and the default value
// is Quorum.
func (s *Session) SetConsistency(cons Consistency) {
	s.mu.Lock()
	s.cons = cons
	s.mu.Unlock()
}

// SetPageSize sets the default page size for this session. A value <= 0 will
// disable paging. This setting can also be changed on a per-query basis.
func (s *Session) SetPageSize(n int) {
	s.mu.Lock()
	s.pageSize = n
	s.mu.Unlock()
}

// SetPrefetch sets the default threshold for pre-fetching new pages. If
// there are only p*pageSize rows remaining, the next page will be requested
// automatically. This value can also be changed on a per-query basis and
// the default value is 0.25.
func (s *Session) SetPrefetch(p float64) {
	s.mu.Lock()
	s.prefetch = p
	s.mu.Unlock()
}

// SetTrace sets the default tracer for this session. This setting can also
// be changed on a per-query basis.
func (s *Session) SetTrace(trace Tracer) {
	s.mu.Lock()
	s.trace = trace
	s.mu.Unlock()
}

// Query generates a new query object for interacting with the database.
// Further details of the query may be tweaked using the resulting query
// value before the query is executed.
func (s *Session) Query(stmt string, values ...interface{}) *Query {
	s.mu.RLock()
	qry := &Query{stmt: stmt, values: values, cons: s.cons,
		session: s, pageSize: s.pageSize, trace: s.trace,
		prefetch: s.prefetch}
	s.mu.RUnlock()
	return qry
}

// Close closes all connections. The session is unusable after this
// operation.
func (s *Session) Close() {
	s.Node.Close()
}

func (s *Session) executeQuery(qry *Query) *Iter {
	conn := s.Node.Pick(nil)
	if conn == nil {
		return &Iter{err: ErrUnavailable}
	}
	return conn.executeQuery(qry)
}

func (s *Session) ExecuteBatch(batch *Batch) error {
	conn := s.Node.Pick(nil)
	if conn == nil {
		return ErrUnavailable
	}
	return conn.executeBatch(batch)
}

// Query represents a CQL statement that can be executed.
type Query struct {
	stmt      string
	values    []interface{}
	cons      Consistency
	pageSize  int
	pageState []byte
	prefetch  float64
	trace     Tracer
	session   *Session
}

// Consistency sets the consistency level for this query. If no consistency
// level have been set, the default consistency level of the cluster
// is used.
func (q *Query) Consistency(c Consistency) *Query {
	q.cons = c
	return q
}

// Trace enables tracing of this query. Look at the documentation of the
// Tracer interface to learn more about tracing.
func (q *Query) Trace(trace Tracer) *Query {
	q.trace = trace
	return q
}

// PageSize will tell the iterator to fetch the result in pages of size n.
// This is useful for iterating over large result sets, but setting the
// page size to low might decrease the performance. This feature is only
// available in Cassandra 2 and onwards.
func (q *Query) PageSize(n int) *Query {
	q.pageSize = n
	return q
}

// SetPrefetch sets the default threshold for pre-fetching new pages. If
// there are only p*pageSize rows remaining, the next page will be requested
// automatically.
func (q *Query) Prefetch(p float64) *Query {
	q.prefetch = p
	return q
}

// Exec executes the query without returning any rows.
func (q *Query) Exec() error {
	iter := q.session.executeQuery(q)
	return iter.err
}

// Iter executes the query and returns an iterator capable of iterating
// over all results.
func (q *Query) Iter() *Iter {
	return q.session.executeQuery(q)
}

// Scan executes the query, copies the columns of the first selected
// row into the values pointed at by dest and discards the rest. If no rows
// were selected, ErrNotFound is returned.
func (q *Query) Scan(dest ...interface{}) error {
	iter := q.Iter()
	if len(iter.rows) == 0 {
		return ErrNotFound
	}
	iter.Scan(dest...)
	return iter.Close()
}

// Iter represents an iterator that can be used to iterate over all rows that
// were returned by a query. The iterator might send additional queries to the
// database during the iteration if paging was enabled.
type Iter struct {
	err     error
	pos     int
	rows    [][][]byte
	columns []ColumnInfo
	next    *nextIter
}

// Columns returns the name and type of the selected columns.
func (iter *Iter) Columns() []ColumnInfo {
	return iter.columns
}

// Scan consumes the next row of the iterator and copies the columns of the
// current row into the values pointed at by dest. Scan might send additional
// queries to the database to retrieve the next set of rows if paging was
// enabled.
//
// Scan returns true if the row was successfully unmarshaled or false if the
// end of the result set was reached or if an error occurred. Close should
// be called afterwards to retrieve any potential errors.
func (iter *Iter) Scan(dest ...interface{}) bool {
	if iter.err != nil {
		return false
	}
	if iter.pos >= len(iter.rows) {
		if iter.next != nil {
			*iter = *iter.next.fetch()
			return iter.Scan(dest...)
		}
		return false
	}
	if iter.next != nil && iter.pos == iter.next.pos {
		go iter.next.fetch()
	}
	if len(dest) != len(iter.columns) {
		iter.err = errors.New("count mismatch")
		return false
	}
	for i := 0; i < len(iter.columns); i++ {
		err := Unmarshal(iter.columns[i].TypeInfo, iter.rows[iter.pos][i], dest[i])
		if err != nil {
			iter.err = err
			return false
		}
	}
	iter.pos++
	return true
}

// Close closes the iterator and returns any errors that happened during
// the query or the iteration.
func (iter *Iter) Close() error {
	return iter.err
}

type nextIter struct {
	qry  Query
	pos  int
	once sync.Once
	next *Iter
}

func (n *nextIter) fetch() *Iter {
	n.once.Do(func() {
		n.next = n.qry.session.executeQuery(&n.qry)
	})
	return n.next
}

type Batch struct {
	Type    BatchType
	Entries []BatchEntry
	Cons    Consistency
}

func NewBatch(typ BatchType) *Batch {
	return &Batch{Type: typ}
}

func (b *Batch) Query(stmt string, args ...interface{}) {
	b.Entries = append(b.Entries, BatchEntry{Stmt: stmt, Args: args})
}

type BatchType int

const (
	LoggedBatch   BatchType = 0
	UnloggedBatch BatchType = 1
	CounterBatch  BatchType = 2
)

type BatchEntry struct {
	Stmt string
	Args []interface{}
}

type Consistency int

const (
	Any Consistency = 1 + iota
	One
	Two
	Three
	Quorum
	All
	LocalQuorum
	EachQuorum
	Serial
	LocalSerial
)

var consinstencyNames = []string{
	0:           "default",
	Any:         "any",
	One:         "one",
	Two:         "two",
	Three:       "three",
	Quorum:      "quorum",
	All:         "all",
	LocalQuorum: "localquorum",
	EachQuorum:  "eachquorum",
	Serial:      "serial",
	LocalSerial: "localserial",
}

func (c Consistency) String() string {
	return consinstencyNames[c]
}

type ColumnInfo struct {
	Keyspace string
	Table    string
	Name     string
	TypeInfo *TypeInfo
}

// Tracer is the interface implemented by query tracers. Tracers have the
// ability to obtain a detailed event log of all events that happened during
// the execution of a query from Cassandra. Gathering this information might
// be essential for debugging and optimizing queries, but this feature should
// not be used on production systems with very high load.
type Tracer interface {
	Trace(traceId []byte)
}

type traceWriter struct {
	session *Session
	w       io.Writer
	mu      sync.Mutex
}

// NewTraceWriter returns a simple Tracer implementation that outputs
// the event log in a textual format.
func NewTraceWriter(session *Session, w io.Writer) Tracer {
	return traceWriter{session: session, w: w}
}

func (t traceWriter) Trace(traceId []byte) {
	var (
		coordinator string
		duration    int
	)
	t.session.Query(`SELECT coordinator, duration
			FROM system_traces.sessions
			WHERE session_id = ?`, traceId).
		Consistency(One).Scan(&coordinator, &duration)

	iter := t.session.Query(`SELECT event_id, activity, source, source_elapsed
			FROM system_traces.events
			WHERE session_id = ?`, traceId).
		Consistency(One).Iter()
	var (
		timestamp time.Time
		activity  string
		source    string
		elapsed   int
	)
	t.mu.Lock()
	defer t.mu.Unlock()
	fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
		traceId, coordinator, time.Duration(duration)*time.Microsecond)
	for iter.Scan(&timestamp, &activity, &source, &elapsed) {
		fmt.Fprintf(t.w, "%s: %s (source: %s, elapsed: %d)\n",
			timestamp.Format("2006/01/02 15:04:05.999999"), activity, source, elapsed)
	}
	if err := iter.Close(); err != nil {
		fmt.Fprintln(t.w, "Error:", err)
	}
}

type Error struct {
	Code    int
	Message string
}

func (e Error) Error() string {
	return e.Message
}

var (
	ErrNotFound    = errors.New("not found")
	ErrUnavailable = errors.New("unavailable")
	ErrProtocol    = errors.New("protocol error")
	ErrUnsupported = errors.New("feature not supported")
)
<|endoftext|>