<commit_before>/*
Package httpsign provides tools for signing and authenticating HTTP requests between
web services. See README.md for more details.
*/
package httpsign

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/mailgun/lemma/random"
	"github.com/mailgun/metrics"
	"github.com/mailgun/timetools"
)

// Modify NonceCacheCapacity and NonceCacheTimeout if your service needs to
// authenticate more than 5,000 requests per second. For example, if you need
// to handle 10,000 requests per second and timeout after one minute,  you may
// want to set NonceCacheTimeout to 60 and NonceCacheCapacity to
// 10000 * cacheTimeout = 600000.
type Config struct {
	Keypath        string   // path to signing key
	HeadersToSign  []string // list of headers to sign
	SignVerbAndURI bool     // include the http verb and uri in request

	NonceCacheCapacity int // capacity of the nonce cache
	NonceCacheTimeout  int // nonce cache timeout

	EmitStats    bool   // toggle emitting metrics or not
	StatsdHost   string // hostname of statsd server
	StatsdPort   int    // port of statsd server
	StatsdPrefix string // prefix to prepend to metrics

	NonceHeaderName            string // default: X-Mailgun-Nonce
	TimestampHeaderName        string // default: X-Mailgun-Timestamp
	SignatureHeaderName        string // default: X-Mailgun-Signature
	SignatureVersionHeaderName string // default: X-Mailgun-Signature-Version
}

// Represents a service that can be used to sign and authenticate requests.
type Service struct {
	config         *Config
	nonceCache     *NonceCache
	randomProvider random.RandomProvider
	timeProvider   timetools.TimeProvider
	secretKey      []byte
	metricsClient  metrics.Client
}

// Return a new Service. Config can not be nil. If you need control over
// setting time and random providers, use NewWithProviders.
func New(config *Config) (*Service, error) {
	return NewWithProviders(
		config,
		&timetools.RealTime{},
		&random.CSPRNG{},
	)
}

// Returns a new Service. Provides control over time and random providers.
func NewWithProviders(config *Config, timeProvider timetools.TimeProvider,
	randomProvider random.RandomProvider) (*Service, error) {

	// config is required!
	if config == nil {
		return nil, fmt.Errorf("config is required.")
	}

	// set defaults if not set
	if config.NonceCacheCapacity < 1 {
		config.NonceCacheCapacity = CacheCapacity
	}
	if config.NonceCacheTimeout < 1 {
		config.NonceCacheTimeout = CacheTimeout
	}
	if config.NonceHeaderName == "" {
		config.NonceHeaderName = XMailgunNonce
	}
	if config.TimestampHeaderName == "" {
		config.TimestampHeaderName = XMailgunTimestamp
	}
	if config.SignatureHeaderName == "" {
		config.SignatureHeaderName = XMailgunSignature
	}
	if config.SignatureVersionHeaderName == "" {
		config.SignatureVersionHeaderName = XMailgunSignatureVersion
	}

	// setup metrics service
	metricsClient := metrics.NewNop()
	if config.EmitStats {
		// get hostname of box
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("failed to obtain hostname: %v", err)
		}

		// build lemma prefix
		prefix := "lemma." + strings.Replace(hostname, ".", "_", -1)
		if config.StatsdPrefix != "" {
			prefix += "." + config.StatsdPrefix
		}

		// build metrics client
		hostport := fmt.Sprintf("%v:%v", config.StatsdHost, config.StatsdPort)
		metricsClient, err = metrics.NewWithOptions(hostport, prefix, metrics.Options{UseBuffering: true})
		if err != nil {
			return nil, err
		}
	}

	// read key from disk, if no key is read that's okay it might be passed in
	keyBytes, err := readKeyFromDisk(config.Keypath)

	// setup nonce cache
	ncache, err := NewNonceCache(config.NonceCacheCapacity, config.NonceCacheTimeout, timeProvider)
	if err != nil {
		return nil, err
	}

	// return service
	return &Service{
		config:         config,
		nonceCache:     ncache,
		secretKey:      keyBytes,
		timeProvider:   timeProvider,
		randomProvider: randomProvider,
		metricsClient:  metricsClient,
	}, nil
}

// Signs a given HTTP request with signature, nonce, and timestamp.
func (s *Service) SignRequest(r *http.Request) error {
	if s.secretKey == nil {
		return fmt.Errorf("service not loaded with key.")
	}
	return s.SignRequestWithKey(r, s.secretKey)
}

// Signs a given HTTP request with signature, nonce, and timestamp. Signs the
// message with the passed in key not the one initialized with.
func (s *Service) SignRequestWithKey(r *http.Request, secretKey []byte) error {
	// extract request body bytes
	var bodyBytes []byte
	var err error
	if r.Body == nil {
		// if we have no body, like a GET request, set it to ""
		bodyBytes = []byte("")
	} else {
		// if we have a body, read it in
		bodyBytes, err = ioutil.ReadAll(r.Body)
		if err != nil {
			return err
		}
		r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))
	}

	// extract any headers if requested
	headerValues, err := extractHeaderValues(r, s.config.HeadersToSign)
	if err != nil {
		return err
	}

	// get 128-bit random number from /dev/urandom and base16 encode it
	nonce, err := s.randomProvider.HexDigest(16)
	if err != nil {
		return fmt.Errorf("unable to get random : %v", err)
	}

	// get current timestamp
	timestamp := strconv.FormatInt(s.timeProvider.UtcNow().Unix(), 10)

	// compute the hmac and base16 encode it
	computedMAC := computeMAC(secretKey, s.config.SignVerbAndURI, r.Method, r.URL.RequestURI(),
		timestamp, nonce, bodyBytes, headerValues)
	signature := hex.EncodeToString(computedMAC)

	// set headers
	r.Header.Set(s.config.NonceHeaderName, nonce)
	r.Header.Set(s.config.TimestampHeaderName, timestamp)
	r.Header.Set(s.config.SignatureHeaderName, signature)
	r.Header.Set(s.config.SignatureVersionHeaderName, "2")

	return nil
}

// Authenticates HTTP request to ensure it was sent by an authorized sender.
func (s *Service) AuthenticateRequest(r *http.Request) error {
	if s.secretKey == nil {
		return fmt.Errorf("service not loaded with key.")
	}
	return s.AuthenticateRequestWithKey(r, s.secretKey)
}

// Authenticates HTTP request to ensure it was sent by an authorized sender.
// Checks message signature with the passed in key, not the one initialized with.
func (s *Service) AuthenticateRequestWithKey(r *http.Request, secretKey []byte) (err error) {
	// Emit a success or failure metric on return.
	defer func() {
		if err == nil {
			s.metricsClient.Inc("success", 1, 1)
		} else {
			s.metricsClient.Inc("failure", 1, 1)
		}
	}()

	// extract parameters
	signature := r.Header.Get(s.config.SignatureHeaderName)
	if signature == "" {
		return fmt.Errorf("header not found: %v", s.config.SignatureHeaderName)
	}
	nonce := r.Header.Get(s.config.NonceHeaderName)
	if nonce == "" {
		return fmt.Errorf("header not found: %v", s.config.NonceHeaderName)
	}
	timestamp := r.Header.Get(s.config.TimestampHeaderName)
	if timestamp == "" {
		return fmt.Errorf("header not found: %v", s.config.TimestampHeaderName)
	}

	// extract request body bytes
	bodyBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))

	// extract any headers if requested
	headerValues, err := extractHeaderValues(r, s.config.HeadersToSign)
	if err != nil {
		return err
	}

	// check the hmac
	isValid, err := checkMAC(secretKey, s.config.SignVerbAndURI, r.Method, r.URL.RequestURI(),
		timestamp, nonce, bodyBytes, headerValues, signature)
	if !isValid {
		return err
	}

	// check timestamp
	isValid, err = s.checkTimestamp(timestamp)
	if !isValid {
		return err
	}

	// check to see if we have seen nonce before
	inCache := s.nonceCache.InCache(nonce)
	if inCache {
		return fmt.Errorf("nonce already in cache: %v", nonce)
	}

	return nil
}

func (s *Service) checkTimestamp(timestampHeader string) (bool, error) {
	// convert unix timestamp string into time struct
	timestamp, err := strconv.ParseInt(timestampHeader, 10, 0)
	if err != nil {
		return false, fmt.Errorf("unable to parse %v: %v", s.config.TimestampHeaderName, timestampHeader)
	}

	now := s.timeProvider.UtcNow().Unix()

	// if timestamp is from the future, it's invalid
	if timestamp >= now+MaxSkewSec {
		return false, fmt.Errorf("timestamp header from the future; now: %v; %v: %v; difference: %v",
			now, s.config.TimestampHeaderName, timestamp, timestamp-now)
	}

	// if the timestamp is older than ttl - skew, it's invalid
	if timestamp <= now-int64(s.nonceCache.cacheTTL-MaxSkewSec) {
		return false, fmt.Errorf("timestamp header too old; now: %v; %v: %v; difference: %v",
			now, s.config.TimestampHeaderName, timestamp, now-timestamp)
	}

	return true, nil
}

func computeMAC(secretKey []byte, signVerbAndUri bool, httpVerb string, httpResourceUri string,
	timestamp string, nonce string, body []byte, headerValues []string) []byte {

	// use hmac-sha256
	mac := hmac.New(sha256.New, secretKey)

	// required parameters (timestamp, nonce, body)
	mac.Write([]byte(fmt.Sprintf("%v|", len(timestamp))))
	mac.Write([]byte(timestamp))
	mac.Write([]byte(fmt.Sprintf("|%v|", len(nonce))))
	mac.Write([]byte(nonce))
	mac.Write([]byte(fmt.Sprintf("|%v|", len(body))))
	mac.Write(body)

	// optional parameters (httpVerb, httpResourceUri)
	if signVerbAndUri {
		mac.Write([]byte(fmt.Sprintf("|%v|", len(httpVerb))))
		mac.Write([]byte(httpVerb))
		mac.Write([]byte(fmt.Sprintf("|%v|", len(httpResourceUri))))
		mac.Write([]byte(httpResourceUri))
	}

	// optional parameters (headers)
	for _, headerValue := range headerValues {
		mac.Write([]byte(fmt.Sprintf("|%v|", len(headerValue))))
		mac.Write([]byte(headerValue))
	}

	return mac.Sum(nil)
}

func checkMAC(secretKey []byte, signVerbAndUri bool, httpVerb string, httpResourceUri string,
	timestamp string, nonce string, body []byte, headerValues []string, signature string) (bool, error) {

	// the hmac we get is a hexdigest (string representation of hex values)
	// which needs to be decoded before before we can use it
	expectedMAC, err := hex.DecodeString(signature)
	if err != nil {
		return false, err
	}

	// compute the hmac
	computedMAC := computeMAC(secretKey, signVerbAndUri, httpVerb, httpResourceUri, timestamp, nonce, body, headerValues)

	// constant time compare
	isEqual := hmac.Equal(expectedMAC, computedMAC)
	if !isEqual {
		return false, fmt.Errorf("signature header value %v does not match computed value %v", expectedMAC, computedMAC)
	}

	return true, nil
}

func extractHeaderValues(r *http.Request, headerNames []string) ([]string, error) {
	if len(headerNames) < 1 {
		return nil, nil
	}

	headerValues := make([]string, len(headerNames))
	for i, headerName := range headerNames {
		_, ok := r.Header[headerName]
		if !ok {
			return nil, fmt.Errorf("header %v not found in request.", headerName)
		}
		headerValues[i] = r.Header.Get(headerName)
	}

	return headerValues, nil
}

func readKeyFromDisk(keypath string) ([]byte, error) {
	// load key from disk
	keyBytes, err := ioutil.ReadFile(keypath)
	if err != nil {
		return nil, err
	}

	// strip newline (\n or 0x0a) if it's at the end
	keyBytes = bytes.TrimSuffix(keyBytes, []byte("\n"))

	return keyBytes, nil
}
<commit_msg>change error message<commit_after>/*
Package httpsign provides tools for signing and authenticating HTTP requests between
web services. See README.md for more details.
*/
package httpsign

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/mailgun/lemma/random"
	"github.com/mailgun/metrics"
	"github.com/mailgun/timetools"
)

// Modify NonceCacheCapacity and NonceCacheTimeout if your service needs to
// authenticate more than 5,000 requests per second. For example, if you need
// to handle 10,000 requests per second and timeout after one minute,  you may
// want to set NonceCacheTimeout to 60 and NonceCacheCapacity to
// 10000 * cacheTimeout = 600000.
type Config struct {
	Keypath        string   // path to signing key
	HeadersToSign  []string // list of headers to sign
	SignVerbAndURI bool     // include the http verb and uri in request

	NonceCacheCapacity int // capacity of the nonce cache
	NonceCacheTimeout  int // nonce cache timeout

	EmitStats    bool   // toggle emitting metrics or not
	StatsdHost   string // hostname of statsd server
	StatsdPort   int    // port of statsd server
	StatsdPrefix string // prefix to prepend to metrics

	NonceHeaderName            string // default: X-Mailgun-Nonce
	TimestampHeaderName        string // default: X-Mailgun-Timestamp
	SignatureHeaderName        string // default: X-Mailgun-Signature
	SignatureVersionHeaderName string // default: X-Mailgun-Signature-Version
}

// Represents a service that can be used to sign and authenticate requests.
type Service struct {
	config         *Config
	nonceCache     *NonceCache
	randomProvider random.RandomProvider
	timeProvider   timetools.TimeProvider
	secretKey      []byte
	metricsClient  metrics.Client
}

// Return a new Service. Config can not be nil. If you need control over
// setting time and random providers, use NewWithProviders.
func New(config *Config) (*Service, error) {
	return NewWithProviders(
		config,
		&timetools.RealTime{},
		&random.CSPRNG{},
	)
}

// Returns a new Service. Provides control over time and random providers.
func NewWithProviders(config *Config, timeProvider timetools.TimeProvider,
	randomProvider random.RandomProvider) (*Service, error) {

	// config is required!
	if config == nil {
		return nil, fmt.Errorf("config is required.")
	}

	// set defaults if not set
	if config.NonceCacheCapacity < 1 {
		config.NonceCacheCapacity = CacheCapacity
	}
	if config.NonceCacheTimeout < 1 {
		config.NonceCacheTimeout = CacheTimeout
	}
	if config.NonceHeaderName == "" {
		config.NonceHeaderName = XMailgunNonce
	}
	if config.TimestampHeaderName == "" {
		config.TimestampHeaderName = XMailgunTimestamp
	}
	if config.SignatureHeaderName == "" {
		config.SignatureHeaderName = XMailgunSignature
	}
	if config.SignatureVersionHeaderName == "" {
		config.SignatureVersionHeaderName = XMailgunSignatureVersion
	}

	// setup metrics service
	metricsClient := metrics.NewNop()
	if config.EmitStats {
		// get hostname of box
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("failed to obtain hostname: %v", err)
		}

		// build lemma prefix
		prefix := "lemma." + strings.Replace(hostname, ".", "_", -1)
		if config.StatsdPrefix != "" {
			prefix += "." + config.StatsdPrefix
		}

		// build metrics client
		hostport := fmt.Sprintf("%v:%v", config.StatsdHost, config.StatsdPort)
		metricsClient, err = metrics.NewWithOptions(hostport, prefix, metrics.Options{UseBuffering: true})
		if err != nil {
			return nil, err
		}
	}

	// read key from disk, if no key is read that's okay it might be passed in
	keyBytes, err := readKeyFromDisk(config.Keypath)

	// setup nonce cache
	ncache, err := NewNonceCache(config.NonceCacheCapacity, config.NonceCacheTimeout, timeProvider)
	if err != nil {
		return nil, err
	}

	// return service
	return &Service{
		config:         config,
		nonceCache:     ncache,
		secretKey:      keyBytes,
		timeProvider:   timeProvider,
		randomProvider: randomProvider,
		metricsClient:  metricsClient,
	}, nil
}

// Signs a given HTTP request with signature, nonce, and timestamp.
func (s *Service) SignRequest(r *http.Request) error {
	if s.secretKey == nil {
		return fmt.Errorf("service not loaded with key.")
	}
	return s.SignRequestWithKey(r, s.secretKey)
}

// Signs a given HTTP request with signature, nonce, and timestamp. Signs the
// message with the passed in key not the one initialized with.
func (s *Service) SignRequestWithKey(r *http.Request, secretKey []byte) error {
	// extract request body bytes
	var bodyBytes []byte
	var err error
	if r.Body == nil {
		// if we have no body, like a GET request, set it to ""
		bodyBytes = []byte("")
	} else {
		// if we have a body, read it in
		bodyBytes, err = ioutil.ReadAll(r.Body)
		if err != nil {
			return err
		}
		r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))
	}

	// extract any headers if requested
	headerValues, err := extractHeaderValues(r, s.config.HeadersToSign)
	if err != nil {
		return err
	}

	// get 128-bit random number from /dev/urandom and base16 encode it
	nonce, err := s.randomProvider.HexDigest(16)
	if err != nil {
		return fmt.Errorf("unable to get random : %v", err)
	}

	// get current timestamp
	timestamp := strconv.FormatInt(s.timeProvider.UtcNow().Unix(), 10)

	// compute the hmac and base16 encode it
	computedMAC := computeMAC(secretKey, s.config.SignVerbAndURI, r.Method, r.URL.RequestURI(),
		timestamp, nonce, bodyBytes, headerValues)
	signature := hex.EncodeToString(computedMAC)

	// set headers
	r.Header.Set(s.config.NonceHeaderName, nonce)
	r.Header.Set(s.config.TimestampHeaderName, timestamp)
	r.Header.Set(s.config.SignatureHeaderName, signature)
	r.Header.Set(s.config.SignatureVersionHeaderName, "2")

	return nil
}

// Authenticates HTTP request to ensure it was sent by an authorized sender.
func (s *Service) AuthenticateRequest(r *http.Request) error {
	if s.secretKey == nil {
		return fmt.Errorf("service not loaded with key.")
	}
	return s.AuthenticateRequestWithKey(r, s.secretKey)
}

// Authenticates HTTP request to ensure it was sent by an authorized sender.
// Checks message signature with the passed in key, not the one initialized with.
func (s *Service) AuthenticateRequestWithKey(r *http.Request, secretKey []byte) (err error) {
	// Emit a success or failure metric on return.
	defer func() {
		if err == nil {
			s.metricsClient.Inc("success", 1, 1)
		} else {
			s.metricsClient.Inc("failure", 1, 1)
		}
	}()

	// extract parameters
	signature := r.Header.Get(s.config.SignatureHeaderName)
	if signature == "" {
		return fmt.Errorf("header not found: %v", s.config.SignatureHeaderName)
	}
	nonce := r.Header.Get(s.config.NonceHeaderName)
	if nonce == "" {
		return fmt.Errorf("header not found: %v", s.config.NonceHeaderName)
	}
	timestamp := r.Header.Get(s.config.TimestampHeaderName)
	if timestamp == "" {
		return fmt.Errorf("header not found: %v", s.config.TimestampHeaderName)
	}

	// extract request body bytes
	bodyBytes, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))

	// extract any headers if requested
	headerValues, err := extractHeaderValues(r, s.config.HeadersToSign)
	if err != nil {
		return err
	}

	// check the hmac
	isValid, err := checkMAC(secretKey, s.config.SignVerbAndURI, r.Method, r.URL.RequestURI(),
		timestamp, nonce, bodyBytes, headerValues, signature)
	if !isValid {
		return err
	}

	// check timestamp
	isValid, err = s.checkTimestamp(timestamp)
	if !isValid {
		return err
	}

	// check to see if we have seen nonce before
	inCache := s.nonceCache.InCache(nonce)
	if inCache {
		return fmt.Errorf("nonce already in cache: %v", nonce)
	}

	return nil
}

func (s *Service) checkTimestamp(timestampHeader string) (bool, error) {
	// convert unix timestamp string into time struct
	timestamp, err := strconv.ParseInt(timestampHeader, 10, 0)
	if err != nil {
		return false, fmt.Errorf("unable to parse %v: %v", s.config.TimestampHeaderName, timestampHeader)
	}

	now := s.timeProvider.UtcNow().Unix()

	// if timestamp is from the future, it's invalid
	if timestamp >= now+MaxSkewSec {
		return false, fmt.Errorf("timestamp header from the future; now: %v; %v: %v; difference: %v",
			now, s.config.TimestampHeaderName, timestamp, timestamp-now)
	}

	// if the timestamp is older than ttl - skew, it's invalid
	if timestamp <= now-int64(s.nonceCache.cacheTTL-MaxSkewSec) {
		return false, fmt.Errorf("timestamp header too old; now: %v; %v: %v; difference: %v",
			now, s.config.TimestampHeaderName, timestamp, now-timestamp)
	}

	return true, nil
}

func computeMAC(secretKey []byte, signVerbAndUri bool, httpVerb string, httpResourceUri string,
	timestamp string, nonce string, body []byte, headerValues []string) []byte {

	// use hmac-sha256
	mac := hmac.New(sha256.New, secretKey)

	// required parameters (timestamp, nonce, body)
	mac.Write([]byte(fmt.Sprintf("%v|", len(timestamp))))
	mac.Write([]byte(timestamp))
	mac.Write([]byte(fmt.Sprintf("|%v|", len(nonce))))
	mac.Write([]byte(nonce))
	mac.Write([]byte(fmt.Sprintf("|%v|", len(body))))
	mac.Write(body)

	// optional parameters (httpVerb, httpResourceUri)
	if signVerbAndUri {
		mac.Write([]byte(fmt.Sprintf("|%v|", len(httpVerb))))
		mac.Write([]byte(httpVerb))
		mac.Write([]byte(fmt.Sprintf("|%v|", len(httpResourceUri))))
		mac.Write([]byte(httpResourceUri))
	}

	// optional parameters (headers)
	for _, headerValue := range headerValues {
		mac.Write([]byte(fmt.Sprintf("|%v|", len(headerValue))))
		mac.Write([]byte(headerValue))
	}

	return mac.Sum(nil)
}

func checkMAC(secretKey []byte, signVerbAndUri bool, httpVerb string, httpResourceUri string,
	timestamp string, nonce string, body []byte, headerValues []string, signature string) (bool, error) {

	// the hmac we get is a hexdigest (string representation of hex values)
	// which needs to be decoded before before we can use it
	expectedMAC, err := hex.DecodeString(signature)
	if err != nil {
		return false, err
	}

	// compute the hmac
	computedMAC := computeMAC(secretKey, signVerbAndUri, httpVerb, httpResourceUri, timestamp, nonce, body, headerValues)

	// constant time compare
	isEqual := hmac.Equal(expectedMAC, computedMAC)
	if !isEqual {
		return false, fmt.Errorf("signature header value %v does not match computed value", expectedMAC)
	}

	return true, nil
}

func extractHeaderValues(r *http.Request, headerNames []string) ([]string, error) {
	if len(headerNames) < 1 {
		return nil, nil
	}

	headerValues := make([]string, len(headerNames))
	for i, headerName := range headerNames {
		_, ok := r.Header[headerName]
		if !ok {
			return nil, fmt.Errorf("header %v not found in request.", headerName)
		}
		headerValues[i] = r.Header.Get(headerName)
	}

	return headerValues, nil
}

func readKeyFromDisk(keypath string) ([]byte, error) {
	// load key from disk
	keyBytes, err := ioutil.ReadFile(keypath)
	if err != nil {
		return nil, err
	}

	// strip newline (\n or 0x0a) if it's at the end
	keyBytes = bytes.TrimSuffix(keyBytes, []byte("\n"))

	return keyBytes, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package leaderelection implements leader election of a set of endpoints.
// It uses an annotation in the endpoints object to store the record of the
// election state.
//
// This implementation does not guarantee that only one client is acting as a
// leader (a.k.a. fencing). A client observes timestamps captured locally to
// infer the state of the leader election. Thus the implementation is tolerant
// to arbitrary clock skew, but is not tolerant to arbitrary clock skew rate.
//
// However the level of tolerance to skew rate can be configured by setting
// RenewDeadline and LeaseDuration appropriately. The tolerance expressed as a
// maximum tolerated ratio of time passed on the fastest node to time passed on
// the slowest node can be approximately achieved with a configuration that sets
// the same ratio of LeaseDuration to RenewDeadline. For example if a user wanted
// to tolerate some nodes progressing forward in time twice as fast as other nodes,
// the user could set LeaseDuration to 60 seconds and RenewDeadline to 30 seconds.
//
// While not required, some method of clock synchronization between nodes in the
// cluster is highly recommended. It's important to keep in mind when configuring
// this client that the tolerance to skew rate varies inversely to master
// availability.
//
// Larger clusters often have a more lenient SLA for API latency. This should be
// taken into account when configuring the client. The rate of leader transitions
// should be monitored and RetryPeriod and LeaseDuration should be increased
// until the rate is stable and acceptably low. It's important to keep in mind
// when configuring this client that the tolerance to API latency varies inversely
// to master availability.
//
// DISCLAIMER: this is an alpha API. This library will likely change significantly
// or even be removed entirely in subsequent releases. Depend on this API at
// your own risk.
package leaderelection

import (
	"fmt"
	"reflect"
	"time"

	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	rl "k8s.io/client-go/tools/leaderelection/resourcelock"

	"github.com/golang/glog"
)

const (
	JitterFactor = 1.2
)

// NewLeaderElector creates a LeaderElector from a LeaderElectionConfig
func NewLeaderElector(lec LeaderElectionConfig) (*LeaderElector, error) {
	if lec.LeaseDuration <= lec.RenewDeadline {
		return nil, fmt.Errorf("leaseDuration must be greater than renewDeadline")
	}
	if lec.RenewDeadline <= time.Duration(JitterFactor*float64(lec.RetryPeriod)) {
		return nil, fmt.Errorf("renewDeadline must be greater than retryPeriod*JitterFactor")
	}
	if lec.Lock == nil {
		return nil, fmt.Errorf("Lock must not be nil.")
	}
	return &LeaderElector{
		config: lec,
	}, nil
}

type LeaderElectionConfig struct {
	// Lock is the resource that will be used for locking
	Lock rl.Interface

	// LeaseDuration is the duration that non-leader candidates will
	// wait to force acquire leadership. This is measured against time of
	// last observed ack.
	LeaseDuration time.Duration
	// RenewDeadline is the duration that the acting master will retry
	// refreshing leadership before giving up.
	RenewDeadline time.Duration
	// RetryPeriod is the duration the LeaderElector clients should wait
	// between tries of actions.
	RetryPeriod time.Duration

	// Callbacks are callbacks that are triggered during certain lifecycle
	// events of the LeaderElector
	Callbacks LeaderCallbacks
}

// LeaderCallbacks are callbacks that are triggered during certain
// lifecycle events of the LeaderElector. These are invoked asynchronously.
//
// possible future callbacks:
//  * OnChallenge()
type LeaderCallbacks struct {
	// OnStartedLeading is called when a LeaderElector client starts leading
	OnStartedLeading func(stop <-chan struct{})
	// OnStoppedLeading is called when a LeaderElector client stops leading
	OnStoppedLeading func()
	// OnNewLeader is called when the client observes a leader that is
	// not the previously observed leader. This includes the first observed
	// leader when the client starts.
	OnNewLeader func(identity string)
}

// LeaderElector is a leader election client.
//
// possible future methods:
//  * (le *LeaderElector) IsLeader()
//  * (le *LeaderElector) GetLeader()
type LeaderElector struct {
	config LeaderElectionConfig
	// internal bookkeeping
	observedRecord rl.LeaderElectionRecord
	observedTime   time.Time
	// used to implement OnNewLeader(), may lag slightly from the
	// value observedRecord.HolderIdentity if the transition has
	// not yet been reported.
	reportedLeader string
}

// Run starts the leader election loop
func (le *LeaderElector) Run() {
	defer func() {
		runtime.HandleCrash()
		le.config.Callbacks.OnStoppedLeading()
	}()
	le.acquire()
	stop := make(chan struct{})
	go le.config.Callbacks.OnStartedLeading(stop)
	le.renew()
	close(stop)
}

// RunOrDie starts a client with the provided config or panics if the config
// fails to validate.
func RunOrDie(lec LeaderElectionConfig) {
	le, err := NewLeaderElector(lec)
	if err != nil {
		panic(err)
	}
	le.Run()
}

// GetLeader returns the identity of the last observed leader or returns the empty string if
// no leader has yet been observed.
func (le *LeaderElector) GetLeader() string {
	return le.observedRecord.HolderIdentity
}

// IsLeader returns true if the last observed leader was this client else returns false.
func (le *LeaderElector) IsLeader() bool {
	return le.observedRecord.HolderIdentity == le.config.Lock.Identity()
}

// acquire loops calling tryAcquireOrRenew and returns immediately when tryAcquireOrRenew succeeds.
func (le *LeaderElector) acquire() {
	stop := make(chan struct{})
	glog.Infof("attempting to acquire leader lease...")
	wait.JitterUntil(func() {
		succeeded := le.tryAcquireOrRenew()
		le.maybeReportTransition()
		desc := le.config.Lock.Describe()
		if !succeeded {
			glog.V(4).Infof("failed to acquire lease %v", desc)
			return
		}
		le.config.Lock.RecordEvent("became leader")
		glog.Infof("successfully acquired lease %v", desc)
		close(stop)
	}, le.config.RetryPeriod, JitterFactor, true, stop)
}

// renew loops calling tryAcquireOrRenew and returns immediately when tryAcquireOrRenew fails.
func (le *LeaderElector) renew() {
	stop := make(chan struct{})
	wait.Until(func() {
		err := wait.Poll(le.config.RetryPeriod, le.config.RenewDeadline, func() (bool, error) {
			return le.tryAcquireOrRenew(), nil
		})
		le.maybeReportTransition()
		desc := le.config.Lock.Describe()
		if err == nil {
			glog.V(4).Infof("successfully renewed lease %v", desc)
			return
		}
		le.config.Lock.RecordEvent("stopped leading")
		glog.Infof("failed to renew lease %v: %v", desc, err)
		close(stop)
	}, 0, stop)
}

// tryAcquireOrRenew tries to acquire a leader lease if it is not already acquired,
// else it tries to renew the lease if it has already been acquired. Returns true
// on success else returns false.
func (le *LeaderElector) tryAcquireOrRenew() bool {
	now := metav1.Now()
	leaderElectionRecord := rl.LeaderElectionRecord{
		HolderIdentity:       le.config.Lock.Identity(),
		LeaseDurationSeconds: int(le.config.LeaseDuration / time.Second),
		RenewTime:            now,
		AcquireTime:          now,
	}

	// 1. obtain or create the ElectionRecord
	oldLeaderElectionRecord, err := le.config.Lock.Get()
	if err != nil {
		if !errors.IsNotFound(err) {
			glog.Errorf("error retrieving resource lock %v: %v", le.config.Lock.Describe(), err)
			return false
		}
		if err = le.config.Lock.Create(leaderElectionRecord); err != nil {
			glog.Errorf("error initially creating leader election record: %v", err)
			return false
		}
		le.observedRecord = leaderElectionRecord
		le.observedTime = time.Now()
		return true
	}

	// 2. Record obtained, check the Identity & Time
	if !reflect.DeepEqual(le.observedRecord, *oldLeaderElectionRecord) {
		le.observedRecord = *oldLeaderElectionRecord
		le.observedTime = time.Now()
	}
	if le.observedTime.Add(le.config.LeaseDuration).After(now.Time) &&
		oldLeaderElectionRecord.HolderIdentity != le.config.Lock.Identity() {
		glog.V(4).Infof("lock is held by %v and has not yet expired", oldLeaderElectionRecord.HolderIdentity)
		return false
	}

	// 3. We're going to try to update. The leaderElectionRecord is set to it's default
	// here. Let's correct it before updating.
	if oldLeaderElectionRecord.HolderIdentity == le.config.Lock.Identity() {
		leaderElectionRecord.AcquireTime = oldLeaderElectionRecord.AcquireTime
		leaderElectionRecord.LeaderTransitions = oldLeaderElectionRecord.LeaderTransitions
	} else {
		leaderElectionRecord.LeaderTransitions = oldLeaderElectionRecord.LeaderTransitions + 1
	}

	// update the lock itself
	if err = le.config.Lock.Update(leaderElectionRecord); err != nil {
		glog.Errorf("Failed to update lock: %v", err)
		return false
	}
	le.observedRecord = leaderElectionRecord
	le.observedTime = time.Now()
	return true
}

func (l *LeaderElector) maybeReportTransition() {
	if l.observedRecord.HolderIdentity == l.reportedLeader {
		return
	}
	l.reportedLeader = l.observedRecord.HolderIdentity
	if l.config.Callbacks.OnNewLeader != nil {
		go l.config.Callbacks.OnNewLeader(l.reportedLeader)
	}
}
<commit_msg>say which lease is being acquired<commit_after>/*
Copyright 2015 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package leaderelection implements leader election of a set of endpoints.
// It uses an annotation in the endpoints object to store the record of the
// election state.
//
// This implementation does not guarantee that only one client is acting as a
// leader (a.k.a. fencing). A client observes timestamps captured locally to
// infer the state of the leader election. Thus the implementation is tolerant
// to arbitrary clock skew, but is not tolerant to arbitrary clock skew rate.
//
// However the level of tolerance to skew rate can be configured by setting
// RenewDeadline and LeaseDuration appropriately. The tolerance expressed as a
// maximum tolerated ratio of time passed on the fastest node to time passed on
// the slowest node can be approximately achieved with a configuration that sets
// the same ratio of LeaseDuration to RenewDeadline. For example if a user wanted
// to tolerate some nodes progressing forward in time twice as fast as other nodes,
// the user could set LeaseDuration to 60 seconds and RenewDeadline to 30 seconds.
//
// While not required, some method of clock synchronization between nodes in the
// cluster is highly recommended. It's important to keep in mind when configuring
// this client that the tolerance to skew rate varies inversely to master
// availability.
//
// Larger clusters often have a more lenient SLA for API latency. This should be
// taken into account when configuring the client. The rate of leader transitions
// should be monitored and RetryPeriod and LeaseDuration should be increased
// until the rate is stable and acceptably low. It's important to keep in mind
// when configuring this client that the tolerance to API latency varies inversely
// to master availability.
//
// DISCLAIMER: this is an alpha API. This library will likely change significantly
// or even be removed entirely in subsequent releases. Depend on this API at
// your own risk.
package leaderelection

import (
	"fmt"
	"reflect"
	"time"

	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	rl "k8s.io/client-go/tools/leaderelection/resourcelock"

	"github.com/golang/glog"
)

const (
	JitterFactor = 1.2
)

// NewLeaderElector creates a LeaderElector from a LeaderElectionConfig
func NewLeaderElector(lec LeaderElectionConfig) (*LeaderElector, error) {
	if lec.LeaseDuration <= lec.RenewDeadline {
		return nil, fmt.Errorf("leaseDuration must be greater than renewDeadline")
	}
	if lec.RenewDeadline <= time.Duration(JitterFactor*float64(lec.RetryPeriod)) {
		return nil, fmt.Errorf("renewDeadline must be greater than retryPeriod*JitterFactor")
	}
	if lec.Lock == nil {
		return nil, fmt.Errorf("Lock must not be nil.")
	}
	return &LeaderElector{
		config: lec,
	}, nil
}

type LeaderElectionConfig struct {
	// Lock is the resource that will be used for locking
	Lock rl.Interface

	// LeaseDuration is the duration that non-leader candidates will
	// wait to force acquire leadership. This is measured against time of
	// last observed ack.
	LeaseDuration time.Duration
	// RenewDeadline is the duration that the acting master will retry
	// refreshing leadership before giving up.
	RenewDeadline time.Duration
	// RetryPeriod is the duration the LeaderElector clients should wait
	// between tries of actions.
	RetryPeriod time.Duration

	// Callbacks are callbacks that are triggered during certain lifecycle
	// events of the LeaderElector
	Callbacks LeaderCallbacks
}

// LeaderCallbacks are callbacks that are triggered during certain
// lifecycle events of the LeaderElector. These are invoked asynchronously.
//
// possible future callbacks:
//  * OnChallenge()
type LeaderCallbacks struct {
	// OnStartedLeading is called when a LeaderElector client starts leading
	OnStartedLeading func(stop <-chan struct{})
	// OnStoppedLeading is called when a LeaderElector client stops leading
	OnStoppedLeading func()
	// OnNewLeader is called when the client observes a leader that is
	// not the previously observed leader. This includes the first observed
	// leader when the client starts.
	OnNewLeader func(identity string)
}

// LeaderElector is a leader election client.
//
// possible future methods:
//  * (le *LeaderElector) IsLeader()
//  * (le *LeaderElector) GetLeader()
type LeaderElector struct {
	config LeaderElectionConfig
	// internal bookkeeping
	observedRecord rl.LeaderElectionRecord
	observedTime   time.Time
	// used to implement OnNewLeader(), may lag slightly from the
	// value observedRecord.HolderIdentity if the transition has
	// not yet been reported.
	reportedLeader string
}

// Run starts the leader election loop
func (le *LeaderElector) Run() {
	defer func() {
		runtime.HandleCrash()
		le.config.Callbacks.OnStoppedLeading()
	}()
	le.acquire()
	stop := make(chan struct{})
	go le.config.Callbacks.OnStartedLeading(stop)
	le.renew()
	close(stop)
}

// RunOrDie starts a client with the provided config or panics if the config
// fails to validate.
func RunOrDie(lec LeaderElectionConfig) {
	le, err := NewLeaderElector(lec)
	if err != nil {
		panic(err)
	}
	le.Run()
}

// GetLeader returns the identity of the last observed leader or returns the empty string if
// no leader has yet been observed.
func (le *LeaderElector) GetLeader() string {
	return le.observedRecord.HolderIdentity
}

// IsLeader returns true if the last observed leader was this client else returns false.
func (le *LeaderElector) IsLeader() bool {
	return le.observedRecord.HolderIdentity == le.config.Lock.Identity()
}

// acquire loops calling tryAcquireOrRenew and returns immediately when tryAcquireOrRenew succeeds.
func (le *LeaderElector) acquire() {
	stop := make(chan struct{})
	desc := le.config.Lock.Describe()
	glog.Infof("attempting to acquire leader lease  %v...", desc)
	wait.JitterUntil(func() {
		succeeded := le.tryAcquireOrRenew()
		le.maybeReportTransition()
		if !succeeded {
			glog.V(4).Infof("failed to acquire lease %v", desc)
			return
		}
		le.config.Lock.RecordEvent("became leader")
		glog.Infof("successfully acquired lease %v", desc)
		close(stop)
	}, le.config.RetryPeriod, JitterFactor, true, stop)
}

// renew loops calling tryAcquireOrRenew and returns immediately when tryAcquireOrRenew fails.
func (le *LeaderElector) renew() {
	stop := make(chan struct{})
	wait.Until(func() {
		err := wait.Poll(le.config.RetryPeriod, le.config.RenewDeadline, func() (bool, error) {
			return le.tryAcquireOrRenew(), nil
		})
		le.maybeReportTransition()
		desc := le.config.Lock.Describe()
		if err == nil {
			glog.V(4).Infof("successfully renewed lease %v", desc)
			return
		}
		le.config.Lock.RecordEvent("stopped leading")
		glog.Infof("failed to renew lease %v: %v", desc, err)
		close(stop)
	}, 0, stop)
}

// tryAcquireOrRenew tries to acquire a leader lease if it is not already acquired,
// else it tries to renew the lease if it has already been acquired. Returns true
// on success else returns false.
func (le *LeaderElector) tryAcquireOrRenew() bool {
	now := metav1.Now()
	leaderElectionRecord := rl.LeaderElectionRecord{
		HolderIdentity:       le.config.Lock.Identity(),
		LeaseDurationSeconds: int(le.config.LeaseDuration / time.Second),
		RenewTime:            now,
		AcquireTime:          now,
	}

	// 1. obtain or create the ElectionRecord
	oldLeaderElectionRecord, err := le.config.Lock.Get()
	if err != nil {
		if !errors.IsNotFound(err) {
			glog.Errorf("error retrieving resource lock %v: %v", le.config.Lock.Describe(), err)
			return false
		}
		if err = le.config.Lock.Create(leaderElectionRecord); err != nil {
			glog.Errorf("error initially creating leader election record: %v", err)
			return false
		}
		le.observedRecord = leaderElectionRecord
		le.observedTime = time.Now()
		return true
	}

	// 2. Record obtained, check the Identity & Time
	if !reflect.DeepEqual(le.observedRecord, *oldLeaderElectionRecord) {
		le.observedRecord = *oldLeaderElectionRecord
		le.observedTime = time.Now()
	}
	if le.observedTime.Add(le.config.LeaseDuration).After(now.Time) &&
		oldLeaderElectionRecord.HolderIdentity != le.config.Lock.Identity() {
		glog.V(4).Infof("lock is held by %v and has not yet expired", oldLeaderElectionRecord.HolderIdentity)
		return false
	}

	// 3. We're going to try to update. The leaderElectionRecord is set to it's default
	// here. Let's correct it before updating.
	if oldLeaderElectionRecord.HolderIdentity == le.config.Lock.Identity() {
		leaderElectionRecord.AcquireTime = oldLeaderElectionRecord.AcquireTime
		leaderElectionRecord.LeaderTransitions = oldLeaderElectionRecord.LeaderTransitions
	} else {
		leaderElectionRecord.LeaderTransitions = oldLeaderElectionRecord.LeaderTransitions + 1
	}

	// update the lock itself
	if err = le.config.Lock.Update(leaderElectionRecord); err != nil {
		glog.Errorf("Failed to update lock: %v", err)
		return false
	}
	le.observedRecord = leaderElectionRecord
	le.observedTime = time.Now()
	return true
}

func (l *LeaderElector) maybeReportTransition() {
	if l.observedRecord.HolderIdentity == l.reportedLeader {
		return
	}
	l.reportedLeader = l.observedRecord.HolderIdentity
	if l.config.Callbacks.OnNewLeader != nil {
		go l.config.Callbacks.OnNewLeader(l.reportedLeader)
	}
}
<|endoftext|>


<commit_before>// +build !providerless

/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package azure

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest/to"

	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog"
	azcache "k8s.io/legacy-cloud-providers/azure/cache"
)

var (
	vmssNameSeparator = "_"

	vmssKey                 = "k8svmssKey"
	vmssVirtualMachinesKey  = "k8svmssVirtualMachinesKey"
	availabilitySetNodesKey = "k8sAvailabilitySetNodesKey"

	availabilitySetNodesCacheTTLDefaultInSeconds = 900
	vmssCacheTTLDefaultInSeconds                 = 600
	vmssVirtualMachinesCacheTTLDefaultInSeconds  = 600
)

type vmssVirtualMachinesEntry struct {
	resourceGroup  string
	vmssName       string
	instanceID     string
	virtualMachine *compute.VirtualMachineScaleSetVM
	lastUpdate     time.Time
}

type vmssEntry struct {
	vmss       *compute.VirtualMachineScaleSet
	lastUpdate time.Time
}

func (ss *scaleSet) newVMSSCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := &sync.Map{} // [vmssName]*vmssEntry

		allResourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range allResourceGroups.List() {
			allScaleSets, rerr := ss.VirtualMachineScaleSetsClient.List(context.Background(), resourceGroup)
			if rerr != nil {
				klog.Errorf("VirtualMachineScaleSetsClient.List failed: %v", rerr)
				return nil, rerr.Error()
			}

			for i := range allScaleSets {
				scaleSet := allScaleSets[i]
				if scaleSet.Name == nil || *scaleSet.Name == "" {
					klog.Warning("failed to get the name of VMSS")
					continue
				}
				localCache.Store(*scaleSet.Name, &vmssEntry{
					vmss:       &scaleSet,
					lastUpdate: time.Now().UTC(),
				})
			}
		}

		return localCache, nil
	}

	if ss.Config.VmssCacheTTLInSeconds == 0 {
		ss.Config.VmssCacheTTLInSeconds = vmssCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.VmssCacheTTLInSeconds)*time.Second, getter)
}

func extractVmssVMName(name string) (string, string, error) {
	split := strings.SplitAfter(name, vmssNameSeparator)
	if len(split) < 2 {
		klog.V(3).Infof("Failed to extract vmssVMName %q", name)
		return "", "", ErrorNotVmssInstance
	}

	ssName := strings.Join(split[0:len(split)-1], "")
	// removing the trailing `vmssNameSeparator` since we used SplitAfter
	ssName = ssName[:len(ssName)-1]
	instanceID := split[len(split)-1]
	return ssName, instanceID, nil
}

func (ss *scaleSet) newVMSSVirtualMachinesCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := &sync.Map{} // [nodeName]*vmssVirtualMachinesEntry

		oldCache := make(map[string]vmssVirtualMachinesEntry)

		if ss.vmssVMCache != nil {
			// get old cache before refreshing the cache
			entry, exists, err := ss.vmssVMCache.Store.GetByKey(vmssVirtualMachinesKey)
			if err != nil {
				return nil, err
			}
			if exists {
				cached := entry.(*azcache.AzureCacheEntry).Data
				if cached != nil {
					virtualMachines := cached.(*sync.Map)
					virtualMachines.Range(func(key, value interface{}) bool {
						oldCache[key.(string)] = *value.(*vmssVirtualMachinesEntry)
						return true
					})
				}
			}
		}

		allResourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range allResourceGroups.List() {
			scaleSetNames, err := ss.listScaleSets(resourceGroup)
			if err != nil {
				return nil, err
			}

			for _, ssName := range scaleSetNames {
				vms, err := ss.listScaleSetVMs(ssName, resourceGroup)
				if err != nil {
					return nil, err
				}

				for i := range vms {
					vm := vms[i]
					if vm.OsProfile == nil || vm.OsProfile.ComputerName == nil {
						klog.Warningf("failed to get computerName for vmssVM (%q)", ssName)
						continue
					}

					computerName := strings.ToLower(*vm.OsProfile.ComputerName)
					vmssVMCacheEntry := &vmssVirtualMachinesEntry{
						resourceGroup:  resourceGroup,
						vmssName:       ssName,
						instanceID:     to.String(vm.InstanceID),
						virtualMachine: &vm,
						lastUpdate:     time.Now().UTC(),
					}
					// set cache entry to nil when the VM is under deleting.
					if vm.VirtualMachineScaleSetVMProperties != nil &&
						strings.EqualFold(to.String(vm.VirtualMachineScaleSetVMProperties.ProvisioningState), string(compute.ProvisioningStateDeleting)) {
						klog.V(4).Infof("VMSS virtualMachine %q is under deleting, setting its cache to nil", computerName)
						vmssVMCacheEntry.virtualMachine = nil
					}
					localCache.Store(computerName, vmssVMCacheEntry)

					if _, exists := oldCache[computerName]; exists {
						delete(oldCache, computerName)
					}
				}
			}

			// add old missing cache data with nil entries to prevent aggressive
			// ARM calls during cache invalidation
			for name, vmEntry := range oldCache {
				// if the nil cache entry has existed for 15 minutes in the cache
				// then it should not be added back to the cache
				if vmEntry.virtualMachine == nil || time.Since(vmEntry.lastUpdate) > 15*time.Minute {
					klog.V(5).Infof("ignoring expired entries from old cache for %s", name)
					continue
				}
				lastUpdate := time.Now().UTC()
				if vmEntry.virtualMachine == nil {
					// if this is already a nil entry then keep the time the nil
					// entry was first created, so we can cleanup unwanted entries
					lastUpdate = vmEntry.lastUpdate
				}

				klog.V(5).Infof("adding old entries to new cache for %s", name)
				localCache.Store(name, &vmssVirtualMachinesEntry{
					resourceGroup:  vmEntry.resourceGroup,
					vmssName:       vmEntry.vmssName,
					instanceID:     vmEntry.instanceID,
					virtualMachine: nil,
					lastUpdate:     lastUpdate,
				})
			}
		}

		return localCache, nil
	}

	if ss.Config.VmssVirtualMachinesCacheTTLInSeconds == 0 {
		ss.Config.VmssVirtualMachinesCacheTTLInSeconds = vmssVirtualMachinesCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.VmssVirtualMachinesCacheTTLInSeconds)*time.Second, getter)
}

func (ss *scaleSet) deleteCacheForNode(nodeName string) error {
	cached, err := ss.vmssVMCache.Get(vmssVirtualMachinesKey, azcache.CacheReadTypeUnsafe)
	if err != nil {
		klog.Errorf("deleteCacheForNode(%s) failed with error: %v", nodeName, err)
		return err
	}

	virtualMachines := cached.(*sync.Map)
	virtualMachines.Delete(nodeName)
	return nil
}

func (ss *scaleSet) newAvailabilitySetNodesCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := sets.NewString()
		resourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range resourceGroups.List() {
			vmList, err := ss.Cloud.ListVirtualMachines(resourceGroup)
			if err != nil {
				return nil, err
			}

			for _, vm := range vmList {
				if vm.Name != nil {
					localCache.Insert(*vm.Name)
				}
			}
		}

		return localCache, nil
	}

	if ss.Config.AvailabilitySetNodesCacheTTLInSeconds == 0 {
		ss.Config.AvailabilitySetNodesCacheTTLInSeconds = availabilitySetNodesCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.AvailabilitySetNodesCacheTTLInSeconds)*time.Second, getter)
}

func (ss *scaleSet) isNodeManagedByAvailabilitySet(nodeName string, crt azcache.AzureCacheReadType) (bool, error) {
	// Assume all nodes are managed by VMSS when DisableAvailabilitySetNodes is enabled.
	if ss.DisableAvailabilitySetNodes {
		klog.V(2).Infof("Assuming node %q is managed by VMSS since DisableAvailabilitySetNodes is set to true", nodeName)
		return false, nil
	}

	cached, err := ss.availabilitySetNodesCache.Get(availabilitySetNodesKey, crt)
	if err != nil {
		return false, err
	}

	availabilitySetNodes := cached.(sets.String)
	return availabilitySetNodes.Has(nodeName), nil
}
<commit_msg>[Azure] cache map remove exists check<commit_after>// +build !providerless

/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package azure

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
	"github.com/Azure/go-autorest/autorest/to"

	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog"
	azcache "k8s.io/legacy-cloud-providers/azure/cache"
)

var (
	vmssNameSeparator = "_"

	vmssKey                 = "k8svmssKey"
	vmssVirtualMachinesKey  = "k8svmssVirtualMachinesKey"
	availabilitySetNodesKey = "k8sAvailabilitySetNodesKey"

	availabilitySetNodesCacheTTLDefaultInSeconds = 900
	vmssCacheTTLDefaultInSeconds                 = 600
	vmssVirtualMachinesCacheTTLDefaultInSeconds  = 600
)

type vmssVirtualMachinesEntry struct {
	resourceGroup  string
	vmssName       string
	instanceID     string
	virtualMachine *compute.VirtualMachineScaleSetVM
	lastUpdate     time.Time
}

type vmssEntry struct {
	vmss       *compute.VirtualMachineScaleSet
	lastUpdate time.Time
}

func (ss *scaleSet) newVMSSCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := &sync.Map{} // [vmssName]*vmssEntry

		allResourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range allResourceGroups.List() {
			allScaleSets, rerr := ss.VirtualMachineScaleSetsClient.List(context.Background(), resourceGroup)
			if rerr != nil {
				klog.Errorf("VirtualMachineScaleSetsClient.List failed: %v", rerr)
				return nil, rerr.Error()
			}

			for i := range allScaleSets {
				scaleSet := allScaleSets[i]
				if scaleSet.Name == nil || *scaleSet.Name == "" {
					klog.Warning("failed to get the name of VMSS")
					continue
				}
				localCache.Store(*scaleSet.Name, &vmssEntry{
					vmss:       &scaleSet,
					lastUpdate: time.Now().UTC(),
				})
			}
		}

		return localCache, nil
	}

	if ss.Config.VmssCacheTTLInSeconds == 0 {
		ss.Config.VmssCacheTTLInSeconds = vmssCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.VmssCacheTTLInSeconds)*time.Second, getter)
}

func extractVmssVMName(name string) (string, string, error) {
	split := strings.SplitAfter(name, vmssNameSeparator)
	if len(split) < 2 {
		klog.V(3).Infof("Failed to extract vmssVMName %q", name)
		return "", "", ErrorNotVmssInstance
	}

	ssName := strings.Join(split[0:len(split)-1], "")
	// removing the trailing `vmssNameSeparator` since we used SplitAfter
	ssName = ssName[:len(ssName)-1]
	instanceID := split[len(split)-1]
	return ssName, instanceID, nil
}

func (ss *scaleSet) newVMSSVirtualMachinesCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := &sync.Map{} // [nodeName]*vmssVirtualMachinesEntry

		oldCache := make(map[string]vmssVirtualMachinesEntry)

		if ss.vmssVMCache != nil {
			// get old cache before refreshing the cache
			entry, exists, err := ss.vmssVMCache.Store.GetByKey(vmssVirtualMachinesKey)
			if err != nil {
				return nil, err
			}
			if exists {
				cached := entry.(*azcache.AzureCacheEntry).Data
				if cached != nil {
					virtualMachines := cached.(*sync.Map)
					virtualMachines.Range(func(key, value interface{}) bool {
						oldCache[key.(string)] = *value.(*vmssVirtualMachinesEntry)
						return true
					})
				}
			}
		}

		allResourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range allResourceGroups.List() {
			scaleSetNames, err := ss.listScaleSets(resourceGroup)
			if err != nil {
				return nil, err
			}

			for _, ssName := range scaleSetNames {
				vms, err := ss.listScaleSetVMs(ssName, resourceGroup)
				if err != nil {
					return nil, err
				}

				for i := range vms {
					vm := vms[i]
					if vm.OsProfile == nil || vm.OsProfile.ComputerName == nil {
						klog.Warningf("failed to get computerName for vmssVM (%q)", ssName)
						continue
					}

					computerName := strings.ToLower(*vm.OsProfile.ComputerName)
					vmssVMCacheEntry := &vmssVirtualMachinesEntry{
						resourceGroup:  resourceGroup,
						vmssName:       ssName,
						instanceID:     to.String(vm.InstanceID),
						virtualMachine: &vm,
						lastUpdate:     time.Now().UTC(),
					}
					// set cache entry to nil when the VM is under deleting.
					if vm.VirtualMachineScaleSetVMProperties != nil &&
						strings.EqualFold(to.String(vm.VirtualMachineScaleSetVMProperties.ProvisioningState), string(compute.ProvisioningStateDeleting)) {
						klog.V(4).Infof("VMSS virtualMachine %q is under deleting, setting its cache to nil", computerName)
						vmssVMCacheEntry.virtualMachine = nil
					}
					localCache.Store(computerName, vmssVMCacheEntry)

					delete(oldCache, computerName)
				}
			}

			// add old missing cache data with nil entries to prevent aggressive
			// ARM calls during cache invalidation
			for name, vmEntry := range oldCache {
				// if the nil cache entry has existed for 15 minutes in the cache
				// then it should not be added back to the cache
				if vmEntry.virtualMachine == nil || time.Since(vmEntry.lastUpdate) > 15*time.Minute {
					klog.V(5).Infof("ignoring expired entries from old cache for %s", name)
					continue
				}
				lastUpdate := time.Now().UTC()
				if vmEntry.virtualMachine == nil {
					// if this is already a nil entry then keep the time the nil
					// entry was first created, so we can cleanup unwanted entries
					lastUpdate = vmEntry.lastUpdate
				}

				klog.V(5).Infof("adding old entries to new cache for %s", name)
				localCache.Store(name, &vmssVirtualMachinesEntry{
					resourceGroup:  vmEntry.resourceGroup,
					vmssName:       vmEntry.vmssName,
					instanceID:     vmEntry.instanceID,
					virtualMachine: nil,
					lastUpdate:     lastUpdate,
				})
			}
		}

		return localCache, nil
	}

	if ss.Config.VmssVirtualMachinesCacheTTLInSeconds == 0 {
		ss.Config.VmssVirtualMachinesCacheTTLInSeconds = vmssVirtualMachinesCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.VmssVirtualMachinesCacheTTLInSeconds)*time.Second, getter)
}

func (ss *scaleSet) deleteCacheForNode(nodeName string) error {
	cached, err := ss.vmssVMCache.Get(vmssVirtualMachinesKey, azcache.CacheReadTypeUnsafe)
	if err != nil {
		klog.Errorf("deleteCacheForNode(%s) failed with error: %v", nodeName, err)
		return err
	}

	virtualMachines := cached.(*sync.Map)
	virtualMachines.Delete(nodeName)
	return nil
}

func (ss *scaleSet) newAvailabilitySetNodesCache() (*azcache.TimedCache, error) {
	getter := func(key string) (interface{}, error) {
		localCache := sets.NewString()
		resourceGroups, err := ss.GetResourceGroups()
		if err != nil {
			return nil, err
		}

		for _, resourceGroup := range resourceGroups.List() {
			vmList, err := ss.Cloud.ListVirtualMachines(resourceGroup)
			if err != nil {
				return nil, err
			}

			for _, vm := range vmList {
				if vm.Name != nil {
					localCache.Insert(*vm.Name)
				}
			}
		}

		return localCache, nil
	}

	if ss.Config.AvailabilitySetNodesCacheTTLInSeconds == 0 {
		ss.Config.AvailabilitySetNodesCacheTTLInSeconds = availabilitySetNodesCacheTTLDefaultInSeconds
	}
	return azcache.NewTimedcache(time.Duration(ss.Config.AvailabilitySetNodesCacheTTLInSeconds)*time.Second, getter)
}

func (ss *scaleSet) isNodeManagedByAvailabilitySet(nodeName string, crt azcache.AzureCacheReadType) (bool, error) {
	// Assume all nodes are managed by VMSS when DisableAvailabilitySetNodes is enabled.
	if ss.DisableAvailabilitySetNodes {
		klog.V(2).Infof("Assuming node %q is managed by VMSS since DisableAvailabilitySetNodes is set to true", nodeName)
		return false, nil
	}

	cached, err := ss.availabilitySetNodesCache.Get(availabilitySetNodesKey, crt)
	if err != nil {
		return false, err
	}

	availabilitySetNodes := cached.(sets.String)
	return availabilitySetNodes.Has(nodeName), nil
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"encoding/json"
	"errors"
	"sync"

	"github.com/OpenBazaar/openbazaar-go/repo"
)

var SettingsNotSetError error = errors.New("Settings not set")

type SettingsDB struct {
	modelStore
}

func NewConfigurationStore(db *sql.DB, lock *sync.Mutex) repo.ConfigurationStore {
	return &SettingsDB{modelStore{db, lock}}
}

func (s *SettingsDB) Put(settings repo.SettingsData) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	b, err := json.MarshalIndent(&settings, "", "    ")
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare("insert or replace into config(key, value) values(?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec("settings", string(b))
	if err != nil {
		tx.Rollback()
		return err
	}
	tx.Commit()
	return nil
}

func (s *SettingsDB) Get() (repo.SettingsData, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	settings := repo.SettingsData{}
	stmt, err := s.db.Prepare("select value from config where key=?")
	if err != nil {
		return settings, err
	}
	defer stmt.Close()
	var settingsBytes []byte
	err = stmt.QueryRow("settings").Scan(&settingsBytes)
	if err != nil {
		return settings, SettingsNotSetError
	}
	err = json.Unmarshal(settingsBytes, &settings)
	if err != nil {
		return settings, err
	}
	return settings, nil
}

func (s *SettingsDB) Update(settings repo.SettingsData) error {
	current, err := s.Get()
	if err != nil {
		return errors.New("Not Found")
	}
	if settings.PaymentDataInQR == nil {
		settings.PaymentDataInQR = current.PaymentDataInQR
	}
	if settings.ShowNotifications == nil {
		settings.ShowNotifications = current.ShowNotifications
	}
	if settings.ShowNsfw == nil {
		settings.ShowNsfw = current.ShowNsfw
	}
	if settings.ShippingAddresses == nil {
		settings.ShippingAddresses = current.ShippingAddresses
	}
	if settings.LocalCurrency == nil {
		settings.LocalCurrency = current.LocalCurrency
	}
	if settings.Country == nil {
		settings.Country = current.Country
	}
	if settings.TermsAndConditions == nil {
		settings.TermsAndConditions = current.TermsAndConditions
	}
	if settings.RefundPolicy == nil {
		settings.RefundPolicy = current.RefundPolicy
	}
	if settings.BlockedNodes == nil {
		settings.BlockedNodes = current.BlockedNodes
	}
	if settings.StoreModerators == nil {
		settings.StoreModerators = current.StoreModerators
	}
	if settings.MisPaymentBuffer == nil {
		settings.MisPaymentBuffer = current.MisPaymentBuffer
	}
	if settings.SMTPSettings == nil {
		settings.SMTPSettings = current.SMTPSettings
	}
	if settings.Version == nil {
		settings.Version = current.Version
	}
	err = s.Put(settings)
	if err != nil {
		return err
	}
	return nil
}

// Delete removes all settings from the database. It's a destructive action that should be used with care.
func (s *SettingsDB) Delete() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	stmt, err := s.db.Prepare("delete from config where key = ?")
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec("settings")

	return err
}
<commit_msg>Capture rollback error reporting in (*SettingsDB).Put<commit_after>package db

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"sync"

	"github.com/OpenBazaar/openbazaar-go/repo"
)

var SettingsNotSetError error = errors.New("Settings not set")

type SettingsDB struct {
	modelStore
}

func NewConfigurationStore(db *sql.DB, lock *sync.Mutex) repo.ConfigurationStore {
	return &SettingsDB{modelStore{db, lock}}
}

func (s *SettingsDB) Put(settings repo.SettingsData) error {
	s.lock.Lock()
	defer s.lock.Unlock()
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	b, err := json.MarshalIndent(&settings, "", "    ")
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare("insert or replace into config(key, value) values(?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec("settings", string(b))
	if err != nil {
		if errRollback := tx.Rollback(); errRollback != nil {
			return fmt.Errorf("rollback: %s\n also: %s", errRollback.Error(), err.Error())
		}
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	return nil
}

func (s *SettingsDB) Get() (repo.SettingsData, error) {
	s.lock.Lock()
	defer s.lock.Unlock()
	settings := repo.SettingsData{}
	stmt, err := s.db.Prepare("select value from config where key=?")
	if err != nil {
		return settings, err
	}
	defer stmt.Close()
	var settingsBytes []byte
	err = stmt.QueryRow("settings").Scan(&settingsBytes)
	if err != nil {
		return settings, SettingsNotSetError
	}
	err = json.Unmarshal(settingsBytes, &settings)
	if err != nil {
		return settings, err
	}
	return settings, nil
}

func (s *SettingsDB) Update(settings repo.SettingsData) error {
	current, err := s.Get()
	if err != nil {
		return errors.New("Not Found")
	}
	if settings.PaymentDataInQR == nil {
		settings.PaymentDataInQR = current.PaymentDataInQR
	}
	if settings.ShowNotifications == nil {
		settings.ShowNotifications = current.ShowNotifications
	}
	if settings.ShowNsfw == nil {
		settings.ShowNsfw = current.ShowNsfw
	}
	if settings.ShippingAddresses == nil {
		settings.ShippingAddresses = current.ShippingAddresses
	}
	if settings.LocalCurrency == nil {
		settings.LocalCurrency = current.LocalCurrency
	}
	if settings.Country == nil {
		settings.Country = current.Country
	}
	if settings.TermsAndConditions == nil {
		settings.TermsAndConditions = current.TermsAndConditions
	}
	if settings.RefundPolicy == nil {
		settings.RefundPolicy = current.RefundPolicy
	}
	if settings.BlockedNodes == nil {
		settings.BlockedNodes = current.BlockedNodes
	}
	if settings.StoreModerators == nil {
		settings.StoreModerators = current.StoreModerators
	}
	if settings.MisPaymentBuffer == nil {
		settings.MisPaymentBuffer = current.MisPaymentBuffer
	}
	if settings.SMTPSettings == nil {
		settings.SMTPSettings = current.SMTPSettings
	}
	if settings.Version == nil {
		settings.Version = current.Version
	}
	err = s.Put(settings)
	if err != nil {
		return err
	}
	return nil
}

// Delete removes all settings from the database. It's a destructive action that should be used with care.
func (s *SettingsDB) Delete() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	stmt, err := s.db.Prepare("delete from config where key = ?")
	if err != nil {
		return err
	}
	defer stmt.Close()

	_, err = stmt.Exec("settings")

	return err
}
<|endoftext|>


<commit_before>package memdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/t3rm1n4l/memdb/skiplist"
	"io"
	"math/rand"
	"os"
	"path"
	"runtime"
	"sync"
	"sync/atomic"
)

type KeyCompare func([]byte, []byte) int
type ItemCallback func(*Item)

type FileType int

const (
	encodeBufSize = 4
	readerBufSize = 10000
)

const (
	ForestdbFile FileType = iota
	RawdbFile
)

func DefaultConfig() Config {
	var cfg Config
	cfg.SetKeyComparator(defaultKeyCmp)
	cfg.SetFileType(RawdbFile)
	return cfg
}

type Item struct {
	bornSn, deadSn uint32
	data           []byte
}

func (itm *Item) Encode(buf []byte, w io.Writer) error {
	l := 2
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint16(buf[0:2], uint16(len(itm.data)))
	if _, err := w.Write(buf[0:2]); err != nil {
		return err
	}
	if _, err := w.Write(itm.data); err != nil {
		return err
	}

	return nil
}

func (itm *Item) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:2]); err != nil {
		return err
	}
	l := binary.BigEndian.Uint16(buf[0:2])
	itm.data = make([]byte, int(l))
	_, err := io.ReadFull(r, itm.data)

	return err
}

func (itm *Item) Bytes() []byte {
	return itm.data
}

func NewItem(data []byte) *Item {
	return &Item{
		data: data,
	}
}

func newInsertCompare(keyCmp KeyCompare) skiplist.CompareFn {
	return func(this skiplist.Item, that skiplist.Item) int {
		var v int
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		if v = keyCmp(thisItem.data, thatItem.data); v == 0 {
			v = int(thisItem.bornSn) - int(thatItem.bornSn)
		}

		return v
	}
}

func newIterCompare(keyCmp KeyCompare) skiplist.CompareFn {
	return func(this skiplist.Item, that skiplist.Item) int {
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		return keyCmp(thisItem.data, thatItem.data)
	}
}

func defaultKeyCmp(this []byte, that []byte) int {
	var l int

	l1 := len(this)
	l2 := len(that)
	if l1 < l2 {
		l = l1
	} else {
		l = l2
	}

	return bytes.Compare(this[:l], that[:l])
}

//
//compare item,sn
type Writer struct {
	rand *rand.Rand
	buf  *skiplist.ActionBuffer
	iter *skiplist.Iterator
	*MemDB
}

func (w *Writer) Put(x *Item) {
	x.bornSn = w.getCurrSn()
	w.store.Insert2(x, w.insCmp, w.buf, w.rand.Float32)
	atomic.AddInt64(&w.count, 1)
}

// Find first item, seek until dead=0, mark dead=sn
func (w *Writer) Delete(x *Item) (success bool) {
	defer func() {
		if success {
			atomic.AddInt64(&w.count, -1)
		}
	}()

	gotNode := w.GetNode(x)
	if gotNode != nil {
		sn := w.getCurrSn()
		gotItem := gotNode.Item().(*Item)
		if gotItem.bornSn == sn {
			success = w.store.DeleteNode(gotNode, w.insCmp, w.buf)
			return
		}

		success = atomic.CompareAndSwapUint32(&gotItem.deadSn, 0, sn)
		return
	}

	return
}

func (w *Writer) Get(x *Item) *Item {
	n := w.GetNode(x)
	if n != nil {
		return n.Item().(*Item)
	}
	return nil
}

func (w *Writer) GetNode(x *Item) *skiplist.Node {
	var curr *skiplist.Node
	found := w.iter.Seek(x)
	if !found {
		return nil
	}

	// Seek until most recent item for key is found
	curr = w.iter.GetNode()
	for {
		w.iter.Next()
		if !w.iter.Valid() {
			break
		}
		next := w.iter.GetNode()
		nxtItm := next.Item().(*Item)
		currItm := curr.Item().(*Item)
		if w.iterCmp(nxtItm, currItm) != 0 {
			break
		}

		curr = next
	}

	currItm := curr.Item().(*Item)
	if currItm.deadSn != 0 {
		return nil
	}

	return curr
}

type Config struct {
	keyCmp  KeyCompare
	insCmp  skiplist.CompareFn
	iterCmp skiplist.CompareFn

	fileType FileType
}

func (cfg *Config) SetKeyComparator(cmp KeyCompare) {
	cfg.keyCmp = cmp
	cfg.insCmp = newInsertCompare(cmp)
	cfg.iterCmp = newIterCompare(cmp)
}

func (cfg *Config) SetFileType(t FileType) error {
	switch t {
	case ForestdbFile, RawdbFile:
	default:
		return errors.New("Invalid format")
	}

	cfg.fileType = t
	return nil
}

type MemDB struct {
	store        *skiplist.Skiplist
	currSn       uint32
	snapshots    *skiplist.Skiplist
	isGCRunning  int32
	lastGCSn     uint32
	leastUnrefSn uint32
	count        int64

	Config
}

func NewWithConfig(cfg Config) *MemDB {
	m := &MemDB{
		store:     skiplist.New(),
		snapshots: skiplist.New(),
		currSn:    1,
		Config:    cfg,
	}

	return m

}

func New() *MemDB {
	return NewWithConfig(DefaultConfig())
}

func (m *MemDB) Reset() {
	m.store = skiplist.New()
	m.snapshots = skiplist.New()
	m.currSn = 1
}

func (m *MemDB) getCurrSn() uint32 {
	return atomic.LoadUint32(&m.currSn)
}

func (m *MemDB) setLeastUnrefSn() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	if iter.Valid() {
		snap := iter.Get().(*Snapshot)
		atomic.StoreUint32(&m.leastUnrefSn, snap.sn-1)
	}
}

func (m *MemDB) getLeastUnrefSn() uint32 {
	return atomic.LoadUint32(&m.leastUnrefSn)
}

func (m *MemDB) NewWriter() *Writer {
	buf := m.store.MakeBuf()

	return &Writer{
		rand:  rand.New(rand.NewSource(int64(rand.Int()))),
		buf:   buf,
		iter:  m.store.NewIterator(m.iterCmp, buf),
		MemDB: m,
	}
}

type Snapshot struct {
	sn       uint32
	refCount int32
	db       *MemDB
	count    int64
}

func (s Snapshot) Count() int64 {
	return s.count
}

func (s *Snapshot) Encode(buf []byte, w io.Writer) error {
	l := 4
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint32(buf[0:4], s.sn)
	if _, err := w.Write(buf[0:4]); err != nil {
		return err
	}

	return nil

}

func (s *Snapshot) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:4]); err != nil {
		return err
	}
	s.sn = binary.BigEndian.Uint32(buf[0:4])
	return nil
}

func (s *Snapshot) Open() bool {
	if atomic.LoadInt32(&s.refCount) == 0 {
		return false
	}
	atomic.AddInt32(&s.refCount, 1)
	return true
}

func (s *Snapshot) Close() {
	newRefcount := atomic.AddInt32(&s.refCount, -1)
	if newRefcount == 0 {
		buf := s.db.snapshots.MakeBuf()
		defer s.db.snapshots.FreeBuf(buf)
		s.db.snapshots.Delete(s, CompareSnapshot, buf)
		s.db.setLeastUnrefSn()
		if atomic.CompareAndSwapInt32(&s.db.isGCRunning, 0, 1) {
			go s.db.GC()
		}
	}
}

func (s *Snapshot) NewIterator() *Iterator {
	return s.db.NewIterator(s)
}

func CompareSnapshot(this skiplist.Item, that skiplist.Item) int {
	thisItem := this.(*Snapshot)
	thatItem := that.(*Snapshot)

	return int(thisItem.sn) - int(thatItem.sn)
}

func (m *MemDB) NewSnapshot() *Snapshot {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	snap := &Snapshot{db: m, sn: m.getCurrSn(), refCount: 1, count: m.ItemsCount()}
	m.snapshots.Insert(snap, CompareSnapshot, buf)
	atomic.AddUint32(&m.currSn, 1)
	return snap
}

type Iterator struct {
	snap *Snapshot
	iter *skiplist.Iterator
	buf  *skiplist.ActionBuffer
}

func (it *Iterator) skipUnwanted() {
loop:
	if !it.iter.Valid() {
		return
	}
	itm := it.iter.Get().(*Item)
	if itm.bornSn > it.snap.sn || (itm.deadSn > 0 && itm.deadSn <= it.snap.sn) {
		it.iter.Next()
		goto loop
	}
}

func (it *Iterator) SeekFirst() {
	it.iter.SeekFirst()
	it.skipUnwanted()
}

func (it *Iterator) Seek(itm *Item) {
	it.iter.Seek(itm)
	it.skipUnwanted()
}

func (it *Iterator) Valid() bool {
	return it.iter.Valid()
}

func (it *Iterator) Get() *Item {
	return it.iter.Get().(*Item)
}

func (it *Iterator) Next() {
	it.iter.Next()
	it.skipUnwanted()
}

func (it *Iterator) Close() {
	it.snap.Close()
	it.snap.db.store.FreeBuf(it.buf)
}

func (m *MemDB) NewIterator(snap *Snapshot) *Iterator {
	if !snap.Open() {
		return nil
	}
	buf := snap.db.store.MakeBuf()
	return &Iterator{
		snap: snap,
		iter: m.store.NewIterator(m.iterCmp, buf),
		buf:  buf,
	}
}

func (m *MemDB) ItemsCount() int64 {
	return atomic.LoadInt64(&m.count)
}

func (m *MemDB) collectDead(sn uint32) {
	buf1 := m.snapshots.MakeBuf()
	buf2 := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf1)
	defer m.snapshots.FreeBuf(buf2)
	iter := m.store.NewIterator(m.iterCmp, buf1)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		itm := iter.Get().(*Item)
		if itm.deadSn > 0 && itm.deadSn <= m.getLeastUnrefSn() {
			m.store.DeleteNode(iter.GetNode(), m.insCmp, buf2)
		}
	}
}

func (m *MemDB) GC() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	if iter.Valid() {
		snap := iter.Get().(*Snapshot)
		if snap.sn != m.lastGCSn && snap.sn > 1 {
			m.lastGCSn = snap.sn - 1
			m.collectDead(m.lastGCSn)
		}
	}

	atomic.CompareAndSwapInt32(&m.isGCRunning, 1, 0)
}

func (m *MemDB) GetSnapshots() []*Snapshot {
	var snaps []*Snapshot
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		snaps = append(snaps, iter.Get().(*Snapshot))
	}

	return snaps
}

func (m *MemDB) StoreToDisk(dir string, snap *Snapshot, callb ItemCallback) error {
	os.MkdirAll(dir, 0755)
	datafile := path.Join(dir, "records.data")
	w := newFileWriter(m.fileType)
	if err := w.Open(datafile); err != nil {
		return err
	}
	defer w.Close()

	itr := m.NewIterator(snap)
	if itr == nil {
		return errors.New("Invalid snapshot")
	}
	defer itr.Close()

	for itr.SeekFirst(); itr.Valid(); itr.Next() {
		itm := itr.Get()
		if err := w.WriteItem(itm); err != nil {
			return err
		}

		if callb != nil {
			callb(itm)
		}
	}

	return nil
}

func (m *MemDB) LoadFromDisk(dir string, callb ItemCallback) (*Snapshot, error) {
	var wg sync.WaitGroup
	datafile := path.Join(dir, "records.data")
	r := newFileReader(m.fileType)
	if err := r.Open(datafile); err != nil {
		return nil, err
	}

	defer r.Close()

	ch := make(chan *Item, readerBufSize)
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go func(wg *sync.WaitGroup) {
			defer wg.Done()
			w := m.NewWriter()
			for itm := range ch {
				w.Put(itm)
				if callb != nil {
					callb(itm)
				}
			}
		}(&wg)
	}

loop:
	for {
		itm, err := r.ReadItem()
		if err != nil {
			return nil, err
		}

		if itm == nil {
			break loop
		}
		ch <- itm
	}

	close(ch)
	wg.Wait()

	snap := m.NewSnapshot()
	return snap, nil
}

func (m *MemDB) DumpStats() string {
	return m.store.GetStats().String()
}
<commit_msg>memdb: minor code cleanup<commit_after>package memdb

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/t3rm1n4l/memdb/skiplist"
	"io"
	"math/rand"
	"os"
	"path"
	"runtime"
	"sync"
	"sync/atomic"
)

type KeyCompare func([]byte, []byte) int
type ItemCallback func(*Item)

type FileType int

const (
	encodeBufSize = 4
	readerBufSize = 10000
)

const (
	ForestdbFile FileType = iota
	RawdbFile
)

func DefaultConfig() Config {
	var cfg Config
	cfg.SetKeyComparator(defaultKeyCmp)
	cfg.SetFileType(RawdbFile)
	return cfg
}

type Item struct {
	bornSn, deadSn uint32
	data           []byte
}

func (itm *Item) Encode(buf []byte, w io.Writer) error {
	l := 2
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint16(buf[0:2], uint16(len(itm.data)))
	if _, err := w.Write(buf[0:2]); err != nil {
		return err
	}
	if _, err := w.Write(itm.data); err != nil {
		return err
	}

	return nil
}

func (itm *Item) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:2]); err != nil {
		return err
	}
	l := binary.BigEndian.Uint16(buf[0:2])
	itm.data = make([]byte, int(l))
	_, err := io.ReadFull(r, itm.data)

	return err
}

func (itm *Item) Bytes() []byte {
	return itm.data
}

func NewItem(data []byte) *Item {
	return &Item{
		data: data,
	}
}

func newInsertCompare(keyCmp KeyCompare) skiplist.CompareFn {
	return func(this skiplist.Item, that skiplist.Item) int {
		var v int
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		if v = keyCmp(thisItem.data, thatItem.data); v == 0 {
			v = int(thisItem.bornSn) - int(thatItem.bornSn)
		}

		return v
	}
}

func newIterCompare(keyCmp KeyCompare) skiplist.CompareFn {
	return func(this skiplist.Item, that skiplist.Item) int {
		thisItem := this.(*Item)
		thatItem := that.(*Item)
		return keyCmp(thisItem.data, thatItem.data)
	}
}

func defaultKeyCmp(this []byte, that []byte) int {
	var l int

	l1 := len(this)
	l2 := len(that)
	if l1 < l2 {
		l = l1
	} else {
		l = l2
	}

	return bytes.Compare(this[:l], that[:l])
}

//
//compare item,sn
type Writer struct {
	rand *rand.Rand
	buf  *skiplist.ActionBuffer
	iter *skiplist.Iterator
	*MemDB
}

func (w *Writer) Put(x *Item) {
	x.bornSn = w.getCurrSn()
	w.store.Insert2(x, w.insCmp, w.buf, w.rand.Float32)
	atomic.AddInt64(&w.count, 1)
}

// Find first item, seek until dead=0, mark dead=sn
func (w *Writer) Delete(x *Item) (success bool) {
	defer func() {
		if success {
			atomic.AddInt64(&w.count, -1)
		}
	}()

	gotNode := w.GetNode(x)
	if gotNode != nil {
		sn := w.getCurrSn()
		gotItem := gotNode.Item().(*Item)
		if gotItem.bornSn == sn {
			success = w.store.DeleteNode(gotNode, w.insCmp, w.buf)
			return
		}

		success = atomic.CompareAndSwapUint32(&gotItem.deadSn, 0, sn)
		return
	}

	return
}

func (w *Writer) Get(x *Item) *Item {
	n := w.GetNode(x)
	if n != nil {
		return n.Item().(*Item)
	}
	return nil
}

func (w *Writer) GetNode(x *Item) *skiplist.Node {
	var curr *skiplist.Node
	found := w.iter.Seek(x)
	if !found {
		return nil
	}

	// Seek until most recent item for key is found
	curr = w.iter.GetNode()
	for {
		w.iter.Next()
		if !w.iter.Valid() {
			break
		}
		next := w.iter.GetNode()
		nxtItm := next.Item().(*Item)
		currItm := curr.Item().(*Item)
		if w.iterCmp(nxtItm, currItm) != 0 {
			break
		}

		curr = next
	}

	currItm := curr.Item().(*Item)
	if currItm.deadSn != 0 {
		return nil
	}

	return curr
}

type Config struct {
	keyCmp  KeyCompare
	insCmp  skiplist.CompareFn
	iterCmp skiplist.CompareFn

	fileType FileType
}

func (cfg *Config) SetKeyComparator(cmp KeyCompare) {
	cfg.keyCmp = cmp
	cfg.insCmp = newInsertCompare(cmp)
	cfg.iterCmp = newIterCompare(cmp)
}

func (cfg *Config) SetFileType(t FileType) error {
	switch t {
	case ForestdbFile, RawdbFile:
	default:
		return errors.New("Invalid format")
	}

	cfg.fileType = t
	return nil
}

type MemDB struct {
	store        *skiplist.Skiplist
	currSn       uint32
	snapshots    *skiplist.Skiplist
	isGCRunning  int32
	lastGCSn     uint32
	leastUnrefSn uint32
	count        int64

	Config
}

func NewWithConfig(cfg Config) *MemDB {
	m := &MemDB{
		store:     skiplist.New(),
		snapshots: skiplist.New(),
		currSn:    1,
		Config:    cfg,
	}

	return m

}

func New() *MemDB {
	return NewWithConfig(DefaultConfig())
}

func (m *MemDB) Reset() {
	m.store = skiplist.New()
	m.snapshots = skiplist.New()
	m.currSn = 1
}

func (m *MemDB) getCurrSn() uint32 {
	return atomic.LoadUint32(&m.currSn)
}

func (m *MemDB) setLeastUnrefSn() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	if iter.Valid() {
		snap := iter.Get().(*Snapshot)
		atomic.StoreUint32(&m.leastUnrefSn, snap.sn-1)
	}
}

func (m *MemDB) getLeastUnrefSn() uint32 {
	return atomic.LoadUint32(&m.leastUnrefSn)
}

func (m *MemDB) NewWriter() *Writer {
	buf := m.store.MakeBuf()

	return &Writer{
		rand:  rand.New(rand.NewSource(int64(rand.Int()))),
		buf:   buf,
		iter:  m.store.NewIterator(m.iterCmp, buf),
		MemDB: m,
	}
}

type Snapshot struct {
	sn       uint32
	refCount int32
	db       *MemDB
	count    int64
}

func (s Snapshot) Count() int64 {
	return s.count
}

func (s *Snapshot) Encode(buf []byte, w io.Writer) error {
	l := 4
	if len(buf) < l {
		return ErrNotEnoughSpace
	}

	binary.BigEndian.PutUint32(buf[0:4], s.sn)
	if _, err := w.Write(buf[0:4]); err != nil {
		return err
	}

	return nil

}

func (s *Snapshot) Decode(buf []byte, r io.Reader) error {
	if _, err := io.ReadFull(r, buf[0:4]); err != nil {
		return err
	}
	s.sn = binary.BigEndian.Uint32(buf[0:4])
	return nil
}

func (s *Snapshot) Open() bool {
	if atomic.LoadInt32(&s.refCount) == 0 {
		return false
	}
	atomic.AddInt32(&s.refCount, 1)
	return true
}

func (s *Snapshot) Close() {
	newRefcount := atomic.AddInt32(&s.refCount, -1)
	if newRefcount == 0 {
		buf := s.db.snapshots.MakeBuf()
		defer s.db.snapshots.FreeBuf(buf)
		s.db.snapshots.Delete(s, CompareSnapshot, buf)
		s.db.setLeastUnrefSn()
		if atomic.CompareAndSwapInt32(&s.db.isGCRunning, 0, 1) {
			go s.db.GC()
		}
	}
}

func (s *Snapshot) NewIterator() *Iterator {
	return s.db.NewIterator(s)
}

func CompareSnapshot(this skiplist.Item, that skiplist.Item) int {
	thisItem := this.(*Snapshot)
	thatItem := that.(*Snapshot)

	return int(thisItem.sn) - int(thatItem.sn)
}

func (m *MemDB) NewSnapshot() *Snapshot {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	snap := &Snapshot{db: m, sn: m.getCurrSn(), refCount: 1, count: m.ItemsCount()}
	m.snapshots.Insert(snap, CompareSnapshot, buf)
	atomic.AddUint32(&m.currSn, 1)
	return snap
}

type Iterator struct {
	snap *Snapshot
	iter *skiplist.Iterator
	buf  *skiplist.ActionBuffer
}

func (it *Iterator) skipUnwanted() {
loop:
	if !it.iter.Valid() {
		return
	}
	itm := it.iter.Get().(*Item)
	if itm.bornSn > it.snap.sn || (itm.deadSn > 0 && itm.deadSn <= it.snap.sn) {
		it.iter.Next()
		goto loop
	}
}

func (it *Iterator) SeekFirst() {
	it.iter.SeekFirst()
	it.skipUnwanted()
}

func (it *Iterator) Seek(itm *Item) {
	it.iter.Seek(itm)
	it.skipUnwanted()
}

func (it *Iterator) Valid() bool {
	return it.iter.Valid()
}

func (it *Iterator) Get() *Item {
	return it.iter.Get().(*Item)
}

func (it *Iterator) Next() {
	it.iter.Next()
	it.skipUnwanted()
}

func (it *Iterator) Close() {
	it.snap.Close()
	it.snap.db.store.FreeBuf(it.buf)
}

func (m *MemDB) NewIterator(snap *Snapshot) *Iterator {
	if !snap.Open() {
		return nil
	}
	buf := snap.db.store.MakeBuf()
	return &Iterator{
		snap: snap,
		iter: m.store.NewIterator(m.iterCmp, buf),
		buf:  buf,
	}
}

func (m *MemDB) ItemsCount() int64 {
	return atomic.LoadInt64(&m.count)
}

func (m *MemDB) collectDead(sn uint32) {
	buf1 := m.snapshots.MakeBuf()
	buf2 := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf1)
	defer m.snapshots.FreeBuf(buf2)
	iter := m.store.NewIterator(m.iterCmp, buf1)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		itm := iter.Get().(*Item)
		if itm.deadSn > 0 && itm.deadSn <= m.getLeastUnrefSn() {
			m.store.DeleteNode(iter.GetNode(), m.insCmp, buf2)
		}
	}
}

func (m *MemDB) GC() {
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)

	sn := m.getLeastUnrefSn()
	if sn != m.lastGCSn && sn > 0 {
		m.lastGCSn = sn
		m.collectDead(m.lastGCSn)
	}

	atomic.CompareAndSwapInt32(&m.isGCRunning, 1, 0)
}

func (m *MemDB) GetSnapshots() []*Snapshot {
	var snaps []*Snapshot
	buf := m.snapshots.MakeBuf()
	defer m.snapshots.FreeBuf(buf)
	iter := m.snapshots.NewIterator(CompareSnapshot, buf)
	iter.SeekFirst()
	for ; iter.Valid(); iter.Next() {
		snaps = append(snaps, iter.Get().(*Snapshot))
	}

	return snaps
}

func (m *MemDB) StoreToDisk(dir string, snap *Snapshot, callb ItemCallback) error {
	os.MkdirAll(dir, 0755)
	datafile := path.Join(dir, "records.data")
	w := newFileWriter(m.fileType)
	if err := w.Open(datafile); err != nil {
		return err
	}
	defer w.Close()

	itr := m.NewIterator(snap)
	if itr == nil {
		return errors.New("Invalid snapshot")
	}
	defer itr.Close()

	for itr.SeekFirst(); itr.Valid(); itr.Next() {
		itm := itr.Get()
		if err := w.WriteItem(itm); err != nil {
			return err
		}

		if callb != nil {
			callb(itm)
		}
	}

	return nil
}

func (m *MemDB) LoadFromDisk(dir string, callb ItemCallback) (*Snapshot, error) {
	var wg sync.WaitGroup
	datafile := path.Join(dir, "records.data")
	r := newFileReader(m.fileType)
	if err := r.Open(datafile); err != nil {
		return nil, err
	}

	defer r.Close()

	ch := make(chan *Item, readerBufSize)
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go func(wg *sync.WaitGroup) {
			defer wg.Done()
			w := m.NewWriter()
			for itm := range ch {
				w.Put(itm)
				if callb != nil {
					callb(itm)
				}
			}
		}(&wg)
	}

loop:
	for {
		itm, err := r.ReadItem()
		if err != nil {
			return nil, err
		}

		if itm == nil {
			break loop
		}
		ch <- itm
	}

	close(ch)
	wg.Wait()

	snap := m.NewSnapshot()
	return snap, nil
}

func (m *MemDB) DumpStats() string {
	return m.store.GetStats().String()
}
<|endoftext|>


<commit_before>// Licensed to The Moov Authors under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. The Moov Authors licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package ach

import (
	"bufio"
	"bytes"
	"fmt"
	"time"
)

const NACHAFileLineLimit = 10000

// MergeFiles is a helper function for consolidating an array of ACH Files into as few files
// as possible. This is useful for optimizing cost and network efficiency.
// This operation will override batch numbers in each file to ensure they do not collide. The ascending batch numbers will start at 1.
//
// Per NACHA rules files must remain under 10,000 lines (when rendered in their ASCII encoding)
//
// File Batches can only be merged if they are unique and routed to and from the same ABA routing numbers.
func MergeFiles(files []*File) ([]*File, error) {
	fs := &mergableFiles{infiles: files}
	for i := range fs.infiles {
		outf := fs.lookupByHeader(fs.infiles[i])
		for j := range fs.infiles[i].Batches {
			batchExistsInMerged := false
			for k := range outf.Batches {
				if fs.infiles[i].Batches[j].Equal(outf.Batches[k]) {
					batchExistsInMerged = true
				}
			}
			if !batchExistsInMerged {
				outf.AddBatch(fs.infiles[i].Batches[j])
				if err := outf.Create(); err != nil {
					return nil, err
				}
				n, err := lineCount(outf)
				if n == 0 || err != nil {
					return nil, fmt.Errorf("problem getting line count of File (header: %#v): %v", outf.Header, err)
				}
				if n > NACHAFileLineLimit {
					outf.RemoveBatch(fs.infiles[i].Batches[j])
					if err := outf.Create(); err != nil { // rebalance ACH file after removing the Batch
						return nil, err
					}
					f := *outf
					fs.locMaxed = append(fs.locMaxed, &f)

					outf = fs.create(outf) // replace output file with the one we just created

					outf.AddBatch(fs.infiles[i].Batches[j])
					if err := outf.Create(); err != nil {
						return nil, err
					}
				}
			}
		}
	}

	// Return LOC-maxed files and merged files
	out := append(fs.locMaxed, fs.outfiles...)

	// Override batch numbers to ensure they don't collide
	for _, f := range out {
		for i, b := range f.Batches {
			b.GetHeader().BatchNumber = i + 1
			b.GetControl().BatchNumber = i + 1
		}
	}

	return out, nil
}

type mergableFiles struct {
	infiles  []*File
	outfiles []*File
	locMaxed []*File
}

// create returns the index of a newly created file in fs.outfiles given the details from f.Header
func (fs *mergableFiles) create(f *File) *File { // returns the outfiles index of the created file
	now := time.Now()

	// remove the current file from outfiles
	for i := range fs.outfiles {
		if fs.outfiles[i].Header.ImmediateDestination == f.Header.ImmediateDestination &&
			fs.outfiles[i].Header.ImmediateOrigin == f.Header.ImmediateOrigin {
			// found a matching file, so remove it from fs.outfiles
			fs.outfiles = append(fs.outfiles[:i], fs.outfiles[i+1:]...)
			goto next
		}
	}
next:
	out := NewFile()
	out.Header = f.Header
	out.Header.FileCreationDate = now.Format("060102") // YYMMDD
	out.Header.FileCreationTime = now.Format("1504")   // HHmm
	out.Create()
	fs.outfiles = append(fs.outfiles, out) // add the new outfile

	return out
}

// lookupByHeader optionally returns a File from fs.files if the FileHeaders match.
// This is done because we append batches into files to minimize the count of output files.
//
// lookupByHeader will return the existing file (stored in outfiles) if no matching file exists.
func (fs *mergableFiles) lookupByHeader(f *File) *File {
	for i := range fs.outfiles {
		if fs.outfiles[i].Header.ImmediateDestination == f.Header.ImmediateDestination &&
			fs.outfiles[i].Header.ImmediateOrigin == f.Header.ImmediateOrigin {
			// found a matching file, so return it
			return fs.outfiles[i]
		}
	}
	fs.outfiles = append(fs.outfiles, f)
	return f
}

func lineCount(f *File) (int, error) {
	if len(f.Batches) < 100 {
		// Ignore Files with low batch counts by returning a valid count.
		// Calling Writer.Write() is costly and so we're going to ignore it in easy cases.
		return 1, nil
	}

	var buf bytes.Buffer
	if err := NewWriter(&buf).Write(f); err != nil {
		return 0, err
	}
	lines := 0
	s := bufio.NewScanner(&buf)
	for s.Scan() {
		if v := s.Text(); v != "" {
			lines++
		}
	}
	return lines, nil
}
<commit_msg>gofmt file to fix breaking build<commit_after>// Licensed to The Moov Authors under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. The Moov Authors licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package ach

import (
	"bufio"
	"bytes"
	"fmt"
	"time"
)

const NACHAFileLineLimit = 10000

// MergeFiles is a helper function for consolidating an array of ACH Files into as few files
// as possible. This is useful for optimizing cost and network efficiency.
// This operation will override batch numbers in each file to ensure they do not collide. The ascending batch numbers will start at 1.
//
// Per NACHA rules files must remain under 10,000 lines (when rendered in their ASCII encoding)
//
// File Batches can only be merged if they are unique and routed to and from the same ABA routing numbers.
func MergeFiles(files []*File) ([]*File, error) {
	fs := &mergableFiles{infiles: files}
	for i := range fs.infiles {
		outf := fs.lookupByHeader(fs.infiles[i])
		for j := range fs.infiles[i].Batches {
			batchExistsInMerged := false
			for k := range outf.Batches {
				if fs.infiles[i].Batches[j].Equal(outf.Batches[k]) {
					batchExistsInMerged = true
				}
			}
			if !batchExistsInMerged {
				outf.AddBatch(fs.infiles[i].Batches[j])
				if err := outf.Create(); err != nil {
					return nil, err
				}
				n, err := lineCount(outf)
				if n == 0 || err != nil {
					return nil, fmt.Errorf("problem getting line count of File (header: %#v): %v", outf.Header, err)
				}
				if n > NACHAFileLineLimit {
					outf.RemoveBatch(fs.infiles[i].Batches[j])
					if err := outf.Create(); err != nil { // rebalance ACH file after removing the Batch
						return nil, err
					}
					f := *outf
					fs.locMaxed = append(fs.locMaxed, &f)

					outf = fs.create(outf) // replace output file with the one we just created

					outf.AddBatch(fs.infiles[i].Batches[j])
					if err := outf.Create(); err != nil {
						return nil, err
					}
				}
			}
		}
	}

	// Return LOC-maxed files and merged files
	out := append(fs.locMaxed, fs.outfiles...)

	// Override batch numbers to ensure they don't collide
	for _, f := range out {
		for i, b := range f.Batches {
			b.GetHeader().BatchNumber = i + 1
			b.GetControl().BatchNumber = i + 1
		}
	}

	return out, nil
}

type mergableFiles struct {
	infiles  []*File
	outfiles []*File
	locMaxed []*File
}

// create returns the index of a newly created file in fs.outfiles given the details from f.Header
func (fs *mergableFiles) create(f *File) *File { // returns the outfiles index of the created file
	now := time.Now()

	// remove the current file from outfiles
	for i := range fs.outfiles {
		if fs.outfiles[i].Header.ImmediateDestination == f.Header.ImmediateDestination &&
			fs.outfiles[i].Header.ImmediateOrigin == f.Header.ImmediateOrigin {
			// found a matching file, so remove it from fs.outfiles
			fs.outfiles = append(fs.outfiles[:i], fs.outfiles[i+1:]...)
			goto next
		}
	}
next:
	out := NewFile()
	out.Header = f.Header
	out.Header.FileCreationDate = now.Format("060102") // YYMMDD
	out.Header.FileCreationTime = now.Format("1504")   // HHmm
	out.Create()
	fs.outfiles = append(fs.outfiles, out) // add the new outfile

	return out
}

// lookupByHeader optionally returns a File from fs.files if the FileHeaders match.
// This is done because we append batches into files to minimize the count of output files.
//
// lookupByHeader will return the existing file (stored in outfiles) if no matching file exists.
func (fs *mergableFiles) lookupByHeader(f *File) *File {
	for i := range fs.outfiles {
		if fs.outfiles[i].Header.ImmediateDestination == f.Header.ImmediateDestination &&
			fs.outfiles[i].Header.ImmediateOrigin == f.Header.ImmediateOrigin {
			// found a matching file, so return it
			return fs.outfiles[i]
		}
	}
	fs.outfiles = append(fs.outfiles, f)
	return f
}

func lineCount(f *File) (int, error) {
	if len(f.Batches) < 100 {
		// Ignore Files with low batch counts by returning a valid count.
		// Calling Writer.Write() is costly and so we're going to ignore it in easy cases.
		return 1, nil
	}

	var buf bytes.Buffer
	if err := NewWriter(&buf).Write(f); err != nil {
		return 0, err
	}
	lines := 0
	s := bufio.NewScanner(&buf)
	for s.Scan() {
		if v := s.Text(); v != "" {
			lines++
		}
	}
	return lines, nil
}
<|endoftext|>


<commit_before>// Package templates implements template execution for files to be dynamically rendered for the client.
package templates

import (
	"bytes"
	"net/http"
	"path"
	"path/filepath"
	"text/template"

	"github.com/mholt/caddy/middleware"
)

// ServeHTTP implements the middleware.Handler interface.
func (t Templates) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	for _, rule := range t.Rules {
		if !middleware.Path(r.URL.Path).Matches(rule.Path) {
			continue
		}

		// Check for index files
		fpath := r.URL.Path
		if idx, ok := middleware.IndexFile(t.FileSys, fpath, rule.IndexFiles); ok {
			fpath = idx
		}

		// Check the extension
		reqExt := path.Ext(fpath)

		for _, ext := range rule.Extensions {
			if reqExt == ext {
				// Create execution context
				ctx := context{root: t.FileSys, req: r, URL: r.URL}

				// Build the template
				tpl, err := template.ParseFiles(filepath.Join(t.Root, fpath))
				if err != nil {
					return http.StatusInternalServerError, err
				}

				// Execute it
				var buf bytes.Buffer
				err = tpl.Execute(&buf, ctx)
				if err != nil {
					return http.StatusInternalServerError, err
				}
				buf.WriteTo(w)

				return http.StatusOK, nil
			}
		}
	}

	return t.Next.ServeHTTP(w, r)
}

// Templates is middleware to render templated files as the HTTP response.
type Templates struct {
	Next    middleware.Handler
	Rules   []Rule
	Root    string
	FileSys http.FileSystem
}

// Rule represents a template rule. A template will only execute
// with this rule if the request path matches the Path specified
// and requests a resource with one of the extensions specified.
type Rule struct {
	Path       string
	Extensions []string
	IndexFiles []string
}
<commit_msg>templates: Better error handling for missing files<commit_after>// Package templates implements template execution for files to be dynamically rendered for the client.
package templates

import (
	"bytes"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"text/template"

	"github.com/mholt/caddy/middleware"
)

// ServeHTTP implements the middleware.Handler interface.
func (t Templates) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	for _, rule := range t.Rules {
		if !middleware.Path(r.URL.Path).Matches(rule.Path) {
			continue
		}

		// Check for index files
		fpath := r.URL.Path
		if idx, ok := middleware.IndexFile(t.FileSys, fpath, rule.IndexFiles); ok {
			fpath = idx
		}

		// Check the extension
		reqExt := path.Ext(fpath)

		for _, ext := range rule.Extensions {
			if reqExt == ext {
				// Create execution context
				ctx := context{root: t.FileSys, req: r, URL: r.URL}

				// Build the template
				tpl, err := template.ParseFiles(filepath.Join(t.Root, fpath))
				if err != nil {
					if os.IsNotExist(err) {
						return http.StatusNotFound, nil
					} else if os.IsPermission(err) {
						return http.StatusForbidden, nil
					}
					return http.StatusInternalServerError, err
				}

				// Execute it
				var buf bytes.Buffer
				err = tpl.Execute(&buf, ctx)
				if err != nil {
					return http.StatusInternalServerError, err
				}
				buf.WriteTo(w)

				return http.StatusOK, nil
			}
		}
	}

	return t.Next.ServeHTTP(w, r)
}

// Templates is middleware to render templated files as the HTTP response.
type Templates struct {
	Next    middleware.Handler
	Rules   []Rule
	Root    string
	FileSys http.FileSystem
}

// Rule represents a template rule. A template will only execute
// with this rule if the request path matches the Path specified
// and requests a resource with one of the extensions specified.
type Rule struct {
	Path       string
	Extensions []string
	IndexFiles []string
}
<|endoftext|>