<commit_before>package watch

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"

	"log"

	"reflect"

	"github.com/foomo/petze/config"
)

var typeDNSConfigErr = reflect.TypeOf(&net.DNSConfigError{})
var typeDNSErr = reflect.TypeOf(&net.DNSError{})
var typeOpErr = reflect.TypeOf(&net.OpError{})
var typeX509CertificateInvalidError = reflect.TypeOf(x509.CertificateInvalidError{})
var typeX509HostnameError = reflect.TypeOf(x509.HostnameError{})
var typeX509SystemRootsError = reflect.TypeOf(x509.SystemRootsError{})
var typeX509UnknownAuthorityError = reflect.TypeOf(x509.UnknownAuthorityError{})

type ErrorType string

const (
	ErrorInvalidEndpoint           ErrorType = "endpointInvalid"
	ErrorTypeServerTooSlow                   = "serverTooSlow"
	ErrorTypeUnknownError                    = "unknownError"
	ErrorTypeClientError                     = "clientError"
	ErrorTypeDNS                             = "dns"
	ErrorTypeDNSConfig                       = "dnsConfig"
	ErrorTypeTLSCertificateInvalid           = "tlsCertificateInvalid"
	ErrorTypeTLSHostNameError                = "tlsHostNameError"
	ErrorTypeTLSSystemRootsError             = "tlsSystemRootsError"
	ErrorTypeTLSUnknownAuthority             = "tlsUnknownAutority"
	ErrorTypeWrongHTTPStatusCode             = "wrongHTTPStatus"
	ErrorTypeCertificateIsExpiring           = "certificateIsExpiring"
)

type WarningType string

const (
	WarningTypeCertificateWillExpire WarningType = "certificateWillExpire"
)

type Error struct {
	Error string
	Type  ErrorType
}

type Warning struct {
	Warning string
	Type    WarningType
}

type Result struct {
	ID         string        `json:"id"`
	Errors     []Error       `json:"errors"`
	Warnings   []Warning     `json:"warnings"`
	Timeout    bool          `json:"timeout"`
	Timestamp  time.Time     `json:"timestamp"`
	RunTime    time.Duration `json:"runTime"`
	StatusCode int           `json:"statusCode"`
}

func NewResult(id string) *Result {
	return &Result{
		ID:        id,
		Errors:    []Error{},
		Timestamp: time.Now(),
	}
}

func (r *Result) addError(e error, t ErrorType) {
	r.Errors = addError(r.Errors, e, t)
}

func addError(errors []Error, err error, t ErrorType) []Error {
	return append(errors, Error{
		Error: err.Error(),
		Type:  t,
	})
}

func addWarning(warnings []Warning, warning string, t WarningType) []Warning {
	return append(warnings, Warning{
		Warning: warning,
		Type:    t,
	})
}

type dialerErrRecorder struct {
	warnings                   []Warning
	errors                     []Error
	err                        net.Error
	dnsError                   net.Error
	dnsConfigError             net.Error
	tlsCertificateInvalidError *x509.CertificateInvalidError
	tlsHostnameError           *x509.HostnameError
	tlsSystemRootsError        *x509.SystemRootsError
	tlsUnknownAuthorityError   *x509.UnknownAuthorityError
}

type Watcher struct {
	active  bool
	service *config.Service
}

// Watch create a watcher and start watching
func Watch(service *config.Service, chanResult chan *Result) *Watcher {
	w := &Watcher{
		active:  true,
		service: service,
	}
	go w.watchLoop(chanResult)
	return w
}

// Stop watching - beware this is async
func (w *Watcher) Stop() {
	w.active = false
}

func (w *Watcher) watchLoop(chanResult chan *Result) {
	for w.active {
		r := watch(w.service)
		if w.active {
			chanResult <- r
			time.Sleep(time.Second * time.Duration(w.service.Interval))
		}
	}
}

func getClientAndDialErrRecorder() (client *http.Client, errRecorder *dialerErrRecorder) {
	errRecorder = &dialerErrRecorder{
		warnings: []Warning{},
		errors:   []Error{},
	}
	tlsConfig := &tls.Config{}
	dialer := &net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 0 * time.Second,
	}
	dialTLS := func(network, address string) (conn net.Conn, err error) {
		tlsConn, tlsErr := tls.DialWithDialer(dialer, network, address, tlsConfig)
		if tlsErr == nil {
			//conn = tlsConn.(net.Conn)
			connectionState := tlsConn.ConnectionState()
			for _, cert := range connectionState.PeerCertificates {
				durationUntilExpiry := cert.NotAfter.Sub(time.Now())
				if durationUntilExpiry < time.Hour*24 {
					// well in less than 24 h is worth an error
					errRecorder.errors = addError(errRecorder.errors, errors.New(fmt.Sprint("cert CN=\"", cert.Subject.CommonName, "\" is expiring in less than 24h: ", cert.NotAfter, ", left: ", durationUntilExpiry)), ErrorTypeCertificateIsExpiring)
				} else if durationUntilExpiry < time.Hour*24*14 {
					// we will start issuing warnings two weeks before a cert expires
					errRecorder.warnings = addWarning(errRecorder.warnings, fmt.Sprint("cert CN=\"", cert.Subject.CommonName, "\" expires on: ", cert.NotAfter, ", left: ", durationUntilExpiry), WarningTypeCertificateWillExpire)
				}
			}
			conn = tlsConn
		} else {
			switch reflect.TypeOf(tlsErr) {
			case typeX509UnknownAuthorityError:
				unknownAuthorityError := tlsErr.(x509.UnknownAuthorityError)
				errRecorder.tlsUnknownAuthorityError = &unknownAuthorityError
			case typeX509HostnameError:
				hostnameErr := tlsErr.(x509.HostnameError)
				errRecorder.tlsHostnameError = &hostnameErr
			case typeX509CertificateInvalidError:
				tlsCertificateInvalidError := tlsErr.(x509.CertificateInvalidError)
				errRecorder.tlsCertificateInvalidError = &tlsCertificateInvalidError
			case typeX509SystemRootsError:
				systemRootsError := tlsErr.(x509.SystemRootsError)
				errRecorder.tlsSystemRootsError = &systemRootsError
			default:
				log.Println("unknown tls error", reflect.TypeOf(tlsErr), tlsErr)
			}
		}
		return conn, tlsErr
	}
	dial := func(network, address string) (conn net.Conn, err error) {
		conn, err = dialer.Dial(network, address)
		if err != nil {
			switch reflect.TypeOf(err) {
			case typeOpErr:
				opError := reflect.ValueOf(err).Elem().Interface().(net.OpError)
				switch reflect.TypeOf(opError.Err) {
				case typeDNSConfigErr:
					log.Println("dns config error")
					errRecorder.dnsConfigError = opError.Err.(net.Error)
				case typeDNSErr:
					log.Println("dns error")
					errRecorder.dnsError = opError.Err.(net.Error)
				default:
					log.Println("misc error", reflect.TypeOf(opError.Err), opError.Err)
					errRecorder.err = opError.Err.(net.Error)
				}
			default:
				log.Println("again some general bullshit", err)
				errRecorder.err = err.(net.Error)
			}
		}
		return
	}
	client = &http.Client{
		Transport: &http.Transport{
			Proxy:               http.ProxyFromEnvironment,
			Dial:                dial,
			DialTLS:             dialTLS,
			TLSHandshakeTimeout: 10 * time.Second,
			TLSClientConfig:     tlsConfig,
		},
	}
	return
}

// actual watch
func watch(service *config.Service) (r *Result) {
	r = NewResult(service.ID)
	// parsing, the endpoint
	request, err := http.NewRequest("GET", service.Endpoint, nil)
	if err != nil {
		r.addError(err, ErrorInvalidEndpoint)
		return r
	}
	// my opersonal dns error check
	if len(request.Host) > 0 {
		host := request.Host
		parts := strings.Split(host, ":")
		if len(parts) > 1 {
			host, _, err = net.SplitHostPort(request.Host)
			if err != nil {
				r.addError(err, ErrorInvalidEndpoint)
				return
			}
		}
		_, lookupErr := net.LookupIP(host)
		if lookupErr != nil {
			r.addError(lookupErr, ErrorTypeDNS)
			return
		}
	}

	// i am explicitly not calling http.Get, because it does 30x handling, that
	// I do not want
	client, errRecorder := getClientAndDialErrRecorder()
	response, err := client.Do(request)
	r.Errors = append(r.Errors, errRecorder.errors...)
	r.Warnings = append(r.Warnings, errRecorder.warnings...)
	r.RunTime = time.Since(r.Timestamp)

	if response != nil && response.Body != nil {
		// always close the body
		response.Body.Close()
	}
	if err != nil {

		var netErr net.Error
		switch true {
		case errRecorder.tlsHostnameError != nil:
			r.addError(errRecorder.tlsHostnameError, ErrorTypeTLSHostNameError)
		case errRecorder.tlsSystemRootsError != nil:
			r.addError(errRecorder.tlsSystemRootsError, ErrorTypeTLSSystemRootsError)
		case errRecorder.tlsUnknownAuthorityError != nil:
			r.addError(errRecorder.tlsUnknownAuthorityError, ErrorTypeTLSUnknownAuthority)
		case errRecorder.tlsCertificateInvalidError != nil:
			r.addError(errRecorder.tlsCertificateInvalidError, ErrorTypeTLSCertificateInvalid)
		case errRecorder.dnsConfigError != nil:
			netErr = errRecorder.dnsConfigError
			r.addError(errRecorder.dnsConfigError, ErrorTypeDNSConfig)
		case errRecorder.dnsError != nil:
			netErr = errRecorder.dnsError
			r.addError(errRecorder.dnsError, ErrorTypeDNS)
		case errRecorder.err != nil:
			netErr = errRecorder.err
			r.addError(errRecorder.err, ErrorTypeUnknownError)
		}
		//log.Println("service", service.ID, err, "errRecorder dns:", errRecorder.dnsError, ", dnsConfig", errRecorder.dnsConfigError, ", err:", errRecorder.err, ", tls cert invalid:", errRecorder.tlsCertificateInvalidError)
		if netErr != nil {
			r.Timeout = netErr.Timeout()
			//r.ErrorIsTemporary = netErr.Temporary()
		}
		if err != nil {
			r.addError(err, ErrorTypeClientError)
		}
		return
	}
	runTimeMilliseconds := r.RunTime.Nanoseconds() / 1000000
	maxRuntime := (int64(service.MaxRuntime) * time.Millisecond.Nanoseconds()) / 1000000
	//log.Println("run time >>>>>>>>>>>", service.ID, runTimeMilliseconds, maxRuntime)
	if runTimeMilliseconds > maxRuntime {
		r.addError(errors.New(fmt.Sprint("response time too slow:", runTimeMilliseconds, ", should not be more than:", maxRuntime)), ErrorTypeServerTooSlow)
	}
	r.StatusCode = response.StatusCode
	if response.StatusCode != http.StatusOK {
		r.addError(errors.New(fmt.Sprint("unexpected status code: ", response.StatusCode)), ErrorTypeWrongHTTPStatusCode)
	}
	return
}
<commit_msg>more errors and initial session support<commit_after>package watch

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/http/cookiejar"
	"strings"
	"time"

	"log"

	"reflect"

	"github.com/foomo/petze/config"
)

var typeDNSConfigErr = reflect.TypeOf(&net.DNSConfigError{})
var typeDNSErr = reflect.TypeOf(&net.DNSError{})
var typeOpErr = reflect.TypeOf(&net.OpError{})
var typeX509CertificateInvalidError = reflect.TypeOf(x509.CertificateInvalidError{})
var typeX509HostnameError = reflect.TypeOf(x509.HostnameError{})
var typeX509SystemRootsError = reflect.TypeOf(x509.SystemRootsError{})
var typeX509UnknownAuthorityError = reflect.TypeOf(x509.UnknownAuthorityError{})

type ErrorType string

const (
	ErrorInvalidEndpoint           ErrorType = "endpointInvalid"
	ErrorTypeServerTooSlow                   = "serverTooSlow"
	ErrorTypeUnknownError                    = "unknownError"
	ErrorTypeClientError                     = "clientError"
	ErrorTypeDNS                             = "dns"
	ErrorTypeDNSConfig                       = "dnsConfig"
	ErrorTypeTLSCertificateInvalid           = "tlsCertificateInvalid"
	ErrorTypeTLSHostNameError                = "tlsHostNameError"
	ErrorTypeTLSSystemRootsError             = "tlsSystemRootsError"
	ErrorTypeTLSUnknownAuthority             = "tlsUnknownAutority"
	ErrorTypeWrongHTTPStatusCode             = "wrongHTTPStatus"
	ErrorTypeCertificateIsExpiring           = "certificateIsExpiring"
	ErrorTypeUnexpectedContentType           = "unexpectedContentType"
	ErrorTypeSessionFail                     = "sessionFail"
	ErrorTypeGoQueryMismatch                 = "goqueryMismatch"
)

type WarningType string

const (
	WarningTypeCertificateWillExpire WarningType = "certificateWillExpire"
)

type Error struct {
	Error   string    `json:"error"`
	Type    ErrorType `json:"type"`
	Comment string    `json:"comment,omitempty"`
}

type Warning struct {
	Warning string      `json:"warning"`
	Type    WarningType `json:"type"`
}

type Result struct {
	ID        string    `json:"id"`
	Errors    []Error   `json:"errors"`
	Warnings  []Warning `json:"warnings"`
	Timeout   bool      `json:"timeout"`
	Timestamp time.Time `json:"timestamp"`
}

func NewResult(id string) *Result {
	return &Result{
		ID:        id,
		Errors:    []Error{},
		Timestamp: time.Now(),
	}
}

func (r *Result) addError(e error, t ErrorType, comment string) {
	r.Errors = addError(r.Errors, e, t, comment)
}

func addError(errors []Error, err error, t ErrorType, comment string) []Error {
	return append(errors, Error{
		Error:   err.Error(),
		Type:    t,
		Comment: comment,
	})
}

func addWarning(warnings []Warning, warning string, t WarningType) []Warning {
	return append(warnings, Warning{
		Warning: warning,
		Type:    t,
	})
}

type dialerErrRecorder struct {
	warnings                   []Warning
	errors                     []Error
	unknownErr                 error
	err                        net.Error
	dnsError                   net.Error
	dnsConfigError             net.Error
	tlsCertificateInvalidError *x509.CertificateInvalidError
	tlsHostnameError           *x509.HostnameError
	tlsSystemRootsError        *x509.SystemRootsError
	tlsUnknownAuthorityError   *x509.UnknownAuthorityError
}

type Watcher struct {
	active  bool
	service *config.Service
}

// Watch create a watcher and start watching
func Watch(service *config.Service, chanResult chan Result) *Watcher {
	w := &Watcher{
		active:  true,
		service: service,
	}
	go w.watchLoop(chanResult)
	return w
}

// Stop watching - beware this is async
func (w *Watcher) Stop() {
	w.active = false
}

func (w *Watcher) watchLoop(chanResult chan Result) {
	for w.active {
		r := watch(w.service)
		if w.active {
			chanResult <- *r
			time.Sleep(time.Second * time.Duration(w.service.Interval))
		}
	}
}

func getClientAndDialErrRecorder() (client *http.Client, errRecorder *dialerErrRecorder) {
	errRecorder = &dialerErrRecorder{
		warnings: []Warning{},
		errors:   []Error{},
	}
	tlsConfig := &tls.Config{}
	dialer := &net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 0 * time.Second,
	}
	dialTLS := func(network, address string) (conn net.Conn, err error) {
		tlsConn, tlsErr := tls.DialWithDialer(dialer, network, address, tlsConfig)
		if tlsErr == nil {
			//conn = tlsConn.(net.Conn)
			connectionState := tlsConn.ConnectionState()
			for _, cert := range connectionState.PeerCertificates {
				durationUntilExpiry := cert.NotAfter.Sub(time.Now())
				if durationUntilExpiry < time.Hour*24 {
					// well in less than 24 h is worth an error
					errRecorder.errors = addError(errRecorder.errors, errors.New(fmt.Sprint("cert CN=\"", cert.Subject.CommonName, "\" is expiring in less than 24h: ", cert.NotAfter, ", left: ", durationUntilExpiry)), ErrorTypeCertificateIsExpiring, "")
				} else if durationUntilExpiry < time.Hour*24*14 {
					// we will start issuing warnings two weeks before a cert expires
					errRecorder.warnings = addWarning(errRecorder.warnings, fmt.Sprint("cert CN=\"", cert.Subject.CommonName, "\" expires on: ", cert.NotAfter, ", left: ", durationUntilExpiry), WarningTypeCertificateWillExpire)
				}
			}
			conn = tlsConn
		} else {

			switch reflect.TypeOf(tlsErr) {
			case typeX509UnknownAuthorityError:
				unknownAuthorityError := tlsErr.(x509.UnknownAuthorityError)
				errRecorder.tlsUnknownAuthorityError = &unknownAuthorityError
			case typeX509HostnameError:
				hostnameErr := tlsErr.(x509.HostnameError)
				errRecorder.tlsHostnameError = &hostnameErr
			case typeX509CertificateInvalidError:
				tlsCertificateInvalidError := tlsErr.(x509.CertificateInvalidError)
				errRecorder.tlsCertificateInvalidError = &tlsCertificateInvalidError
			case typeX509SystemRootsError:
				systemRootsError := tlsErr.(x509.SystemRootsError)
				errRecorder.tlsSystemRootsError = &systemRootsError
			default:
				log.Println("unknown tls error", reflect.TypeOf(tlsErr), tlsErr)
			}
		}
		return conn, tlsErr
	}
	dial := func(network, address string) (conn net.Conn, err error) {
		conn, err = dialer.Dial(network, address)
		if err != nil {
			switch reflect.TypeOf(err) {
			case typeOpErr:
				opError := reflect.ValueOf(err).Elem().Interface().(net.OpError)
				switch reflect.TypeOf(opError.Err) {
				case typeDNSConfigErr:
					log.Println("dns config error")
					errRecorder.dnsConfigError = opError.Err.(net.Error)
				case typeDNSErr:
					log.Println("dns error")
					errRecorder.dnsError = opError.Err.(net.Error)
				default:
					errRecorder.unknownErr = opError.Err
				}
			default:
				log.Println("again some general bullshit", err)
				errRecorder.err = err.(net.Error)
			}
		}
		return
	}
	client = &http.Client{
		Transport: &http.Transport{
			Proxy:               http.ProxyFromEnvironment,
			Dial:                dial,
			DialTLS:             dialTLS,
			TLSHandshakeTimeout: 10 * time.Second,
			TLSClientConfig:     tlsConfig,
		},
	}
	return
}

// actual watch
func watch(service *config.Service) (r *Result) {
	r = NewResult(service.ID)
	// parsing, the endpoint
	request, err := http.NewRequest("GET", service.Endpoint, nil)
	if err != nil {
		r.addError(err, ErrorInvalidEndpoint, "")
		return r
	}
	// my personal dns error check
	if len(request.Host) > 0 {
		host := request.Host
		parts := strings.Split(host, ":")
		if len(parts) > 1 {
			host, _, err = net.SplitHostPort(request.Host)
			if err != nil {
				r.addError(err, ErrorInvalidEndpoint, "")
				return
			}
		}
		_, lookupErr := net.LookupIP(host)
		if lookupErr != nil {
			r.addError(lookupErr, ErrorTypeDNS, "")
			return
		}
	}

	// i am explicitly not calling http.Get, because it does 30x handling and i do not want that
	client, errRecorder := getClientAndDialErrRecorder()
	response, err := client.Do(request)
	r.Errors = append(r.Errors, errRecorder.errors...)
	r.Warnings = append(r.Warnings, errRecorder.warnings...)
	//r.RunTime = time.Since(r.Timestamp)

	if response != nil && response.Body != nil {
		// always close the body
		response.Body.Close()
	}

	if err != nil {
		// sth. went wrong
		r.addError(err, ErrorTypeClientError, "")
		var netErr net.Error
		switch true {
		case errRecorder.tlsHostnameError != nil:
			r.addError(errRecorder.tlsHostnameError, ErrorTypeTLSHostNameError, "")
		case errRecorder.tlsSystemRootsError != nil:
			r.addError(errRecorder.tlsSystemRootsError, ErrorTypeTLSSystemRootsError, "")
		case errRecorder.tlsUnknownAuthorityError != nil:
			r.addError(errRecorder.tlsUnknownAuthorityError, ErrorTypeTLSUnknownAuthority, "")
		case errRecorder.tlsCertificateInvalidError != nil:
			r.addError(errRecorder.tlsCertificateInvalidError, ErrorTypeTLSCertificateInvalid, "")
		case errRecorder.unknownErr != nil:
			r.addError(errRecorder.unknownErr, ErrorTypeUnknownError, "")
		case errRecorder.dnsConfigError != nil:
			netErr = errRecorder.dnsConfigError
			r.addError(errRecorder.dnsConfigError, ErrorTypeDNSConfig, "")
		case errRecorder.dnsError != nil:
			netErr = errRecorder.dnsError
			r.addError(errRecorder.dnsError, ErrorTypeDNS, "")
		case errRecorder.err != nil:
			netErr = errRecorder.err
			r.addError(errRecorder.err, ErrorTypeUnknownError, "")
		}
		//log.Println("service", service.ID, err, "errRecorder dns:", errRecorder.dnsError, ", dnsConfig", errRecorder.dnsConfigError, ", err:", errRecorder.err, ", tls cert invalid:", errRecorder.tlsCertificateInvalidError)
		if netErr != nil {
			r.Timeout = netErr.Timeout()
			//r.ErrorIsTemporary = netErr.Temporary()
		}
		return
	}

	cookieJar, _ := cookiejar.New(nil)
	client.Jar = cookieJar
	errSession := runSession(service, r, client)
	if errSession != nil {
		log.Println("session error", errSession)
		r.addError(errSession, ErrorTypeSessionFail, "")
	}

	// r.addError(errors.New(fmt.Sprint("response time too slow:", runTimeMilliseconds, ", should not be more than:", maxRuntime)), ErrorTypeServerTooSlow)
	//r.StatusCode = response.StatusCode
	if response.StatusCode != http.StatusOK {
		r.addError(errors.New(fmt.Sprint("unexpected status code: ", response.StatusCode)), ErrorTypeWrongHTTPStatusCode, "")
	}
	return
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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 go_kafka_client

import (
	"time"
	"math"
	"fmt"
	"sync"
)

type WorkerManager struct {
	Config *ConsumerConfig
	Strategy         WorkerStrategy
	FailureHook FailedCallback
	FailedAttemptHook FailedAttemptCallback
	Workers          []*Worker
	AvailableWorkers chan *Worker
	CurrentBatch map[TaskId]*Task
	InputChannel     chan [] *Message
	OutputChannel    chan map[TopicAndPartition]int64
	LargestOffsets map[TopicAndPartition]int64
	FailCounter *FailureCounter
}

func (wm *WorkerManager) String() string {
	return fmt.Sprintf("%s-workerManager", wm.Config.ConsumerId)
}

func (wm *WorkerManager) Start() {
	batchProcessed := make(chan bool)
	go wm.processBatch(batchProcessed)
	for {
		batch := <-wm.InputChannel
		for _, message := range batch {
			wm.CurrentBatch[TaskId{ TopicAndPartition{ message.Topic, message.Partition }, message.Offset }] = &Task{
				Msg: message,
			}
		}

		for _, task := range wm.CurrentBatch {
			worker := <-wm.AvailableWorkers
			worker.Start(task, wm.Strategy)
		}

		<-batchProcessed

		wm.OutputChannel <- wm.LargestOffsets
		wm.LargestOffsets = make(map[TopicAndPartition]int64)
	}
}

func (wm *WorkerManager) Stop() chan bool {
	finished := make(chan bool)
	go func() {
		for _, worker := range wm.Workers {
			worker.Close()
		}
		finished <- true
	}()

	return finished
}

func (wm *WorkerManager) IsBatchProcessed() bool {
	return len(wm.CurrentBatch) == 0
}

func (wm *WorkerManager) processBatch(batchProcessed chan bool) {
	outputChannels := make([]*chan WorkerResult, wm.Config.NumWorkers)
	for i, worker := range wm.Workers {
		outputChannels[i] = &worker.OutputChannel
	}

	resultsChannel := make(chan WorkerResult)
	for {
		stopRedirecting := RedirectChannelsTo(outputChannels, resultsChannel)
		result := <- resultsChannel
		task := wm.CurrentBatch[result.Id()]
		if result.Success() {
			wm.taskIsDone(result)
		} else {
			if _, ok := result.(*TimedOutResult); ok {
				task.Callee.OutputChannel = make(chan WorkerResult)
			}

			Warnf(wm, "Worker task %s has failed", result.Id())
			task.Retries++
			if task.Retries >= wm.Config.MaxWorkerRetries {
				Errorf(wm, "Worker task %s has failed after %d retries", result.Id(), wm.Config.MaxWorkerRetries)

				wm.taskIsDone(result)
				if wm.FailCounter.Failed() {
					wm.FailureHook(wm)
				} else {
					wm.FailedAttemptHook(task, result)
				}
			} else {
				Warnf(wm, "Retrying worker task %s %dth time", result.Id(), task.Retries)
				time.Sleep(wm.Config.WorkerBackoff)
				go task.Callee.Start(task, wm.Strategy)
			}
		}

		if wm.IsBatchProcessed() {
			stopRedirecting <- true
			batchProcessed <- true
		}
	}
}

func (wm *WorkerManager) taskIsDone(result WorkerResult) {
	key := result.Id().TopicPartition
	wm.LargestOffsets[key] = int64(math.Max(float64(wm.LargestOffsets[key]), float64(result.Id().Offset)))
	wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
	delete(wm.CurrentBatch, result.Id())
}

func NewWorkerManager(config *ConsumerConfig, batchOutputChannel chan map[TopicAndPartition]int64) *WorkerManager {
	workers := make([]*Worker, config.NumWorkers)
	availableWorkers := make(chan *Worker, config.NumWorkers)
	for i := 0; i < config.NumWorkers; i++ {
		workers[i] = &Worker{
			OutputChannel: make(chan WorkerResult),
			CloseTimeout: config.WorkerCloseTimeout,
			TaskTimeout: config.WorkerTaskTimeout,
		}
		availableWorkers <- workers[i]
	}

	return &WorkerManager {
		Config: config,
		Strategy: config.Strategy,
		FailureHook: config.WorkerFailureCallback,
		FailedAttemptHook: config.WorkerFailedAttemptCallback,
		AvailableWorkers: availableWorkers,
		Workers: workers,
		InputChannel: make(chan [] *Message),
		OutputChannel: batchOutputChannel,
		CurrentBatch: make(map[TaskId]*Task),
		LargestOffsets: make(map[TopicAndPartition]int64),
		FailCounter: NewFailureCounter(config.WorkerRetryThreshold, config.WorkerConsideredFailedTimeWindow),
	}
}

type Worker struct {
	OutputChannel chan WorkerResult
	CloseTimeout time.Duration
	TaskTimeout time.Duration
	Closed bool
}

func (w *Worker) Start(task *Task, strategy WorkerStrategy) {
	task.Callee = w
	go func() {
		resultChannel := make(chan WorkerResult)
		go func() {resultChannel <- strategy(w, task.Msg, task.Id())}()
		select {
		case result := <-resultChannel: {
			w.OutputChannel <- result
		}
		case <-time.After(w.TaskTimeout): {
			w.OutputChannel <- &TimedOutResult{task.Id()}
		}
		}
	}()
}

func (w *Worker) Close() WorkerResult {
	if !w.Closed {
		w.Closed = true
		select {
		case result := <- w.OutputChannel: {
			return result
		}
		case <- time.After(w.CloseTimeout): {
			Warnf(w, "Worker failed to close within %f seconds", w.CloseTimeout.Seconds())
			return nil
		}
		}
	}

	return nil
}

type WorkerStrategy func(*Worker, *Message, TaskId) WorkerResult

type FailedCallback func(*WorkerManager)

type FailedAttemptCallback func(*Task, WorkerResult)

type FailureCounter struct {
	count int32
	failed bool
	countLock sync.Mutex
	failedThreshold int32
}

func NewFailureCounter(FailedThreshold int32, WorkerConsideredFailedTimeWindow time.Duration) *FailureCounter {
	counter := &FailureCounter {
		failedThreshold : FailedThreshold,
	}
	go func() {
		for {
			select {
			case <- time.After(WorkerConsideredFailedTimeWindow): {
				if counter.count >= FailedThreshold {
					counter.failed = true
					return
				} else {
					InLock(&counter.countLock, func() {
						counter.count = 0
					})
				}
			}
			}
		}
	}()
	return counter
}

func (f *FailureCounter) Failed() bool {
	InLock(&f.countLock, func() { f.count++ })
	return f.count >= f.failedThreshold || f.failed
}

type Task struct {
	Msg *Message
	Retries int
	Callee *Worker
}

func (t *Task) Id() TaskId {
	return TaskId{TopicAndPartition{t.Msg.Topic, t.Msg.Partition}, t.Msg.Offset}
}

type WorkerResult interface {
	Id() TaskId
	Success() bool
}

type SuccessfulResult struct {
	id TaskId
}

func NewSuccessfulResult(id TaskId) *SuccessfulResult {
	return &SuccessfulResult{id}
}

func (sr *SuccessfulResult) String() string {
	return fmt.Sprintf("{Success: %s}", sr.Id())
}

func (wr *SuccessfulResult) Id() TaskId {
	return wr.id
}

func (wr *SuccessfulResult) Success() bool {
	return true
}

type ProcessingFailedResult struct {
	id TaskId
}

func (sr *ProcessingFailedResult) String() string {
	return fmt.Sprintf("{Failed: %s}", sr.Id())
}

func (wr *ProcessingFailedResult) Id() TaskId {
	return wr.id
}

func (wr *ProcessingFailedResult) Success() bool {
	return false
}

type TimedOutResult struct {
	id TaskId
}

func (sr *TimedOutResult) String() string {
	return fmt.Sprintf("{Timed out: %s}", sr.Id())
}

func (wr *TimedOutResult) Id() TaskId {
	return wr.id
}

func (wr *TimedOutResult) Success() bool {
	return false
}

type TaskId struct {
	TopicPartition TopicAndPartition
	Offset         int64
}

func (tid TaskId) String() string {
	return fmt.Sprintf("%s, Offset: %d", &tid.TopicPartition, tid.Offset)
}
<commit_msg>Flexible failure management<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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 go_kafka_client

import (
	"time"
	"math"
	"fmt"
	"sync"
)

type WorkerManager struct {
	Config *ConsumerConfig
	Strategy         WorkerStrategy
	FailureHook FailedCallback
	FailedAttemptHook FailedAttemptCallback
	Workers          []*Worker
	AvailableWorkers chan *Worker
	CurrentBatch map[TaskId]*Task
	InputChannel     chan [] *Message
	OutputChannel    chan map[TopicAndPartition]int64
	LargestOffsets map[TopicAndPartition]int64
	FailCounter *FailureCounter
}

func (wm *WorkerManager) String() string {
	return fmt.Sprintf("%s-workerManager", wm.Config.ConsumerId)
}

func (wm *WorkerManager) Start() {
	batchProcessed := make(chan bool)
	go wm.processBatch(batchProcessed)
	for {
		batch := <-wm.InputChannel
		for _, message := range batch {
			wm.CurrentBatch[TaskId{ TopicAndPartition{ message.Topic, message.Partition }, message.Offset }] = &Task{
				Msg: message,
			}
		}

		for _, task := range wm.CurrentBatch {
			worker := <-wm.AvailableWorkers
			worker.Start(task, wm.Strategy)
		}

		<-batchProcessed

		wm.OutputChannel <- wm.LargestOffsets
		wm.LargestOffsets = make(map[TopicAndPartition]int64)
	}
}

func (wm *WorkerManager) Stop() chan bool {
	finished := make(chan bool)
	go func() {
		for _, worker := range wm.Workers {
			worker.Close()
		}
		wm.CurrentBatch = make(map[TaskId]*Task)
		finished <- true
	}()

	return finished
}

func (wm *WorkerManager) IsBatchProcessed() bool {
	return len(wm.CurrentBatch) == 0
}

func (wm *WorkerManager) processBatch(batchProcessed chan bool) {
	outputChannels := make([]*chan WorkerResult, wm.Config.NumWorkers)
	for i, worker := range wm.Workers {
		outputChannels[i] = &worker.OutputChannel
	}

	resultsChannel := make(chan WorkerResult)
	for {
		stopRedirecting := RedirectChannelsTo(outputChannels, resultsChannel)
		result := <- resultsChannel
		task := wm.CurrentBatch[result.Id()]
		if result.Success() {
			wm.taskIsDone(result)
		} else {
			if _, ok := result.(*TimedOutResult); ok {
				task.Callee.OutputChannel = make(chan WorkerResult)
			}

			Warnf(wm, "Worker task %s has failed", result.Id())
			task.Retries++
			if task.Retries >= wm.Config.MaxWorkerRetries {
				Errorf(wm, "Worker task %s has failed after %d retries", result.Id(), wm.Config.MaxWorkerRetries)

				var decision FailedDecision
				if wm.FailCounter.Failed() {
					decision = wm.FailureHook(wm)
				} else {
					decision = wm.FailedAttemptHook(task, result)
				}
				switch decision {
				case CommitOffsetAndContinue: {
					wm.taskIsDone(result)
				}
				case DoNotCommitOffsetAndContinue: {
					wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
					delete(wm.CurrentBatch, result.Id())
				}
				case CommitOffsetAndStop: {
					wm.taskIsDone(result)
					wm.stopBatch()
				}
				case DoNotCommitOffsetAndStop: {
					wm.stopBatch()
				}
				}
			} else {
				Warnf(wm, "Retrying worker task %s %dth time", result.Id(), task.Retries)
				time.Sleep(wm.Config.WorkerBackoff)
				go task.Callee.Start(task, wm.Strategy)
			}
		}

		if wm.IsBatchProcessed() {
			stopRedirecting <- true
			batchProcessed <- true
		}
	}
}

func (wm *WorkerManager) stopBatch() {
	wm.CurrentBatch = make(map[TaskId]*Task)
	for _, worker := range wm.Workers {
		worker.OutputChannel = make(chan WorkerResult)
	}
}

func (wm *WorkerManager) taskIsDone(result WorkerResult) {
	key := result.Id().TopicPartition
	wm.LargestOffsets[key] = int64(math.Max(float64(wm.LargestOffsets[key]), float64(result.Id().Offset)))
	wm.AvailableWorkers <- wm.CurrentBatch[result.Id()].Callee
	delete(wm.CurrentBatch, result.Id())
}

func NewWorkerManager(config *ConsumerConfig, batchOutputChannel chan map[TopicAndPartition]int64) *WorkerManager {
	workers := make([]*Worker, config.NumWorkers)
	availableWorkers := make(chan *Worker, config.NumWorkers)
	for i := 0; i < config.NumWorkers; i++ {
		workers[i] = &Worker{
			OutputChannel: make(chan WorkerResult),
			CloseTimeout: config.WorkerCloseTimeout,
			TaskTimeout: config.WorkerTaskTimeout,
		}
		availableWorkers <- workers[i]
	}

	return &WorkerManager {
		Config: config,
		Strategy: config.Strategy,
		FailureHook: config.WorkerFailureCallback,
		FailedAttemptHook: config.WorkerFailedAttemptCallback,
		AvailableWorkers: availableWorkers,
		Workers: workers,
		InputChannel: make(chan [] *Message),
		OutputChannel: batchOutputChannel,
		CurrentBatch: make(map[TaskId]*Task),
		LargestOffsets: make(map[TopicAndPartition]int64),
		FailCounter: NewFailureCounter(config.WorkerRetryThreshold, config.WorkerConsideredFailedTimeWindow),
	}
}

type Worker struct {
	OutputChannel chan WorkerResult
	CloseTimeout time.Duration
	TaskTimeout time.Duration
	Closed bool
}

func (w *Worker) Start(task *Task, strategy WorkerStrategy) {
	task.Callee = w
	go func() {
		resultChannel := make(chan WorkerResult)
		go func() {resultChannel <- strategy(w, task.Msg, task.Id())}()
		select {
		case result := <-resultChannel: {
			w.OutputChannel <- result
		}
		case <-time.After(w.TaskTimeout): {
			w.OutputChannel <- &TimedOutResult{task.Id()}
		}
		}
	}()
}

func (w *Worker) Close() WorkerResult {
	if !w.Closed {
		w.Closed = true
		select {
		case result := <- w.OutputChannel: {
			return result
		}
		case <- time.After(w.CloseTimeout): {
			Warnf(w, "Worker failed to close within %f seconds", w.CloseTimeout.Seconds())
			return nil
		}
		}
	}

	return nil
}

type WorkerStrategy func(*Worker, *Message, TaskId) WorkerResult

type FailedCallback func(*WorkerManager) FailedDecision

type FailedAttemptCallback func(*Task, WorkerResult) FailedDecision

type FailureCounter struct {
	count int32
	failed bool
	countLock sync.Mutex
	failedThreshold int32
}

func NewFailureCounter(FailedThreshold int32, WorkerConsideredFailedTimeWindow time.Duration) *FailureCounter {
	counter := &FailureCounter {
		failedThreshold : FailedThreshold,
	}
	go func() {
		for {
			select {
			case <- time.After(WorkerConsideredFailedTimeWindow): {
				if counter.count >= FailedThreshold {
					counter.failed = true
					return
				} else {
					InLock(&counter.countLock, func() {
						counter.count = 0
					})
				}
			}
			}
		}
	}()
	return counter
}

func (f *FailureCounter) Failed() bool {
	InLock(&f.countLock, func() { f.count++ })
	return f.count >= f.failedThreshold || f.failed
}

type Task struct {
	Msg *Message
	Retries int
	Callee *Worker
}

func (t *Task) Id() TaskId {
	return TaskId{TopicAndPartition{t.Msg.Topic, t.Msg.Partition}, t.Msg.Offset}
}

type WorkerResult interface {
	Id() TaskId
	Success() bool
}

type SuccessfulResult struct {
	id TaskId
}

func NewSuccessfulResult(id TaskId) *SuccessfulResult {
	return &SuccessfulResult{id}
}

func (sr *SuccessfulResult) String() string {
	return fmt.Sprintf("{Success: %s}", sr.Id())
}

func (wr *SuccessfulResult) Id() TaskId {
	return wr.id
}

func (wr *SuccessfulResult) Success() bool {
	return true
}

type ProcessingFailedResult struct {
	id TaskId
}

func (sr *ProcessingFailedResult) String() string {
	return fmt.Sprintf("{Failed: %s}", sr.Id())
}

func (wr *ProcessingFailedResult) Id() TaskId {
	return wr.id
}

func (wr *ProcessingFailedResult) Success() bool {
	return false
}

type TimedOutResult struct {
	id TaskId
}

func (sr *TimedOutResult) String() string {
	return fmt.Sprintf("{Timed out: %s}", sr.Id())
}

func (wr *TimedOutResult) Id() TaskId {
	return wr.id
}

func (wr *TimedOutResult) Success() bool {
	return false
}

type TaskId struct {
	TopicPartition TopicAndPartition
	Offset         int64
}

func (tid TaskId) String() string {
	return fmt.Sprintf("%s, Offset: %d", &tid.TopicPartition, tid.Offset)
}

type FailedDecision int32

const (
	CommitOffsetAndContinue FailedDecision = iota
	DoNotCommitOffsetAndContinue
	CommitOffsetAndStop
	DoNotCommitOffsetAndStop
)
<|endoftext|>


<commit_before>/*
Copyright 2020 The cert-manager 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 certificates

import (
	"crypto"
	"reflect"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

func mustGenerateRSA(t *testing.T, keySize int) crypto.PrivateKey {
	pk, err := pki.GenerateRSAPrivateKey(keySize)
	if err != nil {
		t.Fatal(err)
	}
	return pk
}

func mustGenerateECDSA(t *testing.T, keySize int) crypto.PrivateKey {
	pk, err := pki.GenerateECPrivateKey(keySize)
	if err != nil {
		t.Fatal(err)
	}
	return pk
}

func TestPrivateKeyMatchesSpec(t *testing.T) {
	tests := map[string]struct {
		key          crypto.PrivateKey
		expectedAlgo cmapi.PrivateKeyAlgorithm
		expectedSize int
		violations   []string
		err          string
	}{
		"should match if keySize and algorithm are correct (RSA)": {
			key:          mustGenerateRSA(t, 2048),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 2048,
		},
		"should not match if RSA keySize is incorrect": {
			key:          mustGenerateRSA(t, 2048),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 4096,
			violations:   []string{"spec.keySize"},
		},
		"should match if keySize and algorithm are correct (ECDSA)": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.ECDSAKeyAlgorithm,
			expectedSize: 256,
		},
		"should not match if ECDSA keySize is incorrect": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.ECDSAKeyAlgorithm,
			expectedSize: pki.ECCurve521,
			violations:   []string{"spec.keySize"},
		},
		"should not match if keyAlgorithm is incorrect": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 2048,
			violations:   []string{"spec.keyAlgorithm"},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			violations, err := PrivateKeyMatchesSpec(
				test.key,
				cmapi.CertificateSpec{
					PrivateKey: &cmapi.CertificatePrivateKey{
						Algorithm: test.expectedAlgo,
						Size:      test.expectedSize,
					},
				},
			)
			switch {
			case err != nil:
				if test.err != err.Error() {
					t.Errorf("error text did not match, got=%s, exp=%s", err.Error(), test.err)
				}
			default:
				if test.err != "" {
					t.Errorf("got no error but expected: %s", test.err)
				}
			}
			if !reflect.DeepEqual(violations, test.violations) {
				t.Errorf("violations did not match, got=%s, exp=%s", violations, test.violations)
			}
		})
	}
}

func TestSecretDataAltNamesMatchSpec(t *testing.T) {
	tests := map[string]struct {
		data       []byte
		spec       cmapi.CertificateSpec
		err        string
		violations []string
	}{
		"should match if common name and dns names exactly equal": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			}),
		},
		"should match if commonName is missing but is present in dnsNames": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"cn", "at", "least", "one"},
			}),
		},
		"should match if commonName is missing but is present in dnsNames (not first)": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one", "cn"},
			}),
		},
		"should match if commonName is one of the requested requested dnsNames": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"least", "one"},
			}),
		},
		"should not match if commonName is not present on certificate": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			}),
			violations: []string{"spec.commonName"},
		},
		"should report violation for both commonName and dnsNames if both are missing": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			}),
			violations: []string{"spec.commonName", "spec.dnsNames"},
		},
		"should report violation for both commonName and dnsNames if not requested": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			}),
			violations: []string{"spec.commonName", "spec.dnsNames"},
		},
		"should not match if certificate has more dnsNames than spec": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			}),
			violations: []string{"spec.dnsNames"},
		},
		"should match if commonName is a duplicated dnsName (but not requested)": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"at", "least", "one"},
			}),
		},
		"should match if commonName is a duplicated dnsName": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"at", "least", "one", "cn"},
			}),
		},
		"should match if ipAddresses are equal": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			}),
		},
		"should not match if ipAddresses are not equal": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.2.1"},
			}),
			violations: []string{"spec.ipAddresses"},
		},
		"should not match if ipAddresses has been made the commonName": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName:  "127.0.0.1",
				IPAddresses: []string{"127.0.0.1"},
			}),
			violations: []string{"spec.commonName"},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			violations, err := SecretDataAltNamesMatchSpec(&corev1.Secret{Data: map[string][]byte{corev1.TLSCertKey: test.data}}, test.spec)
			switch {
			case err != nil:
				if test.err != err.Error() {
					t.Errorf("error text did not match, got=%s, exp=%s", err.Error(), test.err)
				}
			default:
				if test.err != "" {
					t.Errorf("got no error but expected: %s", test.err)
				}
			}
			if !reflect.DeepEqual(violations, test.violations) {
				t.Errorf("violations did not match, got=%s, exp=%s", violations, test.violations)
			}
		})
	}
}

func selfSignCertificate(t *testing.T, spec cmapi.CertificateSpec) []byte {
	pk, err := pki.GenerateRSAPrivateKey(2048)
	if err != nil {
		t.Fatal(err)
	}

	template, err := pki.GenerateTemplate(&cmapi.Certificate{Spec: spec})
	if err != nil {
		t.Fatal(err)
	}

	pemData, _, err := pki.SignCertificate(template, template, pk.Public(), pk)
	if err != nil {
		t.Fatal(err)
	}

	return pemData
}

func TestRenewalTimeWrapper(t *testing.T) {
	type scenario struct {
		notBefore           time.Time
		notAfter            time.Time
		renewBeforeHint     *metav1.Duration
		defaultRenewBefore  time.Duration
		expectedRenewalTime *metav1.Time
	}
	now := time.Now()
	tests := map[string]scenario{
		"no renewBeforeHint, defaultRenewBefore < (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 23)},
		},
		"renewBeforeHint < (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 2},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 22)},
		},
		"no renewBeforeHint, defaultRenewBefore > (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour * 24 * 7,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"renewBeforeHint > (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 24 * 7},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"no renewBeforeHint, defaultRenewBefore == cert duration": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour * 24,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"renewBeforeHint == cert duration": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 24},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},

		// The following two test cases would catch the bug reported in
		// https://github.com/jetstack/cert-manager/issues/3897
		"cert duration very slightly less than defaultRenewBefore": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour*24 + time.Minute,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"cert duration very slightly less than renewBeforeHint": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour*24 + time.Minute,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
	}
	for n, s := range tests {
		t.Run(n, func(t *testing.T) {
			f := RenewalTimeWrapper(s.defaultRenewBefore)
			renewalTime := f(s.notBefore, s.notAfter, s.renewBeforeHint)
			assert.Equal(t, s.expectedRenewalTime, renewalTime)

		})
	}
}
<commit_msg>Feedback from code review<commit_after>/*
Copyright 2020 The cert-manager 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 certificates

import (
	"crypto"
	"fmt"
	"reflect"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	"github.com/jetstack/cert-manager/pkg/util/pki"
)

func mustGenerateRSA(t *testing.T, keySize int) crypto.PrivateKey {
	pk, err := pki.GenerateRSAPrivateKey(keySize)
	if err != nil {
		t.Fatal(err)
	}
	return pk
}

func mustGenerateECDSA(t *testing.T, keySize int) crypto.PrivateKey {
	pk, err := pki.GenerateECPrivateKey(keySize)
	if err != nil {
		t.Fatal(err)
	}
	return pk
}

func TestPrivateKeyMatchesSpec(t *testing.T) {
	tests := map[string]struct {
		key          crypto.PrivateKey
		expectedAlgo cmapi.PrivateKeyAlgorithm
		expectedSize int
		violations   []string
		err          string
	}{
		"should match if keySize and algorithm are correct (RSA)": {
			key:          mustGenerateRSA(t, 2048),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 2048,
		},
		"should not match if RSA keySize is incorrect": {
			key:          mustGenerateRSA(t, 2048),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 4096,
			violations:   []string{"spec.keySize"},
		},
		"should match if keySize and algorithm are correct (ECDSA)": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.ECDSAKeyAlgorithm,
			expectedSize: 256,
		},
		"should not match if ECDSA keySize is incorrect": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.ECDSAKeyAlgorithm,
			expectedSize: pki.ECCurve521,
			violations:   []string{"spec.keySize"},
		},
		"should not match if keyAlgorithm is incorrect": {
			key:          mustGenerateECDSA(t, pki.ECCurve256),
			expectedAlgo: cmapi.RSAKeyAlgorithm,
			expectedSize: 2048,
			violations:   []string{"spec.keyAlgorithm"},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			violations, err := PrivateKeyMatchesSpec(
				test.key,
				cmapi.CertificateSpec{
					PrivateKey: &cmapi.CertificatePrivateKey{
						Algorithm: test.expectedAlgo,
						Size:      test.expectedSize,
					},
				},
			)
			switch {
			case err != nil:
				if test.err != err.Error() {
					t.Errorf("error text did not match, got=%s, exp=%s", err.Error(), test.err)
				}
			default:
				if test.err != "" {
					t.Errorf("got no error but expected: %s", test.err)
				}
			}
			if !reflect.DeepEqual(violations, test.violations) {
				t.Errorf("violations did not match, got=%s, exp=%s", violations, test.violations)
			}
		})
	}
}

func TestSecretDataAltNamesMatchSpec(t *testing.T) {
	tests := map[string]struct {
		data       []byte
		spec       cmapi.CertificateSpec
		err        string
		violations []string
	}{
		"should match if common name and dns names exactly equal": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			}),
		},
		"should match if commonName is missing but is present in dnsNames": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"cn", "at", "least", "one"},
			}),
		},
		"should match if commonName is missing but is present in dnsNames (not first)": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one", "cn"},
			}),
		},
		"should match if commonName is one of the requested requested dnsNames": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"least", "one"},
			}),
		},
		"should not match if commonName is not present on certificate": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			}),
			violations: []string{"spec.commonName"},
		},
		"should report violation for both commonName and dnsNames if both are missing": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			}),
			violations: []string{"spec.commonName", "spec.dnsNames"},
		},
		"should report violation for both commonName and dnsNames if not requested": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			}),
			violations: []string{"spec.commonName", "spec.dnsNames"},
		},
		"should not match if certificate has more dnsNames than spec": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one", "other"},
			}),
			violations: []string{"spec.dnsNames"},
		},
		"should match if commonName is a duplicated dnsName (but not requested)": {
			spec: cmapi.CertificateSpec{
				DNSNames: []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"at", "least", "one"},
			}),
		},
		"should match if commonName is a duplicated dnsName": {
			spec: cmapi.CertificateSpec{
				CommonName: "cn",
				DNSNames:   []string{"at", "least", "one"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName: "at",
				DNSNames:   []string{"at", "least", "one", "cn"},
			}),
		},
		"should match if ipAddresses are equal": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			}),
		},
		"should not match if ipAddresses are not equal": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.2.1"},
			}),
			violations: []string{"spec.ipAddresses"},
		},
		"should not match if ipAddresses has been made the commonName": {
			spec: cmapi.CertificateSpec{
				IPAddresses: []string{"127.0.0.1"},
			},
			data: selfSignCertificate(t, cmapi.CertificateSpec{
				CommonName:  "127.0.0.1",
				IPAddresses: []string{"127.0.0.1"},
			}),
			violations: []string{"spec.commonName"},
		},
	}
	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			violations, err := SecretDataAltNamesMatchSpec(&corev1.Secret{Data: map[string][]byte{corev1.TLSCertKey: test.data}}, test.spec)
			switch {
			case err != nil:
				if test.err != err.Error() {
					t.Errorf("error text did not match, got=%s, exp=%s", err.Error(), test.err)
				}
			default:
				if test.err != "" {
					t.Errorf("got no error but expected: %s", test.err)
				}
			}
			if !reflect.DeepEqual(violations, test.violations) {
				t.Errorf("violations did not match, got=%s, exp=%s", violations, test.violations)
			}
		})
	}
}

func selfSignCertificate(t *testing.T, spec cmapi.CertificateSpec) []byte {
	pk, err := pki.GenerateRSAPrivateKey(2048)
	if err != nil {
		t.Fatal(err)
	}

	template, err := pki.GenerateTemplate(&cmapi.Certificate{Spec: spec})
	if err != nil {
		t.Fatal(err)
	}

	pemData, _, err := pki.SignCertificate(template, template, pk.Public(), pk)
	if err != nil {
		t.Fatal(err)
	}

	return pemData
}

func TestRenewalTimeWrapper(t *testing.T) {
	type scenario struct {
		notBefore           time.Time
		notAfter            time.Time
		renewBeforeHint     *metav1.Duration
		defaultRenewBefore  time.Duration
		expectedRenewalTime *metav1.Time
	}
	now := time.Now()
	tests := map[string]scenario{
		"no renewBeforeHint, defaultRenewBefore < (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 23)},
		},
		"renewBeforeHint < (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 2},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 22)},
		},
		"no renewBeforeHint, defaultRenewBefore > (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour * 24 * 7,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"renewBeforeHint > (cert duration / 3)": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 24 * 7},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"no renewBeforeHint, defaultRenewBefore == cert duration": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour * 24,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},
		"renewBeforeHint == cert duration": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour * 24),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 24},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour * 16)},
		},

		// The following two test cases would catch the bug reported in
		// https://github.com/jetstack/cert-manager/issues/3897
		"cert duration very slightly more than defaultRenewBefore": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour*24 + time.Minute*3),
			renewBeforeHint:     nil,
			defaultRenewBefore:  time.Hour * 24,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour*16 + time.Minute*2)},
		},
		"cert duration very slightly more than renewBeforeHint": {
			notBefore:           now,
			notAfter:            now.Add(time.Hour*24 + time.Minute*3),
			renewBeforeHint:     &metav1.Duration{Duration: time.Hour * 24},
			defaultRenewBefore:  time.Hour,
			expectedRenewalTime: &metav1.Time{Time: now.Add(time.Hour*16 + time.Minute*2)},
		},
	}
	for n, s := range tests {
		t.Run(n, func(t *testing.T) {
			f := RenewalTimeWrapper(s.defaultRenewBefore)
			renewalTime := f(s.notBefore, s.notAfter, s.renewBeforeHint)
			assert.Equal(t, s.expectedRenewalTime, renewalTime, fmt.Sprintf("Expected renewal time: %v got: %v", s.expectedRenewalTime, renewalTime))

		})
	}
}
<|endoftext|>


<commit_before>package chevalier

import (
	"testing"
)

func TestEmptyDataSource(t *testing.T) {
	ds := new(DataSource)
	addr := uint64(42);
	ds.Address = &addr;
	if !ds.Empty() {
		t.Errorf("I think an empty datasource is non-empty.")
	}
	tg_f := "foo"
	tg_v := "bar"
	tg := new(DataSource_Tag)
	tg.Field = &tg_f
	tg.Value = &tg_v
	tgs := make([]*DataSource_Tag, 1)
	tgs[0] = tg
	ds.Source = tgs
	if ds.Empty() {
		t.Errorf("I think a non-empty datasource is empty.")
	}
}

<commit_msg>This is a better test of the empty check<commit_after>package chevalier

import (
	"testing"
)

func TestEmptyDataSource(t *testing.T) {
	ds := new(DataSource)
	addr := uint64(42);
	ds.Address = &addr;
	if !ds.Empty() {
		t.Errorf("I think an empty datasource is non-empty.")
	}
	tg_f := "foo"
	tg_v := "bar"
	tg := new(DataSource_Tag)
	tg.Field = &tg_f
	tg.Value = &tg_v
	tgs := make([]*DataSource_Tag, 1)
	tgs[0] = tg
	ds.Source = tgs
	if ds.Empty() {
		t.Errorf("I think a non-empty datasource is empty.")
	}
	ds.Source = ds.Source[:0]
	if !ds.Empty() {
		t.Errorf("I think a cleared datasource is non-empty.")
	}
}

<|endoftext|>


<commit_before>package flac

import (
	"fmt"
	"io"

	"github.com/faiface/beep"
	"github.com/mewkiz/flac"
	"github.com/pkg/errors"
)

// Decode takes a ReadCloser containing audio data in FLAC format and returns a StreamSeekCloser,
// which streams that audio. The Seek method will panic if rc is not io.Seeker.
//
// Do not close the supplied ReadSeekCloser, instead, use the Close method of the returned
// StreamSeekCloser when you want to release the resources.
func Decode(rc io.ReadCloser) (s beep.StreamSeekCloser, format beep.Format, err error) {
	d := decoder{rc: rc}
	defer func() { // hacky way to always close rc if an error occured
		if err != nil {
			d.rc.Close()
		}
	}()
	rsc, ok := rc.(io.ReadSeeker)
	if !ok {
		panic(fmt.Errorf("%T does not implement io.Seeker", rc))
	}
	d.stream, err = flac.New(rsc)
	if err != nil {
		return nil, beep.Format{}, errors.Wrap(err, "flac")
	}
	format = beep.Format{
		SampleRate:  beep.SampleRate(d.stream.Info.SampleRate),
		NumChannels: int(d.stream.Info.NChannels),
		Precision:   int(d.stream.Info.BitsPerSample / 8),
	}
	return &d, format, nil
}

type decoder struct {
	rc     io.ReadCloser
	stream *flac.Stream
	buf    [][2]float64
	pos    int
	err    error
}

func (d *decoder) Stream(samples [][2]float64) (n int, ok bool) {
	if d.err != nil {
		return 0, false
	}
	// Copy samples from buffer.
	j := 0
	for i := range samples {
		if j >= len(d.buf) {
			// refill buffer.
			if err := d.refill(); err != nil {
				d.err = err
				return n, n > 0
			}
			j = 0
		}
		samples[i] = d.buf[j]
		j++
		n++
	}
	d.buf = d.buf[j:]
	d.pos += n
	return n, true
}

// refill decodes audio samples to fill the decode buffer.
func (d *decoder) refill() error {
	// Empty buffer.
	d.buf = d.buf[:0]
	// Parse audio frame.
	frame, err := d.stream.ParseNext()
	if err != nil {
		return err
	}
	// Expand buffer size if needed.
	n := len(frame.Subframes[0].Samples)
	if cap(d.buf) < n {
		d.buf = make([][2]float64, n)
	} else {
		d.buf = d.buf[:n]
	}
	// Decode audio samples.
	bps := d.stream.Info.BitsPerSample
	nchannels := d.stream.Info.NChannels
	switch {
	case bps == 8 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
			d.buf[i][1] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
		}
	case bps == 16 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
			d.buf[i][1] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
		}
	case bps == 24 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int32(frame.Subframes[0].Samples[i])) / (1<<23 - 1)
			d.buf[i][1] = float64(int32(frame.Subframes[1].Samples[i])) / (1<<23 - 1)
		}
	case bps == 8 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
			d.buf[i][1] = float64(int8(frame.Subframes[1].Samples[i])) / (1<<7 - 1)
		}
	case bps == 16 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
			d.buf[i][1] = float64(int16(frame.Subframes[1].Samples[i])) / (1<<15 - 1)
		}
	case bps == 24 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int32(frame.Subframes[0].Samples[i])) / (1<<23 - 1)
			d.buf[i][1] = float64(int32(frame.Subframes[1].Samples[i])) / (1<<23 - 1)
		}
	default:
		panic(fmt.Errorf("support for %d bits-per-sample and %d channels combination not yet implemented", bps, nchannels))
	}
	return nil
}

func (d *decoder) Err() error {
	return d.err
}

func (d *decoder) Len() int {
	return int(d.stream.Info.NSamples)
}

func (d *decoder) Position() int {
	return d.pos
}

func (d *decoder) Seek(p int) error {
	panic("not yet implemented")
}

func (d *decoder) Close() error {
	err := d.rc.Close()
	if err != nil {
		return errors.Wrap(err, "flac")
	}
	return nil
}
<commit_msg>update flac.pos on EOF error<commit_after>package flac

import (
	"fmt"
	"io"

	"github.com/faiface/beep"
	"github.com/mewkiz/flac"
	"github.com/pkg/errors"
)

// Decode takes a ReadCloser containing audio data in FLAC format and returns a StreamSeekCloser,
// which streams that audio. The Seek method will panic if rc is not io.Seeker.
//
// Do not close the supplied ReadSeekCloser, instead, use the Close method of the returned
// StreamSeekCloser when you want to release the resources.
func Decode(rc io.ReadCloser) (s beep.StreamSeekCloser, format beep.Format, err error) {
	d := decoder{rc: rc}
	defer func() { // hacky way to always close rc if an error occured
		if err != nil {
			d.rc.Close()
		}
	}()
	rsc, ok := rc.(io.ReadSeeker)
	if !ok {
		panic(fmt.Errorf("%T does not implement io.Seeker", rc))
	}
	d.stream, err = flac.New(rsc)
	if err != nil {
		return nil, beep.Format{}, errors.Wrap(err, "flac")
	}
	format = beep.Format{
		SampleRate:  beep.SampleRate(d.stream.Info.SampleRate),
		NumChannels: int(d.stream.Info.NChannels),
		Precision:   int(d.stream.Info.BitsPerSample / 8),
	}
	return &d, format, nil
}

type decoder struct {
	rc     io.ReadCloser
	stream *flac.Stream
	buf    [][2]float64
	pos    int
	err    error
}

func (d *decoder) Stream(samples [][2]float64) (n int, ok bool) {
	if d.err != nil {
		return 0, false
	}
	// Copy samples from buffer.
	j := 0
	for i := range samples {
		if j >= len(d.buf) {
			// refill buffer.
			if err := d.refill(); err != nil {
				d.err = err
				d.pos += n
				return n, n > 0
			}
			j = 0
		}
		samples[i] = d.buf[j]
		j++
		n++
	}
	d.buf = d.buf[j:]
	d.pos += n
	return n, true
}

// refill decodes audio samples to fill the decode buffer.
func (d *decoder) refill() error {
	// Empty buffer.
	d.buf = d.buf[:0]
	// Parse audio frame.
	frame, err := d.stream.ParseNext()
	if err != nil {
		return err
	}
	// Expand buffer size if needed.
	n := len(frame.Subframes[0].Samples)
	if cap(d.buf) < n {
		d.buf = make([][2]float64, n)
	} else {
		d.buf = d.buf[:n]
	}
	// Decode audio samples.
	bps := d.stream.Info.BitsPerSample
	nchannels := d.stream.Info.NChannels
	switch {
	case bps == 8 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
			d.buf[i][1] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
		}
	case bps == 16 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
			d.buf[i][1] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
		}
	case bps == 24 && nchannels == 1:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int32(frame.Subframes[0].Samples[i])) / (1<<23 - 1)
			d.buf[i][1] = float64(int32(frame.Subframes[1].Samples[i])) / (1<<23 - 1)
		}
	case bps == 8 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int8(frame.Subframes[0].Samples[i])) / (1<<7 - 1)
			d.buf[i][1] = float64(int8(frame.Subframes[1].Samples[i])) / (1<<7 - 1)
		}
	case bps == 16 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int16(frame.Subframes[0].Samples[i])) / (1<<15 - 1)
			d.buf[i][1] = float64(int16(frame.Subframes[1].Samples[i])) / (1<<15 - 1)
		}
	case bps == 24 && nchannels >= 2:
		for i := 0; i < n; i++ {
			d.buf[i][0] = float64(int32(frame.Subframes[0].Samples[i])) / (1<<23 - 1)
			d.buf[i][1] = float64(int32(frame.Subframes[1].Samples[i])) / (1<<23 - 1)
		}
	default:
		panic(fmt.Errorf("support for %d bits-per-sample and %d channels combination not yet implemented", bps, nchannels))
	}
	return nil
}

func (d *decoder) Err() error {
	return d.err
}

func (d *decoder) Len() int {
	return int(d.stream.Info.NSamples)
}

func (d *decoder) Position() int {
	return d.pos
}

func (d *decoder) Seek(p int) error {
	panic("not yet implemented")
}

func (d *decoder) Close() error {
	err := d.rc.Close()
	if err != nil {
		return errors.Wrap(err, "flac")
	}
	return nil
}
<|endoftext|>


<commit_before>package cli

import "flag"

// StringFlag is a flag with type string
type StringFlag struct {
	Name        string
	Aliases     []string
	Usage       string
	EnvVars     []string
	FilePath    string
	Required    bool
	Hidden      bool
	TakesFile   bool
	Value       string
	DefaultText string
	Destination *string
	HasBeenSet  bool
}

// IsSet returns whether or not the flag has been set through env or file
func (f *StringFlag) IsSet() bool {
	return f.HasBeenSet
}

// String returns a readable representation of this value
// (for usage defaults)
func (f *StringFlag) String() string {
	return FlagStringer(f)
}

// Names returns the names of the flag
func (f *StringFlag) Names() []string {
	return flagNames(f.Name, f.Aliases)
}

// IsRequired returns whether or not the flag is required
func (f *StringFlag) IsRequired() bool {
	return f.Required
}

// TakesValue returns true of the flag takes a value, otherwise false
func (f *StringFlag) TakesValue() bool {
	return true
}

// GetUsage returns the usage string for the flag
func (f *StringFlag) GetUsage() string {
	return f.Usage
}

// GetValue returns the flags value as string representation and an empty
// string if the flag takes no value at all.
func (f *StringFlag) GetValue() string {
	return f.Value
}

// IsVisible returns true if the flag is not hidden, otherwise false
func (f *StringFlag) IsVisible() bool {
	return !f.Hidden
}

// Apply populates the flag given the flag set and environment
func (f *StringFlag) Apply(set *flag.FlagSet) error {
	if val, ok := flagFromEnvOrFile(f.EnvVars, f.FilePath); ok {
		f.Value = val
		f.HasBeenSet = true
	}

	for _, name := range f.Names() {
		if f.Destination != nil {
			set.StringVar(f.Destination, name, f.Value, f.Usage)
			continue
		}
		set.String(name, f.Value, f.Usage)
	}

	return nil
}

// String looks up the value of a local StringFlag, returns
// "" if not found
func (c *Context) String(name string) string {
	if fs := c.lookupFlagSet(name); fs != nil {
		return lookupString(name, fs)
	}
	return ""
}

func lookupString(name string, set *flag.FlagSet) string {
	f := set.Lookup(name)
	if f != nil {
		parsed, err := f.Value.String(), error(nil)
		if err != nil {
			return ""
		}
		return parsed
	}
	return ""
}
<commit_msg>remove useless variable declarations (#1330)<commit_after>package cli

import "flag"

// StringFlag is a flag with type string
type StringFlag struct {
	Name        string
	Aliases     []string
	Usage       string
	EnvVars     []string
	FilePath    string
	Required    bool
	Hidden      bool
	TakesFile   bool
	Value       string
	DefaultText string
	Destination *string
	HasBeenSet  bool
}

// IsSet returns whether or not the flag has been set through env or file
func (f *StringFlag) IsSet() bool {
	return f.HasBeenSet
}

// String returns a readable representation of this value
// (for usage defaults)
func (f *StringFlag) String() string {
	return FlagStringer(f)
}

// Names returns the names of the flag
func (f *StringFlag) Names() []string {
	return flagNames(f.Name, f.Aliases)
}

// IsRequired returns whether or not the flag is required
func (f *StringFlag) IsRequired() bool {
	return f.Required
}

// TakesValue returns true of the flag takes a value, otherwise false
func (f *StringFlag) TakesValue() bool {
	return true
}

// GetUsage returns the usage string for the flag
func (f *StringFlag) GetUsage() string {
	return f.Usage
}

// GetValue returns the flags value as string representation and an empty
// string if the flag takes no value at all.
func (f *StringFlag) GetValue() string {
	return f.Value
}

// IsVisible returns true if the flag is not hidden, otherwise false
func (f *StringFlag) IsVisible() bool {
	return !f.Hidden
}

// Apply populates the flag given the flag set and environment
func (f *StringFlag) Apply(set *flag.FlagSet) error {
	if val, ok := flagFromEnvOrFile(f.EnvVars, f.FilePath); ok {
		f.Value = val
		f.HasBeenSet = true
	}

	for _, name := range f.Names() {
		if f.Destination != nil {
			set.StringVar(f.Destination, name, f.Value, f.Usage)
			continue
		}
		set.String(name, f.Value, f.Usage)
	}

	return nil
}

// String looks up the value of a local StringFlag, returns
// "" if not found
func (c *Context) String(name string) string {
	if fs := c.lookupFlagSet(name); fs != nil {
		return lookupString(name, fs)
	}
	return ""
}

func lookupString(name string, set *flag.FlagSet) string {
	f := set.Lookup(name)
	if f != nil {
		parsed := f.Value.String()
		return parsed
	}
	return ""
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative 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 kpa

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"knative.dev/pkg/apis/duck"
	"knative.dev/pkg/injection/clients/dynamicclient"
	"knative.dev/pkg/logging"

	network "knative.dev/networking/pkg"
	"knative.dev/networking/pkg/apis/networking"
	nv1a1 "knative.dev/networking/pkg/apis/networking/v1alpha1"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/network/prober"
	"knative.dev/serving/pkg/activator"
	pav1alpha1 "knative.dev/serving/pkg/apis/autoscaling/v1alpha1"
	"knative.dev/serving/pkg/autoscaler/config/autoscalerconfig"
	"knative.dev/serving/pkg/reconciler/autoscaling/config"
	kparesources "knative.dev/serving/pkg/reconciler/autoscaling/kpa/resources"
	aresources "knative.dev/serving/pkg/reconciler/autoscaling/resources"
	"knative.dev/serving/pkg/resources"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/dynamic"
)

const (
	scaleUnknown = -1
	probePeriod  = 1 * time.Second
	probeTimeout = 45 * time.Second

	// The time after which the PA will be re-enqueued.
	// This number is small, since `handleScaleToZero` below will
	// re-enqueue for the configured grace period.
	reenqeuePeriod = 1 * time.Second

	// TODO(#3456): Remove this buffer once KPA does pod failure diagnostics.
	//
	// KPA will scale the Deployment down to zero if it fails to activate after ProgressDeadlineSeconds,
	// however, after ProgressDeadlineSeconds, the Deployment itself updates its status, which causes
	// the Revision to re-reconcile and diagnose pod failures. If we use the same timeout here, we will
	// race the Revision reconciler and scale down the pods before it can actually surface the pod errors.
	// We should instead do pod failure diagnostics here immediately before scaling down the Deployment.
	activationTimeoutBuffer = 30 * time.Second
)

var probeOptions = []interface{}{
	prober.WithHeader(network.UserAgentKey, network.AutoscalingUserAgent),
	prober.WithHeader(network.ProbeHeaderName, activator.Name),
	prober.ExpectsBody(activator.Name),
	prober.ExpectsStatusCodes([]int{http.StatusOK}),
}

// for mocking in tests
type asyncProber interface {
	Offer(context.Context, string, interface{}, time.Duration, time.Duration, ...interface{}) bool
}

// scaler scales the target of a kpa-class PA up or down including scaling to zero.
type scaler struct {
	psInformerFactory duck.InformerFactory
	dynamicClient     dynamic.Interface
	transport         http.RoundTripper

	// For sync probes.
	activatorProbe func(pa *pav1alpha1.PodAutoscaler, transport http.RoundTripper) (bool, error)

	// For async probes.
	probeManager asyncProber
	enqueueCB    func(interface{}, time.Duration)
}

// newScaler creates a scaler.
func newScaler(ctx context.Context, psInformerFactory duck.InformerFactory, enqueueCB func(interface{}, time.Duration)) *scaler {
	logger := logging.FromContext(ctx)
	transport := pkgnet.NewProberTransport()
	ks := &scaler{
		// Wrap it in a cache, so that we don't stamp out a new
		// informer/lister each time.
		psInformerFactory: psInformerFactory,
		dynamicClient:     dynamicclient.Get(ctx),
		transport:         transport,

		// Production setup uses the default probe implementation.
		activatorProbe: activatorProbe,
		probeManager: prober.New(func(arg interface{}, success bool, err error) {
			logger.Infof("Async prober is done for %v: success?: %v error: %v", arg, success, err)
			// Re-enqueue the PA in any case. If the probe timed out to retry again, if succeeded to scale to 0.
			enqueueCB(arg, reenqeuePeriod)
		}, transport),
		enqueueCB: enqueueCB,
	}
	return ks
}

// Resolves the pa to the probing endpoint Eg. http://hostname:port/healthz
func paToProbeTarget(pa *pav1alpha1.PodAutoscaler) string {
	svc := pkgnet.GetServiceHostname(pa.Status.ServiceName, pa.Namespace)
	port := networking.ServicePort(pa.Spec.ProtocolType)

	return fmt.Sprintf("http://%s:%d/%s", svc, port, network.ProbePath)
}

// activatorProbe returns true if via probe it determines that the
// PA is backed by the Activator.
func activatorProbe(pa *pav1alpha1.PodAutoscaler, transport http.RoundTripper) (bool, error) {
	// No service name -- no probe.
	if pa.Status.ServiceName == "" {
		return false, nil
	}
	return prober.Do(context.Background(), transport, paToProbeTarget(pa), probeOptions...)
}

func lastPodRetention(pa *pav1alpha1.PodAutoscaler, cfg *autoscalerconfig.Config) time.Duration {
	d, ok := pa.ScaleToZeroPodRetention()
	if ok {
		return d
	}
	return cfg.ScaleToZeroPodRetentionPeriod
}

// pre: 0 <= min <= max && 0 <= x
func applyBounds(min, max, x int32) int32 {
	if x < min {
		return min
	}
	if max != 0 && x > max {
		return max
	}
	return x
}

func durationMax(d1, d2 time.Duration) time.Duration {
	if d1 < d2 {
		return d2
	}
	return d1
}

func (ks *scaler) handleScaleToZero(ctx context.Context, pa *pav1alpha1.PodAutoscaler,
	sks *nv1a1.ServerlessService, desiredScale int32) (int32, bool) {
	if desiredScale != 0 {
		return desiredScale, true
	}

	// We should only scale to zero when three of the following conditions are true:
	//   a) enable-scale-to-zero from configmap is true
	//   b) The PA has been active for at least the stable window, after which it
	//			gets marked inactive, and
	//   c) the PA has been backed by the Activator for at least the grace period
	//      of time.
	//  Alternatively, if (a) and the revision did not succeed to activate in
	//  `activationTimeout` time -- also scale it to 0.
	cfgs := config.FromContext(ctx)
	cfgAS := cfgs.Autoscaler

	if !cfgAS.EnableScaleToZero {
		return 1, true
	}
	cfgD := cfgs.Deployment
	activationTimeout := cfgD.ProgressDeadline + activationTimeoutBuffer

	now := time.Now()
	logger := logging.FromContext(ctx)
	switch {
	case pa.Status.IsActivating(): // Active=Unknown
		// If we are stuck activating for longer than our progress deadline, presume we cannot succeed and scale to 0.
		if pa.Status.CanFailActivation(now, activationTimeout) {
			logger.Info("Activation has timed out after ", activationTimeout)
			return desiredScale, true
		}
		ks.enqueueCB(pa, activationTimeout)
		return scaleUnknown, false
	case pa.Status.IsActive(): // Active=True
		// Don't scale-to-zero if the PA is active
		// but return `(0, false)` to mark PA inactive, instead.
		sw := aresources.StableWindow(pa, cfgAS)
		af := pa.Status.ActiveFor(now)
		if af >= sw {
			// If SKS is in proxy mode, then there is high probability
			// of SKS not changing its spec/status and thus not triggering
			// a new reconciliation of PA.
			if sks.Spec.Mode == nv1a1.SKSOperationModeProxy {
				logger.Debug("SKS is already in proxy mode, auto-re-enqueue PA")
				// Long enough to ensure current iteration is finished.
				ks.enqueueCB(pa, 3*time.Second)
			}
			logger.Info("Can deactivate PA, was active for ", af)
			return desiredScale, false
		}
		// Otherwise, scale down to at most 1 for the remainder of the idle period and then
		// reconcile PA again.
		logger.Infof("Sleeping additionally for %v before can scale to 0", sw-af)
		ks.enqueueCB(pa, sw-af)
		return 1, true
	default: // Active=False
		var (
			err error
			r   bool = true
		)

		if resolveTBC(ctx, pa) != -1 {
			// if TBC is -1 activator is guaranteed to already be in the path.
			// Otherwise, probe to make sure Activator is in path.
			r, err = ks.activatorProbe(pa, ks.transport)
			logger.Infof("Probing activator = %v, err = %v", r, err)
		}

		if r {
			// This enforces that the revision has been backed by the Activator for at least
			// ScaleToZeroGracePeriod time.
			// And at least ScaleToZeroPodRetentionPeriod since PA became inactive.

			// Most conservative check, if it passes we're good.
			lastPodTimeout := lastPodRetention(pa, cfgAS)
			lastPodMaxTimeout := durationMax(cfgAS.ScaleToZeroGracePeriod, lastPodTimeout)
			// If we have been inactive for this long, we can scale to 0!
			if pa.Status.InactiveFor(now) >= lastPodMaxTimeout {
				return desiredScale, true
			}

			// Now check last pod retention timeout. Since it's a hard deadline, regardless
			// of network programming state we should circle back after that time period.
			if lastPodTimeout > 0 {
				if inactiveTime := pa.Status.InactiveFor(now); inactiveTime < lastPodTimeout {
					logger.Infof("Can't scale to 0; InactiveFor %v < ScaleToZeroPodRetentionPeriod = %v",
						inactiveTime, lastPodTimeout)
					ks.enqueueCB(pa, lastPodTimeout-inactiveTime)
					return desiredScale, false
				}
				logger.Debug("Last pod timeout satisfied")
			}

			// Otherwise check how long SKS was in proxy mode.
			// Compute the difference between time we've been proxying with the timeout.
			// If it's positive, that's the time we need to sleep, if negative -- we
			// can scale to zero.
			pf := sks.Status.ProxyFor()
			to := cfgAS.ScaleToZeroGracePeriod - pf
			if to <= 0 {
				logger.Info("Fast path scaling to 0, in proxy mode for: ", pf)
				return desiredScale, true
			}

			// Re-enqueue the PA for reconciliation with timeout of `to` to make sure we wait
			// long enough.
			logger.Info("Enqueueing PA after ", to)
			ks.enqueueCB(pa, to)
			return desiredScale, false
		}

		// Otherwise (any prober failure) start the async probe.
		logger.Info("PA is not yet backed by activator, cannot scale to zero")
		if !ks.probeManager.Offer(context.Background(), paToProbeTarget(pa), pa, probePeriod, probeTimeout, probeOptions...) {
			logger.Info("Probe for revision is already in flight")
		}
		return desiredScale, false
	}
}

func (ks *scaler) applyScale(ctx context.Context, pa *pav1alpha1.PodAutoscaler, desiredScale int32,
	ps *pav1alpha1.PodScalable) error {
	logger := logging.FromContext(ctx)

	gvr, name, err := resources.ScaleResourceArguments(pa.Spec.ScaleTargetRef)
	if err != nil {
		return err
	}

	psNew := ps.DeepCopy()
	psNew.Spec.Replicas = &desiredScale
	patch, err := duck.CreatePatch(ps, psNew)
	if err != nil {
		return err
	}
	patchBytes, err := patch.MarshalJSON()
	if err != nil {
		return err
	}

	_, err = ks.dynamicClient.Resource(*gvr).Namespace(pa.Namespace).Patch(ctx, ps.Name, types.JSONPatchType,
		patchBytes, metav1.PatchOptions{})
	if err != nil {
		return fmt.Errorf("failed to apply scale %d to scale target %s: %w", desiredScale, name, err)
	}

	logger.Debug("Successfully scaled to ", desiredScale)
	return nil
}

// scale attempts to scale the given PA's target reference to the desired scale.
func (ks *scaler) scale(ctx context.Context, pa *pav1alpha1.PodAutoscaler, sks *nv1a1.ServerlessService, desiredScale int32) (int32, error) {
	asConfig := config.FromContext(ctx).Autoscaler
	logger := logging.FromContext(ctx)

	if desiredScale < 0 && !pa.Status.IsActivating() {
		logger.Debug("Metrics are not yet being collected.")
		return desiredScale, nil
	}

	min, max := pa.ScaleBounds(asConfig)
	initialScale := kparesources.GetInitialScale(asConfig, pa)
	// If initial scale has been attained, ignore the initialScale altogether.
	if initialScale > 1 && !pa.Status.IsScaleTargetInitialized() {
		// Ignore initial scale if minScale >= initialScale.
		if min < initialScale {
			logger.Debugf("Adjusting min to meet the initial scale: %d -> %d", min, initialScale)
		}
		min = intMax(initialScale, min)
	}
	if newScale := applyBounds(min, max, desiredScale); newScale != desiredScale {
		logger.Debugf("Adjusting desiredScale to meet the min and max bounds before applying: %d -> %d", desiredScale, newScale)
		desiredScale = newScale
	}

	desiredScale, shouldApplyScale := ks.handleScaleToZero(ctx, pa, sks, desiredScale)
	if !shouldApplyScale {
		return desiredScale, nil
	}

	ps, err := resources.GetScaleResource(ctx, pa.Namespace, pa.Spec.ScaleTargetRef, ks.psInformerFactory)
	if err != nil {
		return desiredScale, fmt.Errorf("failed to get scale target %v: %w", pa.Spec.ScaleTargetRef, err)
	}

	currentScale := int32(1)
	if ps.Spec.Replicas != nil {
		currentScale = *ps.Spec.Replicas
	}
	if desiredScale == currentScale {
		return desiredScale, nil
	}

	logger.Infof("Scaling from %d to %d", currentScale, desiredScale)
	return desiredScale, ks.applyScale(ctx, pa, desiredScale, ps)
}
<commit_msg>Add logging to aid with debugging (#9698)<commit_after>/*
Copyright 2019 The Knative 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 kpa

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"knative.dev/pkg/apis/duck"
	"knative.dev/pkg/injection/clients/dynamicclient"
	"knative.dev/pkg/logging"

	network "knative.dev/networking/pkg"
	"knative.dev/networking/pkg/apis/networking"
	nv1a1 "knative.dev/networking/pkg/apis/networking/v1alpha1"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/network/prober"
	"knative.dev/serving/pkg/activator"
	pav1alpha1 "knative.dev/serving/pkg/apis/autoscaling/v1alpha1"
	"knative.dev/serving/pkg/autoscaler/config/autoscalerconfig"
	"knative.dev/serving/pkg/reconciler/autoscaling/config"
	kparesources "knative.dev/serving/pkg/reconciler/autoscaling/kpa/resources"
	aresources "knative.dev/serving/pkg/reconciler/autoscaling/resources"
	"knative.dev/serving/pkg/resources"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/dynamic"
)

const (
	scaleUnknown = -1
	probePeriod  = 1 * time.Second
	probeTimeout = 45 * time.Second

	// The time after which the PA will be re-enqueued.
	// This number is small, since `handleScaleToZero` below will
	// re-enqueue for the configured grace period.
	reenqeuePeriod = 1 * time.Second

	// TODO(#3456): Remove this buffer once KPA does pod failure diagnostics.
	//
	// KPA will scale the Deployment down to zero if it fails to activate after ProgressDeadlineSeconds,
	// however, after ProgressDeadlineSeconds, the Deployment itself updates its status, which causes
	// the Revision to re-reconcile and diagnose pod failures. If we use the same timeout here, we will
	// race the Revision reconciler and scale down the pods before it can actually surface the pod errors.
	// We should instead do pod failure diagnostics here immediately before scaling down the Deployment.
	activationTimeoutBuffer = 30 * time.Second
)

var probeOptions = []interface{}{
	prober.WithHeader(network.UserAgentKey, network.AutoscalingUserAgent),
	prober.WithHeader(network.ProbeHeaderName, activator.Name),
	prober.ExpectsBody(activator.Name),
	prober.ExpectsStatusCodes([]int{http.StatusOK}),
}

// for mocking in tests
type asyncProber interface {
	Offer(context.Context, string, interface{}, time.Duration, time.Duration, ...interface{}) bool
}

// scaler scales the target of a kpa-class PA up or down including scaling to zero.
type scaler struct {
	psInformerFactory duck.InformerFactory
	dynamicClient     dynamic.Interface
	transport         http.RoundTripper

	// For sync probes.
	activatorProbe func(pa *pav1alpha1.PodAutoscaler, transport http.RoundTripper) (bool, error)

	// For async probes.
	probeManager asyncProber
	enqueueCB    func(interface{}, time.Duration)
}

// newScaler creates a scaler.
func newScaler(ctx context.Context, psInformerFactory duck.InformerFactory, enqueueCB func(interface{}, time.Duration)) *scaler {
	logger := logging.FromContext(ctx)
	transport := pkgnet.NewProberTransport()
	ks := &scaler{
		// Wrap it in a cache, so that we don't stamp out a new
		// informer/lister each time.
		psInformerFactory: psInformerFactory,
		dynamicClient:     dynamicclient.Get(ctx),
		transport:         transport,

		// Production setup uses the default probe implementation.
		activatorProbe: activatorProbe,
		probeManager: prober.New(func(arg interface{}, success bool, err error) {
			logger.Infof("Async prober is done for %v: success?: %v error: %v", arg, success, err)
			// Re-enqueue the PA in any case. If the probe timed out to retry again, if succeeded to scale to 0.
			enqueueCB(arg, reenqeuePeriod)
		}, transport),
		enqueueCB: enqueueCB,
	}
	return ks
}

// Resolves the pa to the probing endpoint Eg. http://hostname:port/healthz
func paToProbeTarget(pa *pav1alpha1.PodAutoscaler) string {
	svc := pkgnet.GetServiceHostname(pa.Status.ServiceName, pa.Namespace)
	port := networking.ServicePort(pa.Spec.ProtocolType)

	return fmt.Sprintf("http://%s:%d/%s", svc, port, network.ProbePath)
}

// activatorProbe returns true if via probe it determines that the
// PA is backed by the Activator.
func activatorProbe(pa *pav1alpha1.PodAutoscaler, transport http.RoundTripper) (bool, error) {
	// No service name -- no probe.
	if pa.Status.ServiceName == "" {
		return false, nil
	}
	return prober.Do(context.Background(), transport, paToProbeTarget(pa), probeOptions...)
}

func lastPodRetention(pa *pav1alpha1.PodAutoscaler, cfg *autoscalerconfig.Config) time.Duration {
	d, ok := pa.ScaleToZeroPodRetention()
	if ok {
		return d
	}
	return cfg.ScaleToZeroPodRetentionPeriod
}

// pre: 0 <= min <= max && 0 <= x
func applyBounds(min, max, x int32) int32 {
	if x < min {
		return min
	}
	if max != 0 && x > max {
		return max
	}
	return x
}

func durationMax(d1, d2 time.Duration) time.Duration {
	if d1 < d2 {
		return d2
	}
	return d1
}

func (ks *scaler) handleScaleToZero(ctx context.Context, pa *pav1alpha1.PodAutoscaler,
	sks *nv1a1.ServerlessService, desiredScale int32) (int32, bool) {
	if desiredScale != 0 {
		return desiredScale, true
	}

	// We should only scale to zero when three of the following conditions are true:
	//   a) enable-scale-to-zero from configmap is true
	//   b) The PA has been active for at least the stable window, after which it
	//			gets marked inactive, and
	//   c) the PA has been backed by the Activator for at least the grace period
	//      of time.
	//  Alternatively, if (a) and the revision did not succeed to activate in
	//  `activationTimeout` time -- also scale it to 0.
	cfgs := config.FromContext(ctx)
	cfgAS := cfgs.Autoscaler

	if !cfgAS.EnableScaleToZero {
		return 1, true
	}
	cfgD := cfgs.Deployment
	activationTimeout := cfgD.ProgressDeadline + activationTimeoutBuffer

	now := time.Now()
	logger := logging.FromContext(ctx)
	switch {
	case pa.Status.IsActivating(): // Active=Unknown
		// If we are stuck activating for longer than our progress deadline, presume we cannot succeed and scale to 0.
		if pa.Status.CanFailActivation(now, activationTimeout) {
			logger.Info("Activation has timed out after ", activationTimeout)
			return desiredScale, true
		}
		ks.enqueueCB(pa, activationTimeout)
		return scaleUnknown, false
	case pa.Status.IsActive(): // Active=True
		// Don't scale-to-zero if the PA is active
		// but return `(0, false)` to mark PA inactive, instead.
		sw := aresources.StableWindow(pa, cfgAS)
		af := pa.Status.ActiveFor(now)
		if af >= sw {
			// If SKS is in proxy mode, then there is high probability
			// of SKS not changing its spec/status and thus not triggering
			// a new reconciliation of PA.
			if sks.Spec.Mode == nv1a1.SKSOperationModeProxy {
				logger.Debug("SKS is already in proxy mode, auto-re-enqueue PA")
				// Long enough to ensure current iteration is finished.
				ks.enqueueCB(pa, 3*time.Second)
			}
			logger.Info("Can deactivate PA, was active for ", af)
			return desiredScale, false
		}
		// Otherwise, scale down to at most 1 for the remainder of the idle period and then
		// reconcile PA again.
		logger.Infof("Sleeping additionally for %v before can scale to 0", sw-af)
		ks.enqueueCB(pa, sw-af)
		return 1, true
	default: // Active=False
		var (
			err error
			r   bool = true
		)

		if resolveTBC(ctx, pa) != -1 {
			// if TBC is -1 activator is guaranteed to already be in the path.
			// Otherwise, probe to make sure Activator is in path.
			r, err = ks.activatorProbe(pa, ks.transport)
			logger.Infof("Probing activator = %v, err = %v", r, err)
		}

		if r {
			// This enforces that the revision has been backed by the Activator for at least
			// ScaleToZeroGracePeriod time.
			// And at least ScaleToZeroPodRetentionPeriod since PA became inactive.

			// Most conservative check, if it passes we're good.
			lastPodTimeout := lastPodRetention(pa, cfgAS)
			lastPodMaxTimeout := durationMax(cfgAS.ScaleToZeroGracePeriod, lastPodTimeout)
			// If we have been inactive for this long, we can scale to 0!
			if pa.Status.InactiveFor(now) >= lastPodMaxTimeout {
				return desiredScale, true
			}

			// Now check last pod retention timeout. Since it's a hard deadline, regardless
			// of network programming state we should circle back after that time period.
			if lastPodTimeout > 0 {
				if inactiveTime := pa.Status.InactiveFor(now); inactiveTime < lastPodTimeout {
					logger.Infof("Can't scale to 0; InactiveFor %v < ScaleToZeroPodRetentionPeriod = %v",
						inactiveTime, lastPodTimeout)
					ks.enqueueCB(pa, lastPodTimeout-inactiveTime)
					return desiredScale, false
				}
				logger.Debug("Last pod timeout satisfied")
			}

			// Otherwise check how long SKS was in proxy mode.
			// Compute the difference between time we've been proxying with the timeout.
			// If it's positive, that's the time we need to sleep, if negative -- we
			// can scale to zero.
			pf := sks.Status.ProxyFor()
			to := cfgAS.ScaleToZeroGracePeriod - pf
			if to <= 0 {
				logger.Info("Fast path scaling to 0, in proxy mode for: ", pf)
				return desiredScale, true
			}

			// Re-enqueue the PA for reconciliation with timeout of `to` to make sure we wait
			// long enough.
			logger.Info("Enqueueing PA after ", to)
			ks.enqueueCB(pa, to)
			return desiredScale, false
		}

		// Otherwise (any prober failure) start the async probe.
		logger.Info("PA is not yet backed by activator, cannot scale to zero")
		if !ks.probeManager.Offer(context.Background(), paToProbeTarget(pa), pa, probePeriod, probeTimeout, probeOptions...) {
			logger.Info("Probe for revision is already in flight")
		}
		return desiredScale, false
	}
}

func (ks *scaler) applyScale(ctx context.Context, pa *pav1alpha1.PodAutoscaler, desiredScale int32,
	ps *pav1alpha1.PodScalable) error {
	logger := logging.FromContext(ctx)

	gvr, name, err := resources.ScaleResourceArguments(pa.Spec.ScaleTargetRef)
	if err != nil {
		return err
	}

	psNew := ps.DeepCopy()
	psNew.Spec.Replicas = &desiredScale
	patch, err := duck.CreatePatch(ps, psNew)
	if err != nil {
		return err
	}
	patchBytes, err := patch.MarshalJSON()
	if err != nil {
		return err
	}

	_, err = ks.dynamicClient.Resource(*gvr).Namespace(pa.Namespace).Patch(ctx, ps.Name, types.JSONPatchType,
		patchBytes, metav1.PatchOptions{})
	if err != nil {
		return fmt.Errorf("failed to apply scale %d to scale target %s: %w", desiredScale, name, err)
	}

	logger.Debug("Successfully scaled to ", desiredScale)
	return nil
}

// scale attempts to scale the given PA's target reference to the desired scale.
func (ks *scaler) scale(ctx context.Context, pa *pav1alpha1.PodAutoscaler, sks *nv1a1.ServerlessService, desiredScale int32) (int32, error) {
	asConfig := config.FromContext(ctx).Autoscaler
	logger := logging.FromContext(ctx)

	if desiredScale < 0 && !pa.Status.IsActivating() {
		logger.Debug("Metrics are not yet being collected.")
		return desiredScale, nil
	}

	min, max := pa.ScaleBounds(asConfig)
	initialScale := kparesources.GetInitialScale(asConfig, pa)
	logger.Debugf("MinScale = %d, MaxScale = %d, InitialScale = %d, DesiredScale = %d",
		min, max, initialScale, desiredScale)
	// If initial scale has been attained, ignore the initialScale altogether.
	if initialScale > 1 && !pa.Status.IsScaleTargetInitialized() {
		// Ignore initial scale if minScale >= initialScale.
		if min < initialScale {
			logger.Debugf("Adjusting min to meet the initial scale: %d -> %d", min, initialScale)
		}
		min = intMax(initialScale, min)
	}
	if newScale := applyBounds(min, max, desiredScale); newScale != desiredScale {
		logger.Debugf("Adjusting desiredScale to meet the min and max bounds before applying: %d -> %d", desiredScale, newScale)
		desiredScale = newScale
	}

	desiredScale, shouldApplyScale := ks.handleScaleToZero(ctx, pa, sks, desiredScale)
	if !shouldApplyScale {
		return desiredScale, nil
	}

	ps, err := resources.GetScaleResource(ctx, pa.Namespace, pa.Spec.ScaleTargetRef, ks.psInformerFactory)
	if err != nil {
		return desiredScale, fmt.Errorf("failed to get scale target %v: %w", pa.Spec.ScaleTargetRef, err)
	}

	currentScale := int32(1)
	if ps.Spec.Replicas != nil {
		currentScale = *ps.Spec.Replicas
	}
	if desiredScale == currentScale {
		return desiredScale, nil
	}

	logger.Infof("Scaling from %d to %d", currentScale, desiredScale)
	return desiredScale, ks.applyScale(ctx, pa, desiredScale, ps)
}
<|endoftext|>