<commit_before>package auroraconfig

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"

	"github.com/skatteetaten/ao/pkg/configuration"
	"github.com/skatteetaten/ao/pkg/fileutil"
	"github.com/skatteetaten/ao/pkg/jsonutil"
	"github.com/skatteetaten/ao/pkg/serverapi"
	"fmt"
	"strings"
)

const InvalidConfigurationError = "Invalid configuration"

func GetContent(filename string, auroraConfig *serverapi.AuroraConfig) (content string, version string, err error) {
	var fileFound bool = false

	_, fileFound = auroraConfig.Files[filename]
	if fileFound {
		content = string(auroraConfig.Files[filename])
	}

	version = auroraConfig.Versions[filename]

	if !fileFound {
		return "", "", errors.New("Illegal file/folder")
	}
	return content, version, nil

}

func GetAllContent(outputFolder string, configuration *configuration.ConfigurationClass) (output string, err error) {
	auroraConfig, err := GetAuroraConfig(configuration)
	if err != nil {
		return
	}

	if outputFolder != "" {
		if fileutil.IsLegalFileFolder(outputFolder) == fileutil.SpecIllegal {
			err = errors.New("Illegal file/folder")
			return "", err

		}
		var content string
		for filename := range auroraConfig.Files {
			content = jsonutil.PrettyPrintJson(string(auroraConfig.Files[filename]))
			err = fileutil.WriteFile(outputFolder, filename, content)
			if err != nil {
				return "", err
			}
		}
	}

	outputBytes, err := json.Marshal(auroraConfig)
	output = jsonutil.PrettyPrintJson(string(outputBytes))
	return output, err

}

func GetFileList(configuration *configuration.ConfigurationClass) (filenames []string, err error) {
	auroraConfig, err := GetAuroraConfig(configuration)
	if err != nil {
		return
	}
	filenames = make([]string, len(auroraConfig.Files))

	var filenameIndex = 0
	for filename := range auroraConfig.Files {
		filenames[filenameIndex] = filename
		filenameIndex++
	}
	return filenames, nil
}

func GetAuroraConfigRequest(configuration *configuration.ConfigurationClass) (request *serverapi.Request) {
	request = new(serverapi.Request)
	request.Method = http.MethodGet
	request.ApiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"
	return request
}

func Response2AuroraConfig(response serverapi.Response) (auroraConfig serverapi.AuroraConfig, err error) {

	if !response.Success {
		output, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return auroraConfig, err
		}
		err = errors.New(output)
		return auroraConfig, err

	}

	auroraConfig, err = serverapi.ResponseItems2AuroraConfig(response)

	return auroraConfig, nil

}

func ValidateAuroraConfig(auroraConfig *serverapi.AuroraConfig, config *configuration.ConfigurationClass) error {
	payload, err := json.Marshal(auroraConfig)
	if err != nil {
		return err
	}

	endpoint := fmt.Sprintf("/affiliation/%s/auroraconfig/validate", config.GetAffiliation())
	response, err := serverapi.CallApi(http.MethodPut, endpoint, string(payload), config)
	if err != nil {
		return err
	}

	printValidationResponse(response)

	return nil
}

type Validation struct {
	IllegalFieldErrors []string
	MissingFieldErrors []string
	InvalidFieldErrors []string
	UniqueErrors       map[string]bool
}

func (v *Validation) GetAllErrors() []string {
	errorMessages := append(v.IllegalFieldErrors, v.InvalidFieldErrors...)
	return append(errorMessages, v.MissingFieldErrors...)
}

func (v *Validation) Contains(key string) bool {
	return v.UniqueErrors[key]
}

func printValidationResponse(res serverapi.Response) {
	fmt.Println(res.Message)

	if res.Success {
		return
	}

	validation := &Validation{
		UniqueErrors: make(map[string]bool),
	}

	for _, item := range res.Items {
		var res serverapi.ResponseItemError
		json.Unmarshal(item, &res)
		validation.formatValidationError(&res)
	}

	for _, e := range validation.GetAllErrors() {
		fmt.Print(e)
	}
}

func (v *Validation) formatValidationError(res *serverapi.ResponseItemError) {
	illegalFieldFormat := `
Filename:    %s
Path:        %s
Value:       %s
Message:     %s
`
	missingFieldFormat := `
Application: %s/%s
Path:        %s
Message:     %s
`

	invalidFieldFormat := `
Filename:    %s
Path:        %s
Message:     %s
`

	for _, message := range res.Messages {
		k := []string{
			message.Field.Source,
			message.Field.Path,
			message.Field.Value,
		}
		key := strings.Join(k, "|")

		if v.Contains(key) {
			continue
		}

		if message.Type != "MISSING" {
			v.UniqueErrors[key] = true
		}

		switch message.Type {
		case "ILLEGAL":
			{
				validationError := fmt.Sprintf(illegalFieldFormat,
					message.Field.Source,
					message.Field.Path,
					message.Field.Value,
					message.Message,
				)
				v.IllegalFieldErrors = append(v.IllegalFieldErrors, validationError)
			}

		case "INVALID":
			{
				invalid := fmt.Sprintf(invalidFieldFormat,
					message.Field.Source,
					message.Field.Path,
					message.Message,
				)
				v.InvalidFieldErrors = append(v.InvalidFieldErrors, invalid);
			}

		case "MISSING":
			{
				missing := fmt.Sprintf(missingFieldFormat,
					res.Environment,
					res.Application,
					message.Field.Path,
					message.Message,
				)
				v.MissingFieldErrors = append(v.MissingFieldErrors, missing);
			}
		}
	}
}

func GetAuroraConfig(configuration *configuration.ConfigurationClass) (auroraConfig serverapi.AuroraConfig, err error) {
	var apiEndpoint string = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"

	response, err := serverapi.CallApi(http.MethodGet, apiEndpoint, "", configuration)
	if err != nil {
		return auroraConfig, err
	}

	if !response.Success {
		output, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return auroraConfig, err
		}
		err = errors.New(output)
		return auroraConfig, err

	}

	auroraConfig, err = serverapi.ResponseItems2AuroraConfig(response)

	return auroraConfig, nil
}

func PutAuroraConfig(auroraConfig serverapi.AuroraConfig, configuration *configuration.ConfigurationClass) (err error) {
	content, err := json.Marshal(auroraConfig)
	if err != nil {
		return err
	}

	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"

	_, err = putContent(apiEndpoint, string(content), "", configuration)
	if err != nil {
		return err
	}
	return
}

func putContent(apiEndpoint string, content string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {

	var versionHeader = make(map[string]string)
	versionHeader["AuroraConfigFileVersion"] = version

	response, err := serverapi.CallApiWithHeaders(versionHeader, http.MethodPut, apiEndpoint, content, configuration)

	if err != nil {
		validationMessages, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return "", err
		}
		return validationMessages, errors.New(InvalidConfigurationError + "\n" + validationMessages)
	}

	if !response.Success {
		validationMessages, _ := serverapi.ResponsItems2MessageString(response)
		return validationMessages, errors.New(InvalidConfigurationError + "\n" + validationMessages)
	}

	return
}

func PutFile(filename string, content string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfigfile/" + filename

	return putContent(apiEndpoint, content, version, configuration)
}

func PutSecret(vaultname string, secretname string, secret string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/vault/" + vaultname + "/secret/" + secretname

	encodedSecret := base64.StdEncoding.EncodeToString([]byte(secret))
	return putContent(apiEndpoint, encodedSecret, version, configuration)
}

func deleteContent(apiEndpoint string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {

	var versionHeader = make(map[string]string)
	versionHeader["AuroraConfigFileVersion"] = version

	response, err := serverapi.CallApiWithHeaders(versionHeader, http.MethodDelete, apiEndpoint, "", configuration)

	if err != nil {
		return "", err
	}

	if !response.Success {
		validationMessages, _ := serverapi.ResponsItems2MessageString(response)
		return validationMessages, errors.New(validationMessages)
	}

	return
}

func DeleteVault(vaultname string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/vault/" + vaultname

	return deleteContent(apiEndpoint, "", configuration)

}
<commit_msg>naming<commit_after>package auroraconfig

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"

	"github.com/skatteetaten/ao/pkg/configuration"
	"github.com/skatteetaten/ao/pkg/fileutil"
	"github.com/skatteetaten/ao/pkg/jsonutil"
	"github.com/skatteetaten/ao/pkg/serverapi"
	"fmt"
	"strings"
)

const InvalidConfigurationError = "Invalid configuration"

func GetContent(filename string, auroraConfig *serverapi.AuroraConfig) (content string, version string, err error) {
	var fileFound bool = false

	_, fileFound = auroraConfig.Files[filename]
	if fileFound {
		content = string(auroraConfig.Files[filename])
	}

	version = auroraConfig.Versions[filename]

	if !fileFound {
		return "", "", errors.New("Illegal file/folder")
	}
	return content, version, nil

}

func GetAllContent(outputFolder string, configuration *configuration.ConfigurationClass) (output string, err error) {
	auroraConfig, err := GetAuroraConfig(configuration)
	if err != nil {
		return
	}

	if outputFolder != "" {
		if fileutil.IsLegalFileFolder(outputFolder) == fileutil.SpecIllegal {
			err = errors.New("Illegal file/folder")
			return "", err

		}
		var content string
		for filename := range auroraConfig.Files {
			content = jsonutil.PrettyPrintJson(string(auroraConfig.Files[filename]))
			err = fileutil.WriteFile(outputFolder, filename, content)
			if err != nil {
				return "", err
			}
		}
	}

	outputBytes, err := json.Marshal(auroraConfig)
	output = jsonutil.PrettyPrintJson(string(outputBytes))
	return output, err

}

func GetFileList(configuration *configuration.ConfigurationClass) (filenames []string, err error) {
	auroraConfig, err := GetAuroraConfig(configuration)
	if err != nil {
		return
	}
	filenames = make([]string, len(auroraConfig.Files))

	var filenameIndex = 0
	for filename := range auroraConfig.Files {
		filenames[filenameIndex] = filename
		filenameIndex++
	}
	return filenames, nil
}

func GetAuroraConfigRequest(configuration *configuration.ConfigurationClass) (request *serverapi.Request) {
	request = new(serverapi.Request)
	request.Method = http.MethodGet
	request.ApiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"
	return request
}

func Response2AuroraConfig(response serverapi.Response) (auroraConfig serverapi.AuroraConfig, err error) {

	if !response.Success {
		output, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return auroraConfig, err
		}
		err = errors.New(output)
		return auroraConfig, err

	}

	auroraConfig, err = serverapi.ResponseItems2AuroraConfig(response)

	return auroraConfig, nil

}

func ValidateAuroraConfig(auroraConfig *serverapi.AuroraConfig, config *configuration.ConfigurationClass) error {
	payload, err := json.Marshal(auroraConfig)
	if err != nil {
		return err
	}

	endpoint := fmt.Sprintf("/affiliation/%s/auroraconfig/validate", config.GetAffiliation())
	response, err := serverapi.CallApi(http.MethodPut, endpoint, string(payload), config)
	if err != nil {
		return err
	}

	printValidationResponse(response)

	return nil
}

type Validation struct {
	IllegalFieldErrors []string
	MissingFieldErrors []string
	InvalidFieldErrors []string
	UniqueErrors       map[string]bool
}

func (v *Validation) GetAllErrors() []string {
	errorMessages := append(v.IllegalFieldErrors, v.InvalidFieldErrors...)
	return append(errorMessages, v.MissingFieldErrors...)
}

func (v *Validation) Contains(key string) bool {
	return v.UniqueErrors[key]
}

func printValidationResponse(res serverapi.Response) {
	fmt.Println(res.Message)

	if res.Success {
		return
	}

	validation := &Validation{
		UniqueErrors: make(map[string]bool),
	}

	for _, item := range res.Items {
		var res serverapi.ResponseItemError
		json.Unmarshal(item, &res)
		validation.formatValidationError(&res)
	}

	for _, e := range validation.GetAllErrors() {
		fmt.Print(e)
	}
}

func (v *Validation) formatValidationError(res *serverapi.ResponseItemError) {
	illegalFieldFormat := `
Filename:    %s
Path:        %s
Value:       %s
Message:     %s
`
	missingFieldFormat := `
Application: %s/%s
Path:        %s (Missing)
Message:     %s
`

	invalidFieldFormat := `
Filename:    %s
Path:        %s
Message:     %s
`

	for _, message := range res.Messages {
		k := []string{
			message.Field.Source,
			message.Field.Path,
			message.Field.Value,
		}
		key := strings.Join(k, "|")

		if v.Contains(key) {
			continue
		}

		if message.Type != "MISSING" {
			v.UniqueErrors[key] = true
		}

		switch message.Type {
		case "ILLEGAL":
			{
				illegal := fmt.Sprintf(illegalFieldFormat,
					message.Field.Source,
					message.Field.Path,
					message.Field.Value,
					message.Message,
				)
				v.IllegalFieldErrors = append(v.IllegalFieldErrors, illegal)
			}

		case "INVALID":
			{
				invalid := fmt.Sprintf(invalidFieldFormat,
					message.Field.Source,
					message.Field.Path,
					message.Message,
				)
				v.InvalidFieldErrors = append(v.InvalidFieldErrors, invalid);
			}

		case "MISSING":
			{
				missing := fmt.Sprintf(missingFieldFormat,
					res.Environment,
					res.Application,
					message.Field.Path,
					message.Message,
				)
				v.MissingFieldErrors = append(v.MissingFieldErrors, missing);
			}
		}
	}
}

func GetAuroraConfig(configuration *configuration.ConfigurationClass) (auroraConfig serverapi.AuroraConfig, err error) {
	var apiEndpoint string = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"

	response, err := serverapi.CallApi(http.MethodGet, apiEndpoint, "", configuration)
	if err != nil {
		return auroraConfig, err
	}

	if !response.Success {
		output, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return auroraConfig, err
		}
		err = errors.New(output)
		return auroraConfig, err

	}

	auroraConfig, err = serverapi.ResponseItems2AuroraConfig(response)

	return auroraConfig, nil
}

func PutAuroraConfig(auroraConfig serverapi.AuroraConfig, configuration *configuration.ConfigurationClass) (err error) {
	content, err := json.Marshal(auroraConfig)
	if err != nil {
		return err
	}

	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfig"

	_, err = putContent(apiEndpoint, string(content), "", configuration)
	if err != nil {
		return err
	}
	return
}

func putContent(apiEndpoint string, content string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {

	var versionHeader = make(map[string]string)
	versionHeader["AuroraConfigFileVersion"] = version

	response, err := serverapi.CallApiWithHeaders(versionHeader, http.MethodPut, apiEndpoint, content, configuration)

	if err != nil {
		validationMessages, err := serverapi.ResponsItems2MessageString(response)
		if err != nil {
			return "", err
		}
		return validationMessages, errors.New(InvalidConfigurationError + "\n" + validationMessages)
	}

	if !response.Success {
		validationMessages, _ := serverapi.ResponsItems2MessageString(response)
		return validationMessages, errors.New(InvalidConfigurationError + "\n" + validationMessages)
	}

	return
}

func PutFile(filename string, content string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/auroraconfigfile/" + filename

	return putContent(apiEndpoint, content, version, configuration)
}

func PutSecret(vaultname string, secretname string, secret string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/vault/" + vaultname + "/secret/" + secretname

	encodedSecret := base64.StdEncoding.EncodeToString([]byte(secret))
	return putContent(apiEndpoint, encodedSecret, version, configuration)
}

func deleteContent(apiEndpoint string, version string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {

	var versionHeader = make(map[string]string)
	versionHeader["AuroraConfigFileVersion"] = version

	response, err := serverapi.CallApiWithHeaders(versionHeader, http.MethodDelete, apiEndpoint, "", configuration)

	if err != nil {
		return "", err
	}

	if !response.Success {
		validationMessages, _ := serverapi.ResponsItems2MessageString(response)
		return validationMessages, errors.New(validationMessages)
	}

	return
}

func DeleteVault(vaultname string, configuration *configuration.ConfigurationClass) (validationMessages string, err error) {
	var apiEndpoint = "/affiliation/" + configuration.GetAffiliation() + "/vault/" + vaultname

	return deleteContent(apiEndpoint, "", configuration)

}
<|endoftext|>


<commit_before>package proxy

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/datawire/dlib/dlog"
	"golang.org/x/net/proxy"
)

// A Proxy listens to a port and forwards incoming connections to a router
type Proxy struct {
	listener net.Listener
	router   func(*net.TCPConn) (string, error)
}

// NewProxy returns a new Proxy instance that is listening to the given tcp address
func NewProxy(c context.Context, address string, router func(*net.TCPConn) (string, error)) (proxy *Proxy, err error) {
	setRlimit(c)
	ln, err := net.Listen("tcp", address)
	if err == nil {
		proxy = &Proxy{ln, router}
	}
	return
}

func setRlimit(c context.Context) {
	var rLimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "error getting rlimit: %s", err.Error())
	} else {
		dlog.Debugf(c, "initial rlimit: %d", rLimit)
	}

	rLimit.Max = 999999
	rLimit.Cur = 999999
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "Error setting rlimit: %s", err.Error())
	}

	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "Error getting rlimit: %s", err.Error())
	} else {
		dlog.Debugf(c, "Final rlimit: %d", rLimit)
	}
}

// Run starts the proxy accept loop and runs it until the context is cancelled
func (pxy *Proxy) Run(c context.Context, limit int32) {
	dlog.Debugf(c, "listening limit=%v", limit)
	connQueue := make(chan net.Conn, limit)
	capacity := limit
	closing := false
	go func() {
		for {
			conn, err := pxy.listener.Accept()
			if err != nil {
				if closing {
					return
				}
				dlog.Error(c, err.Error())
			} else {
				atomic.AddInt32(&capacity, -1)
				connQueue <- conn
			}
		}
	}()
	for {
		select {
		case <-c.Done():
			closing = true
			_ = pxy.listener.Close()
			return
		case conn := <-connQueue:
			cpy := atomic.AddInt32(&capacity, 1)
			switch conn := conn.(type) {
			case *net.TCPConn:
				dlog.Debugf(c, "CAPACITY: %v", cpy)
				pxy.handleConnection(c, conn)
			default:
				dlog.Errorf(c, "unknown connection type: %v", conn)
			}
		}
	}
}

func (pxy *Proxy) handleConnection(c context.Context, conn *net.TCPConn) {
	host, err := pxy.router(conn)
	if err != nil {
		dlog.Errorf(c, "router error: %v", err)
		return
	}

	dlog.Debugf(c, "CONNECT %s %s", conn.RemoteAddr(), host)

	// setting up an ssh tunnel with dynamic socks proxy at this end
	// seems faster than connecting directly to a socks proxy
	dialer, err := proxy.SOCKS5("tcp", "localhost:1080", nil, proxy.Direct)
	//	dialer, err := proxy.SOCKS5("tcp", "localhost:9050", nil, proxy.Direct)
	if err != nil {
		dlog.Error(c, err.Error())
		conn.Close()
		return
	}

	tc, _ := context.WithTimeout(c, 5*time.Second)
	_proxy, err := dialer.(proxy.ContextDialer).DialContext(tc, "tcp", host)
	if err != nil {
		dlog.Error(c, err.Error())
		conn.Close()
		return
	}
	px := _proxy.(*net.TCPConn)

	done := sync.WaitGroup{}
	done.Add(2)
	go pxy.pipe(c, conn, px, &done)
	go pxy.pipe(c, px, conn, &done)

	done.Wait()
}

func (pxy *Proxy) pipe(c context.Context, from, to *net.TCPConn, done *sync.WaitGroup) {
	defer done.Done()

	closed := int32(0)
	closePipe := func() {
		if atomic.CompareAndSwapInt32(&closed, 0, 1) {
			dlog.Debugf(c, "CLOSED WRITE %v", to.RemoteAddr())
			_ = to.CloseWrite()
			dlog.Debugf(c, "CLOSED READ %v", from.RemoteAddr())
			_ = from.CloseRead()
		}
	}
	defer closePipe()

	// Close pipes when context is done
	eop := make(chan bool)
	go func() {
		select {
		case <-eop:
			// just end this goroutine
		case <-c.Done():
			// close the pipe
			closePipe()
		}
	}()

	const size = 64 * 1024
	var buf [size]byte
	for {
		n, err := from.Read(buf[0:size])
		if err != nil {
			if err != io.EOF {
				dlog.Error(c, err.Error())
			}
			break
		}
		_, err = to.Write(buf[0:n])
		if err != nil {
			dlog.Error(c, err.Error())
			break
		}
	}
	close(eop)
}
<commit_msg>Fix interface confusion when dealing with proxy connection.<commit_after>package proxy

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/datawire/dlib/dlog"
	"golang.org/x/net/proxy"
)

// A Proxy listens to a port and forwards incoming connections to a router
type Proxy struct {
	listener net.Listener
	router   func(*net.TCPConn) (string, error)
}

// NewProxy returns a new Proxy instance that is listening to the given tcp address
func NewProxy(c context.Context, address string, router func(*net.TCPConn) (string, error)) (proxy *Proxy, err error) {
	setRlimit(c)
	ln, err := net.Listen("tcp", address)
	if err == nil {
		proxy = &Proxy{ln, router}
	}
	return
}

func setRlimit(c context.Context) {
	var rLimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "error getting rlimit: %s", err.Error())
	} else {
		dlog.Debugf(c, "initial rlimit: %d", rLimit)
	}

	rLimit.Max = 999999
	rLimit.Cur = 999999
	err = syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "Error setting rlimit: %s", err.Error())
	}

	err = syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit)
	if err != nil {
		dlog.Errorf(c, "Error getting rlimit: %s", err.Error())
	} else {
		dlog.Debugf(c, "Final rlimit: %d", rLimit)
	}
}

// Run starts the proxy accept loop and runs it until the context is cancelled
func (pxy *Proxy) Run(c context.Context, limit int32) {
	dlog.Debugf(c, "listening limit=%v", limit)
	connQueue := make(chan net.Conn, limit)
	capacity := limit
	closing := false
	go func() {
		for {
			conn, err := pxy.listener.Accept()
			if err != nil {
				if closing {
					return
				}
				dlog.Error(c, err.Error())
			} else {
				atomic.AddInt32(&capacity, -1)
				connQueue <- conn
			}
		}
	}()
	for {
		select {
		case <-c.Done():
			closing = true
			_ = pxy.listener.Close()
			return
		case conn := <-connQueue:
			cpy := atomic.AddInt32(&capacity, 1)
			switch conn := conn.(type) {
			case *net.TCPConn:
				dlog.Debugf(c, "CAPACITY: %v", cpy)
				pxy.handleConnection(c, conn)
			default:
				dlog.Errorf(c, "unknown connection type: %v", conn)
			}
		}
	}
}

func (pxy *Proxy) handleConnection(c context.Context, conn *net.TCPConn) {
	host, err := pxy.router(conn)
	if err != nil {
		dlog.Errorf(c, "router error: %v", err)
		return
	}

	dlog.Debugf(c, "CONNECT %s %s", conn.RemoteAddr(), host)

	// setting up an ssh tunnel with dynamic socks proxy at this end
	// seems faster than connecting directly to a socks proxy
	dialer, err := proxy.SOCKS5("tcp", "localhost:1080", nil, proxy.Direct)
	//	dialer, err := proxy.SOCKS5("tcp", "localhost:9050", nil, proxy.Direct)
	if err != nil {
		dlog.Error(c, err.Error())
		conn.Close()
		return
	}

	tc, _ := context.WithTimeout(c, 5*time.Second)
	px, err := dialer.(proxy.ContextDialer).DialContext(tc, "tcp", host)
	if err != nil {
		dlog.Error(c, err.Error())
		conn.Close()
		return
	}

	done := sync.WaitGroup{}
	done.Add(2)
	go pxy.pipe(c, conn, px, &done)
	go pxy.pipe(c, px, conn, &done)

	done.Wait()
}

func (pxy *Proxy) pipe(c context.Context, from, to net.Conn, done *sync.WaitGroup) {
	defer done.Done()

	closed := int32(0)
	closePipe := func() {
		if atomic.CompareAndSwapInt32(&closed, 0, 1) {
			dlog.Debugf(c, "CLOSED %v -> %v", from.LocalAddr(), from.RemoteAddr())
			_ = from.Close()
		}
	}
	defer closePipe()

	// Close pipes when context is done
	eop := make(chan bool)
	go func() {
		select {
		case <-eop:
			// just end this goroutine
		case <-c.Done():
			// close the pipe
			closePipe()
		}
	}()

	const size = 64 * 1024
	var buf [size]byte
	for {
		n, err := from.Read(buf[0:size])
		if err != nil {
			if err != io.EOF {
				dlog.Error(c, err.Error())
			}
			break
		}
		_, err = to.Write(buf[0:n])
		if err != nil {
			dlog.Error(c, err.Error())
			break
		}
	}
	close(eop)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/facebookgo/inject"
	"github.com/patrickmn/go-cache"
	"golang.org/x/sync/errgroup"

	"github.com/grafana/grafana/pkg/api"
	"github.com/grafana/grafana/pkg/api/routing"
	"github.com/grafana/grafana/pkg/bus"
	_ "github.com/grafana/grafana/pkg/extensions"
	"github.com/grafana/grafana/pkg/infra/log"
	_ "github.com/grafana/grafana/pkg/infra/metrics"
	_ "github.com/grafana/grafana/pkg/infra/remotecache"
	_ "github.com/grafana/grafana/pkg/infra/serverlock"
	_ "github.com/grafana/grafana/pkg/infra/tracing"
	_ "github.com/grafana/grafana/pkg/infra/usagestats"
	"github.com/grafana/grafana/pkg/login"
	"github.com/grafana/grafana/pkg/login/social"
	"github.com/grafana/grafana/pkg/middleware"
	_ "github.com/grafana/grafana/pkg/plugins"
	"github.com/grafana/grafana/pkg/registry"
	_ "github.com/grafana/grafana/pkg/services/alerting"
	_ "github.com/grafana/grafana/pkg/services/auth"
	_ "github.com/grafana/grafana/pkg/services/cleanup"
	_ "github.com/grafana/grafana/pkg/services/notifications"
	_ "github.com/grafana/grafana/pkg/services/provisioning"
	_ "github.com/grafana/grafana/pkg/services/rendering"
	_ "github.com/grafana/grafana/pkg/services/search"
	_ "github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
)

func NewGrafanaServer() *GrafanaServerImpl {
	rootCtx, shutdownFn := context.WithCancel(context.Background())
	childRoutines, childCtx := errgroup.WithContext(rootCtx)

	return &GrafanaServerImpl{
		context:       childCtx,
		shutdownFn:    shutdownFn,
		childRoutines: childRoutines,
		log:           log.New("server"),
		cfg:           setting.NewCfg(),
	}
}

type GrafanaServerImpl struct {
	context            context.Context
	shutdownFn         context.CancelFunc
	childRoutines      *errgroup.Group
	log                log.Logger
	cfg                *setting.Cfg
	shutdownReason     string
	shutdownInProgress bool

	RouteRegister routing.RouteRegister `inject:""`
	HttpServer    *api.HTTPServer       `inject:""`
}

func (g *GrafanaServerImpl) Run() error {
	var err error
	g.loadConfiguration()
	g.writePIDFile()

	login.Init()
	social.NewOAuthService()

	serviceGraph := inject.Graph{}
	err = serviceGraph.Provide(&inject.Object{Value: bus.GetBus()})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: g.cfg})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: routing.NewRouteRegister(middleware.RequestMetrics, middleware.RequestTracing)})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: cache.New(5*time.Minute, 10*time.Minute)})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}

	// self registered services
	services := registry.GetServices()

	// Add all services to dependency graph
	for _, service := range services {
		err = serviceGraph.Provide(&inject.Object{Value: service.Instance})
		if err != nil {
			return fmt.Errorf("Failed to provide object to the graph: %v", err)
		}
	}

	err = serviceGraph.Provide(&inject.Object{Value: g})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}

	// Inject dependencies to services
	if err := serviceGraph.Populate(); err != nil {
		return fmt.Errorf("Failed to populate service dependency: %v", err)
	}

	// Init & start services
	for _, service := range services {
		if registry.IsDisabled(service.Instance) {
			continue
		}

		g.log.Info("Initializing " + service.Name)

		if err := service.Instance.Init(); err != nil {
			return fmt.Errorf("Service init failed: %v", err)
		}
	}

	// Start background services
	for _, srv := range services {
		// variable needed for accessing loop variable in function callback
		descriptor := srv
		service, ok := srv.Instance.(registry.BackgroundService)
		if !ok {
			continue
		}

		if registry.IsDisabled(descriptor.Instance) {
			continue
		}

		g.childRoutines.Go(func() error {
			// Skip starting new service when shutting down
			// Can happen when service stop/return during startup
			if g.shutdownInProgress {
				return nil
			}

			err := service.Run(g.context)

			// If error is not canceled then the service crashed
			if err != context.Canceled && err != nil {
				g.log.Error("Stopped "+descriptor.Name, "reason", err)
			} else {
				g.log.Info("Stopped "+descriptor.Name, "reason", err)
			}

			// Mark that we are in shutdown mode
			// So more services are not started
			g.shutdownInProgress = true
			return err
		})
	}

	sendSystemdNotification("READY=1")

	return g.childRoutines.Wait()
}

func (g *GrafanaServerImpl) loadConfiguration() {
	err := g.cfg.Load(&setting.CommandLineArgs{
		Config:   *configFile,
		HomePath: *homePath,
		Args:     flag.Args(),
	})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start grafana. error: %s\n", err.Error())
		os.Exit(1)
	}

	g.log.Info("Starting "+setting.ApplicationName, "version", version, "commit", commit, "branch", buildBranch, "compiled", time.Unix(setting.BuildStamp, 0))
	g.cfg.LogConfigSources()
}

func (g *GrafanaServerImpl) Shutdown(reason string) {
	g.log.Info("Shutdown started", "reason", reason)
	g.shutdownReason = reason
	g.shutdownInProgress = true

	// call cancel func on root context
	g.shutdownFn()

	// wait for child routines
	g.childRoutines.Wait()
}

func (g *GrafanaServerImpl) Exit(reason error) int {
	// default exit code is 1
	code := 1

	if reason == context.Canceled && g.shutdownReason != "" {
		reason = fmt.Errorf(g.shutdownReason)
		code = 0
	}

	g.log.Error("Server shutdown", "reason", reason)
	return code
}

func (g *GrafanaServerImpl) writePIDFile() {
	if *pidFile == "" {
		return
	}

	// Ensure the required directory structure exists.
	err := os.MkdirAll(filepath.Dir(*pidFile), 0700)
	if err != nil {
		g.log.Error("Failed to verify pid directory", "error", err)
		os.Exit(1)
	}

	// Retrieve the PID and write it.
	pid := strconv.Itoa(os.Getpid())
	if err := ioutil.WriteFile(*pidFile, []byte(pid), 0644); err != nil {
		g.log.Error("Failed to write pidfile", "error", err)
		os.Exit(1)
	}

	g.log.Info("Writing PID file", "path", *pidFile, "pid", pid)
}

func sendSystemdNotification(state string) error {
	notifySocket := os.Getenv("NOTIFY_SOCKET")

	if notifySocket == "" {
		return fmt.Errorf("NOTIFY_SOCKET environment variable empty or unset")
	}

	socketAddr := &net.UnixAddr{
		Name: notifySocket,
		Net:  "unixgram",
	}

	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)

	if err != nil {
		return err
	}

	_, err = conn.Write([]byte(state))

	conn.Close()

	return err
}
<commit_msg>Fix so that correct cache is provided to di registry (#17566)<commit_after>package main

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/facebookgo/inject"
	"golang.org/x/sync/errgroup"

	"github.com/grafana/grafana/pkg/api"
	"github.com/grafana/grafana/pkg/api/routing"
	"github.com/grafana/grafana/pkg/bus"
	_ "github.com/grafana/grafana/pkg/extensions"
	"github.com/grafana/grafana/pkg/infra/localcache"
	"github.com/grafana/grafana/pkg/infra/log"
	_ "github.com/grafana/grafana/pkg/infra/metrics"
	_ "github.com/grafana/grafana/pkg/infra/remotecache"
	_ "github.com/grafana/grafana/pkg/infra/serverlock"
	_ "github.com/grafana/grafana/pkg/infra/tracing"
	_ "github.com/grafana/grafana/pkg/infra/usagestats"
	"github.com/grafana/grafana/pkg/login"
	"github.com/grafana/grafana/pkg/login/social"
	"github.com/grafana/grafana/pkg/middleware"
	_ "github.com/grafana/grafana/pkg/plugins"
	"github.com/grafana/grafana/pkg/registry"
	_ "github.com/grafana/grafana/pkg/services/alerting"
	_ "github.com/grafana/grafana/pkg/services/auth"
	_ "github.com/grafana/grafana/pkg/services/cleanup"
	_ "github.com/grafana/grafana/pkg/services/notifications"
	_ "github.com/grafana/grafana/pkg/services/provisioning"
	_ "github.com/grafana/grafana/pkg/services/rendering"
	_ "github.com/grafana/grafana/pkg/services/search"
	_ "github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
)

func NewGrafanaServer() *GrafanaServerImpl {
	rootCtx, shutdownFn := context.WithCancel(context.Background())
	childRoutines, childCtx := errgroup.WithContext(rootCtx)

	return &GrafanaServerImpl{
		context:       childCtx,
		shutdownFn:    shutdownFn,
		childRoutines: childRoutines,
		log:           log.New("server"),
		cfg:           setting.NewCfg(),
	}
}

type GrafanaServerImpl struct {
	context            context.Context
	shutdownFn         context.CancelFunc
	childRoutines      *errgroup.Group
	log                log.Logger
	cfg                *setting.Cfg
	shutdownReason     string
	shutdownInProgress bool

	RouteRegister routing.RouteRegister `inject:""`
	HttpServer    *api.HTTPServer       `inject:""`
}

func (g *GrafanaServerImpl) Run() error {
	var err error
	g.loadConfiguration()
	g.writePIDFile()

	login.Init()
	social.NewOAuthService()

	serviceGraph := inject.Graph{}
	err = serviceGraph.Provide(&inject.Object{Value: bus.GetBus()})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: g.cfg})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: routing.NewRouteRegister(middleware.RequestMetrics, middleware.RequestTracing)})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}
	err = serviceGraph.Provide(&inject.Object{Value: localcache.New(5*time.Minute, 10*time.Minute)})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}

	// self registered services
	services := registry.GetServices()

	// Add all services to dependency graph
	for _, service := range services {
		err = serviceGraph.Provide(&inject.Object{Value: service.Instance})
		if err != nil {
			return fmt.Errorf("Failed to provide object to the graph: %v", err)
		}
	}

	err = serviceGraph.Provide(&inject.Object{Value: g})
	if err != nil {
		return fmt.Errorf("Failed to provide object to the graph: %v", err)
	}

	// Inject dependencies to services
	if err := serviceGraph.Populate(); err != nil {
		return fmt.Errorf("Failed to populate service dependency: %v", err)
	}

	// Init & start services
	for _, service := range services {
		if registry.IsDisabled(service.Instance) {
			continue
		}

		g.log.Info("Initializing " + service.Name)

		if err := service.Instance.Init(); err != nil {
			return fmt.Errorf("Service init failed: %v", err)
		}
	}

	// Start background services
	for _, srv := range services {
		// variable needed for accessing loop variable in function callback
		descriptor := srv
		service, ok := srv.Instance.(registry.BackgroundService)
		if !ok {
			continue
		}

		if registry.IsDisabled(descriptor.Instance) {
			continue
		}

		g.childRoutines.Go(func() error {
			// Skip starting new service when shutting down
			// Can happen when service stop/return during startup
			if g.shutdownInProgress {
				return nil
			}

			err := service.Run(g.context)

			// If error is not canceled then the service crashed
			if err != context.Canceled && err != nil {
				g.log.Error("Stopped "+descriptor.Name, "reason", err)
			} else {
				g.log.Info("Stopped "+descriptor.Name, "reason", err)
			}

			// Mark that we are in shutdown mode
			// So more services are not started
			g.shutdownInProgress = true
			return err
		})
	}

	sendSystemdNotification("READY=1")

	return g.childRoutines.Wait()
}

func (g *GrafanaServerImpl) loadConfiguration() {
	err := g.cfg.Load(&setting.CommandLineArgs{
		Config:   *configFile,
		HomePath: *homePath,
		Args:     flag.Args(),
	})

	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to start grafana. error: %s\n", err.Error())
		os.Exit(1)
	}

	g.log.Info("Starting "+setting.ApplicationName, "version", version, "commit", commit, "branch", buildBranch, "compiled", time.Unix(setting.BuildStamp, 0))
	g.cfg.LogConfigSources()
}

func (g *GrafanaServerImpl) Shutdown(reason string) {
	g.log.Info("Shutdown started", "reason", reason)
	g.shutdownReason = reason
	g.shutdownInProgress = true

	// call cancel func on root context
	g.shutdownFn()

	// wait for child routines
	g.childRoutines.Wait()
}

func (g *GrafanaServerImpl) Exit(reason error) int {
	// default exit code is 1
	code := 1

	if reason == context.Canceled && g.shutdownReason != "" {
		reason = fmt.Errorf(g.shutdownReason)
		code = 0
	}

	g.log.Error("Server shutdown", "reason", reason)
	return code
}

func (g *GrafanaServerImpl) writePIDFile() {
	if *pidFile == "" {
		return
	}

	// Ensure the required directory structure exists.
	err := os.MkdirAll(filepath.Dir(*pidFile), 0700)
	if err != nil {
		g.log.Error("Failed to verify pid directory", "error", err)
		os.Exit(1)
	}

	// Retrieve the PID and write it.
	pid := strconv.Itoa(os.Getpid())
	if err := ioutil.WriteFile(*pidFile, []byte(pid), 0644); err != nil {
		g.log.Error("Failed to write pidfile", "error", err)
		os.Exit(1)
	}

	g.log.Info("Writing PID file", "path", *pidFile, "pid", pid)
}

func sendSystemdNotification(state string) error {
	notifySocket := os.Getenv("NOTIFY_SOCKET")

	if notifySocket == "" {
		return fmt.Errorf("NOTIFY_SOCKET environment variable empty or unset")
	}

	socketAddr := &net.UnixAddr{
		Name: notifySocket,
		Net:  "unixgram",
	}

	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)

	if err != nil {
		return err
	}

	_, err = conn.Write([]byte(state))

	conn.Close()

	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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 download

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"

	"cloud.google.com/go/storage"
	"google.golang.org/api/option"

	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/style"
)

const (
	// PreloadVersion is the current version of the preloaded tarball
	//
	// NOTE: You may need to bump this version up when upgrading auxiliary docker images
	PreloadVersion = "v11"
	// PreloadBucket is the name of the GCS bucket where preloaded volume tarballs exist
	PreloadBucket = "minikube-preloaded-volume-tarballs"
)

// Enumeration for preload existence cache.
const (
	preloadUnknown = iota // Value when preload status has not been checked.
	preloadMissing = iota // Value when preload has been checked and is missing.
	preloadPresent = iota // Value when preload has been checked and is present.
)

var (
	preloadState int = preloadUnknown
)

// TarballName returns name of the tarball
func TarballName(k8sVersion, containerRuntime string) string {
	if containerRuntime == "crio" {
		containerRuntime = "cri-o"
	}
	var storageDriver string
	if containerRuntime == "cri-o" {
		storageDriver = "overlay"
	} else {
		storageDriver = "overlay2"
	}
	return fmt.Sprintf("preloaded-images-k8s-%s-%s-%s-%s-%s.tar.lz4", PreloadVersion, k8sVersion, containerRuntime, storageDriver, runtime.GOARCH)
}

// returns the name of the checksum file
func checksumName(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("%s.checksum", TarballName(k8sVersion, containerRuntime))
}

// returns target dir for all cached items related to preloading
func targetDir() string {
	return localpath.MakeMiniPath("cache", "preloaded-tarball")
}

// PreloadChecksumPath returns the local path to the cached checksum file
func PreloadChecksumPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), checksumName(k8sVersion, containerRuntime))
}

// TarballPath returns the local path to the cached preload tarball
func TarballPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), TarballName(k8sVersion, containerRuntime))
}

// remoteTarballURL returns the URL for the remote tarball in GCS
func remoteTarballURL(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("https://storage.googleapis.com/%s/%s", PreloadBucket, TarballName(k8sVersion, containerRuntime))
}

// PreloadExists returns true if there is a preloaded tarball that can be used
func PreloadExists(k8sVersion, containerRuntime string, forcePreload ...bool) bool {
	// TODO (#8166): Get rid of the need for this and viper at all
	force := false
	if len(forcePreload) > 0 {
		force = forcePreload[0]
	}

	// TODO: debug why this func is being called two times
	klog.Infof("Checking if preload exists for k8s version %s and runtime %s", k8sVersion, containerRuntime)
	if !viper.GetBool("preload") && !force {
		return false
	}

	// If the preload existence is cached, just return that value.
	if preloadState != preloadUnknown {
		return preloadState == preloadPresent
	}

	// Omit remote check if tarball exists locally
	targetPath := TarballPath(k8sVersion, containerRuntime)
	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found local preload: %s", targetPath)
		preloadState = preloadPresent
		return true
	}

	url := remoteTarballURL(k8sVersion, containerRuntime)
	resp, err := http.Head(url)
	if err != nil {
		klog.Warningf("%s fetch error: %v", url, err)
		preloadState = preloadMissing
		return false
	}

	// note: err won't be set if it's a 404
	if resp.StatusCode != 200 {
		klog.Warningf("%s status code: %d", url, resp.StatusCode)
		preloadState = preloadMissing
		return false
	}

	klog.Infof("Found remote preload: %s", url)
	preloadState = preloadPresent
	return true
}

var checkPreloadExists = PreloadExists

// Preload caches the preloaded images tarball on the host machine
func Preload(k8sVersion, containerRuntime string) error {
	targetPath := TarballPath(k8sVersion, containerRuntime)
	targetLock := targetPath + ".lock"

	releaser, err := lockDownload(targetLock)
	if releaser != nil {
		defer releaser.Release()
	}
	if err != nil {
		return err
	}

	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found %s in cache, skipping download", targetPath)
		return nil
	}

	// Make sure we support this k8s version
	if !checkPreloadExists(k8sVersion, containerRuntime) {
		klog.Infof("Preloaded tarball for k8s version %s does not exist", k8sVersion)
		return nil
	}

	out.Step(style.FileDownload, "Downloading Kubernetes {{.version}} preload ...", out.V{"version": k8sVersion})
	url := remoteTarballURL(k8sVersion, containerRuntime)

	checksum, err := getChecksum(k8sVersion, containerRuntime)
	var realPath string
	if err != nil {
		klog.Warningf("No checksum for preloaded tarball for k8s version %s: %v", k8sVersion, err)
		realPath = targetPath
		tmp, err := ioutil.TempFile(targetDir(), TarballName(k8sVersion, containerRuntime)+".*")
		if err != nil {
			return errors.Wrap(err, "tempfile")
		}
		targetPath = tmp.Name()
	} else if checksum != "" {
		url += "?checksum=" + checksum
	}

	if err := download(url, targetPath); err != nil {
		return errors.Wrapf(err, "download failed: %s", url)
	}

	if err := ensureChecksumValid(k8sVersion, containerRuntime, targetPath); err != nil {
		return err
	}

	if realPath != "" {
		klog.Infof("renaming tempfile to %s ...", TarballName(k8sVersion, containerRuntime))
		err := os.Rename(targetPath, realPath)
		if err != nil {
			return errors.Wrap(err, "rename")
		}
	}

	// If the download was successful, mark off that the preload exists in the cache.
	preloadState = preloadPresent
	return nil
}

func getStorageAttrs(name string) (*storage.ObjectAttrs, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	attrs, err := client.Bucket(PreloadBucket).Object(name).Attrs(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage object")
	}
	return attrs, nil
}

var getChecksum = func(k8sVersion, containerRuntime string) (string, error) {
	klog.Infof("getting checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return "", err
	}
	md5 := hex.EncodeToString(attrs.MD5)
	return fmt.Sprintf("md5:%s", md5), nil
}

func saveChecksumFile(k8sVersion, containerRuntime string) error {
	klog.Infof("saving checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return err
	}
	checksum := attrs.MD5
	return ioutil.WriteFile(PreloadChecksumPath(k8sVersion, containerRuntime), checksum, 0o644)
}

// verifyChecksum returns true if the checksum of the local binary matches
// the checksum of the remote binary
func verifyChecksum(k8sVersion, containerRuntime, path string) error {
	klog.Infof("verifying checksumm of %s ...", path)
	// get md5 checksum of tarball path
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		return errors.Wrap(err, "reading tarball")
	}
	checksum := md5.Sum(contents)

	remoteChecksum, err := ioutil.ReadFile(PreloadChecksumPath(k8sVersion, containerRuntime))
	if err != nil {
		return errors.Wrap(err, "reading checksum file")
	}

	// create a slice of checksum, which is [16]byte
	if string(remoteChecksum) != string(checksum[:]) {
		return fmt.Errorf("checksum of %s does not match remote checksum (%s != %s)", path, string(remoteChecksum), string(checksum[:]))
	}
	return nil
}

// ensureChecksumValid saves and verifies local binary checksum matches remote binary checksum
var ensureChecksumValid = func(k8sVersion, containerRuntime, targetPath string) error {
	if err := saveChecksumFile(k8sVersion, containerRuntime); err != nil {
		return errors.Wrap(err, "saving checksum file")
	}

	if err := verifyChecksum(k8sVersion, containerRuntime, targetPath); err != nil {
		return errors.Wrap(err, "verify")
	}

	return nil
}
<commit_msg>Remove iota from all except first preload enum.<commit_after>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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 download

import (
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"

	"cloud.google.com/go/storage"
	"google.golang.org/api/option"

	"github.com/pkg/errors"
	"github.com/spf13/viper"
	"k8s.io/klog/v2"
	"k8s.io/minikube/pkg/minikube/localpath"
	"k8s.io/minikube/pkg/minikube/out"
	"k8s.io/minikube/pkg/minikube/style"
)

const (
	// PreloadVersion is the current version of the preloaded tarball
	//
	// NOTE: You may need to bump this version up when upgrading auxiliary docker images
	PreloadVersion = "v11"
	// PreloadBucket is the name of the GCS bucket where preloaded volume tarballs exist
	PreloadBucket = "minikube-preloaded-volume-tarballs"
)

// Enumeration for preload existence cache.
const (
	preloadUnknown = iota // Value when preload status has not been checked.
	preloadMissing        // Value when preload has been checked and is missing.
	preloadPresent        // Value when preload has been checked and is present.
)

var (
	preloadState int = preloadUnknown
)

// TarballName returns name of the tarball
func TarballName(k8sVersion, containerRuntime string) string {
	if containerRuntime == "crio" {
		containerRuntime = "cri-o"
	}
	var storageDriver string
	if containerRuntime == "cri-o" {
		storageDriver = "overlay"
	} else {
		storageDriver = "overlay2"
	}
	return fmt.Sprintf("preloaded-images-k8s-%s-%s-%s-%s-%s.tar.lz4", PreloadVersion, k8sVersion, containerRuntime, storageDriver, runtime.GOARCH)
}

// returns the name of the checksum file
func checksumName(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("%s.checksum", TarballName(k8sVersion, containerRuntime))
}

// returns target dir for all cached items related to preloading
func targetDir() string {
	return localpath.MakeMiniPath("cache", "preloaded-tarball")
}

// PreloadChecksumPath returns the local path to the cached checksum file
func PreloadChecksumPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), checksumName(k8sVersion, containerRuntime))
}

// TarballPath returns the local path to the cached preload tarball
func TarballPath(k8sVersion, containerRuntime string) string {
	return filepath.Join(targetDir(), TarballName(k8sVersion, containerRuntime))
}

// remoteTarballURL returns the URL for the remote tarball in GCS
func remoteTarballURL(k8sVersion, containerRuntime string) string {
	return fmt.Sprintf("https://storage.googleapis.com/%s/%s", PreloadBucket, TarballName(k8sVersion, containerRuntime))
}

// PreloadExists returns true if there is a preloaded tarball that can be used
func PreloadExists(k8sVersion, containerRuntime string, forcePreload ...bool) bool {
	// TODO (#8166): Get rid of the need for this and viper at all
	force := false
	if len(forcePreload) > 0 {
		force = forcePreload[0]
	}

	// TODO: debug why this func is being called two times
	klog.Infof("Checking if preload exists for k8s version %s and runtime %s", k8sVersion, containerRuntime)
	if !viper.GetBool("preload") && !force {
		return false
	}

	// If the preload existence is cached, just return that value.
	if preloadState != preloadUnknown {
		return preloadState == preloadPresent
	}

	// Omit remote check if tarball exists locally
	targetPath := TarballPath(k8sVersion, containerRuntime)
	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found local preload: %s", targetPath)
		preloadState = preloadPresent
		return true
	}

	url := remoteTarballURL(k8sVersion, containerRuntime)
	resp, err := http.Head(url)
	if err != nil {
		klog.Warningf("%s fetch error: %v", url, err)
		preloadState = preloadMissing
		return false
	}

	// note: err won't be set if it's a 404
	if resp.StatusCode != 200 {
		klog.Warningf("%s status code: %d", url, resp.StatusCode)
		preloadState = preloadMissing
		return false
	}

	klog.Infof("Found remote preload: %s", url)
	preloadState = preloadPresent
	return true
}

var checkPreloadExists = PreloadExists

// Preload caches the preloaded images tarball on the host machine
func Preload(k8sVersion, containerRuntime string) error {
	targetPath := TarballPath(k8sVersion, containerRuntime)
	targetLock := targetPath + ".lock"

	releaser, err := lockDownload(targetLock)
	if releaser != nil {
		defer releaser.Release()
	}
	if err != nil {
		return err
	}

	if _, err := checkCache(targetPath); err == nil {
		klog.Infof("Found %s in cache, skipping download", targetPath)
		return nil
	}

	// Make sure we support this k8s version
	if !checkPreloadExists(k8sVersion, containerRuntime) {
		klog.Infof("Preloaded tarball for k8s version %s does not exist", k8sVersion)
		return nil
	}

	out.Step(style.FileDownload, "Downloading Kubernetes {{.version}} preload ...", out.V{"version": k8sVersion})
	url := remoteTarballURL(k8sVersion, containerRuntime)

	checksum, err := getChecksum(k8sVersion, containerRuntime)
	var realPath string
	if err != nil {
		klog.Warningf("No checksum for preloaded tarball for k8s version %s: %v", k8sVersion, err)
		realPath = targetPath
		tmp, err := ioutil.TempFile(targetDir(), TarballName(k8sVersion, containerRuntime)+".*")
		if err != nil {
			return errors.Wrap(err, "tempfile")
		}
		targetPath = tmp.Name()
	} else if checksum != "" {
		url += "?checksum=" + checksum
	}

	if err := download(url, targetPath); err != nil {
		return errors.Wrapf(err, "download failed: %s", url)
	}

	if err := ensureChecksumValid(k8sVersion, containerRuntime, targetPath); err != nil {
		return err
	}

	if realPath != "" {
		klog.Infof("renaming tempfile to %s ...", TarballName(k8sVersion, containerRuntime))
		err := os.Rename(targetPath, realPath)
		if err != nil {
			return errors.Wrap(err, "rename")
		}
	}

	// If the download was successful, mark off that the preload exists in the cache.
	preloadState = preloadPresent
	return nil
}

func getStorageAttrs(name string) (*storage.ObjectAttrs, error) {
	ctx := context.Background()
	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	attrs, err := client.Bucket(PreloadBucket).Object(name).Attrs(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage object")
	}
	return attrs, nil
}

var getChecksum = func(k8sVersion, containerRuntime string) (string, error) {
	klog.Infof("getting checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return "", err
	}
	md5 := hex.EncodeToString(attrs.MD5)
	return fmt.Sprintf("md5:%s", md5), nil
}

func saveChecksumFile(k8sVersion, containerRuntime string) error {
	klog.Infof("saving checksum for %s ...", TarballName(k8sVersion, containerRuntime))
	attrs, err := getStorageAttrs(TarballName(k8sVersion, containerRuntime))
	if err != nil {
		return err
	}
	checksum := attrs.MD5
	return ioutil.WriteFile(PreloadChecksumPath(k8sVersion, containerRuntime), checksum, 0o644)
}

// verifyChecksum returns true if the checksum of the local binary matches
// the checksum of the remote binary
func verifyChecksum(k8sVersion, containerRuntime, path string) error {
	klog.Infof("verifying checksumm of %s ...", path)
	// get md5 checksum of tarball path
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		return errors.Wrap(err, "reading tarball")
	}
	checksum := md5.Sum(contents)

	remoteChecksum, err := ioutil.ReadFile(PreloadChecksumPath(k8sVersion, containerRuntime))
	if err != nil {
		return errors.Wrap(err, "reading checksum file")
	}

	// create a slice of checksum, which is [16]byte
	if string(remoteChecksum) != string(checksum[:]) {
		return fmt.Errorf("checksum of %s does not match remote checksum (%s != %s)", path, string(remoteChecksum), string(checksum[:]))
	}
	return nil
}

// ensureChecksumValid saves and verifies local binary checksum matches remote binary checksum
var ensureChecksumValid = func(k8sVersion, containerRuntime, targetPath string) error {
	if err := saveChecksumFile(k8sVersion, containerRuntime); err != nil {
		return errors.Wrap(err, "saving checksum file")
	}

	if err := verifyChecksum(k8sVersion, containerRuntime, targetPath); err != nil {
		return errors.Wrap(err, "verify")
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 observability sets up and configures observability tools.
package observability

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

	"github.com/google/exposure-notifications-server/pkg/logging"
	"go.uber.org/zap"
	"google.golang.org/api/option"

	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
)

var _ Exporter = (*stackdriverExporter)(nil)

type stackdriverExporter struct {
	exporter *stackdriver.Exporter
	config   *StackdriverConfig
}

type debugRoundTripper struct {
	logger *zap.SugaredLogger
	http.RoundTripper
}

func (r *debugRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := r.RoundTripper.RoundTrip(req)
	r.logger.Debugw("http debug", "request", req, "response", resp, "err", err)
	return resp, err
}

// NewStackdriver creates a new metrics and trace exporter for Stackdriver.
func NewStackdriver(ctx context.Context, config *StackdriverConfig) (Exporter, error) {
	logger := logging.FromContext(ctx).Named("stackdriver")

	projectID := config.ProjectID
	if projectID == "" {
		return nil, fmt.Errorf("missing PROJECT_ID in Stackdriver exporter")
	}

	monitoredResource := NewStackdriverMonitoredResource(ctx, config)
	logger.Debugw("monitored resource", "resource", monitoredResource)

	var opts []option.ClientOption

	if config.DebugLogging {
		c := &http.Client{
			Transport: &debugRoundTripper{
				RoundTripper: http.DefaultTransport,
				logger:       logging.FromContext(ctx).Named("debugRoundTripper"),
			},
		}
		opts = append(opts, option.WithHTTPClient(c))
	}

	exporter, err := stackdriver.NewExporter(stackdriver.Options{
		Context:                 ctx,
		ProjectID:               projectID,
		ReportingInterval:       config.ReportingInterval,
		BundleDelayThreshold:    config.BundleDelayThreshold,
		Timeout:                 config.Timeout,
		BundleCountThreshold:    int(config.BundleCountThreshold),
		MonitoredResource:       monitoredResource,
		DefaultMonitoringLabels: &stackdriver.Labels{},
		MonitoringClientOptions: opts,
		OnError: func(err error) {
			logger.Errorw("failed to export metric", "error", err, "resource", monitoredResource)
		},
	})

	if err != nil {
		return nil, fmt.Errorf("failed to create Stackdriver exporter: %w", err)
	}
	return &stackdriverExporter{exporter, config}, nil
}

// StartExporter starts the exporter.
func (e *stackdriverExporter) StartExporter() error {
	if err := registerViews(); err != nil {
		return fmt.Errorf("failed to register views: %w", err)
	}

	if err := e.exporter.StartMetricsExporter(); err != nil {
		return fmt.Errorf("failed to start stackdriver exporter: %w", err)
	}

	trace.ApplyConfig(trace.Config{
		DefaultSampler: trace.ProbabilitySampler(e.config.SampleRate),
	})
	trace.RegisterExporter(e.exporter)

	view.RegisterExporter(e.exporter)

	return nil
}

// Close halts the exporter.
func (e *stackdriverExporter) Close() error {
	// Flush any existing metrics
	e.exporter.Flush()

	// Stop the exporter
	e.exporter.StopMetricsExporter()

	// Unregister the exporter
	trace.UnregisterExporter(e.exporter)
	view.UnregisterExporter(e.exporter)

	return nil
}
<commit_msg>Fix the metric missing issue caused by repeated metric point upload. (#1039)<commit_after>// Copyright 2020 Google LLC
//
// 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 observability sets up and configures observability tools.
package observability

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

	"github.com/google/exposure-notifications-server/pkg/logging"
	"go.uber.org/zap"
	"google.golang.org/api/option"

	"contrib.go.opencensus.io/exporter/stackdriver"
	"go.opencensus.io/trace"
)

var _ Exporter = (*stackdriverExporter)(nil)

type stackdriverExporter struct {
	exporter *stackdriver.Exporter
	config   *StackdriverConfig
}

type debugRoundTripper struct {
	logger *zap.SugaredLogger
	http.RoundTripper
}

func (r *debugRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	resp, err := r.RoundTripper.RoundTrip(req)
	r.logger.Debugw("http debug", "request", req, "response", resp, "err", err)
	return resp, err
}

// NewStackdriver creates a new metrics and trace exporter for Stackdriver.
func NewStackdriver(ctx context.Context, config *StackdriverConfig) (Exporter, error) {
	logger := logging.FromContext(ctx).Named("stackdriver")

	projectID := config.ProjectID
	if projectID == "" {
		return nil, fmt.Errorf("missing PROJECT_ID in Stackdriver exporter")
	}

	monitoredResource := NewStackdriverMonitoredResource(ctx, config)
	logger.Debugw("monitored resource", "resource", monitoredResource)

	var opts []option.ClientOption

	if config.DebugLogging {
		c := &http.Client{
			Transport: &debugRoundTripper{
				RoundTripper: http.DefaultTransport,
				logger:       logging.FromContext(ctx).Named("debugRoundTripper"),
			},
		}
		opts = append(opts, option.WithHTTPClient(c))
	}

	exporter, err := stackdriver.NewExporter(stackdriver.Options{
		Context:                 ctx,
		ProjectID:               projectID,
		ReportingInterval:       config.ReportingInterval,
		BundleDelayThreshold:    config.BundleDelayThreshold,
		Timeout:                 config.Timeout,
		BundleCountThreshold:    int(config.BundleCountThreshold),
		MonitoredResource:       monitoredResource,
		DefaultMonitoringLabels: &stackdriver.Labels{},
		MonitoringClientOptions: opts,
		OnError: func(err error) {
			logger.Errorw("failed to export metric", "error", err, "resource", monitoredResource)
		},
	})

	if err != nil {
		return nil, fmt.Errorf("failed to create Stackdriver exporter: %w", err)
	}
	return &stackdriverExporter{exporter, config}, nil
}

// StartExporter starts the exporter.
func (e *stackdriverExporter) StartExporter() error {
	if err := registerViews(); err != nil {
		return fmt.Errorf("failed to register views: %w", err)
	}

	if err := e.exporter.StartMetricsExporter(); err != nil {
		return fmt.Errorf("failed to start stackdriver exporter: %w", err)
	}

	trace.ApplyConfig(trace.Config{
		DefaultSampler: trace.ProbabilitySampler(e.config.SampleRate),
	})
	trace.RegisterExporter(e.exporter)

	return nil
}

// Close halts the exporter.
func (e *stackdriverExporter) Close() error {
	// Flush any existing metrics
	e.exporter.Flush()

	// Stop the exporter
	e.exporter.StopMetricsExporter()

	// Unregister the exporter
	trace.UnregisterExporter(e.exporter)

	return nil
}
<|endoftext|>


<commit_before>package scheduler

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  *redis.Client
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	if err := s.ImportFromMemStorage(); err != nil {
		s.log.Errorln("Something went wrong while importing old storage", err)
	}
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				if d, err := time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			_, err = s.broker.PushJob(et.Trigger(event))
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			s.client.ZRem(SchedKey, results[0])
			return err
		}
		switch t := t.(type) {
		case *AtTrigger:
			job := t.Trigger()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Trigger()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		case *EventTrigger: // Debounced
			job := t.Trigger(nil)
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		return s.client.HSet(hKey, t.ID(), t.Infos().Arguments).Err()
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	// TODO(pagination): use a sort of couchdb.ForeachDocs function when available.
	req := &couchdb.AllDocsRequest{Limit: 1000}
	err := couchdb.GetAllDocs(db, consts.Triggers, req, &infos)
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

var _ Scheduler = &RedisScheduler{}
<commit_msg>Fix gometalinter<commit_after>package scheduler

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  *redis.Client
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	if err := s.ImportFromMemStorage(); err != nil {
		s.log.Errorln("Something went wrong while importing old storage", err)
	}
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				var d time.Duration
				if d, err = time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			_, err = s.broker.PushJob(et.Trigger(event))
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			s.client.ZRem(SchedKey, results[0])
			return err
		}
		switch t := t.(type) {
		case *EventTrigger: // Debounced
			job := t.Trigger(nil)
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		case *AtTrigger:
			job := t.Trigger()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Trigger()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		return s.client.HSet(hKey, t.ID(), t.Infos().Arguments).Err()
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	// TODO(pagination): use a sort of couchdb.ForeachDocs function when available.
	req := &couchdb.AllDocsRequest{Limit: 1000}
	err := couchdb.GetAllDocs(db, consts.Triggers, req, &infos)
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

var _ Scheduler = &RedisScheduler{}
<|endoftext|>


<commit_before>//
// Copyright (c) 2014 The foocsim 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 caches

import (
	"fmt"
	"github.com/boltdb/bolt"
	"github.com/lpabon/godbc"
)

var valueset = []byte{1}
var valueunset = []byte{0}

type BoltDBCache struct {
	stats        CacheStats
	buckets      map[string]int
	db           *bolt.DB
	cachesize    uint64
	writethrough bool
	keyn         uint64
}

func NewBoltDBCache(cachesize uint64, writethrough bool) *BoltDBCache {
	var err error
	bdc := &BoltDBCache{}
	bdc.buckets = make(map[string]int)
	bdc.db, err = bolt.Open("cache.db", 0600, nil)
	godbc.Check(err == nil)

	return bdc
}

func (c *BoltDBCache) boltput(bucket string, key string, val []byte) (err error) {
	err = c.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			b, _ = tx.CreateBucket([]byte(bucket))
			c.buckets[bucket] = 1
		}
		return b.Put([]byte(key), val)
	})
	c.keyn++
	return
}

func (c *BoltDBCache) boltget(bucket string, key string) (val []byte, err error) {
	err = c.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			b, _ = tx.CreateBucket([]byte(bucket))
			c.buckets[bucket] = 1
		}
		val = b.Get([]byte(key))
		return nil
	})
	return
}

func (c *BoltDBCache) boltdelete(bucket string, key string) (err error) {
	err = c.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			return nil
		}
		return b.Delete([]byte(key))
	})
	c.keyn--
	return
}

func (c *BoltDBCache) Close() {
	c.db.Close()
}

func (c *BoltDBCache) Invalidate(obj, chunk string) {
	if val, _ := c.boltget(obj, chunk); val != nil {
		c.stats.writehits++
		c.stats.invalidations++

		c.boltdelete(obj, chunk)
	}
}

func (c *BoltDBCache) Evict() {
	c.stats.evictions++

	// BIG ASSUMPTION! I have no idea
	// if Go keeps track of the iteration
	// through a map
	for evicted := false; evicted; {
		for bucket, _ := range c.buckets {
			c.db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte(bucket))
				c := b.Cursor()

				for k, v := c.First(); k != nil; k, v = c.Next() {
					if v[0] == 1 {
						b.Put([]byte(k), valueunset)
					} else {
						b.Delete(k)
						evicted = true
						return nil
					}
				}

				return nil
			})
		}
	}
}

func (c *BoltDBCache) Insert(obj, chunk string) {
	c.stats.insertions++

	if c.keyn >= c.cachesize {
		c.Evict()
	}

	c.boltput(obj, chunk, valueset)
}

func (c *BoltDBCache) Write(obj string, chunk string) {
	c.stats.writes++

	// Invalidate
	c.Invalidate(obj, chunk)

	// We would do back end IO here

	// Insert
	if c.writethrough {
		c.Insert(obj, chunk)
	}
}

func (c *BoltDBCache) Read(obj, chunk string) {
	c.stats.reads++

	if val, _ := c.boltget(obj, chunk); val != nil {
		// Read Hit
		c.stats.readhits++

		// Clock Algorithm: Set that we looked
		// at it
		c.boltput(obj, chunk, valueset)
	} else {
		// Read miss
		// We would do IO here
		c.Insert(obj, chunk)
	}

}

func (c *BoltDBCache) Delete(obj string) {
	c.stats.deletions++

	c.db.Update(func(tx *bolt.Tx) (err error) {
		err = tx.DeleteBucket([]byte(obj))
		if err == nil {
			c.stats.deletionhits++
			delete(c.buckets, obj)
		}

		return err
	})
}

func (c *BoltDBCache) String() string {
	return fmt.Sprintf(
		"== Cache Information ==\n"+
			"Cache Utilization: 0\n") +
		c.stats.String() +
		c.db.String()
}

func (c *BoltDBCache) Stats() *CacheStats {
	return c.stats.Copy()
}
<commit_msg>Builds and runs.. although super slow<commit_after>//
// Copyright (c) 2014 The foocsim 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 caches

import (
	"errors"
	"fmt"
	"github.com/boltdb/bolt"
	"github.com/lpabon/godbc"
)

var valueset = []byte{1}
var valueunset = []byte{0}
var ErrKeyMissing = errors.New("No Key Found")

type BoltDBCache struct {
	stats        CacheStats
	buckets      map[string]int
	db           *bolt.DB
	cachesize    uint64
	writethrough bool
	keyn         uint64
}

func NewBoltDBCache(cachesize uint64, writethrough bool) *BoltDBCache {
	var err error
	bdc := &BoltDBCache{}
	bdc.buckets = make(map[string]int)
	bdc.db, err = bolt.Open("cache.db", 0600, nil)
	godbc.Check(err == nil)

	return bdc
}

func (c *BoltDBCache) boltput(bucket string, key string, val []byte) (err error) {
	err = c.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			b, _ = tx.CreateBucket([]byte(bucket))
			c.buckets[bucket] = 1
		}
		return b.Put([]byte(key), val)
	})
	c.keyn++
	return
}

func (c *BoltDBCache) boltget(bucket string, key string) (val []byte, err error) {
	err = c.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			return ErrKeyMissing
		}
		val = b.Get([]byte(key))
		return nil
	})
	return
}

func (c *BoltDBCache) boltdelete(bucket string, key string) (err error) {
	err = c.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucket))
		if nil == b {
			return nil
		}
		return b.Delete([]byte(key))
	})
	c.keyn--
	return
}

func (c *BoltDBCache) Close() {
	c.db.Close()
}

func (c *BoltDBCache) Invalidate(obj, chunk string) {
	if _, err := c.boltget(obj, chunk); err == nil {
		c.stats.writehits++
		c.stats.invalidations++

		c.boltdelete(obj, chunk)
	}
}

func (c *BoltDBCache) Evict() {
	c.stats.evictions++

	// BIG ASSUMPTION! I have no idea
	// if Go keeps track of the iteration
	// through a map
	for evicted := false; evicted; {
		for bucket, _ := range c.buckets {
			c.db.Update(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte(bucket))
				c := b.Cursor()

				for k, v := c.First(); k != nil; k, v = c.Next() {
					if v[0] == 1 {
						b.Put([]byte(k), valueunset)
					} else {
						b.Delete(k)
						evicted = true
						return nil
					}
				}

				return nil
			})
		}
	}
}

func (c *BoltDBCache) Insert(obj, chunk string) {
	c.stats.insertions++

	if c.keyn >= c.cachesize {
		c.Evict()
	}

	c.boltput(obj, chunk, valueset)
}

func (c *BoltDBCache) Write(obj string, chunk string) {
	c.stats.writes++

	// Invalidate
	c.Invalidate(obj, chunk)

	// We would do back end IO here

	// Insert
	if c.writethrough {
		c.Insert(obj, chunk)
	}
}

func (c *BoltDBCache) Read(obj, chunk string) {
	c.stats.reads++

	if _, err := c.boltget(obj, chunk); err == nil {
		// Read Hit
		c.stats.readhits++

		// Clock Algorithm: Set that we looked
		// at it
		c.boltput(obj, chunk, valueset)
	} else {
		// Read miss
		// We would do IO here
		c.Insert(obj, chunk)
	}

}

func (c *BoltDBCache) Delete(obj string) {
	c.stats.deletions++

	c.db.Update(func(tx *bolt.Tx) (err error) {
		err = tx.DeleteBucket([]byte(obj))
		if err == nil {
			c.stats.deletionhits++
			delete(c.buckets, obj)
		}

		return err
	})
}

func (c *BoltDBCache) String() string {
	return fmt.Sprintf(
		"== Cache Information ==\n"+
			"Cache Utilization: 0\n") +
		c.stats.String() +
		c.db.String()
}

func (c *BoltDBCache) Stats() *CacheStats {
	return c.stats.Copy()
}
<|endoftext|>