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

// Waiting for FDs via kqueue/kevent.

package net

import (
	"os"
	"syscall"
)

type pollster struct {
	kq       int
	eventbuf [10]syscall.Kevent_t
	events   []syscall.Kevent_t
}

func newpollster() (p *pollster, err os.Error) {
	p = new(pollster)
	var e int
	if p.kq, e = syscall.Kqueue(); e != 0 {
		return nil, os.NewSyscallError("kqueue", e)
	}
	p.events = p.eventbuf[0:0]
	return p, nil
}

func (p *pollster) AddFD(fd int, mode int, repeat bool) os.Error {
	var kmode int
	if mode == 'r' {
		kmode = syscall.EVFILT_READ
	} else {
		kmode = syscall.EVFILT_WRITE
	}
	var events [1]syscall.Kevent_t
	ev := &events[0]
	// EV_ADD - add event to kqueue list
	// EV_ONESHOT - delete the event the first time it triggers
	flags := syscall.EV_ADD
	if !repeat {
		flags |= syscall.EV_ONESHOT
	}
	syscall.SetKevent(ev, fd, kmode, flags)

	n, e := syscall.Kevent(p.kq, &events, nil, nil)
	if e != 0 {
		return os.NewSyscallError("kevent", e)
	}
	if n != 1 || (ev.Flags&syscall.EV_ERROR) == 0 || int(ev.Ident) != fd || int(ev.Filter) != kmode {
		return os.NewSyscallError("kqueue phase error", e)
	}
	if ev.Data != 0 {
		return os.Errno(int(ev.Data))
	}
	return nil
}

func (p *pollster) DelFD(fd int, mode int) {
	var kmode int
	if mode == 'r' {
		kmode = syscall.EVFILT_READ
	} else {
		kmode = syscall.EVFILT_WRITE
	}
	var events [1]syscall.Kevent_t
	ev := &events[0]
	// EV_DELETE - delete event from kqueue list
	syscall.SetKevent(ev, fd, kmode, syscall.EV_DELETE)
	syscall.Kevent(p.kq, &events, nil, nil)
}

func (p *pollster) WaitFD(nsec int64) (fd int, mode int, err os.Error) {
	var t *syscall.Timespec
	for len(p.events) == 0 {
		if nsec > 0 {
			if t == nil {
				t = new(syscall.Timespec)
			}
			*t = syscall.NsecToTimespec(nsec)
		}
		nn, e := syscall.Kevent(p.kq, nil, &p.eventbuf, t)
		if e != 0 {
			if e == syscall.EINTR {
				continue
			}
			return -1, 0, os.NewSyscallError("kevent", e)
		}
		if nn == 0 {
			return -1, 0, nil
		}
		p.events = p.eventbuf[0:nn]
	}
	ev := &p.events[0]
	p.events = p.events[1:]
	fd = int(ev.Ident)
	if ev.Filter == syscall.EVFILT_READ {
		mode = 'r'
	} else {
		mode = 'w'
	}
	return fd, mode, nil
}

func (p *pollster) Close() os.Error { return os.NewSyscallError("close", syscall.Close(p.kq)) }
<commit_msg>net: fix freebsd build<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Waiting for FDs via kqueue/kevent.

package net

import (
	"os"
	"syscall"
)

type pollster struct {
	kq       int
	eventbuf [10]syscall.Kevent_t
	events   []syscall.Kevent_t
}

func newpollster() (p *pollster, err os.Error) {
	p = new(pollster)
	var e int
	if p.kq, e = syscall.Kqueue(); e != 0 {
		return nil, os.NewSyscallError("kqueue", e)
	}
	p.events = p.eventbuf[0:0]
	return p, nil
}

func (p *pollster) AddFD(fd int, mode int, repeat bool) os.Error {
	var kmode int
	if mode == 'r' {
		kmode = syscall.EVFILT_READ
	} else {
		kmode = syscall.EVFILT_WRITE
	}
	var events [1]syscall.Kevent_t
	ev := &events[0]
	// EV_ADD - add event to kqueue list
	// EV_ONESHOT - delete the event the first time it triggers
	flags := syscall.EV_ADD
	if !repeat {
		flags |= syscall.EV_ONESHOT
	}
	syscall.SetKevent(ev, fd, kmode, flags)

	n, e := syscall.Kevent(p.kq, events[:], nil, nil)
	if e != 0 {
		return os.NewSyscallError("kevent", e)
	}
	if n != 1 || (ev.Flags&syscall.EV_ERROR) == 0 || int(ev.Ident) != fd || int(ev.Filter) != kmode {
		return os.NewSyscallError("kqueue phase error", e)
	}
	if ev.Data != 0 {
		return os.Errno(int(ev.Data))
	}
	return nil
}

func (p *pollster) DelFD(fd int, mode int) {
	var kmode int
	if mode == 'r' {
		kmode = syscall.EVFILT_READ
	} else {
		kmode = syscall.EVFILT_WRITE
	}
	var events [1]syscall.Kevent_t
	ev := &events[0]
	// EV_DELETE - delete event from kqueue list
	syscall.SetKevent(ev, fd, kmode, syscall.EV_DELETE)
	syscall.Kevent(p.kq, events[:], nil, nil)
}

func (p *pollster) WaitFD(nsec int64) (fd int, mode int, err os.Error) {
	var t *syscall.Timespec
	for len(p.events) == 0 {
		if nsec > 0 {
			if t == nil {
				t = new(syscall.Timespec)
			}
			*t = syscall.NsecToTimespec(nsec)
		}
		nn, e := syscall.Kevent(p.kq, nil, p.eventbuf[:], t)
		if e != 0 {
			if e == syscall.EINTR {
				continue
			}
			return -1, 0, os.NewSyscallError("kevent", e)
		}
		if nn == 0 {
			return -1, 0, nil
		}
		p.events = p.eventbuf[0:nn]
	}
	ev := &p.events[0]
	p.events = p.events[1:]
	fd = int(ev.Ident)
	if ev.Filter == syscall.EVFILT_READ {
		mode = 'r'
	} else {
		mode = 'w'
	}
	return fd, mode, nil
}

func (p *pollster) Close() os.Error { return os.NewSyscallError("close", syscall.Close(p.kq)) }
<|endoftext|>


<commit_before>package asset

import (
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"path"
	"path/filepath"

	"github.com/kubernetes-incubator/bootkube/pkg/tlsutil"
)

const (
	AssetPathSecrets                     = "tls"
	AssetPathCAKey                       = "tls/ca.key"
	AssetPathCACert                      = "tls/ca.crt"
	AssetPathAPIServerKey                = "tls/apiserver.key"
	AssetPathAPIServerCert               = "tls/apiserver.crt"
	AssetPathEtcdClientCA                = "tls/etcd-client-ca.crt"
	AssetPathEtcdClientCert              = "tls/etcd-client.crt"
	AssetPathEtcdClientKey               = "tls/etcd-client.key"
	AssetPathEtcdServerCA                = "tls/etcd/server-ca.crt"
	AssetPathEtcdServerCert              = "tls/etcd/server.crt"
	AssetPathEtcdServerKey               = "tls/etcd/server.key"
	AssetPathEtcdPeerCA                  = "tls/etcd/peer-ca.crt"
	AssetPathEtcdPeerCert                = "tls/etcd/peer.crt"
	AssetPathEtcdPeerKey                 = "tls/etcd/peer.key"
	AssetPathServiceAccountPrivKey       = "tls/service-account.key"
	AssetPathServiceAccountPubKey        = "tls/service-account.pub"
	AssetPathAdminKey                    = "tls/admin.key"
	AssetPathAdminCert                   = "tls/admin.crt"
	AssetPathAdminKubeConfig             = "auth/kubeconfig"
	AssetPathKubeletKubeConfig           = "auth/kubeconfig-kubelet"
	AssetPathManifests                   = "manifests"
	AssetPathKubeConfigInCluster         = "manifests/kubeconfig-in-cluster.yaml"
	AssetPathKubeletBootstrapToken       = "manifests/kubelet-bootstrap-token.yaml"
	AssetPathProxy                       = "manifests/kube-proxy.yaml"
	AssetPathProxySA                     = "manifests/kube-proxy-sa.yaml"
	AssetPathProxyRoleBinding            = "manifests/kube-proxy-role-binding.yaml"
	AssetPathFlannel                     = "manifests/flannel.yaml"
	AssetPathFlannelCfg                  = "manifests/flannel-cfg.yaml"
	AssetPathFlannelClusterRole          = "manifests/flannel-cluster-role.yaml"
	AssetPathFlannelClusterRoleBinding   = "manifests/flannel-cluster-role-binding.yaml"
	AssetPathFlannelSA                   = "manifests/flannel-sa.yaml"
	AssetPathCalico                      = "manifests/calico.yaml"
	AssetPathCalicoPolicyOnly            = "manifests/calico-policy-only.yaml"
	AssetPathCalicoCfg                   = "manifests/calico-config.yaml"
	AssetPathCalicoSA                    = "manifests/calico-service-account.yaml"
	AssetPathCalicoRole                  = "manifests/calico-role.yaml"
	AssetPathCalicoRoleBinding           = "manifests/calico-role-binding.yaml"
	AssetPathCalicoBGPConfigsCRD         = "manifests/calico-bgp-configs-crd.yaml"
	AssetPathCalicoFelixConfigsCRD       = "manifests/calico-felix-configs-crd.yaml"
	AssetPathCalicoNetworkPoliciesCRD    = "manifests/calico-network-policies-crd.yaml"
	AssetPathCalicoIPPoolsCRD            = "manifests/calico-ip-pools-crd.yaml"
	AssetPathAPIServerSecret             = "manifests/kube-apiserver-secret.yaml"
	AssetPathAPIServer                   = "manifests/kube-apiserver.yaml"
	AssetPathControllerManager           = "manifests/kube-controller-manager.yaml"
	AssetPathControllerManagerSA         = "manifests/kube-controller-manager-service-account.yaml"
	AssetPathControllerManagerRB         = "manifests/kube-controller-manager-role-binding.yaml"
	AssetPathControllerManagerSecret     = "manifests/kube-controller-manager-secret.yaml"
	AssetPathControllerManagerDisruption = "manifests/kube-controller-manager-disruption.yaml"
	AssetPathScheduler                   = "manifests/kube-scheduler.yaml"
	AssetPathSchedulerDisruption         = "manifests/kube-scheduler-disruption.yaml"
	AssetPathKubeDNSDeployment           = "manifests/kube-dns-deployment.yaml"
	AssetPathKubeDNSSvc                  = "manifests/kube-dns-svc.yaml"
	AssetPathSystemNamespace             = "manifests/kube-system-ns.yaml"
	AssetPathCheckpointer                = "manifests/pod-checkpointer.yaml"
	AssetPathCheckpointerSA              = "manifests/pod-checkpointer-sa.yaml"
	AssetPathCheckpointerRole            = "manifests/pod-checkpointer-role.yaml"
	AssetPathCheckpointerRoleBinding     = "manifests/pod-checkpointer-role-binding.yaml"
	AssetPathEtcdClientSecret            = "manifests/etcd-client-tls.yaml"
	AssetPathEtcdPeerSecret              = "manifests/etcd-peer-tls.yaml"
	AssetPathEtcdServerSecret            = "manifests/etcd-server-tls.yaml"
	AssetPathCSRBootstrapRoleBinding     = "manifests/csr-bootstrap-role-binding.yaml"
	AssetPathCSRApproverRoleBinding      = "manifests/csr-approver-role-binding.yaml"
	AssetPathCSRRenewalRoleBinding       = "manifests/csr-renewal-role-binding.yaml"
	AssetPathKubeSystemSARoleBinding     = "manifests/kube-system-rbac-role-binding.yaml"
	AssetPathBootstrapManifests          = "bootstrap-manifests"
	AssetPathBootstrapAPIServer          = "bootstrap-manifests/bootstrap-apiserver.yaml"
	AssetPathBootstrapControllerManager  = "bootstrap-manifests/bootstrap-controller-manager.yaml"
	AssetPathBootstrapScheduler          = "bootstrap-manifests/bootstrap-scheduler.yaml"
	AssetPathBootstrapEtcd               = "bootstrap-manifests/bootstrap-etcd.yaml"
	AssetPathBootstrapEtcdService        = "etcd/bootstrap-etcd-service.json"
	AssetPathMigrateEtcdCluster          = "etcd/migrate-etcd-cluster.json"
)

var (
	BootstrapSecretsDir = "/etc/kubernetes/bootstrap-secrets" // Overridden for testing.
)

// AssetConfig holds all configuration needed when generating
// the default set of assets.
type Config struct {
	EtcdCACert             *x509.Certificate
	EtcdClientCert         *x509.Certificate
	EtcdClientKey          *rsa.PrivateKey
	EtcdServers            []*url.URL
	EtcdUseTLS             bool
	APIServers             []*url.URL
	CACert                 *x509.Certificate
	CAPrivKey              *rsa.PrivateKey
	AltNames               *tlsutil.AltNames
	PodCIDR                *net.IPNet
	ServiceCIDR            *net.IPNet
	APIServiceIP           net.IP
	BootEtcdServiceIP      net.IP
	DNSServiceIP           net.IP
	EtcdServiceIP          net.IP
	EtcdServiceName        string
	CloudProvider          string
	NetworkProvider        string
	BootstrapSecretsSubdir string
	Images                 ImageVersions
}

// ImageVersions holds all the images (and their versions) that are rendered into the templates.
type ImageVersions struct {
	Etcd            string
	EtcdOperator    string
	Flannel         string
	FlannelCNI      string
	Calico          string
	CalicoCNI       string
	Hyperkube       string
	Kenc            string
	KubeDNS         string
	KubeDNSMasq     string
	KubeDNSSidecar  string
	PodCheckpointer string
}

// NewDefaultAssets returns a list of default assets, optionally
// configured via a user provided AssetConfig. Default assets include
// TLS assets (certs, keys and secrets), and k8s component manifests.
func NewDefaultAssets(conf Config) (Assets, error) {
	conf.BootstrapSecretsSubdir = path.Base(BootstrapSecretsDir)

	as := newStaticAssets(conf.Images)
	as = append(as, newDynamicAssets(conf)...)

	// Add kube-apiserver service IP
	conf.AltNames.IPs = append(conf.AltNames.IPs, conf.APIServiceIP)

	// Create a CA if none was provided.
	if conf.CACert == nil {
		var err error
		conf.CAPrivKey, conf.CACert, err = newCACert()
		if err != nil {
			return Assets{}, err
		}
	}

	// TLS assets
	tlsAssets, err := newTLSAssets(conf.CACert, conf.CAPrivKey, *conf.AltNames)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, tlsAssets...)

	// etcd TLS assets.
	if conf.EtcdUseTLS {
		etcdTLSAssets, err := newEtcdTLSAssets(conf.EtcdCACert, conf.EtcdClientCert, conf.EtcdClientKey, conf.CACert, conf.CAPrivKey, conf.EtcdServers)
		if err != nil {
			return Assets{}, err
		}
		as = append(as, etcdTLSAssets...)
	}

	kubeConfigAssets, err := newKubeConfigAssets(as, conf)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, kubeConfigAssets...)

	// K8S APIServer secret
	apiSecret, err := newAPIServerSecretAsset(as, conf.EtcdUseTLS)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, apiSecret)

	// K8S ControllerManager secret
	cmSecret, err := newControllerManagerSecretAsset(as)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, cmSecret)

	return as, nil
}

type Asset struct {
	Name string
	Data []byte
}

type Assets []Asset

func (as Assets) Get(name string) (Asset, error) {
	for _, asset := range as {
		if asset.Name == name {
			return asset, nil
		}
	}
	return Asset{}, fmt.Errorf("asset %q does not exist", name)
}

func (as Assets) WriteFiles(path string) error {
	if err := os.Mkdir(path, 0755); err != nil {
		return err
	}
	for _, asset := range as {
		if err := asset.WriteFile(path); err != nil {
			return err
		}
	}
	return nil
}

func (a Asset) WriteFile(path string) error {
	f := filepath.Join(path, a.Name)
	if err := os.MkdirAll(filepath.Dir(f), 0755); err != nil {
		return err
	}
	fmt.Printf("Writing asset: %s\n", f)
	return ioutil.WriteFile(f, a.Data, 0600)
}
<commit_msg>assets: use MkdirAll to not error on existing dir (#886)<commit_after>package asset

import (
	"crypto/rsa"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"path"
	"path/filepath"

	"github.com/kubernetes-incubator/bootkube/pkg/tlsutil"
)

const (
	AssetPathSecrets                     = "tls"
	AssetPathCAKey                       = "tls/ca.key"
	AssetPathCACert                      = "tls/ca.crt"
	AssetPathAPIServerKey                = "tls/apiserver.key"
	AssetPathAPIServerCert               = "tls/apiserver.crt"
	AssetPathEtcdClientCA                = "tls/etcd-client-ca.crt"
	AssetPathEtcdClientCert              = "tls/etcd-client.crt"
	AssetPathEtcdClientKey               = "tls/etcd-client.key"
	AssetPathEtcdServerCA                = "tls/etcd/server-ca.crt"
	AssetPathEtcdServerCert              = "tls/etcd/server.crt"
	AssetPathEtcdServerKey               = "tls/etcd/server.key"
	AssetPathEtcdPeerCA                  = "tls/etcd/peer-ca.crt"
	AssetPathEtcdPeerCert                = "tls/etcd/peer.crt"
	AssetPathEtcdPeerKey                 = "tls/etcd/peer.key"
	AssetPathServiceAccountPrivKey       = "tls/service-account.key"
	AssetPathServiceAccountPubKey        = "tls/service-account.pub"
	AssetPathAdminKey                    = "tls/admin.key"
	AssetPathAdminCert                   = "tls/admin.crt"
	AssetPathAdminKubeConfig             = "auth/kubeconfig"
	AssetPathKubeletKubeConfig           = "auth/kubeconfig-kubelet"
	AssetPathManifests                   = "manifests"
	AssetPathKubeConfigInCluster         = "manifests/kubeconfig-in-cluster.yaml"
	AssetPathKubeletBootstrapToken       = "manifests/kubelet-bootstrap-token.yaml"
	AssetPathProxy                       = "manifests/kube-proxy.yaml"
	AssetPathProxySA                     = "manifests/kube-proxy-sa.yaml"
	AssetPathProxyRoleBinding            = "manifests/kube-proxy-role-binding.yaml"
	AssetPathFlannel                     = "manifests/flannel.yaml"
	AssetPathFlannelCfg                  = "manifests/flannel-cfg.yaml"
	AssetPathFlannelClusterRole          = "manifests/flannel-cluster-role.yaml"
	AssetPathFlannelClusterRoleBinding   = "manifests/flannel-cluster-role-binding.yaml"
	AssetPathFlannelSA                   = "manifests/flannel-sa.yaml"
	AssetPathCalico                      = "manifests/calico.yaml"
	AssetPathCalicoPolicyOnly            = "manifests/calico-policy-only.yaml"
	AssetPathCalicoCfg                   = "manifests/calico-config.yaml"
	AssetPathCalicoSA                    = "manifests/calico-service-account.yaml"
	AssetPathCalicoRole                  = "manifests/calico-role.yaml"
	AssetPathCalicoRoleBinding           = "manifests/calico-role-binding.yaml"
	AssetPathCalicoBGPConfigsCRD         = "manifests/calico-bgp-configs-crd.yaml"
	AssetPathCalicoFelixConfigsCRD       = "manifests/calico-felix-configs-crd.yaml"
	AssetPathCalicoNetworkPoliciesCRD    = "manifests/calico-network-policies-crd.yaml"
	AssetPathCalicoIPPoolsCRD            = "manifests/calico-ip-pools-crd.yaml"
	AssetPathAPIServerSecret             = "manifests/kube-apiserver-secret.yaml"
	AssetPathAPIServer                   = "manifests/kube-apiserver.yaml"
	AssetPathControllerManager           = "manifests/kube-controller-manager.yaml"
	AssetPathControllerManagerSA         = "manifests/kube-controller-manager-service-account.yaml"
	AssetPathControllerManagerRB         = "manifests/kube-controller-manager-role-binding.yaml"
	AssetPathControllerManagerSecret     = "manifests/kube-controller-manager-secret.yaml"
	AssetPathControllerManagerDisruption = "manifests/kube-controller-manager-disruption.yaml"
	AssetPathScheduler                   = "manifests/kube-scheduler.yaml"
	AssetPathSchedulerDisruption         = "manifests/kube-scheduler-disruption.yaml"
	AssetPathKubeDNSDeployment           = "manifests/kube-dns-deployment.yaml"
	AssetPathKubeDNSSvc                  = "manifests/kube-dns-svc.yaml"
	AssetPathSystemNamespace             = "manifests/kube-system-ns.yaml"
	AssetPathCheckpointer                = "manifests/pod-checkpointer.yaml"
	AssetPathCheckpointerSA              = "manifests/pod-checkpointer-sa.yaml"
	AssetPathCheckpointerRole            = "manifests/pod-checkpointer-role.yaml"
	AssetPathCheckpointerRoleBinding     = "manifests/pod-checkpointer-role-binding.yaml"
	AssetPathEtcdClientSecret            = "manifests/etcd-client-tls.yaml"
	AssetPathEtcdPeerSecret              = "manifests/etcd-peer-tls.yaml"
	AssetPathEtcdServerSecret            = "manifests/etcd-server-tls.yaml"
	AssetPathCSRBootstrapRoleBinding     = "manifests/csr-bootstrap-role-binding.yaml"
	AssetPathCSRApproverRoleBinding      = "manifests/csr-approver-role-binding.yaml"
	AssetPathCSRRenewalRoleBinding       = "manifests/csr-renewal-role-binding.yaml"
	AssetPathKubeSystemSARoleBinding     = "manifests/kube-system-rbac-role-binding.yaml"
	AssetPathBootstrapManifests          = "bootstrap-manifests"
	AssetPathBootstrapAPIServer          = "bootstrap-manifests/bootstrap-apiserver.yaml"
	AssetPathBootstrapControllerManager  = "bootstrap-manifests/bootstrap-controller-manager.yaml"
	AssetPathBootstrapScheduler          = "bootstrap-manifests/bootstrap-scheduler.yaml"
	AssetPathBootstrapEtcd               = "bootstrap-manifests/bootstrap-etcd.yaml"
	AssetPathBootstrapEtcdService        = "etcd/bootstrap-etcd-service.json"
	AssetPathMigrateEtcdCluster          = "etcd/migrate-etcd-cluster.json"
)

var (
	BootstrapSecretsDir = "/etc/kubernetes/bootstrap-secrets" // Overridden for testing.
)

// AssetConfig holds all configuration needed when generating
// the default set of assets.
type Config struct {
	EtcdCACert             *x509.Certificate
	EtcdClientCert         *x509.Certificate
	EtcdClientKey          *rsa.PrivateKey
	EtcdServers            []*url.URL
	EtcdUseTLS             bool
	APIServers             []*url.URL
	CACert                 *x509.Certificate
	CAPrivKey              *rsa.PrivateKey
	AltNames               *tlsutil.AltNames
	PodCIDR                *net.IPNet
	ServiceCIDR            *net.IPNet
	APIServiceIP           net.IP
	BootEtcdServiceIP      net.IP
	DNSServiceIP           net.IP
	EtcdServiceIP          net.IP
	EtcdServiceName        string
	CloudProvider          string
	NetworkProvider        string
	BootstrapSecretsSubdir string
	Images                 ImageVersions
}

// ImageVersions holds all the images (and their versions) that are rendered into the templates.
type ImageVersions struct {
	Etcd            string
	EtcdOperator    string
	Flannel         string
	FlannelCNI      string
	Calico          string
	CalicoCNI       string
	Hyperkube       string
	Kenc            string
	KubeDNS         string
	KubeDNSMasq     string
	KubeDNSSidecar  string
	PodCheckpointer string
}

// NewDefaultAssets returns a list of default assets, optionally
// configured via a user provided AssetConfig. Default assets include
// TLS assets (certs, keys and secrets), and k8s component manifests.
func NewDefaultAssets(conf Config) (Assets, error) {
	conf.BootstrapSecretsSubdir = path.Base(BootstrapSecretsDir)

	as := newStaticAssets(conf.Images)
	as = append(as, newDynamicAssets(conf)...)

	// Add kube-apiserver service IP
	conf.AltNames.IPs = append(conf.AltNames.IPs, conf.APIServiceIP)

	// Create a CA if none was provided.
	if conf.CACert == nil {
		var err error
		conf.CAPrivKey, conf.CACert, err = newCACert()
		if err != nil {
			return Assets{}, err
		}
	}

	// TLS assets
	tlsAssets, err := newTLSAssets(conf.CACert, conf.CAPrivKey, *conf.AltNames)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, tlsAssets...)

	// etcd TLS assets.
	if conf.EtcdUseTLS {
		etcdTLSAssets, err := newEtcdTLSAssets(conf.EtcdCACert, conf.EtcdClientCert, conf.EtcdClientKey, conf.CACert, conf.CAPrivKey, conf.EtcdServers)
		if err != nil {
			return Assets{}, err
		}
		as = append(as, etcdTLSAssets...)
	}

	kubeConfigAssets, err := newKubeConfigAssets(as, conf)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, kubeConfigAssets...)

	// K8S APIServer secret
	apiSecret, err := newAPIServerSecretAsset(as, conf.EtcdUseTLS)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, apiSecret)

	// K8S ControllerManager secret
	cmSecret, err := newControllerManagerSecretAsset(as)
	if err != nil {
		return Assets{}, err
	}
	as = append(as, cmSecret)

	return as, nil
}

type Asset struct {
	Name string
	Data []byte
}

type Assets []Asset

func (as Assets) Get(name string) (Asset, error) {
	for _, asset := range as {
		if asset.Name == name {
			return asset, nil
		}
	}
	return Asset{}, fmt.Errorf("asset %q does not exist", name)
}

func (as Assets) WriteFiles(path string) error {
	if err := os.MkdirAll(path, 0755); err != nil {
		return err
	}
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}
	if len(files) > 0 {
		return errors.New("asset directory must be empty")
	}
	for _, asset := range as {
		if err := asset.WriteFile(path); err != nil {
			return err
		}
	}
	return nil
}

func (a Asset) WriteFile(path string) error {
	f := filepath.Join(path, a.Name)
	if err := os.MkdirAll(filepath.Dir(f), 0755); err != nil {
		return err
	}
	fmt.Printf("Writing asset: %s\n", f)
	return ioutil.WriteFile(f, a.Data, 0600)
}
<|endoftext|>


<commit_before>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	log.Println("Syncing Repositories ...")
	if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
		return err
	}

	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		err := c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://kubernetes-charts.storage.googleapis.com",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/\w+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}
<commit_msg>fix(tool/charts): \w+ -> .+ (#380)<commit_after>package helm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/Masterminds/semver"
	"sigs.k8s.io/yaml"
)

// LoadChartfile opens a Chartfile tree
func LoadChartfile(projectRoot string) (*Charts, error) {
	// make sure project root is valid
	abs, err := filepath.Abs(projectRoot)
	if err != nil {
		return nil, err
	}

	// open chartfile
	chartfile := filepath.Join(abs, Filename)
	data, err := ioutil.ReadFile(chartfile)
	if err != nil {
		return nil, err
	}

	// parse it
	c := Chartfile{
		Version:   Version,
		Directory: DefaultDir,
	}
	if err := yaml.UnmarshalStrict(data, &c); err != nil {
		return nil, err
	}

	for i, r := range c.Requires {
		if r.Chart == "" {
			return nil, fmt.Errorf("requirements[%v]: 'chart' must be set", i)
		}
	}

	// return Charts handle
	charts := &Charts{
		Manifest:    c,
		projectRoot: abs,

		// default to ExecHelm, but allow injecting from the outside
		Helm: ExecHelm{},
	}
	return charts, nil
}

// Charts exposes the central Chartfile management functions
type Charts struct {
	// Manifest are the chartfile.yaml contents. It holds data about the developers intentions
	Manifest Chartfile

	// projectRoot is the enclosing directory of chartfile.yaml
	projectRoot string

	// Helm is the helm implementation underneath. ExecHelm is the default, but
	// any implementation of the Helm interface may be used
	Helm Helm
}

// ChartDir returns the directory pulled charts are saved in
func (c Charts) ChartDir() string {
	return filepath.Join(c.projectRoot, c.Manifest.Directory)
}

// ManifestFile returns the full path to the chartfile.yaml
func (c Charts) ManifestFile() string {
	return filepath.Join(c.projectRoot, Filename)
}

// Vendor pulls all Charts specified in the manifest into the local charts
// directory. It fetches the repository index before doing so.
func (c Charts) Vendor() error {
	dir := c.ChartDir()
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		return err
	}

	log.Println("Syncing Repositories ...")
	if err := c.Helm.RepoUpdate(Opts{Repositories: c.Manifest.Repositories}); err != nil {
		return err
	}

	log.Println("Pulling Charts ...")
	for _, r := range c.Manifest.Requires {
		err := c.Helm.Pull(r.Chart, r.Version.String(), PullOpts{
			Destination: dir,
			Opts:        Opts{Repositories: c.Manifest.Repositories},
		})
		if err != nil {
			return err
		}

		log.Printf(" %s@%s", r.Chart, r.Version.String())
	}

	return nil
}

// Add adds every Chart in reqs to the Manifest after validation, and runs
// Vendor afterwards
func (c Charts) Add(reqs []string) error {
	log.Printf("Adding %v Charts ...", len(reqs))

	skip := func(s string, err error) {
		log.Printf(" Skipping %s: %s.", s, err)
	}

	// parse new charts, append in memory
	added := 0
	for _, s := range reqs {
		r, err := parseReq(s)
		if err != nil {
			skip(s, err)
			continue
		}

		if c.Manifest.Requires.Has(*r) {
			skip(s, fmt.Errorf("already exists"))
			continue
		}

		c.Manifest.Requires = append(c.Manifest.Requires, *r)
		added++
		log.Println(" OK:", s)
	}

	// write out
	if err := write(c.Manifest, c.ManifestFile()); err != nil {
		return err
	}

	// skipped some? fail then
	if added != len(reqs) {
		return fmt.Errorf("%v Charts were skipped. Please check above logs for details", len(reqs)-added)
	}

	// worked fine? vendor it
	log.Printf("Added %v Charts to helmfile.yaml. Vendoring ...", added)
	return c.Vendor()
}

func InitChartfile(path string) (*Charts, error) {
	c := Chartfile{
		Version: Version,
		Repositories: []Repo{{
			Name: "stable",
			URL:  "https://kubernetes-charts.storage.googleapis.com",
		}},
		Requires: make(Requirements, 0),
	}

	if err := write(c, path); err != nil {
		return nil, err
	}

	return LoadChartfile(filepath.Dir(path))
}

// write saves a Chartfile to dest
func write(c Chartfile, dest string) error {
	data, err := yaml.Marshal(c)
	if err != nil {
		return err
	}

	return ioutil.WriteFile(dest, data, 0644)
}

var chartExp = regexp.MustCompile(`\w+\/.+@.+`)

// parseReq parses a requirement from a string of the format `repo/name@version`
func parseReq(s string) (*Requirement, error) {
	if !chartExp.MatchString(s) {
		return nil, fmt.Errorf("not of form 'repo/chart@version'")
	}

	elems := strings.Split(s, "@")
	chart := elems[0]
	ver, err := semver.NewVersion(elems[1])
	if errors.Is(err, semver.ErrInvalidSemVer) {
		return nil, fmt.Errorf("version is invalid")
	} else if err != nil {
		return nil, fmt.Errorf("version is invalid: %s", err)
	}

	return &Requirement{
		Chart:   chart,
		Version: *ver,
	}, nil
}
<|endoftext|>


<commit_before>package tls

import (
	"context"

	"github.com/rancher/dynamiclistener"
	v1 "github.com/rancher/k3s/pkg/apis/k3s.cattle.io/v1"
	k3sclient "github.com/rancher/k3s/pkg/generated/controllers/k3s.cattle.io/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	ns   = "kube-system"
	name = "tls-config"
)

func NewServer(ctx context.Context, listenerConfigs k3sclient.ListenerConfigController, config dynamiclistener.UserConfig) (dynamiclistener.ServerInterface, error) {
	storage := &listenerConfigStorage{
		client: listenerConfigs,
		cache:  listenerConfigs.Cache(),
	}

	server, err := dynamiclistener.NewServer(storage, config)
	if err != nil {
		return nil, err
	}

	listenerConfigs.OnChange(ctx, "listen-config", func(key string, obj *v1.ListenerConfig) (*v1.ListenerConfig, error) {
		if obj == nil {
			return nil, nil
		}
		return obj, server.Update(fromStorage(obj))
	})

	return server, err
}

type listenerConfigStorage struct {
	cache  k3sclient.ListenerConfigCache
	client k3sclient.ListenerConfigClient
}

func (l *listenerConfigStorage) Set(config *dynamiclistener.ListenerStatus) (*dynamiclistener.ListenerStatus, error) {
	if config == nil {
		return nil, nil
	}

	obj, err := l.cache.Get(ns, name)
	if errors.IsNotFound(err) {
		ls := v1.NewListenerConfig(ns, name, v1.ListenerConfig{
			Status: *config,
		})

		ls, err := l.client.Create(ls)
		return fromStorage(ls), err
	} else if err != nil {
		return nil, err
	}

	obj = obj.DeepCopy()
	obj.ResourceVersion = config.Revision
	obj.Status = *config
	obj.Status.Revision = ""

	obj, err = l.client.Update(obj)
	return fromStorage(obj), err
}

func (l *listenerConfigStorage) Get() (*dynamiclistener.ListenerStatus, error) {
	obj, err := l.cache.Get(ns, name)
	if errors.IsNotFound(err) {
		obj, err = l.client.Get(ns, name, metav1.GetOptions{})
	}
	if errors.IsNotFound(err) {
		return &dynamiclistener.ListenerStatus{}, nil
	}
	return fromStorage(obj), err
}

func fromStorage(obj *v1.ListenerConfig) *dynamiclistener.ListenerStatus {
	if obj == nil {
		return nil
	}

	copy := obj.DeepCopy()
	copy.Status.Revision = obj.ResourceVersion
	return &copy.Status
}
<commit_msg>Remove CA Certs/Key from listenerconfig storage<commit_after>package tls

import (
	"context"

	"github.com/rancher/dynamiclistener"
	v1 "github.com/rancher/k3s/pkg/apis/k3s.cattle.io/v1"
	k3sclient "github.com/rancher/k3s/pkg/generated/controllers/k3s.cattle.io/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	ns   = "kube-system"
	name = "tls-config"
)

func NewServer(ctx context.Context, listenerConfigs k3sclient.ListenerConfigController, config dynamiclistener.UserConfig) (dynamiclistener.ServerInterface, error) {
	storage := &listenerConfigStorage{
		client: listenerConfigs,
		cache:  listenerConfigs.Cache(),
		config: config,
	}

	server, err := dynamiclistener.NewServer(storage, config)
	if err != nil {
		return nil, err
	}

	listenerConfigs.OnChange(ctx, "listen-config", func(key string, obj *v1.ListenerConfig) (*v1.ListenerConfig, error) {
		if obj == nil {
			return nil, nil
		}
		return obj, server.Update(storage.fromStorage(obj))
	})

	return server, err
}

type listenerConfigStorage struct {
	cache  k3sclient.ListenerConfigCache
	client k3sclient.ListenerConfigClient
	config dynamiclistener.UserConfig
}

func (l *listenerConfigStorage) Set(config *dynamiclistener.ListenerStatus) (*dynamiclistener.ListenerStatus, error) {
	if config == nil {
		return nil, nil
	}

	obj, err := l.cache.Get(ns, name)
	if errors.IsNotFound(err) {
		ls := v1.NewListenerConfig(ns, name, v1.ListenerConfig{
			Status: *config,
		})

		ls, err := l.client.Create(ls)
		return l.fromStorage(ls), err
	} else if err != nil {
		return nil, err
	}

	obj = obj.DeepCopy()
	obj.ResourceVersion = config.Revision
	obj.Status = *config
	obj.Status.Revision = ""

	if l.config.CACerts != "" && l.config.CAKey != "" {
		obj.Status.CACert = ""
		obj.Status.CAKey = ""
	}

	obj, err = l.client.Update(obj)
	return l.fromStorage(obj), err
}

func (l *listenerConfigStorage) Get() (*dynamiclistener.ListenerStatus, error) {
	obj, err := l.cache.Get(ns, name)
	if errors.IsNotFound(err) {
		obj, err = l.client.Get(ns, name, metav1.GetOptions{})
	}
	if errors.IsNotFound(err) {
		return &dynamiclistener.ListenerStatus{}, nil
	}
	return l.fromStorage(obj), err
}

func (l *listenerConfigStorage) fromStorage(obj *v1.ListenerConfig) *dynamiclistener.ListenerStatus {
	if obj == nil {
		return nil
	}

	copy := obj.DeepCopy()
	copy.Status.Revision = obj.ResourceVersion

	if l.config.CACerts != "" && l.config.CAKey != "" {
		copy.Status.CACert = l.config.CACerts
		copy.Status.CAKey = l.config.CAKey
	}

	return &copy.Status
}
<|endoftext|>


<commit_before>package query

import (
	"core"
	"fmt"
	"sort"
)

// ReverseDeps For each input label, finds all targets which depend upon it.
func ReverseDeps(state *core.BuildState, labels []core.BuildLabel) {

	uniqueTargets := make(map[*core.BuildTarget]struct{})

	graph := state.Graph
	for _, label := range labels {
		//labelTarget := graph.TargetOrDie(label)
		for _, child := range graph.PackageOrDie(label).AllChildren(graph.TargetOrDie(label)) {
			for _, target := range graph.ReverseDependencies(child) {
				if parent := target.Parent(graph); parent != nil {
					uniqueTargets[parent] = struct{}{}
				} else {
					uniqueTargets[target] = struct{}{}
				}
			}
		}
	}
	// Check for anything subincluding this guy too
	for _, pkg := range graph.PackageMap() {
		for _, label := range labels {
			if pkg.HasSubinclude(label) {
				for _, target := range pkg.AllTargets() {
					uniqueTargets[target] = struct{}{}
				}
			}
		}
	}

	targets := make(core.BuildLabels, 0, len(uniqueTargets))
	for _, label := range labels {
		delete(uniqueTargets, graph.TargetOrDie(label))
	}
	for target := range uniqueTargets {
		if state.ShouldInclude(target) {
			targets = append(targets, target.Label)
		}
	}
	sort.Sort(targets)
	for _, target := range targets {
		fmt.Printf("%s\n", target)
	}
}
<commit_msg>commit<commit_after>package query

import (
	"core"
	"fmt"
	"sort"
)

// ReverseDeps For each input label, finds all targets which depend upon it.
func ReverseDeps(state *core.BuildState, labels []core.BuildLabel) {

	uniqueTargets := make(map[*core.BuildTarget]struct{})

	graph := state.Graph
	for _, label := range labels {
		for _, child := range graph.PackageOrDie(label).AllChildren(graph.TargetOrDie(label)) {
			for _, target := range graph.ReverseDependencies(child) {
				if parent := target.Parent(graph); parent != nil {
					uniqueTargets[parent] = struct{}{}
				} else {
					uniqueTargets[target] = struct{}{}
				}
			}
		}
	}
	// Check for anything subincluding this guy too
	for _, pkg := range graph.PackageMap() {
		for _, label := range labels {
			if pkg.HasSubinclude(label) {
				for _, target := range pkg.AllTargets() {
					uniqueTargets[target] = struct{}{}
				}
			}
		}
	}

	targets := make(core.BuildLabels, 0, len(uniqueTargets))
	for _, label := range labels {
		delete(uniqueTargets, graph.TargetOrDie(label))
	}
	for target := range uniqueTargets {
		if state.ShouldInclude(target) {
			targets = append(targets, target.Label)
		}
	}
	sort.Sort(targets)
	for _, target := range targets {
		fmt.Printf("%s\n", target)
	}
}
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/hashicorp/hil/ast"
	"github.com/mitchellh/reflectwalk"
)

func TestInterpolationWalker_detect(t *testing.T) {
	cases := []struct {
		Input  interface{}
		Result []string
	}{
		{
			Input: map[string]interface{}{
				"foo": "$${var.foo}",
			},
			Result: []string{
				"Literal(TypeString, ${var.foo})",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${var.foo}",
			},
			Result: []string{
				"Variable(var.foo)",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${aws_instance.foo.*.num}",
			},
			Result: []string{
				"Variable(aws_instance.foo.*.num)",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${lookup(var.foo)}",
			},
			Result: []string{
				"Call(lookup, Variable(var.foo))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${file("test.txt")}`,
			},
			Result: []string{
				"Call(file, Literal(TypeString, test.txt))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${file("foo/bar.txt")}`,
			},
			Result: []string{
				"Call(file, Literal(TypeString, foo/bar.txt))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${join(",", foo.bar.*.id)}`,
			},
			Result: []string{
				"Call(join, Literal(TypeString, ,), Variable(foo.bar.*.id))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${concat("localhost", ":8080")}`,
			},
			Result: []string{
				"Call(concat, Literal(TypeString, localhost), Literal(TypeString, :8080))",
			},
		},
	}

	for i, tc := range cases {
		var actual []string
		detectFn := func(root ast.Node) (interface{}, error) {
			actual = append(actual, fmt.Sprintf("%s", root))
			return "", nil
		}

		w := &interpolationWalker{F: detectFn}
		if err := reflectwalk.Walk(tc.Input, w); err != nil {
			t.Fatalf("err: %s", err)
		}

		if !reflect.DeepEqual(actual, tc.Result) {
			t.Fatalf("%d: bad:\n\n%#v", i, actual)
		}
	}
}

func TestInterpolationWalker_replace(t *testing.T) {
	cases := []struct {
		Input  interface{}
		Output interface{}
		Value  interface{}
	}{
		{
			Input: map[string]interface{}{
				"foo": "$${var.foo}",
			},
			Output: map[string]interface{}{
				"foo": "bar",
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": "hello, ${var.foo}",
			},
			Output: map[string]interface{}{
				"foo": "bar",
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": map[string]interface{}{
					"${var.foo}": "bar",
				},
			},
			Output: map[string]interface{}{
				"foo": map[string]interface{}{
					"bar": "bar",
				},
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": []interface{}{
					"${var.foo}",
					"bing",
				},
			},
			Output: map[string]interface{}{
				"foo": []interface{}{
					"bar",
					"baz",
					"bing",
				},
			},
			Value: []interface{}{"bar", "baz"},
		},

		{
			Input: map[string]interface{}{
				"foo": []interface{}{
					"${var.foo}",
					"bing",
				},
			},
			Output: map[string]interface{}{},
			Value:  []interface{}{UnknownVariableValue, "baz"},
		},
	}

	for i, tc := range cases {
		fn := func(ast.Node) (interface{}, error) {
			return tc.Value, nil
		}

		w := &interpolationWalker{F: fn, Replace: true}
		if err := reflectwalk.Walk(tc.Input, w); err != nil {
			t.Fatalf("err: %s", err)
		}

		if !reflect.DeepEqual(tc.Input, tc.Output) {
			t.Fatalf("%d: bad:\n\nexpected:%#v\ngot:%#v", i, tc.Output, tc.Input)
		}
	}
}
<commit_msg>Update reflectwalk to fix failing tests<commit_after>package config

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/hashicorp/hil/ast"
	"github.com/mitchellh/reflectwalk"
)

func TestInterpolationWalker_detect(t *testing.T) {
	cases := []struct {
		Input  interface{}
		Result []string
	}{
		{
			Input: map[string]interface{}{
				"foo": "$${var.foo}",
			},
			Result: []string{
				"Literal(TypeString, ${var.foo})",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${var.foo}",
			},
			Result: []string{
				"Variable(var.foo)",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${aws_instance.foo.*.num}",
			},
			Result: []string{
				"Variable(aws_instance.foo.*.num)",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": "${lookup(var.foo)}",
			},
			Result: []string{
				"Call(lookup, Variable(var.foo))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${file("test.txt")}`,
			},
			Result: []string{
				"Call(file, Literal(TypeString, test.txt))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${file("foo/bar.txt")}`,
			},
			Result: []string{
				"Call(file, Literal(TypeString, foo/bar.txt))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${join(",", foo.bar.*.id)}`,
			},
			Result: []string{
				"Call(join, Literal(TypeString, ,), Variable(foo.bar.*.id))",
			},
		},

		{
			Input: map[string]interface{}{
				"foo": `${concat("localhost", ":8080")}`,
			},
			Result: []string{
				"Call(concat, Literal(TypeString, localhost), Literal(TypeString, :8080))",
			},
		},
	}

	for i, tc := range cases {
		var actual []string
		detectFn := func(root ast.Node) (interface{}, error) {
			actual = append(actual, fmt.Sprintf("%s", root))
			return "", nil
		}

		w := &interpolationWalker{F: detectFn}
		if err := reflectwalk.Walk(tc.Input, w); err != nil {
			t.Fatalf("err: %s", err)
		}

		if !reflect.DeepEqual(actual, tc.Result) {
			t.Fatalf("%d: bad:\n\n%#v", i, actual)
		}
	}
}

func TestInterpolationWalker_replace(t *testing.T) {
	cases := []struct {
		Input  interface{}
		Output interface{}
		Value  interface{}
	}{
		{
			Input: map[string]interface{}{
				"foo": "$${var.foo}",
			},
			Output: map[string]interface{}{
				"foo": "bar",
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": "hello, ${var.foo}",
			},
			Output: map[string]interface{}{
				"foo": "bar",
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": map[string]interface{}{
					"${var.foo}": "bar",
				},
			},
			Output: map[string]interface{}{
				"foo": map[string]interface{}{
					"bar": "bar",
				},
			},
			Value: "bar",
		},

		{
			Input: map[string]interface{}{
				"foo": []interface{}{
					"${var.foo}",
					"bing",
				},
			},
			Output: map[string]interface{}{
				"foo": []interface{}{
					"bar",
					"baz",
					"bing",
				},
			},
			Value: []interface{}{"bar", "baz"},
		},

		{
			Input: map[string]interface{}{
				"foo": []interface{}{
					"${var.foo}",
					"bing",
				},
			},
			Output: map[string]interface{}{},
			Value:  []interface{}{UnknownVariableValue, "baz"},
		},
	}

	for i, tc := range cases {
		fn := func(ast.Node) (interface{}, error) {
			return tc.Value, nil
		}

		t.Run(fmt.Sprintf("walk-%d", i), func(t *testing.T) {
			w := &interpolationWalker{F: fn, Replace: true}
			if err := reflectwalk.Walk(tc.Input, w); err != nil {
				t.Fatalf("err: %s", err)
			}

			if !reflect.DeepEqual(tc.Input, tc.Output) {
				t.Fatalf("%d: bad:\n\nexpected:%#v\ngot:%#v", i, tc.Output, tc.Input)
			}
		})
	}
}
<|endoftext|>


<commit_before>package main

import (
	"data"
	"fmt"
	"math/rand"
	"time"

	"github.com/skizzehq/goskizze/skizze"
)

func main() {
	words := data.GetData()
	// shuffle
	rand.Seed(time.Now().UnixNano())
	for i := range words {
		j := rand.Intn(i + 1)
		words[i], words[j] = words[j], words[i]
	}
	client, err := skizze.Dial("127.0.0.1:3596", skizze.Options{Insecure: true})

	if err != nil {
		fmt.Printf("Error connecting to Skizze: %s\n", err)
		return
	}
	domainName := "skizze_stress"
	if _, err := client.CreateDomain(domainName); err != nil {
		fmt.Println(err)
	}

	end := time.Duration(0)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	zipf := rand.NewZipf(r, 1.1, 1.1, 500000)
	totalAdds := uint64(0)
	for i := 0; i < 100; i++ {
		word := words[i]
		n := zipf.Uint64() + 1
		fmt.Printf("%d Push: %s (%d times)\n", i, word, n)
		totalAdds += n
		fill := make([]string, n, n)
		for j := 0; j < len(fill); j++ {
			fill[j] = word
		}
		t := time.Now()
		if err := client.AddToDomain(domainName, fill...); err != nil {
			fmt.Println(err)
			return
		}
		end += time.Since(t)
	}

	client.Close()
	fmt.Printf("Added %d values (%d unique) in %ds\n", totalAdds, len(words), int(end.Seconds()))
}
<commit_msg>update stress<commit_after>package main

import (
	"data"
	"fmt"
	"math/rand"
	"time"

	"github.com/skizzehq/goskizze/skizze"
)

func main() {
	words := data.GetData()
	// shuffle
	rand.Seed(time.Now().UnixNano())
	for i := range words {
		j := rand.Intn(i + 1)
		words[i], words[j] = words[j], words[i]
	}
	client, err := skizze.Dial("127.0.0.1:3596", skizze.Options{Insecure: true})

	if err != nil {
		fmt.Printf("Error connecting to Skizze: %s\n", err)
		return
	}
	domainName := "skizze_stress"
	if _, err := client.CreateDomain(domainName); err != nil {
		fmt.Println(err)
	}

	end := time.Duration(0)
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	zipf := rand.NewZipf(r, 1.1, 1.1, uint64(len(words)-1))
	totalAdds := 0
	for i := 0; i < 100000; i++ {
		fill := make([]string, 1000, 1000)
		for j := 0; j < len(fill); j++ {
			k := zipf.Uint64()
			fill[j] = words[k]
		}
		totalAdds += len(fill)

		t := time.Now()
		if err := client.AddToDomain(domainName, fill...); err != nil {
			fmt.Println(err)
			return
		}
		end += time.Since(t)
		if end.Seconds() > 0 {
			fmt.Printf("Added %d values (%d unique) in %ds (avg. %d v/s)\n", totalAdds, len(words), int(end.Seconds()), totalAdds/int(end.Seconds()+1))
		}
	}

	client.Close()
	fmt.Printf("Added %d values (%d unique) in %ds (avg. %d v/s)\n", totalAdds, len(words), int(end.Seconds()), totalAdds/int(end.Seconds()+1))
}
<|endoftext|>


<commit_before>// Copyright © 2017 ben dewan <benj.dewan@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package connection

import (
	"fmt"

	compose "github.com/benjdewan/gocomposeapi"
)

func update(cxn *Connection, dep Deployment) error {
	existing, ok := cxn.getDeploymentByName(dep.GetName())
	if !ok {
		return fmt.Errorf("Attempting to update '%s', but it doesn't exist",
			dep.GetName())
	}

	id := existing.ID
	timeout := dep.GetTimeout()

	scaling := dep.GetScaling()
	if err := updateScalings(cxn, id, scaling, timeout); err != nil {
		return err
	}

	version := dep.GetVersion()
	if len(version) > 0 {
		if err := updateVersion(cxn, id, version, timeout); err != nil {
			return err
		}
	}

	if err := addTeamRoles(cxn, id, dep.GetTeamRoles()); err != nil {
		return err
	}

	// Changing versions and sizes can change the deployment ID. Ensure
	// we have the latest/live value
	updatedDeployment, errs := cxn.client.GetDeploymentByName(dep.GetName())
	if len(errs) != 0 {
		return fmt.Errorf("Unable to get deployment information for '%s':\n%s",
			dep.GetName(), errsOut(errs))
	}
	cxn.newDeploymentIDs.Store(updatedDeployment.ID, struct{}{})
	return nil
}

func updateScalings(cxn *Connection, ID string, newScale int, timeout float64) error {
	existingScalings, errs := cxn.client.GetScalings(ID)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to get current scaling for '%s':\n%v",
			ID, errsOut(errs))
	}

	if existingScalings.AllocatedUnits == newScale {
		return nil
	}

	params := compose.ScalingsParams{
		DeploymentID: ID,
		Units:        newScale,
	}

	recipe, errs := cxn.client.SetScalings(params)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to resize '%s':\n%v",
			ID, errsOut(errs))
	}

	err := cxn.waitOnRecipe(recipe.ID, timeout)
	if err != nil {
		return err
	}
	return nil
}

func updateVersion(cxn *Connection, ID, newVersion string, timeout float64) error {
	deployment, errs := cxn.client.GetDeployment(ID)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to fetch current deployment information for '%s':\n%v",
			ID, errsOut(errs))
	}

	if deployment.Version == newVersion {
		return nil
	}

	transitions, errs := cxn.client.GetVersionsForDeployment(ID)
	if len(errs) != 0 || transitions == nil {
		return fmt.Errorf("Error fetching upgrade information for '%s':\n%v",
			ID, errsOut(errs))
	}

	validTransition := false
	for _, transition := range *transitions {
		if transition.ToVersion == newVersion {
			validTransition = true
			break
		}
	}
	if !validTransition {
		return fmt.Errorf("Cannot upgrade '%s' to version '%s'.",
			ID, newVersion)
	}

	recipe, errs := cxn.client.UpdateVersion(ID, newVersion)
	if errs != nil {
		return fmt.Errorf("Unable to upgrade '%s' to version '%s':\n%v",
			ID, newVersion, errsOut(errs))
	}

	err := cxn.waitOnRecipe(recipe.ID, timeout)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Remove unnecessary assignment<commit_after>// Copyright © 2017 ben dewan <benj.dewan@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package connection

import (
	"fmt"

	compose "github.com/benjdewan/gocomposeapi"
)

func update(cxn *Connection, dep Deployment) error {
	existing, ok := cxn.getDeploymentByName(dep.GetName())
	if !ok {
		return fmt.Errorf("Attempting to update '%s', but it doesn't exist",
			dep.GetName())
	}

	id := existing.ID
	timeout := dep.GetTimeout()

	if err := updateScalings(cxn, id, dep.GetScaling(), timeout); err != nil {
		return err
	}

	version := dep.GetVersion()
	if len(version) > 0 {
		if err := updateVersion(cxn, id, version, timeout); err != nil {
			return err
		}
	}

	if err := addTeamRoles(cxn, id, dep.GetTeamRoles()); err != nil {
		return err
	}

	// Changing versions and sizes can change the deployment ID. Ensure
	// we have the latest/live value
	updatedDeployment, errs := cxn.client.GetDeploymentByName(dep.GetName())
	if len(errs) != 0 {
		return fmt.Errorf("Unable to get deployment information for '%s':\n%s",
			dep.GetName(), errsOut(errs))
	}
	cxn.newDeploymentIDs.Store(updatedDeployment.ID, struct{}{})
	return nil
}

func updateScalings(cxn *Connection, ID string, newScale int, timeout float64) error {
	existingScalings, errs := cxn.client.GetScalings(ID)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to get current scaling for '%s':\n%v",
			ID, errsOut(errs))
	}

	if existingScalings.AllocatedUnits == newScale {
		return nil
	}

	params := compose.ScalingsParams{
		DeploymentID: ID,
		Units:        newScale,
	}

	recipe, errs := cxn.client.SetScalings(params)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to resize '%s':\n%v",
			ID, errsOut(errs))
	}

	err := cxn.waitOnRecipe(recipe.ID, timeout)
	if err != nil {
		return err
	}
	return nil
}

func updateVersion(cxn *Connection, ID, newVersion string, timeout float64) error {
	deployment, errs := cxn.client.GetDeployment(ID)
	if len(errs) != 0 {
		return fmt.Errorf("Unable to fetch current deployment information for '%s':\n%v",
			ID, errsOut(errs))
	}

	if deployment.Version == newVersion {
		return nil
	}

	transitions, errs := cxn.client.GetVersionsForDeployment(ID)
	if len(errs) != 0 || transitions == nil {
		return fmt.Errorf("Error fetching upgrade information for '%s':\n%v",
			ID, errsOut(errs))
	}

	validTransition := false
	for _, transition := range *transitions {
		if transition.ToVersion == newVersion {
			validTransition = true
			break
		}
	}
	if !validTransition {
		return fmt.Errorf("Cannot upgrade '%s' to version '%s'.",
			ID, newVersion)
	}

	recipe, errs := cxn.client.UpdateVersion(ID, newVersion)
	if errs != nil {
		return fmt.Errorf("Unable to upgrade '%s' to version '%s':\n%v",
			ID, newVersion, errsOut(errs))
	}

	err := cxn.waitOnRecipe(recipe.ID, timeout)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/wvanbergen/kazoo-go"
)

var (
	AlreadyClosing = errors.New("The consumer group is already shutting down.")
)

type Config struct {
	*sarama.Config

	Zookeeper *kazoo.Config

	Offsets struct {
		Initial           int64         // The initial offset method to use if the consumer has no previously stored offset. Must be either sarama.OffsetOldest (default) or sarama.OffsetNewest.
		ProcessingTimeout time.Duration // Time to wait for all the offsets for a partition to be processed after stopping to consume from it. Defaults to 1 minute.
		CommitInterval    time.Duration // The interval between which the prossed offsets are commited.
	}
}

func NewConfig() *Config {
	config := &Config{}
	config.Config = sarama.NewConfig()
	config.Zookeeper = kazoo.NewConfig()
	config.Offsets.Initial = sarama.OffsetOldest
	config.Offsets.ProcessingTimeout = 60 * time.Second
	config.Offsets.CommitInterval = 10 * time.Second

	return config
}

func (cgc *Config) Validate() error {
	if cgc.Zookeeper.Timeout <= 0 {
		return sarama.ConfigurationError("ZookeeperTimeout should have a duration > 0")
	}

	if cgc.Offsets.CommitInterval <= 0 {
		return sarama.ConfigurationError("CommitInterval should have a duration > 0")
	}

	if cgc.Offsets.Initial != sarama.OffsetOldest && cgc.Offsets.Initial != sarama.OffsetNewest {
		return errors.New("Offsets.Initial should be sarama.OffsetOldest or sarama.OffsetNewest.")
	}

	if cgc.Config != nil {
		if err := cgc.Config.Validate(); err != nil {
			return err
		}
	}

	return nil
}

// The ConsumerGroup type holds all the information for a consumer that is part
// of a consumer group. Call JoinConsumerGroup to start a consumer.
type ConsumerGroup struct {
	config *Config

	consumer sarama.Consumer
	kazoo    *kazoo.Kazoo
	group    *kazoo.Consumergroup
	instance *kazoo.ConsumergroupInstance

	wg             sync.WaitGroup
	singleShutdown sync.Once

	messages chan *sarama.ConsumerMessage
	errors   chan *sarama.ConsumerError
	stopper  chan struct{}

	consumers kazoo.ConsumergroupInstanceList

	offsetManager OffsetManager
}

// Connects to a consumer group, using Zookeeper for auto-discovery
func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *Config) (cg *ConsumerGroup, err error) {

	if name == "" {
		return nil, sarama.ConfigurationError("Empty consumergroup name")
	}

	if len(topics) == 0 {
		return nil, sarama.ConfigurationError("No topics provided")
	}

	if len(zookeeper) == 0 {
		return nil, errors.New("You need to provide at least one zookeeper node address!")
	}

	if config == nil {
		config = NewConfig()
	}
	config.ClientID = name

	// Validate configuration
	if err = config.Validate(); err != nil {
		return
	}

	var kz *kazoo.Kazoo
	if kz, err = kazoo.NewKazoo(zookeeper, config.Zookeeper); err != nil {
		return
	}

	brokers, err := kz.BrokerList()
	if err != nil {
		kz.Close()
		return
	}

	group := kz.Consumergroup(name)
	instance := group.NewInstance()

	var consumer sarama.Consumer
	if consumer, err = sarama.NewConsumer(brokers, config.Config); err != nil {
		kz.Close()
		return
	}

	cg = &ConsumerGroup{
		config:   config,
		consumer: consumer,

		kazoo:    kz,
		group:    group,
		instance: instance,

		messages: make(chan *sarama.ConsumerMessage, config.ChannelBufferSize),
		errors:   make(chan *sarama.ConsumerError, config.ChannelBufferSize),
		stopper:  make(chan struct{}),
	}

	// Register consumer group
	if exists, err := cg.group.Exists(); err != nil {
		cg.Logf("FAILED to check for existence of consumergroup: %s!\n", err)
		_ = consumer.Close()
		_ = kz.Close()
		return nil, err
	} else if !exists {
		cg.Logf("Consumergroup `%s` does not yet exists, creating...\n", cg.group.Name)
		if err := cg.group.Create(); err != nil {
			cg.Logf("FAILED to create consumergroup in Zookeeper: %s!\n", err)
			_ = consumer.Close()
			_ = kz.Close()
			return nil, err
		}
	}

	// Register itself with zookeeper
	if err := cg.instance.Register(topics); err != nil {
		cg.Logf("FAILED to register consumer instance: %s!\n", err)
		return nil, err
	} else {
		cg.Logf("Consumer instance registered (%s).", cg.instance.ID)
	}

	offsetConfig := OffsetManagerConfig{CommitInterval: config.Offsets.CommitInterval}
	cg.offsetManager = NewZookeeperOffsetManager(cg, &offsetConfig)

	go cg.topicListConsumer(topics)

	return
}

// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Messages() <-chan *sarama.ConsumerMessage {
	return cg.messages
}

// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Errors() <-chan *sarama.ConsumerError {
	return cg.errors
}

func (cg *ConsumerGroup) Closed() bool {
	return cg.instance == nil
}

func (cg *ConsumerGroup) Close() error {
	shutdownError := AlreadyClosing
	cg.singleShutdown.Do(func() {
		defer cg.kazoo.Close()

		shutdownError = nil

		close(cg.stopper)
		cg.wg.Wait()

		if err := cg.offsetManager.Close(); err != nil {
			cg.Logf("FAILED closing the offset manager: %s!\n", err)
		}

		if shutdownError = cg.instance.Deregister(); shutdownError != nil {
			cg.Logf("FAILED deregistering consumer instance: %s!\n", shutdownError)
		} else {
			cg.Logf("Deregistered consumer instance %s.\n", cg.instance.ID)
		}

		if shutdownError = cg.consumer.Close(); shutdownError != nil {
			cg.Logf("FAILED closing the Sarama client: %s\n", shutdownError)
		}

		close(cg.messages)
		close(cg.errors)
		cg.instance = nil
	})

	return shutdownError
}

func (cg *ConsumerGroup) Logf(format string, args ...interface{}) {
	var identifier string
	if cg.instance == nil {
		identifier = "(defunct)"
	} else {
		identifier = cg.instance.ID[len(cg.instance.ID)-12:]
	}
	sarama.Logger.Printf("[%s/%s] %s", cg.group.Name, identifier, fmt.Sprintf(format, args...))
}

func (cg *ConsumerGroup) InstanceRegistered() (bool, error) {
	return cg.instance.Registered()
}

func (cg *ConsumerGroup) CommitUpto(message *sarama.ConsumerMessage) error {
	cg.offsetManager.MarkAsProcessed(message.Topic, message.Partition, message.Offset)
	return nil
}

func (cg *ConsumerGroup) topicListConsumer(topics []string) {
	for {
		select {
		case <-cg.stopper:
			return
		default:
		}

		consumers, consumerChanges, err := cg.group.WatchInstances()
		if err != nil {
			cg.Logf("FAILED to get list of registered consumer instances: %s\n", err)
			return
		}

		cg.consumers = consumers
		cg.Logf("Currently registered consumers: %d\n", len(cg.consumers))

		stopper := make(chan struct{})

		for _, topic := range topics {
			cg.wg.Add(1)
			go cg.topicConsumer(topic, cg.messages, cg.errors, stopper)
		}

		select {
		case <-cg.stopper:
			close(stopper)
			return

		case <-consumerChanges:
			cg.Logf("Triggering rebalance due to consumer list change\n")
			close(stopper)
			cg.wg.Wait()
		}
	}
}

func (cg *ConsumerGroup) topicConsumer(topic string, messages chan<- *sarama.ConsumerMessage, errors chan<- *sarama.ConsumerError, stopper <-chan struct{}) {
	defer cg.wg.Done()

	select {
	case <-stopper:
		return
	default:
	}

	cg.Logf("%s :: Started topic consumer\n", topic)

	// Fetch a list of partition IDs
	partitions, err := cg.kazoo.Topic(topic).Partitions()
	if err != nil {
		cg.Logf("%s :: FAILED to get list of partitions: %s\n", topic, err)
		cg.errors <- &sarama.ConsumerError{
			Topic:     topic,
			Partition: -1,
			Err:       err,
		}
		return
	}

	partitionLeaders, err := retrievePartitionLeaders(partitions)
	if err != nil {
		cg.Logf("%s :: FAILED to get leaders of partitions: %s\n", topic, err)
		cg.errors <- &sarama.ConsumerError{
			Topic:     topic,
			Partition: -1,
			Err:       err,
		}
		return
	}

	dividedPartitions := dividePartitionsBetweenConsumers(cg.consumers, partitionLeaders)
	myPartitions := dividedPartitions[cg.instance.ID]
	cg.Logf("%s :: Claiming %d of %d partitions", topic, len(myPartitions), len(partitionLeaders))

	// Consume all the assigned partitions
	var wg sync.WaitGroup
	for _, pid := range myPartitions {

		wg.Add(1)
		go cg.partitionConsumer(topic, pid.ID, messages, errors, &wg, stopper)
	}

	wg.Wait()
	cg.Logf("%s :: Stopped topic consumer\n", topic)
}

// Consumes a partition
func (cg *ConsumerGroup) partitionConsumer(topic string, partition int32, messages chan<- *sarama.ConsumerMessage, errors chan<- *sarama.ConsumerError, wg *sync.WaitGroup, stopper <-chan struct{}) {
	defer wg.Done()

	select {
	case <-stopper:
		return
	default:
	}

	for maxRetries, tries := 3, 0; tries < maxRetries; tries++ {
		if err := cg.instance.ClaimPartition(topic, partition); err == nil {
			break
		} else if err == kazoo.ErrPartitionClaimedByOther && tries+1 < maxRetries {
			time.Sleep(1 * time.Second)
		} else {
			cg.Logf("%s/%d :: FAILED to claim the partition: %s\n", topic, partition, err)
			return
		}
	}
	defer cg.instance.ReleasePartition(topic, partition)

	nextOffset, err := cg.offsetManager.InitializePartition(topic, partition)
	if err != nil {
		cg.Logf("%s/%d :: FAILED to determine initial offset: %s\n", topic, partition, err)
		return
	}

	if nextOffset >= 0 {
		cg.Logf("%s/%d :: Partition consumer starting at offset %d.\n", topic, partition, nextOffset)
	} else {
		nextOffset = cg.config.Offsets.Initial
		if nextOffset == sarama.OffsetOldest {
			cg.Logf("%s/%d :: Partition consumer starting at the oldest available offset.\n", topic, partition)
		} else if nextOffset == sarama.OffsetNewest {
			cg.Logf("%s/%d :: Partition consumer listening for new messages only.\n", topic, partition)
		}
	}

	consumer, err := cg.consumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		cg.Logf("%s/%d :: Partition consumer offset out of Range.\n", topic, partition)
		// if the offset is out of range, simplistically decide whether to use OffsetNewest or OffsetOldest
		// if the configuration specified offsetOldest, then switch to the oldest available offset, else
		// switch to the newest available offset.
		if cg.config.Offsets.Initial == sarama.OffsetOldest {
			nextOffset = sarama.OffsetOldest
			cg.Logf("%s/%d :: Partition consumer offset reset to oldest available offset.\n", topic, partition)
		} else {
			nextOffset = sarama.OffsetNewest
			cg.Logf("%s/%d :: Partition consumer offset reset to newest available offset.\n", topic, partition)
		}
		// retry the consumePartition with the adjusted offset
		consumer, err = cg.consumer.ConsumePartition(topic, partition, nextOffset)
	}
	if err != nil {
		cg.Logf("%s/%d :: FAILED to start partition consumer: %s\n", topic, partition, err)
		return
	}
	defer consumer.Close()

	err = nil
	var lastOffset int64 = -1 // aka unknown
partitionConsumerLoop:
	for {
		select {
		case <-stopper:
			break partitionConsumerLoop

		case err := <-consumer.Errors():
			for {
				select {
				case errors <- err:
					continue partitionConsumerLoop

				case <-stopper:
					break partitionConsumerLoop
				}
			}

		case message := <-consumer.Messages():
			for {
				select {
				case <-stopper:
					break partitionConsumerLoop

				case messages <- message:
					lastOffset = message.Offset
					continue partitionConsumerLoop
				}
			}
		}
	}

	cg.Logf("%s/%d :: Stopping partition consumer at offset %d\n", topic, partition, lastOffset)
	if err := cg.offsetManager.FinalizePartition(topic, partition, lastOffset, cg.config.Offsets.ProcessingTimeout); err != nil {
		cg.Logf("%s/%d :: %s\n", topic, partition, err)
	}
}
<commit_msg>Fix typo in config godoc<commit_after>package consumergroup

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/wvanbergen/kazoo-go"
)

var (
	AlreadyClosing = errors.New("The consumer group is already shutting down.")
)

type Config struct {
	*sarama.Config

	Zookeeper *kazoo.Config

	Offsets struct {
		Initial           int64         // The initial offset method to use if the consumer has no previously stored offset. Must be either sarama.OffsetOldest (default) or sarama.OffsetNewest.
		ProcessingTimeout time.Duration // Time to wait for all the offsets for a partition to be processed after stopping to consume from it. Defaults to 1 minute.
		CommitInterval    time.Duration // The interval between which the processed offsets are commited.
	}
}

func NewConfig() *Config {
	config := &Config{}
	config.Config = sarama.NewConfig()
	config.Zookeeper = kazoo.NewConfig()
	config.Offsets.Initial = sarama.OffsetOldest
	config.Offsets.ProcessingTimeout = 60 * time.Second
	config.Offsets.CommitInterval = 10 * time.Second

	return config
}

func (cgc *Config) Validate() error {
	if cgc.Zookeeper.Timeout <= 0 {
		return sarama.ConfigurationError("ZookeeperTimeout should have a duration > 0")
	}

	if cgc.Offsets.CommitInterval <= 0 {
		return sarama.ConfigurationError("CommitInterval should have a duration > 0")
	}

	if cgc.Offsets.Initial != sarama.OffsetOldest && cgc.Offsets.Initial != sarama.OffsetNewest {
		return errors.New("Offsets.Initial should be sarama.OffsetOldest or sarama.OffsetNewest.")
	}

	if cgc.Config != nil {
		if err := cgc.Config.Validate(); err != nil {
			return err
		}
	}

	return nil
}

// The ConsumerGroup type holds all the information for a consumer that is part
// of a consumer group. Call JoinConsumerGroup to start a consumer.
type ConsumerGroup struct {
	config *Config

	consumer sarama.Consumer
	kazoo    *kazoo.Kazoo
	group    *kazoo.Consumergroup
	instance *kazoo.ConsumergroupInstance

	wg             sync.WaitGroup
	singleShutdown sync.Once

	messages chan *sarama.ConsumerMessage
	errors   chan *sarama.ConsumerError
	stopper  chan struct{}

	consumers kazoo.ConsumergroupInstanceList

	offsetManager OffsetManager
}

// Connects to a consumer group, using Zookeeper for auto-discovery
func JoinConsumerGroup(name string, topics []string, zookeeper []string, config *Config) (cg *ConsumerGroup, err error) {

	if name == "" {
		return nil, sarama.ConfigurationError("Empty consumergroup name")
	}

	if len(topics) == 0 {
		return nil, sarama.ConfigurationError("No topics provided")
	}

	if len(zookeeper) == 0 {
		return nil, errors.New("You need to provide at least one zookeeper node address!")
	}

	if config == nil {
		config = NewConfig()
	}
	config.ClientID = name

	// Validate configuration
	if err = config.Validate(); err != nil {
		return
	}

	var kz *kazoo.Kazoo
	if kz, err = kazoo.NewKazoo(zookeeper, config.Zookeeper); err != nil {
		return
	}

	brokers, err := kz.BrokerList()
	if err != nil {
		kz.Close()
		return
	}

	group := kz.Consumergroup(name)
	instance := group.NewInstance()

	var consumer sarama.Consumer
	if consumer, err = sarama.NewConsumer(brokers, config.Config); err != nil {
		kz.Close()
		return
	}

	cg = &ConsumerGroup{
		config:   config,
		consumer: consumer,

		kazoo:    kz,
		group:    group,
		instance: instance,

		messages: make(chan *sarama.ConsumerMessage, config.ChannelBufferSize),
		errors:   make(chan *sarama.ConsumerError, config.ChannelBufferSize),
		stopper:  make(chan struct{}),
	}

	// Register consumer group
	if exists, err := cg.group.Exists(); err != nil {
		cg.Logf("FAILED to check for existence of consumergroup: %s!\n", err)
		_ = consumer.Close()
		_ = kz.Close()
		return nil, err
	} else if !exists {
		cg.Logf("Consumergroup `%s` does not yet exists, creating...\n", cg.group.Name)
		if err := cg.group.Create(); err != nil {
			cg.Logf("FAILED to create consumergroup in Zookeeper: %s!\n", err)
			_ = consumer.Close()
			_ = kz.Close()
			return nil, err
		}
	}

	// Register itself with zookeeper
	if err := cg.instance.Register(topics); err != nil {
		cg.Logf("FAILED to register consumer instance: %s!\n", err)
		return nil, err
	} else {
		cg.Logf("Consumer instance registered (%s).", cg.instance.ID)
	}

	offsetConfig := OffsetManagerConfig{CommitInterval: config.Offsets.CommitInterval}
	cg.offsetManager = NewZookeeperOffsetManager(cg, &offsetConfig)

	go cg.topicListConsumer(topics)

	return
}

// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Messages() <-chan *sarama.ConsumerMessage {
	return cg.messages
}

// Returns a channel that you can read to obtain events from Kafka to process.
func (cg *ConsumerGroup) Errors() <-chan *sarama.ConsumerError {
	return cg.errors
}

func (cg *ConsumerGroup) Closed() bool {
	return cg.instance == nil
}

func (cg *ConsumerGroup) Close() error {
	shutdownError := AlreadyClosing
	cg.singleShutdown.Do(func() {
		defer cg.kazoo.Close()

		shutdownError = nil

		close(cg.stopper)
		cg.wg.Wait()

		if err := cg.offsetManager.Close(); err != nil {
			cg.Logf("FAILED closing the offset manager: %s!\n", err)
		}

		if shutdownError = cg.instance.Deregister(); shutdownError != nil {
			cg.Logf("FAILED deregistering consumer instance: %s!\n", shutdownError)
		} else {
			cg.Logf("Deregistered consumer instance %s.\n", cg.instance.ID)
		}

		if shutdownError = cg.consumer.Close(); shutdownError != nil {
			cg.Logf("FAILED closing the Sarama client: %s\n", shutdownError)
		}

		close(cg.messages)
		close(cg.errors)
		cg.instance = nil
	})

	return shutdownError
}

func (cg *ConsumerGroup) Logf(format string, args ...interface{}) {
	var identifier string
	if cg.instance == nil {
		identifier = "(defunct)"
	} else {
		identifier = cg.instance.ID[len(cg.instance.ID)-12:]
	}
	sarama.Logger.Printf("[%s/%s] %s", cg.group.Name, identifier, fmt.Sprintf(format, args...))
}

func (cg *ConsumerGroup) InstanceRegistered() (bool, error) {
	return cg.instance.Registered()
}

func (cg *ConsumerGroup) CommitUpto(message *sarama.ConsumerMessage) error {
	cg.offsetManager.MarkAsProcessed(message.Topic, message.Partition, message.Offset)
	return nil
}

func (cg *ConsumerGroup) topicListConsumer(topics []string) {
	for {
		select {
		case <-cg.stopper:
			return
		default:
		}

		consumers, consumerChanges, err := cg.group.WatchInstances()
		if err != nil {
			cg.Logf("FAILED to get list of registered consumer instances: %s\n", err)
			return
		}

		cg.consumers = consumers
		cg.Logf("Currently registered consumers: %d\n", len(cg.consumers))

		stopper := make(chan struct{})

		for _, topic := range topics {
			cg.wg.Add(1)
			go cg.topicConsumer(topic, cg.messages, cg.errors, stopper)
		}

		select {
		case <-cg.stopper:
			close(stopper)
			return

		case <-consumerChanges:
			cg.Logf("Triggering rebalance due to consumer list change\n")
			close(stopper)
			cg.wg.Wait()
		}
	}
}

func (cg *ConsumerGroup) topicConsumer(topic string, messages chan<- *sarama.ConsumerMessage, errors chan<- *sarama.ConsumerError, stopper <-chan struct{}) {
	defer cg.wg.Done()

	select {
	case <-stopper:
		return
	default:
	}

	cg.Logf("%s :: Started topic consumer\n", topic)

	// Fetch a list of partition IDs
	partitions, err := cg.kazoo.Topic(topic).Partitions()
	if err != nil {
		cg.Logf("%s :: FAILED to get list of partitions: %s\n", topic, err)
		cg.errors <- &sarama.ConsumerError{
			Topic:     topic,
			Partition: -1,
			Err:       err,
		}
		return
	}

	partitionLeaders, err := retrievePartitionLeaders(partitions)
	if err != nil {
		cg.Logf("%s :: FAILED to get leaders of partitions: %s\n", topic, err)
		cg.errors <- &sarama.ConsumerError{
			Topic:     topic,
			Partition: -1,
			Err:       err,
		}
		return
	}

	dividedPartitions := dividePartitionsBetweenConsumers(cg.consumers, partitionLeaders)
	myPartitions := dividedPartitions[cg.instance.ID]
	cg.Logf("%s :: Claiming %d of %d partitions", topic, len(myPartitions), len(partitionLeaders))

	// Consume all the assigned partitions
	var wg sync.WaitGroup
	for _, pid := range myPartitions {

		wg.Add(1)
		go cg.partitionConsumer(topic, pid.ID, messages, errors, &wg, stopper)
	}

	wg.Wait()
	cg.Logf("%s :: Stopped topic consumer\n", topic)
}

// Consumes a partition
func (cg *ConsumerGroup) partitionConsumer(topic string, partition int32, messages chan<- *sarama.ConsumerMessage, errors chan<- *sarama.ConsumerError, wg *sync.WaitGroup, stopper <-chan struct{}) {
	defer wg.Done()

	select {
	case <-stopper:
		return
	default:
	}

	for maxRetries, tries := 3, 0; tries < maxRetries; tries++ {
		if err := cg.instance.ClaimPartition(topic, partition); err == nil {
			break
		} else if err == kazoo.ErrPartitionClaimedByOther && tries+1 < maxRetries {
			time.Sleep(1 * time.Second)
		} else {
			cg.Logf("%s/%d :: FAILED to claim the partition: %s\n", topic, partition, err)
			return
		}
	}
	defer cg.instance.ReleasePartition(topic, partition)

	nextOffset, err := cg.offsetManager.InitializePartition(topic, partition)
	if err != nil {
		cg.Logf("%s/%d :: FAILED to determine initial offset: %s\n", topic, partition, err)
		return
	}

	if nextOffset >= 0 {
		cg.Logf("%s/%d :: Partition consumer starting at offset %d.\n", topic, partition, nextOffset)
	} else {
		nextOffset = cg.config.Offsets.Initial
		if nextOffset == sarama.OffsetOldest {
			cg.Logf("%s/%d :: Partition consumer starting at the oldest available offset.\n", topic, partition)
		} else if nextOffset == sarama.OffsetNewest {
			cg.Logf("%s/%d :: Partition consumer listening for new messages only.\n", topic, partition)
		}
	}

	consumer, err := cg.consumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		cg.Logf("%s/%d :: Partition consumer offset out of Range.\n", topic, partition)
		// if the offset is out of range, simplistically decide whether to use OffsetNewest or OffsetOldest
		// if the configuration specified offsetOldest, then switch to the oldest available offset, else
		// switch to the newest available offset.
		if cg.config.Offsets.Initial == sarama.OffsetOldest {
			nextOffset = sarama.OffsetOldest
			cg.Logf("%s/%d :: Partition consumer offset reset to oldest available offset.\n", topic, partition)
		} else {
			nextOffset = sarama.OffsetNewest
			cg.Logf("%s/%d :: Partition consumer offset reset to newest available offset.\n", topic, partition)
		}
		// retry the consumePartition with the adjusted offset
		consumer, err = cg.consumer.ConsumePartition(topic, partition, nextOffset)
	}
	if err != nil {
		cg.Logf("%s/%d :: FAILED to start partition consumer: %s\n", topic, partition, err)
		return
	}
	defer consumer.Close()

	err = nil
	var lastOffset int64 = -1 // aka unknown
partitionConsumerLoop:
	for {
		select {
		case <-stopper:
			break partitionConsumerLoop

		case err := <-consumer.Errors():
			for {
				select {
				case errors <- err:
					continue partitionConsumerLoop

				case <-stopper:
					break partitionConsumerLoop
				}
			}

		case message := <-consumer.Messages():
			for {
				select {
				case <-stopper:
					break partitionConsumerLoop

				case messages <- message:
					lastOffset = message.Offset
					continue partitionConsumerLoop
				}
			}
		}
	}

	cg.Logf("%s/%d :: Stopping partition consumer at offset %d\n", topic, partition, lastOffset)
	if err := cg.offsetManager.FinalizePartition(topic, partition, lastOffset, cg.config.Offsets.ProcessingTimeout); err != nil {
		cg.Logf("%s/%d :: %s\n", topic, partition, err)
	}
}
<|endoftext|>