<commit_before>// Copyright 2011 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.

package httptest

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
)

// ResponseRecorder is an implementation of http.ResponseWriter that
// records its mutations for later inspection in tests.
type ResponseRecorder struct {
	// Code is the HTTP response code set by WriteHeader.
	//
	// Note that if a Handler never calls WriteHeader or Write,
	// this might end up being 0, rather than the implicit
	// http.StatusOK. To get the implicit value, use the Result
	// method.
	Code int

	// HeaderMap contains the headers explicitly set by the Handler.
	//
	// To get the implicit headers set by the server (such as
	// automatic Content-Type), use the Result method.
	HeaderMap http.Header

	// Body is the buffer that a Handler's Write calls are sent to.
	// If nil, the Writes are silently discard.
	Body *bytes.Buffer

	// Flushed is whether the Handler called Flush.
	Flushed bool

	result      *http.Response // cache of Result's return value
	snapHeader  http.Header    // snapshot of HeaderMap at first Write
	wroteHeader bool
}

// NewRecorder returns an initialized ResponseRecorder.
func NewRecorder() *ResponseRecorder {
	return &ResponseRecorder{
		HeaderMap: make(http.Header),
		Body:      new(bytes.Buffer),
		Code:      200,
	}
}

// DefaultRemoteAddr is the default remote address to return in RemoteAddr if
// an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
const DefaultRemoteAddr = "1.2.3.4"

// Header returns the response headers.
func (rw *ResponseRecorder) Header() http.Header {
	m := rw.HeaderMap
	if m == nil {
		m = make(http.Header)
		rw.HeaderMap = m
	}
	return m
}

// writeHeader writes a header if it was not written yet and
// detects Content-Type if needed.
//
// bytes or str are the beginning of the response body.
// We pass both to avoid unnecessarily generate garbage
// in rw.WriteString which was created for performance reasons.
// Non-nil bytes win.
func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
	if rw.wroteHeader {
		return
	}
	if len(str) > 512 {
		str = str[:512]
	}

	m := rw.Header()

	_, hasType := m["Content-Type"]
	hasTE := m.Get("Transfer-Encoding") != ""
	if !hasType && !hasTE {
		if b == nil {
			b = []byte(str)
		}
		m.Set("Content-Type", http.DetectContentType(b))
	}

	rw.WriteHeader(200)
}

// Write always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
	rw.writeHeader(buf, "")
	if rw.Body != nil {
		rw.Body.Write(buf)
	}
	return len(buf), nil
}

// WriteString always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) WriteString(str string) (int, error) {
	rw.writeHeader(nil, str)
	if rw.Body != nil {
		rw.Body.WriteString(str)
	}
	return len(str), nil
}

// WriteHeader sets rw.Code. After it is called, changing rw.Header
// will not affect rw.HeaderMap.
func (rw *ResponseRecorder) WriteHeader(code int) {
	if rw.wroteHeader {
		return
	}
	rw.Code = code
	rw.wroteHeader = true
	if rw.HeaderMap == nil {
		rw.HeaderMap = make(http.Header)
	}
	rw.snapHeader = cloneHeader(rw.HeaderMap)
}

func cloneHeader(h http.Header) http.Header {
	h2 := make(http.Header, len(h))
	for k, vv := range h {
		vv2 := make([]string, len(vv))
		copy(vv2, vv)
		h2[k] = vv2
	}
	return h2
}

// Flush sets rw.Flushed to true.
func (rw *ResponseRecorder) Flush() {
	if !rw.wroteHeader {
		rw.WriteHeader(200)
	}
	rw.Flushed = true
}

// Result returns the response generated by the handler.
//
// The returned Response will have at least its StatusCode,
// Header, Body, and optionally Trailer populated.
// More fields may be populated in the future, so callers should
// not DeepEqual the result in tests.
//
// The Response.Header is a snapshot of the headers at the time of the
// first write call, or at the time of this call, if the handler never
// did a write.
//
// The Response.Body is guaranteed to be non-nil and Body.Read call is
// guaranteed to not return any error other than io.EOF.
//
// Result must only be called after the handler has finished running.
func (rw *ResponseRecorder) Result() *http.Response {
	if rw.result != nil {
		return rw.result
	}
	if rw.snapHeader == nil {
		rw.snapHeader = cloneHeader(rw.HeaderMap)
	}
	res := &http.Response{
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		StatusCode: rw.Code,
		Header:     rw.snapHeader,
	}
	rw.result = res
	if res.StatusCode == 0 {
		res.StatusCode = 200
	}
	res.Status = http.StatusText(res.StatusCode)
	if rw.Body != nil {
		res.Body = ioutil.NopCloser(bytes.NewReader(rw.Body.Bytes()))
	}
	res.ContentLength = parseContentLength(res.Header.Get("Content-Length"))

	if trailers, ok := rw.snapHeader["Trailer"]; ok {
		res.Trailer = make(http.Header, len(trailers))
		for _, k := range trailers {
			// TODO: use http2.ValidTrailerHeader, but we can't
			// get at it easily because it's bundled into net/http
			// unexported. This is good enough for now:
			switch k {
			case "Transfer-Encoding", "Content-Length", "Trailer":
				// Ignore since forbidden by RFC 2616 14.40.
				continue
			}
			k = http.CanonicalHeaderKey(k)
			vv, ok := rw.HeaderMap[k]
			if !ok {
				continue
			}
			vv2 := make([]string, len(vv))
			copy(vv2, vv)
			res.Trailer[k] = vv2
		}
	}
	return res
}

// parseContentLength trims whitespace from s and returns -1 if no value
// is set, or the value if it's >= 0.
//
// This a modified version of same function found in net/http/transfer.go. This
// one just ignores an invalid header.
func parseContentLength(cl string) int64 {
	cl = strings.TrimSpace(cl)
	if cl == "" {
		return -1
	}
	n, err := strconv.ParseInt(cl, 10, 64)
	if err != nil {
		return -1
	}
	return n
}
<commit_msg>net/http/httptest: doc tweaks<commit_after>// Copyright 2011 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.

package httptest

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
)

// ResponseRecorder is an implementation of http.ResponseWriter that
// records its mutations for later inspection in tests.
type ResponseRecorder struct {
	// Code is the HTTP response code set by WriteHeader.
	//
	// Note that if a Handler never calls WriteHeader or Write,
	// this might end up being 0, rather than the implicit
	// http.StatusOK. To get the implicit value, use the Result
	// method.
	Code int

	// HeaderMap contains the headers explicitly set by the Handler.
	//
	// To get the implicit headers set by the server (such as
	// automatic Content-Type), use the Result method.
	HeaderMap http.Header

	// Body is the buffer to which the Handler's Write calls are sent.
	// If nil, the Writes are silently discarded.
	Body *bytes.Buffer

	// Flushed is whether the Handler called Flush.
	Flushed bool

	result      *http.Response // cache of Result's return value
	snapHeader  http.Header    // snapshot of HeaderMap at first Write
	wroteHeader bool
}

// NewRecorder returns an initialized ResponseRecorder.
func NewRecorder() *ResponseRecorder {
	return &ResponseRecorder{
		HeaderMap: make(http.Header),
		Body:      new(bytes.Buffer),
		Code:      200,
	}
}

// DefaultRemoteAddr is the default remote address to return in RemoteAddr if
// an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
const DefaultRemoteAddr = "1.2.3.4"

// Header returns the response headers.
func (rw *ResponseRecorder) Header() http.Header {
	m := rw.HeaderMap
	if m == nil {
		m = make(http.Header)
		rw.HeaderMap = m
	}
	return m
}

// writeHeader writes a header if it was not written yet and
// detects Content-Type if needed.
//
// bytes or str are the beginning of the response body.
// We pass both to avoid unnecessarily generate garbage
// in rw.WriteString which was created for performance reasons.
// Non-nil bytes win.
func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
	if rw.wroteHeader {
		return
	}
	if len(str) > 512 {
		str = str[:512]
	}

	m := rw.Header()

	_, hasType := m["Content-Type"]
	hasTE := m.Get("Transfer-Encoding") != ""
	if !hasType && !hasTE {
		if b == nil {
			b = []byte(str)
		}
		m.Set("Content-Type", http.DetectContentType(b))
	}

	rw.WriteHeader(200)
}

// Write always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
	rw.writeHeader(buf, "")
	if rw.Body != nil {
		rw.Body.Write(buf)
	}
	return len(buf), nil
}

// WriteString always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) WriteString(str string) (int, error) {
	rw.writeHeader(nil, str)
	if rw.Body != nil {
		rw.Body.WriteString(str)
	}
	return len(str), nil
}

// WriteHeader sets rw.Code. After it is called, changing rw.Header
// will not affect rw.HeaderMap.
func (rw *ResponseRecorder) WriteHeader(code int) {
	if rw.wroteHeader {
		return
	}
	rw.Code = code
	rw.wroteHeader = true
	if rw.HeaderMap == nil {
		rw.HeaderMap = make(http.Header)
	}
	rw.snapHeader = cloneHeader(rw.HeaderMap)
}

func cloneHeader(h http.Header) http.Header {
	h2 := make(http.Header, len(h))
	for k, vv := range h {
		vv2 := make([]string, len(vv))
		copy(vv2, vv)
		h2[k] = vv2
	}
	return h2
}

// Flush sets rw.Flushed to true.
func (rw *ResponseRecorder) Flush() {
	if !rw.wroteHeader {
		rw.WriteHeader(200)
	}
	rw.Flushed = true
}

// Result returns the response generated by the handler.
//
// The returned Response will have at least its StatusCode,
// Header, Body, and optionally Trailer populated.
// More fields may be populated in the future, so callers should
// not DeepEqual the result in tests.
//
// The Response.Header is a snapshot of the headers at the time of the
// first write call, or at the time of this call, if the handler never
// did a write.
//
// The Response.Body is guaranteed to be non-nil and Body.Read call is
// guaranteed to not return any error other than io.EOF.
//
// Result must only be called after the handler has finished running.
func (rw *ResponseRecorder) Result() *http.Response {
	if rw.result != nil {
		return rw.result
	}
	if rw.snapHeader == nil {
		rw.snapHeader = cloneHeader(rw.HeaderMap)
	}
	res := &http.Response{
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		StatusCode: rw.Code,
		Header:     rw.snapHeader,
	}
	rw.result = res
	if res.StatusCode == 0 {
		res.StatusCode = 200
	}
	res.Status = http.StatusText(res.StatusCode)
	if rw.Body != nil {
		res.Body = ioutil.NopCloser(bytes.NewReader(rw.Body.Bytes()))
	}
	res.ContentLength = parseContentLength(res.Header.Get("Content-Length"))

	if trailers, ok := rw.snapHeader["Trailer"]; ok {
		res.Trailer = make(http.Header, len(trailers))
		for _, k := range trailers {
			// TODO: use http2.ValidTrailerHeader, but we can't
			// get at it easily because it's bundled into net/http
			// unexported. This is good enough for now:
			switch k {
			case "Transfer-Encoding", "Content-Length", "Trailer":
				// Ignore since forbidden by RFC 2616 14.40.
				continue
			}
			k = http.CanonicalHeaderKey(k)
			vv, ok := rw.HeaderMap[k]
			if !ok {
				continue
			}
			vv2 := make([]string, len(vv))
			copy(vv2, vv)
			res.Trailer[k] = vv2
		}
	}
	return res
}

// parseContentLength trims whitespace from s and returns -1 if no value
// is set, or the value if it's >= 0.
//
// This a modified version of same function found in net/http/transfer.go. This
// one just ignores an invalid header.
func parseContentLength(cl string) int64 {
	cl = strings.TrimSpace(cl)
	if cl == "" {
		return -1
	}
	n, err := strconv.ParseInt(cl, 10, 64)
	if err != nil {
		return -1
	}
	return n
}
<|endoftext|>


<commit_before>package tar

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestTarGzFile(t *testing.T) {
	var assert = assert.New(t)

	folder, err := ioutil.TempDir("", "targztest")
	assert.NoError(err)

	file, err := os.Create(folder + "/folder.tar.gz")
	assert.NoError(err)

	empty, err := os.Create(folder + "/empty.txt")
	assert.NoError(err)

	archive := New(file)
	assert.NoError(archive.Add("empty.txt", empty.Name()))
	assert.Error(archive.Add("dont.txt", empty.Name()+"_nope"))
	assert.NoError(archive.Close())
}
<commit_msg>added more tar tests<commit_after>package tar

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestTarGzFile(t *testing.T) {
	var assert = assert.New(t)

	folder, err := ioutil.TempDir("", "targztest")
	assert.NoError(err)

	file, err := os.Create(folder + "/folder.tar.gz")
	assert.NoError(err)

	empty, err := os.Create(folder + "/empty.txt")
	assert.NoError(err)

	empty2, err := os.Create(folder + "/empty2.txt")
	assert.NoError(err)

	archive := New(file)
	assert.NoError(archive.Add("empty.txt", empty.Name()))
	assert.Error(archive.Add("dont.txt", empty.Name()+"_nope"))
	assert.NoError(archive.Close())
	assert.Error(archive.Add("empty2.txt", empty2.Name()))
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Pani Networks
// 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 provides tool for Romana CNI plugin to interact with kubernetes.
package kubernetes

import (
	"fmt"
	"net"

	"github.com/containernetworking/cni/pkg/types"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
)

type PodDescription struct {
	Name        string
	Namespace   string
	Labels      map[string]string
	Annotations map[string]string
}

// K8sArgs is the valid CNI_ARGS used for Kubernetes.
type K8sArgs struct {
	types.CommonArgs
	IP                         net.IP
	K8S_POD_NAME               types.UnmarshallableString
	K8S_POD_NAMESPACE          types.UnmarshallableString
	K8S_POD_INFRA_CONTAINER_ID types.UnmarshallableString
}

// MakePodName returns unique pod name.
func (k8s K8sArgs) MakePodName() string {
	const suffixLength = 8
	var suffix string
	infra := string(k8s.K8S_POD_INFRA_CONTAINER_ID)
	if len(infra) > suffixLength {
		suffix = infra[:suffixLength]
	} else {
		suffix = infra
	}

	return fmt.Sprintf("%s.%s.%s", k8s.K8S_POD_NAME, k8s.K8S_POD_NAMESPACE, suffix)
}

// MakeVethName generates veth name that can be used for external part
// of the veth interface.
func (k8s K8sArgs) MakeVethName() string {
	const suffixLength = 8
	const vethPrefix = "romana"
	var suffix string
	infra := string(k8s.K8S_POD_INFRA_CONTAINER_ID)
	if len(infra) > suffixLength {
		suffix = infra[:suffixLength]
	} else {
		suffix = infra
	}

	return fmt.Sprintf("%s-%s", vethPrefix, suffix)
}

// GetPodDescription retrieves additional information about pod that being created
// or deleted using CNI.
func GetPodDescription(args K8sArgs, configFile string) (*PodDescription, error) {
	// Init kubernetes client. Attempt to load from statically configured k8s config or fallback on in-cluster
	kubeClientConfig, err := clientcmd.BuildConfigFromFlags("", configFile)
	if err != nil {
		return nil, err
	}
	kubeClient, err := kubernetes.NewForConfig(kubeClientConfig)

	pod, err := kubeClient.Core().Pods(string(args.K8S_POD_NAMESPACE)).Get(fmt.Sprintf("%s", args.K8S_POD_NAME))
	if err != nil {
		return nil, fmt.Errorf("Failed to discover a pod %s, err=(%s)", args.K8S_POD_NAME, err)
	}

	res := PodDescription{
		Name:        args.MakePodName(),
		Namespace:   string(args.K8S_POD_NAMESPACE),
		Labels:      pod.Labels,
		Annotations: pod.Annotations,
	}

	return &res, nil
}
<commit_msg>Remove unused function<commit_after>// Copyright (c) 2017 Pani Networks
// 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 provides tool for Romana CNI plugin to interact with kubernetes.
package kubernetes

import (
	"fmt"
	"net"

	"github.com/containernetworking/cni/pkg/types"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
)

type PodDescription struct {
	Name        string
	Namespace   string
	Labels      map[string]string
	Annotations map[string]string
}

// K8sArgs is the valid CNI_ARGS used for Kubernetes.
type K8sArgs struct {
	types.CommonArgs
	IP                         net.IP
	K8S_POD_NAME               types.UnmarshallableString
	K8S_POD_NAMESPACE          types.UnmarshallableString
	K8S_POD_INFRA_CONTAINER_ID types.UnmarshallableString
}

// MakePodName returns unique pod name.
func (k8s K8sArgs) MakePodName() string {
	const suffixLength = 8
	var suffix string
	infra := string(k8s.K8S_POD_INFRA_CONTAINER_ID)
	if len(infra) > suffixLength {
		suffix = infra[:suffixLength]
	} else {
		suffix = infra
	}

	return fmt.Sprintf("%s.%s.%s", k8s.K8S_POD_NAME, k8s.K8S_POD_NAMESPACE, suffix)
}

// GetPodDescription retrieves additional information about pod that being created
// or deleted using CNI.
func GetPodDescription(args K8sArgs, configFile string) (*PodDescription, error) {
	// Init kubernetes client. Attempt to load from statically configured k8s config or fallback on in-cluster
	kubeClientConfig, err := clientcmd.BuildConfigFromFlags("", configFile)
	if err != nil {
		return nil, err
	}
	kubeClient, err := kubernetes.NewForConfig(kubeClientConfig)

	pod, err := kubeClient.Core().Pods(string(args.K8S_POD_NAMESPACE)).Get(fmt.Sprintf("%s", args.K8S_POD_NAME))
	if err != nil {
		return nil, fmt.Errorf("Failed to discover a pod %s, err=(%s)", args.K8S_POD_NAME, err)
	}

	res := PodDescription{
		Name:        args.MakePodName(),
		Namespace:   string(args.K8S_POD_NAMESPACE),
		Labels:      pod.Labels,
		Annotations: pod.Annotations,
	}

	return &res, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package gcemodel

import (
	"fmt"
	"strings"

	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/pkg/model/iam"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/gce"
	"k8s.io/kops/upup/pkg/fi/cloudup/gcetasks"
	"k8s.io/kops/util/pkg/vfs"
)

// StorageAclBuilder configures storage acls
type StorageAclBuilder struct {
	*GCEModelContext
	Cloud     gce.GCECloud
	Lifecycle fi.Lifecycle
}

var _ fi.ModelBuilder = &NetworkModelBuilder{}

// Build creates the tasks that set up storage acls

func (b *StorageAclBuilder) Build(c *fi.ModelBuilderContext) error {
	if featureflag.GoogleCloudBucketACL.Enabled() {
		if b.Cluster.Spec.CloudConfig.GCEServiceAccount == "" {
			return fmt.Errorf("featureflag GoogleCloudBucketACL not supported with per-instancegroup GCEServiceAccount")
		}

		klog.Warningf("featureflag GoogleCloudBucketACL is no longer recommended; use per-instancegroup GCEServiceAccounts instead")

		gceDefaultServiceAccount, err := b.Cloud.ServiceAccount()
		if err != nil {
			return fmt.Errorf("error fetching default ServiceAccount: %w", err)
		}

		clusterPath := b.Cluster.Spec.ConfigBase
		p, err := vfs.Context.BuildVfsPath(clusterPath)
		if err != nil {
			return fmt.Errorf("cannot parse cluster path %q: %w", clusterPath, err)
		}

		switch p := p.(type) {
		case *vfs.GSPath:
			// It's not ideal that we have to do this at the bucket level,
			// but GCS doesn't seem to have a way to do subtrees (like AWS IAM does)
			// Note this permission only lets us list objects, not read them
			c.AddTask(&gcetasks.StorageBucketAcl{
				Name:      s("serviceaccount-statestore-list"),
				Lifecycle: b.Lifecycle,
				Bucket:    s(p.Bucket()),
				Entity:    s("user-" + gceDefaultServiceAccount),
				Role:      s("READER"),
			})
		}

		klog.Warningf("we need to split master / node roles")
		nodeRole, err := iam.BuildNodeRoleSubject(kops.InstanceGroupRoleMaster, false)
		if err != nil {
			return err
		}

		writeablePaths, err := iam.WriteableVFSPaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}

		buckets := sets.NewString()
		for _, p := range writeablePaths {
			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}

			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				continue
			}
			buckets.Insert(bucket)

			klog.Warningf("adding bucket level write ACL to gs://%s to support etcd backup", bucket)

			c.AddTask(&gcetasks.StorageBucketAcl{
				Name:      s("serviceaccount-backup-readwrite-" + bucket),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Entity:    s("user-" + gceDefaultServiceAccount),
				Role:      s("WRITER"),
			})
		}

		return nil
	}

	type serviceAccountRole struct {
		Email string
		Role  kops.InstanceGroupRole
	}
	serviceAccountRoles := make(map[serviceAccountRole]bool)

	for _, ig := range b.InstanceGroups {
		serviceAccount := b.LinkToServiceAccount(ig)

		email := *serviceAccount.Email
		serviceAccountRoles[serviceAccountRole{Email: email, Role: ig.Spec.Role}] = true
	}

	for serviceAccountRole := range serviceAccountRoles {
		role := serviceAccountRole.Role

		nodeRole, err := iam.BuildNodeRoleSubject(role, false)
		if err != nil {
			return err
		}

		buckets := sets.NewString()

		writeablePaths, err := iam.WriteableVFSPaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}
		for _, p := range writeablePaths {
			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}

			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				continue
			}
			buckets.Insert(bucket)

			nameForTask := strings.ToLower(string(role))

			klog.Warningf("adding bucket level write IAM for role %q to gs://%s to support etcd backup", bucket, role)

			c.AddTask(&gcetasks.StorageBucketIAM{
				Name:      s("objectadmin-" + bucket + "-serviceaccount-" + nameForTask),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Member:    s("serviceAccount:" + serviceAccountRole.Email),
				Role:      s("roles/storage.objectAdmin"),
			})
		}

		// Add bucket read permissions if we need to read from the bucket
		readablePaths, err := iam.ReadableStatePaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}
		if len(readablePaths) != 0 {
			p, err := vfs.Context.BuildVfsPath(b.Cluster.Spec.ConfigStore)
			if err != nil {
				return fmt.Errorf("cannot parse VFS path %q: %v", b.Cluster.Spec.ConfigStore, err)
			}

			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}
			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				// Already marked as writeable; we can skip
				continue
			}
			buckets.Insert(bucket)

			nameForTask := strings.ToLower(string(role))

			klog.Warningf("adding bucket level read IAM to gs://%s for role %q", bucket, role)

			c.AddTask(&gcetasks.StorageBucketIAM{
				Name:      s("objectviewer-" + bucket + "-serviceaccount-" + nameForTask),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Member:    s("serviceAccount:" + serviceAccountRole.Email),
				Role:      s("roles/storage.objectViewer"),
			})
		}
	}

	return nil
}
<commit_msg>gce: don't set per-IG permissions when using shared account<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package gcemodel

import (
	"fmt"
	"strings"

	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/klog/v2"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/pkg/model/iam"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/gce"
	"k8s.io/kops/upup/pkg/fi/cloudup/gcetasks"
	"k8s.io/kops/util/pkg/vfs"
)

// StorageAclBuilder configures storage acls
type StorageAclBuilder struct {
	*GCEModelContext
	Cloud     gce.GCECloud
	Lifecycle fi.Lifecycle
}

var _ fi.ModelBuilder = &NetworkModelBuilder{}

// Build creates the tasks that set up storage acls

func (b *StorageAclBuilder) Build(c *fi.ModelBuilderContext) error {
	if featureflag.GoogleCloudBucketACL.Enabled() {
		if b.Cluster.Spec.CloudConfig.GCEServiceAccount == "" {
			return fmt.Errorf("featureflag GoogleCloudBucketACL not supported with per-instancegroup GCEServiceAccount")
		}

		klog.Warningf("featureflag GoogleCloudBucketACL is no longer recommended; use per-instancegroup GCEServiceAccounts instead")

		gceDefaultServiceAccount, err := b.Cloud.ServiceAccount()
		if err != nil {
			return fmt.Errorf("error fetching default ServiceAccount: %w", err)
		}

		clusterPath := b.Cluster.Spec.ConfigBase
		p, err := vfs.Context.BuildVfsPath(clusterPath)
		if err != nil {
			return fmt.Errorf("cannot parse cluster path %q: %w", clusterPath, err)
		}

		switch p := p.(type) {
		case *vfs.GSPath:
			// It's not ideal that we have to do this at the bucket level,
			// but GCS doesn't seem to have a way to do subtrees (like AWS IAM does)
			// Note this permission only lets us list objects, not read them
			c.AddTask(&gcetasks.StorageBucketAcl{
				Name:      s("serviceaccount-statestore-list"),
				Lifecycle: b.Lifecycle,
				Bucket:    s(p.Bucket()),
				Entity:    s("user-" + gceDefaultServiceAccount),
				Role:      s("READER"),
			})
		}

		klog.Warningf("we need to split master / node roles")
		nodeRole, err := iam.BuildNodeRoleSubject(kops.InstanceGroupRoleMaster, false)
		if err != nil {
			return err
		}

		writeablePaths, err := iam.WriteableVFSPaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}

		buckets := sets.NewString()
		for _, p := range writeablePaths {
			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}

			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				continue
			}
			buckets.Insert(bucket)

			klog.Warningf("adding bucket level write ACL to gs://%s to support etcd backup", bucket)

			c.AddTask(&gcetasks.StorageBucketAcl{
				Name:      s("serviceaccount-backup-readwrite-" + bucket),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Entity:    s("user-" + gceDefaultServiceAccount),
				Role:      s("WRITER"),
			})
		}

		return nil
	}

	if b.Cluster.Spec.CloudConfig.GCEServiceAccount != "" {
		// We can't set up per-instancegroup permissions if we're using a cluster-level account
		// Historically, we did not grant the serviceaccount permissions in this case.
		return nil
	}

	type serviceAccountRole struct {
		Email string
		Role  kops.InstanceGroupRole
	}
	serviceAccountRoles := make(map[serviceAccountRole]bool)

	for _, ig := range b.InstanceGroups {
		serviceAccount := b.LinkToServiceAccount(ig)

		email := *serviceAccount.Email
		serviceAccountRoles[serviceAccountRole{Email: email, Role: ig.Spec.Role}] = true
	}

	for serviceAccountRole := range serviceAccountRoles {
		role := serviceAccountRole.Role

		nodeRole, err := iam.BuildNodeRoleSubject(role, false)
		if err != nil {
			return err
		}

		buckets := sets.NewString()

		writeablePaths, err := iam.WriteableVFSPaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}
		for _, p := range writeablePaths {
			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}

			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				continue
			}
			buckets.Insert(bucket)

			nameForTask := strings.ToLower(string(role))

			klog.Warningf("adding bucket level write IAM for role %q to gs://%s to support etcd backup", bucket, role)

			c.AddTask(&gcetasks.StorageBucketIAM{
				Name:      s("objectadmin-" + bucket + "-serviceaccount-" + nameForTask),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Member:    s("serviceAccount:" + serviceAccountRole.Email),
				Role:      s("roles/storage.objectAdmin"),
			})
		}

		// Add bucket read permissions if we need to read from the bucket
		readablePaths, err := iam.ReadableStatePaths(b.Cluster, nodeRole)
		if err != nil {
			return err
		}
		if len(readablePaths) != 0 {
			p, err := vfs.Context.BuildVfsPath(b.Cluster.Spec.ConfigStore)
			if err != nil {
				return fmt.Errorf("cannot parse VFS path %q: %v", b.Cluster.Spec.ConfigStore, err)
			}

			gcsPath, ok := p.(*vfs.GSPath)
			if !ok {
				klog.Warningf("unknown path, can't apply IAM policy: %q", p)
				continue
			}
			bucket := gcsPath.Bucket()
			if buckets.Has(bucket) {
				// Already marked as writeable; we can skip
				continue
			}
			buckets.Insert(bucket)

			nameForTask := strings.ToLower(string(role))

			klog.Warningf("adding bucket level read IAM to gs://%s for role %q", bucket, role)

			c.AddTask(&gcetasks.StorageBucketIAM{
				Name:      s("objectviewer-" + bucket + "-serviceaccount-" + nameForTask),
				Lifecycle: b.Lifecycle,
				Bucket:    s(bucket),
				Member:    s("serviceAccount:" + serviceAccountRole.Email),
				Role:      s("roles/storage.objectViewer"),
			})
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package dynamodb

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	dynamodbBatchLimit = 100

	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	var params *godynamodb.DescribeLimitsInput
	_, err := d.svc.DescribeLimits(params)
	if err != nil {
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
		TableName: aws.String(param.Name),
		TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
			AttributeName: aws.String("TTL"),
			Enabled:       aws.Bool(true),
		},
	})
	if err != nil {
		return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}
	if _, err := d.svc.UpdateItem(params); err != nil {
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<commit_msg>Add dynamodb.Put comment<commit_after>package dynamodb

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"math"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	godynamodb "github.com/aws/aws-sdk-go/service/dynamodb"
	godynamodbiface "github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
	"github.com/pkg/errors"

	"github.com/yuuki/diamondb/pkg/config"
	"github.com/yuuki/diamondb/pkg/model"
	"github.com/yuuki/diamondb/pkg/storage/util"
	"github.com/yuuki/diamondb/pkg/timeparser"
)

//go:generate mockgen -source ../../../vendor/github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface/interface.go -destination dynamodb_mock.go -package dynamodb

// ReadWriter defines the interface for DynamoDB reader and writer.
type ReadWriter interface {
	Ping() error
	Client() godynamodbiface.DynamoDBAPI
	CreateTable(*CreateTableParam) error
	Fetch(string, time.Time, time.Time) (model.SeriesMap, error)
	batchGet(q *query) (model.SeriesMap, error)
	Put(string, string, string, int64, map[int64]float64) error
}

// DynamoDB provides a dynamodb client.
type DynamoDB struct {
	svc godynamodbiface.DynamoDBAPI
}

type timeSlot struct {
	itemEpoch int64
	step      int
}

type query struct {
	names []string
	start time.Time
	end   time.Time
	slot  *timeSlot
	// context
}

const (
	oneYear time.Duration = time.Duration(24*360) * time.Hour
	oneWeek time.Duration = time.Duration(24*7) * time.Hour
	oneDay  time.Duration = time.Duration(24*1) * time.Hour
)

var (
	dynamodbBatchLimit = 100

	oneYearSeconds = int(oneYear.Seconds())
	oneWeekSeconds = int(oneWeek.Seconds())
	oneDaySeconds  = int(oneDay.Seconds())
)

var _ ReadWriter = &DynamoDB{}

// New creates a new DynamoDB.
func New() (*DynamoDB, error) {
	awsConf := aws.NewConfig().WithRegion(config.Config.DynamoDBRegion)
	if config.Config.DynamoDBEndpoint != "" {
		// For dynamodb-local configuration
		awsConf.WithEndpoint(config.Config.DynamoDBEndpoint)
		awsConf.WithCredentials(credentials.NewStaticCredentials("dummy", "dummy", "dummy"))
	}
	sess, err := session.NewSession(awsConf)
	if err != nil {
		return nil, errors.Wrapf(err,
			"failed to create session for dynamodb (%s,%s)",
			config.Config.DynamoDBRegion,
			config.Config.DynamoDBEndpoint,
		)
	}
	return &DynamoDB{
		svc: godynamodb.New(sess),
	}, nil
}

// Client returns the DynamoDB client.
func (d *DynamoDB) Client() godynamodbiface.DynamoDBAPI {
	return d.svc
}

// Ping pings DynamoDB endpoint.
func (d *DynamoDB) Ping() error {
	var params *godynamodb.DescribeLimitsInput
	_, err := d.svc.DescribeLimits(params)
	if err != nil {
		return errors.Wrapf(err, "failed to ping dynamodb")
	}
	return nil
}

// CreateTableParam is parameter set of CreateTable.
type CreateTableParam struct {
	Name string
	RCU  int64 // ReadCapacityUnits
	WCU  int64 // WriteCapacityUnits
}

// CreateTable creates a dynamodb table to store time series data.
// Skip creating table if the table already exists.
func (d *DynamoDB) CreateTable(param *CreateTableParam) error {
	_, err := d.svc.CreateTable(&godynamodb.CreateTableInput{
		TableName: aws.String(param.Name),
		AttributeDefinitions: []*godynamodb.AttributeDefinition{
			{
				AttributeName: aws.String("Name"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
			{
				AttributeName: aws.String("Timestamp"),
				AttributeType: aws.String(godynamodb.ScalarAttributeTypeS),
			},
		},
		KeySchema: []*godynamodb.KeySchemaElement{
			{
				AttributeName: aws.String("Name"),
				KeyType:       aws.String("HASH"),
			},
			{
				AttributeName: aws.String("Timestamp"),
				KeyType:       aws.String("RANGE"),
			},
		},
		ProvisionedThroughput: &godynamodb.ProvisionedThroughput{
			ReadCapacityUnits:  aws.Int64(param.RCU),
			WriteCapacityUnits: aws.Int64(param.WCU),
		},
		// TODO StreamSpecification to export to s3
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceInUseException" {
				// Skip if the table already exists
				log.Printf("Skip creating DynamoDB table because %s already exists\n", param.Name)
				return nil
			}
		}
		return errors.Wrapf(err, "failed to create dynamodb table (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	log.Printf("Creating DynamoDB table (name:%s, rcu:%d, wcu:%d) ...\n",
		param.Name, param.RCU, param.WCU)

	err = d.svc.WaitUntilTableExists(&godynamodb.DescribeTableInput{
		TableName: aws.String(param.Name),
	})
	if err != nil {
		return errors.Wrapf(err, "failed to wait until table exists (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}

	_, err = d.svc.UpdateTimeToLive(&godynamodb.UpdateTimeToLiveInput{
		TableName: aws.String(param.Name),
		TimeToLiveSpecification: &godynamodb.TimeToLiveSpecification{
			AttributeName: aws.String("TTL"),
			Enabled:       aws.Bool(true),
		},
	})
	if err != nil {
		return errors.Wrapf(err, "failed to set TTL to (%s,%d,%d)",
			param.Name, param.RCU, param.WCU)
	}
	return nil
}

// Fetch fetches datapoints by name from start until end.
func (d *DynamoDB) Fetch(name string, start, end time.Time) (model.SeriesMap, error) {
	slots := selectTimeSlots(start, end)
	nameGroups := util.GroupNames(util.SplitName(name), dynamodbBatchLimit)
	numQueries := len(slots) * len(nameGroups)

	type result struct {
		value model.SeriesMap
		err   error
	}
	c := make(chan *result, numQueries)
	for _, slot := range slots {
		for _, names := range nameGroups {
			q := &query{
				names: names,
				start: start,
				end:   end,
				slot:  slot,
			}
			go func(q *query) {
				sm, err := d.batchGet(q)
				c <- &result{value: sm, err: err}
			}(q)
		}
	}
	sm := make(model.SeriesMap, len(nameGroups))
	for i := 0; i < numQueries; i++ {
		ret := <-c
		if ret.err != nil {
			return nil, ret.err
		}
		sm.MergePointsToMap(ret.value)
	}
	return sm, nil
}

func batchGetResultToMap(resp *godynamodb.BatchGetItemOutput, q *query) model.SeriesMap {
	sm := make(model.SeriesMap, len(resp.Responses))
	for _, xs := range resp.Responses {
		for _, x := range xs {
			name := (*x["Name"].S)
			points := make(model.DataPoints, 0, len(x["Values"].BS))
			for _, y := range x["Values"].BS {
				t := int64(binary.BigEndian.Uint64(y[0:8]))
				v := math.Float64frombits(binary.BigEndian.Uint64(y[8:]))
				// Trim datapoints out of [start, end]
				if t < q.start.Unix() || q.end.Unix() < t {
					continue
				}
				points = append(points, model.NewDataPoint(t, v))
			}
			sm[name] = model.NewSeriesPoint(name, points, q.slot.step)
		}
	}
	return sm
}

func (d *DynamoDB) batchGet(q *query) (model.SeriesMap, error) {
	var keys []map[string]*godynamodb.AttributeValue
	for _, name := range q.names {
		keys = append(keys, map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", q.slot.itemEpoch, q.slot.step))},
		})
	}
	items := make(map[string]*godynamodb.KeysAndAttributes)
	items[config.Config.DynamoDBTableName] = &godynamodb.KeysAndAttributes{Keys: keys}
	params := &godynamodb.BatchGetItemInput{
		RequestItems:           items,
		ReturnConsumedCapacity: aws.String("NONE"),
	}
	resp, err := d.svc.BatchGetItem(params)
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				// Don't handle ResourceNotFoundException as error
				// bacause diamondb web return length 0 series as 200.
				return model.SeriesMap{}, nil
			}
		}
		return nil, errors.Wrapf(err,
			"failed to call dynamodb API batchGetItem (%s,%d,%d)",
			config.Config.DynamoDBTableName, q.slot.itemEpoch, q.slot.step,
		)
	}
	return batchGetResultToMap(resp, q), nil
}

// Put writes the datapoints into DynamoDB. It creates item
// if item doesn't exist and updates item if it exists.
func (d *DynamoDB) Put(name, slot, history string, itemEpoch int64, tv map[int64]float64) error {
	stepDuration, err := timeparser.ParseTimeOffset(slot)
	if err != nil {
		return err
	}
	historyDuration, err := timeparser.ParseTimeOffset(history)
	if err != nil {
		return err
	}
	ttl := itemEpoch + int64(historyDuration.Seconds())

	vals := make([][]byte, 0, len(tv))
	for timestamp, value := range tv {
		buf := new(bytes.Buffer)
		binary.Write(buf, binary.BigEndian, timestamp)
		binary.Write(buf, binary.BigEndian, math.Float64bits(value))
		vals = append(vals, buf.Bytes())
	}

	params := &godynamodb.UpdateItemInput{
		TableName: aws.String(config.Config.DynamoDBTableName),
		Key: map[string]*godynamodb.AttributeValue{
			"Name":      {S: aws.String(name)},
			"Timestamp": {S: aws.String(fmt.Sprintf("%d:%d", itemEpoch, int64(stepDuration.Seconds())))},
		},
		UpdateExpression: aws.String(`
			SET #ttl = :new_ttl
			ADD #values_set :new_values
		`),
		ExpressionAttributeNames: map[string]*string{
			"#ttl":        aws.String("TTL"),
			"#values_set": aws.String("Values"),
		},
		ExpressionAttributeValues: map[string]*godynamodb.AttributeValue{
			":new_ttl":    {N: aws.String(fmt.Sprintf("%d", ttl))},
			":new_values": {BS: vals},
		},
		ReturnValues: aws.String("NONE"),
	}
	if _, err := d.svc.UpdateItem(params); err != nil {
		return errors.Wrapf(err, "failed to call dynamodb API putItem (%s,%s,%d)",
			config.Config.DynamoDBTableName, name, itemEpoch)
	}
	return nil
}

func selectTimeSlots(startTime, endTime time.Time) []*timeSlot {
	var (
		step          int
		itemEpochStep int
	)
	diff := endTime.Sub(startTime)
	switch {
	case oneYear <= diff:
		itemEpochStep = oneYearSeconds
		step = 60 * 60 * 24
	case oneWeek <= diff:
		itemEpochStep = oneWeekSeconds
		step = 60 * 60
	case oneDay <= diff:
		itemEpochStep = oneDaySeconds
		step = 5 * 60
	default:
		itemEpochStep = 60 * 60
		step = 60
	}

	var slots []*timeSlot
	startItemEpoch := startTime.Unix() - startTime.Unix()%int64(itemEpochStep)
	endItemEpoch := endTime.Unix()
	for epoch := startItemEpoch; epoch < endItemEpoch; epoch += int64(itemEpochStep) {
		slots = append(slots, &timeSlot{itemEpoch: epoch, step: step})
	}

	return slots
}
<|endoftext|>


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

package mat64

import (
	"sync"

	"gonum.org/v1/gonum/blas"
	"gonum.org/v1/gonum/blas/blas64"
	"gonum.org/v1/gonum/matrix"
)

var tab64 = [64]byte{
	0x3f, 0x00, 0x3a, 0x01, 0x3b, 0x2f, 0x35, 0x02,
	0x3c, 0x27, 0x30, 0x1b, 0x36, 0x21, 0x2a, 0x03,
	0x3d, 0x33, 0x25, 0x28, 0x31, 0x12, 0x1c, 0x14,
	0x37, 0x1e, 0x22, 0x0b, 0x2b, 0x0e, 0x16, 0x04,
	0x3e, 0x39, 0x2e, 0x34, 0x26, 0x1a, 0x20, 0x29,
	0x32, 0x24, 0x11, 0x13, 0x1d, 0x0a, 0x0d, 0x15,
	0x38, 0x2d, 0x19, 0x1f, 0x23, 0x10, 0x09, 0x0c,
	0x2c, 0x18, 0x0f, 0x08, 0x17, 0x07, 0x06, 0x05,
}

// bits returns the ceiling of base 2 log of v.
// Approach based on http://stackoverflow.com/a/11398748.
func bits(v uint64) byte {
	if v == 0 {
		return 0
	}
	v <<= 2
	v--
	v |= v >> 1
	v |= v >> 2
	v |= v >> 4
	v |= v >> 8
	v |= v >> 16
	v |= v >> 32
	return tab64[((v-(v>>1))*0x07EDD5E59A4E28C2)>>58] - 1
}

var (
	// pool contains size stratified workspace Dense pools.
	// Each pool element i returns sized matrices with a data
	// slice capped at 1<<i.
	pool [63]sync.Pool

	// poolSym is the SymDense equivalent of pool.
	poolSym [63]sync.Pool

	// poolTri is the TriDense equivalent of pool.
	poolTri [63]sync.Pool

	// poolVec is the Vector equivalent of pool.
	poolVec [63]sync.Pool

	// poolFloats is the []float64 equivalent of pool.
	poolFloats [63]sync.Pool

	// poolInts is the []int equivalent of pool.
	poolInts [63]sync.Pool
)

func init() {
	for i := range pool {
		l := 1 << uint(i)
		pool[i].New = func() interface{} {
			return &Dense{mat: blas64.General{
				Data: make([]float64, l),
			}}
		}
		poolSym[i].New = func() interface{} {
			return &SymDense{mat: blas64.Symmetric{
				Uplo: blas.Upper,
				Data: make([]float64, l),
			}}
		}
		poolTri[i].New = func() interface{} {
			return &TriDense{mat: blas64.Triangular{
				Data: make([]float64, l),
			}}
		}
		poolVec[i].New = func() interface{} {
			return &Vector{mat: blas64.Vector{
				Inc:  1,
				Data: make([]float64, l),
			}}
		}
		poolFloats[i].New = func() interface{} {
			return make([]float64, l)
		}
		poolInts[i].New = func() interface{} {
			return make([]int, l)
		}
	}
}

// getWorkspace returns a *Dense of size r×c and a data slice
// with a cap that is less than 2*r*c. If clear is true, the
// data slice visible through the Matrix interface is zeroed.
func getWorkspace(r, c int, clear bool) *Dense {
	l := uint64(r * c)
	w := pool[bits(l)].Get().(*Dense)
	w.mat.Data = w.mat.Data[:l]
	if clear {
		zero(w.mat.Data)
	}
	w.mat.Rows = r
	w.mat.Cols = c
	w.mat.Stride = c
	w.capRows = r
	w.capCols = c
	return w
}

// putWorkspace replaces a used *Dense into the appropriate size
// workspace pool. putWorkspace must not be called with a matrix
// where references to the underlying data slice has been kept.
func putWorkspace(w *Dense) {
	pool[bits(uint64(cap(w.mat.Data)))].Put(w)
}

// getWorkspaceSym returns a *SymDense of size n and a cap that is less than 2*n. If clear is true, the
// data slice visible through the Matrix interface is zeroed.
func getWorkspaceSym(n int, clear bool) *SymDense {
	l := uint64(n)
	l *= l
	s := poolSym[bits(l)].Get().(*SymDense)
	s.mat.Data = s.mat.Data[:l]
	if clear {
		zero(s.mat.Data)
	}
	s.mat.N = n
	s.mat.Stride = n
	s.cap = n
	return s
}

// getWorkspaceTri returns a *TriDense of size n and a cap that is less than 2*n.
// If clear is true, the data slice visible through the Matrix interface is zeroed.
func getWorkspaceTri(n int, kind matrix.TriKind, clear bool) *TriDense {
	l := uint64(n)
	l *= l
	t := poolTri[bits(l)].Get().(*TriDense)
	t.mat.Data = t.mat.Data[:l]
	if clear {
		zero(t.mat.Data)
	}
	t.mat.N = n
	t.mat.Stride = n
	if kind == matrix.Upper {
		t.mat.Uplo = blas.Upper
	} else if kind == matrix.Lower {
		t.mat.Uplo = blas.Lower
	} else {
		panic(matrix.ErrTriangle)
	}
	t.mat.Diag = blas.NonUnit
	t.cap = n
	return t
}

// putWorkspaceSym replaces a used *SymDense into the appropriate size
// workspace pool. putWorkspaceSym must not be called with a matrix
// where references to the underlying data slice has been kept.
func putWorkspaceSym(s *SymDense) {
	poolSym[bits(uint64(cap(s.mat.Data)))].Put(s)
}

// putWorkspaceTri replaces a used *TriDense into the appropriate size
// workspace pool. putWorkspaceTri must not be called with a matrix
// where references to the underlying data slice has been kept.
func putWorkspaceTri(t *TriDense) {
	poolTri[bits(uint64(cap(t.mat.Data)))].Put(t)
}

// getWorkspaceVec returns a *Vector of length n and a cap that is less than 2*n. If clear is true, the
// data slice visible through the Matrix interface is zeroed.
func getWorkspaceVec(n int, clear bool) *Vector {
	l := uint64(n)
	v := poolVec[bits(l)].Get().(*Vector)
	v.mat.Data = v.mat.Data[:l]
	if clear {
		zero(v.mat.Data)
	}
	v.n = n
	return v
}

// putWorkspaceVec replaces a used *Vector into the appropriate size
// workspace pool. putWorkspaceVec must not be called with a matrix
// where references to the underlying data slice has been kept.
func putWorkspaceVec(v *Vector) {
	poolVec[bits(uint64(cap(v.mat.Data)))].Put(v)
}

// getFloats returns a []float64 of length l and a cap that is
// less than 2*l. If clear is true, the slice visible is zeroed.
func getFloats(l int, clear bool) []float64 {
	w := poolFloats[bits(uint64(l))].Get().([]float64)
	w = w[:l]
	if clear {
		zero(w)
	}
	return w
}

// putFloats replaces a used []float64 into the appropriate size
// workspace pool. putFloats must not be called with a slice
// where references to the underlying data have been kept.
func putFloats(w []float64) {
	poolFloats[bits(uint64(cap(w)))].Put(w)
}

// getInts returns a []ints of length l and a cap that is
// less than 2*l. If clear is true, the slice visible is zeroed.
func getInts(l int, clear bool) []int {
	w := poolInts[bits(uint64(l))].Get().([]int)
	w = w[:l]
	if clear {
		for i := range w {
			w[i] = 0
		}
	}
	return w
}

// putInts replaces a used []int into the appropriate size
// workspace pool. putInts must not be called with a slice
// where references to the underlying data have been kept.
func putInts(w []int) {
	poolInts[bits(uint64(cap(w)))].Put(w)
}
<commit_msg>matrix/mat64: regroup functions and clean up docs<commit_after>// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mat64

import (
	"sync"

	"gonum.org/v1/gonum/blas"
	"gonum.org/v1/gonum/blas/blas64"
	"gonum.org/v1/gonum/matrix"
)

var tab64 = [64]byte{
	0x3f, 0x00, 0x3a, 0x01, 0x3b, 0x2f, 0x35, 0x02,
	0x3c, 0x27, 0x30, 0x1b, 0x36, 0x21, 0x2a, 0x03,
	0x3d, 0x33, 0x25, 0x28, 0x31, 0x12, 0x1c, 0x14,
	0x37, 0x1e, 0x22, 0x0b, 0x2b, 0x0e, 0x16, 0x04,
	0x3e, 0x39, 0x2e, 0x34, 0x26, 0x1a, 0x20, 0x29,
	0x32, 0x24, 0x11, 0x13, 0x1d, 0x0a, 0x0d, 0x15,
	0x38, 0x2d, 0x19, 0x1f, 0x23, 0x10, 0x09, 0x0c,
	0x2c, 0x18, 0x0f, 0x08, 0x17, 0x07, 0x06, 0x05,
}

// bits returns the ceiling of base 2 log of v.
// Approach based on http://stackoverflow.com/a/11398748.
func bits(v uint64) byte {
	if v == 0 {
		return 0
	}
	v <<= 2
	v--
	v |= v >> 1
	v |= v >> 2
	v |= v >> 4
	v |= v >> 8
	v |= v >> 16
	v |= v >> 32
	return tab64[((v-(v>>1))*0x07EDD5E59A4E28C2)>>58] - 1
}

var (
	// pool contains size stratified workspace Dense pools.
	// Each pool element i returns sized matrices with a data
	// slice capped at 1<<i.
	pool [63]sync.Pool

	// poolSym is the SymDense equivalent of pool.
	poolSym [63]sync.Pool

	// poolTri is the TriDense equivalent of pool.
	poolTri [63]sync.Pool

	// poolVec is the Vector equivalent of pool.
	poolVec [63]sync.Pool

	// poolFloats is the []float64 equivalent of pool.
	poolFloats [63]sync.Pool

	// poolInts is the []int equivalent of pool.
	poolInts [63]sync.Pool
)

func init() {
	for i := range pool {
		l := 1 << uint(i)
		pool[i].New = func() interface{} {
			return &Dense{mat: blas64.General{
				Data: make([]float64, l),
			}}
		}
		poolSym[i].New = func() interface{} {
			return &SymDense{mat: blas64.Symmetric{
				Uplo: blas.Upper,
				Data: make([]float64, l),
			}}
		}
		poolTri[i].New = func() interface{} {
			return &TriDense{mat: blas64.Triangular{
				Data: make([]float64, l),
			}}
		}
		poolVec[i].New = func() interface{} {
			return &Vector{mat: blas64.Vector{
				Inc:  1,
				Data: make([]float64, l),
			}}
		}
		poolFloats[i].New = func() interface{} {
			return make([]float64, l)
		}
		poolInts[i].New = func() interface{} {
			return make([]int, l)
		}
	}
}

// getWorkspace returns a *Dense of size r×c and a data slice
// with a cap that is less than 2*r*c. If clear is true, the
// data slice visible through the Matrix interface is zeroed.
func getWorkspace(r, c int, clear bool) *Dense {
	l := uint64(r * c)
	w := pool[bits(l)].Get().(*Dense)
	w.mat.Data = w.mat.Data[:l]
	if clear {
		zero(w.mat.Data)
	}
	w.mat.Rows = r
	w.mat.Cols = c
	w.mat.Stride = c
	w.capRows = r
	w.capCols = c
	return w
}

// putWorkspace replaces a used *Dense into the appropriate size
// workspace pool. putWorkspace must not be called with a matrix
// where references to the underlying data slice have been kept.
func putWorkspace(w *Dense) {
	pool[bits(uint64(cap(w.mat.Data)))].Put(w)
}

// getWorkspaceSym returns a *SymDense of size n and a cap that
// is less than 2*n. If clear is true, the data slice visible
// through the Matrix interface is zeroed.
func getWorkspaceSym(n int, clear bool) *SymDense {
	l := uint64(n)
	l *= l
	s := poolSym[bits(l)].Get().(*SymDense)
	s.mat.Data = s.mat.Data[:l]
	if clear {
		zero(s.mat.Data)
	}
	s.mat.N = n
	s.mat.Stride = n
	s.cap = n
	return s
}

// putWorkspaceSym replaces a used *SymDense into the appropriate size
// workspace pool. putWorkspaceSym must not be called with a matrix
// where references to the underlying data slice have been kept.
func putWorkspaceSym(s *SymDense) {
	poolSym[bits(uint64(cap(s.mat.Data)))].Put(s)
}

// getWorkspaceTri returns a *TriDense of size n and a cap that
// is less than 2*n. If clear is true, the data slice visible
// through the Matrix interface is zeroed.
func getWorkspaceTri(n int, kind matrix.TriKind, clear bool) *TriDense {
	l := uint64(n)
	l *= l
	t := poolTri[bits(l)].Get().(*TriDense)
	t.mat.Data = t.mat.Data[:l]
	if clear {
		zero(t.mat.Data)
	}
	t.mat.N = n
	t.mat.Stride = n
	if kind == matrix.Upper {
		t.mat.Uplo = blas.Upper
	} else if kind == matrix.Lower {
		t.mat.Uplo = blas.Lower
	} else {
		panic(matrix.ErrTriangle)
	}
	t.mat.Diag = blas.NonUnit
	t.cap = n
	return t
}

// putWorkspaceTri replaces a used *TriDense into the appropriate size
// workspace pool. putWorkspaceTri must not be called with a matrix
// where references to the underlying data slice have been kept.
func putWorkspaceTri(t *TriDense) {
	poolTri[bits(uint64(cap(t.mat.Data)))].Put(t)
}

// getWorkspaceVec returns a *Vector of length n and a cap that
// is less than 2*n. If clear is true, the data slice visible
// through the Matrix interface is zeroed.
func getWorkspaceVec(n int, clear bool) *Vector {
	l := uint64(n)
	v := poolVec[bits(l)].Get().(*Vector)
	v.mat.Data = v.mat.Data[:l]
	if clear {
		zero(v.mat.Data)
	}
	v.n = n
	return v
}

// putWorkspaceVec replaces a used *Vector into the appropriate size
// workspace pool. putWorkspaceVec must not be called with a matrix
// where references to the underlying data slice have been kept.
func putWorkspaceVec(v *Vector) {
	poolVec[bits(uint64(cap(v.mat.Data)))].Put(v)
}

// getFloats returns a []float64 of length l and a cap that is
// less than 2*l. If clear is true, the slice visible is zeroed.
func getFloats(l int, clear bool) []float64 {
	w := poolFloats[bits(uint64(l))].Get().([]float64)
	w = w[:l]
	if clear {
		zero(w)
	}
	return w
}

// putFloats replaces a used []float64 into the appropriate size
// workspace pool. putFloats must not be called with a slice
// where references to the underlying data have been kept.
func putFloats(w []float64) {
	poolFloats[bits(uint64(cap(w)))].Put(w)
}

// getInts returns a []ints of length l and a cap that is
// less than 2*l. If clear is true, the slice visible is zeroed.
func getInts(l int, clear bool) []int {
	w := poolInts[bits(uint64(l))].Get().([]int)
	w = w[:l]
	if clear {
		for i := range w {
			w[i] = 0
		}
	}
	return w
}

// putInts replaces a used []int into the appropriate size
// workspace pool. putInts must not be called with a slice
// where references to the underlying data have been kept.
func putInts(w []int) {
	poolInts[bits(uint64(cap(w)))].Put(w)
}
<|endoftext|>


<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.

package gob

import (
	"reflect"
	"testing"
)

type typeT struct {
	id  typeId
	str string
}

var basicTypes = []typeT{
	{tBool, "bool"},
	{tInt, "int"},
	{tUint, "uint"},
	{tFloat, "float"},
	{tBytes, "bytes"},
	{tString, "string"},
}

func getTypeUnlocked(name string, rt reflect.Type) gobType {
	typeLock.Lock()
	defer typeLock.Unlock()
	t, err := getBaseType(name, rt)
	if err != nil {
		panic("getTypeUnlocked: " + err.Error())
	}
	return t
}

// Sanity checks
func TestBasic(t *testing.T) {
	for _, tt := range basicTypes {
		if tt.id.string() != tt.str {
			t.Errorf("checkType: expected %q got %s", tt.str, tt.id.string())
		}
		if tt.id == 0 {
			t.Errorf("id for %q is zero", tt.str)
		}
	}
}

// Reregister some basic types to check registration is idempotent.
func TestReregistration(t *testing.T) {
	newtyp := getTypeUnlocked("int", reflect.TypeOf(int(0)))
	if newtyp != tInt.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("uint", reflect.TypeOf(uint(0)))
	if newtyp != tUint.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("string", reflect.TypeOf("hello"))
	if newtyp != tString.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
}

func TestArrayType(t *testing.T) {
	var a3 [3]int
	a3int := getTypeUnlocked("foo", reflect.TypeOf(a3))
	newa3int := getTypeUnlocked("bar", reflect.TypeOf(a3))
	if a3int != newa3int {
		t.Errorf("second registration of [3]int creates new type")
	}
	var a4 [4]int
	a4int := getTypeUnlocked("goo", reflect.TypeOf(a4))
	if a3int == a4int {
		t.Errorf("registration of [3]int creates same type as [4]int")
	}
	var b3 [3]bool
	a3bool := getTypeUnlocked("", reflect.TypeOf(b3))
	if a3int == a3bool {
		t.Errorf("registration of [3]bool creates same type as [3]int")
	}
	str := a3bool.string()
	expected := "[3]bool"
	if str != expected {
		t.Errorf("array printed as %q; expected %q", str, expected)
	}
}

func TestSliceType(t *testing.T) {
	var s []int
	sint := getTypeUnlocked("slice", reflect.TypeOf(s))
	var news []int
	newsint := getTypeUnlocked("slice1", reflect.TypeOf(news))
	if sint != newsint {
		t.Errorf("second registration of []int creates new type")
	}
	var b []bool
	sbool := getTypeUnlocked("", reflect.TypeOf(b))
	if sbool == sint {
		t.Errorf("registration of []bool creates same type as []int")
	}
	str := sbool.string()
	expected := "[]bool"
	if str != expected {
		t.Errorf("slice printed as %q; expected %q", str, expected)
	}
}

func TestMapType(t *testing.T) {
	var m map[string]int
	mapStringInt := getTypeUnlocked("map", reflect.TypeOf(m))
	var newm map[string]int
	newMapStringInt := getTypeUnlocked("map1", reflect.TypeOf(newm))
	if mapStringInt != newMapStringInt {
		t.Errorf("second registration of map[string]int creates new type")
	}
	var b map[string]bool
	mapStringBool := getTypeUnlocked("", reflect.TypeOf(b))
	if mapStringBool == mapStringInt {
		t.Errorf("registration of map[string]bool creates same type as map[string]int")
	}
	str := mapStringBool.string()
	expected := "map[string]bool"
	if str != expected {
		t.Errorf("map printed as %q; expected %q", str, expected)
	}
}

type Bar struct {
	X string
}

// This structure has pointers and refers to itself, making it a good test case.
type Foo struct {
	A int
	B int32 // will become int
	C string
	D []byte
	E *float64    // will become float64
	F ****float64 // will become float64
	G *Bar
	H *Bar // should not interpolate the definition of Bar again
	I *Foo // will not explode
}

func TestStructType(t *testing.T) {
	sstruct := getTypeUnlocked("Foo", reflect.TypeOf(Foo{}))
	str := sstruct.string()
	// If we can print it correctly, we built it correctly.
	expected := "Foo = struct { A int; B int; C string; D bytes; E float; F float; G Bar = struct { X string; }; H Bar; I Foo; }"
	if str != expected {
		t.Errorf("struct printed as %q; expected %q", str, expected)
	}
}

// Should be OK to register the same type multiple times, as long as they're
// at the same level of indirection.
func TestRegistration(t *testing.T) {
	type T struct{ a int }
	Register(new(T))
	Register(new(T))
}
<commit_msg>encoding/gob: test for type registration name.<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.

package gob

import (
	"reflect"
	"testing"
)

type typeT struct {
	id  typeId
	str string
}

var basicTypes = []typeT{
	{tBool, "bool"},
	{tInt, "int"},
	{tUint, "uint"},
	{tFloat, "float"},
	{tBytes, "bytes"},
	{tString, "string"},
}

func getTypeUnlocked(name string, rt reflect.Type) gobType {
	typeLock.Lock()
	defer typeLock.Unlock()
	t, err := getBaseType(name, rt)
	if err != nil {
		panic("getTypeUnlocked: " + err.Error())
	}
	return t
}

// Sanity checks
func TestBasic(t *testing.T) {
	for _, tt := range basicTypes {
		if tt.id.string() != tt.str {
			t.Errorf("checkType: expected %q got %s", tt.str, tt.id.string())
		}
		if tt.id == 0 {
			t.Errorf("id for %q is zero", tt.str)
		}
	}
}

// Reregister some basic types to check registration is idempotent.
func TestReregistration(t *testing.T) {
	newtyp := getTypeUnlocked("int", reflect.TypeOf(int(0)))
	if newtyp != tInt.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("uint", reflect.TypeOf(uint(0)))
	if newtyp != tUint.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
	newtyp = getTypeUnlocked("string", reflect.TypeOf("hello"))
	if newtyp != tString.gobType() {
		t.Errorf("reregistration of %s got new type", newtyp.string())
	}
}

func TestArrayType(t *testing.T) {
	var a3 [3]int
	a3int := getTypeUnlocked("foo", reflect.TypeOf(a3))
	newa3int := getTypeUnlocked("bar", reflect.TypeOf(a3))
	if a3int != newa3int {
		t.Errorf("second registration of [3]int creates new type")
	}
	var a4 [4]int
	a4int := getTypeUnlocked("goo", reflect.TypeOf(a4))
	if a3int == a4int {
		t.Errorf("registration of [3]int creates same type as [4]int")
	}
	var b3 [3]bool
	a3bool := getTypeUnlocked("", reflect.TypeOf(b3))
	if a3int == a3bool {
		t.Errorf("registration of [3]bool creates same type as [3]int")
	}
	str := a3bool.string()
	expected := "[3]bool"
	if str != expected {
		t.Errorf("array printed as %q; expected %q", str, expected)
	}
}

func TestSliceType(t *testing.T) {
	var s []int
	sint := getTypeUnlocked("slice", reflect.TypeOf(s))
	var news []int
	newsint := getTypeUnlocked("slice1", reflect.TypeOf(news))
	if sint != newsint {
		t.Errorf("second registration of []int creates new type")
	}
	var b []bool
	sbool := getTypeUnlocked("", reflect.TypeOf(b))
	if sbool == sint {
		t.Errorf("registration of []bool creates same type as []int")
	}
	str := sbool.string()
	expected := "[]bool"
	if str != expected {
		t.Errorf("slice printed as %q; expected %q", str, expected)
	}
}

func TestMapType(t *testing.T) {
	var m map[string]int
	mapStringInt := getTypeUnlocked("map", reflect.TypeOf(m))
	var newm map[string]int
	newMapStringInt := getTypeUnlocked("map1", reflect.TypeOf(newm))
	if mapStringInt != newMapStringInt {
		t.Errorf("second registration of map[string]int creates new type")
	}
	var b map[string]bool
	mapStringBool := getTypeUnlocked("", reflect.TypeOf(b))
	if mapStringBool == mapStringInt {
		t.Errorf("registration of map[string]bool creates same type as map[string]int")
	}
	str := mapStringBool.string()
	expected := "map[string]bool"
	if str != expected {
		t.Errorf("map printed as %q; expected %q", str, expected)
	}
}

type Bar struct {
	X string
}

// This structure has pointers and refers to itself, making it a good test case.
type Foo struct {
	A int
	B int32 // will become int
	C string
	D []byte
	E *float64    // will become float64
	F ****float64 // will become float64
	G *Bar
	H *Bar // should not interpolate the definition of Bar again
	I *Foo // will not explode
}

func TestStructType(t *testing.T) {
	sstruct := getTypeUnlocked("Foo", reflect.TypeOf(Foo{}))
	str := sstruct.string()
	// If we can print it correctly, we built it correctly.
	expected := "Foo = struct { A int; B int; C string; D bytes; E float; F float; G Bar = struct { X string; }; H Bar; I Foo; }"
	if str != expected {
		t.Errorf("struct printed as %q; expected %q", str, expected)
	}
}

// Should be OK to register the same type multiple times, as long as they're
// at the same level of indirection.
func TestRegistration(t *testing.T) {
	type T struct{ a int }
	Register(new(T))
	Register(new(T))
}

type N1 struct{}
type N2 struct{}

// See comment in type.go/Register.
func TestRegistrationNaming(t *testing.T) {
	testCases := []struct {
		t    interface{}
		name string
	}{
		{&N1{}, "*gob.N1"},
		{N2{}, "encoding/gob.N2"},
	}

	for _, tc := range testCases {
		Register(tc.t)

		tct := reflect.TypeOf(tc.t)
		if ct := nameToConcreteType[tc.name]; ct != tct {
			t.Errorf("nameToConcreteType[%q] = %v, want %v", tc.name, ct, tct)
		}
		// concreteTypeToName is keyed off the base type.
		if tct.Kind() == reflect.Ptr {
			tct = tct.Elem()
		}
		if n := concreteTypeToName[tct]; n != tc.name {
			t.Errorf("concreteTypeToName[%v] got %v, want %v", tct, n, tc.name)
		}
	}
}
<|endoftext|>


<commit_before>package message

import "github.com/tychoish/grip/level"

// Composer defines an interface with a "String()" method that
// returns the message in string format. Objects that implement this
// interface, in combination to the Compose[*] operations, the
// String() method is only caled if the priority of the method is
// greater than the threshold priority. This makes it possible to
// defer building log messages (that may be somewhat expensive to
// generate) until it's certain that we're going to be outputting the
// message.
type Composer interface {
	// Returns the content of the message as a string for use in
	// line-printing logging engines.
	String() string

	// A "raw" format of the logging output for use by some Sender
	// implementations that write logged items to interfaces that
	// accept JSON or another structured format.
	Raw() interface{}

	// Returns "true" when the message has content and should be
	// logged, and false otherwise. When false, the sender can
	// (and should!) ignore messages even if they are otherwise
	// above the logging threshold.
	Loggable() bool

	// Priority returns the priority of the message.
	Priority() level.Priority
	SetPriority(level.Priority) error
}

// ConvertToComposer can coerce unknown objects into Composer
// instances, as possible.
func ConvertToComposer(p level.Priority, message interface{}) Composer {
	switch message := message.(type) {
	case Composer:
		_ = message.SetPriority(p)
		return message
	case string:
		return NewDefaultMessage(p, message)
	case error:
		return NewErrorMessage(p, message)
	case []string, []interface{}:
		return NewLineMessage(p, message...)
	case []byte:
		return NewBytesMessage(p, message)
	case map[string]interface{}:
		return NewFields(p, Fields(message))
	case Fields:
		return NewFields(p, message)
	case nil:
		return NewLineMessage(p)
	default:
		return NewFormattedMessage(p, "%+v", message)
	}
}
<commit_msg>Revert "fix formatting of string slices"<commit_after>package message

import "github.com/tychoish/grip/level"

// Composer defines an interface with a "String()" method that
// returns the message in string format. Objects that implement this
// interface, in combination to the Compose[*] operations, the
// String() method is only caled if the priority of the method is
// greater than the threshold priority. This makes it possible to
// defer building log messages (that may be somewhat expensive to
// generate) until it's certain that we're going to be outputting the
// message.
type Composer interface {
	// Returns the content of the message as a string for use in
	// line-printing logging engines.
	String() string

	// A "raw" format of the logging output for use by some Sender
	// implementations that write logged items to interfaces that
	// accept JSON or another structured format.
	Raw() interface{}

	// Returns "true" when the message has content and should be
	// logged, and false otherwise. When false, the sender can
	// (and should!) ignore messages even if they are otherwise
	// above the logging threshold.
	Loggable() bool

	// Priority returns the priority of the message.
	Priority() level.Priority
	SetPriority(level.Priority) error
}

// ConvertToComposer can coerce unknown objects into Composer
// instances, as possible.
func ConvertToComposer(p level.Priority, message interface{}) Composer {
	switch message := message.(type) {
	case Composer:
		_ = message.SetPriority(p)
		return message
	case string:
		return NewDefaultMessage(p, message)
	case error:
		return NewErrorMessage(p, message)
	case []string, []interface{}:
		return NewLineMessage(p, message)
	case []byte:
		return NewBytesMessage(p, message)
	case map[string]interface{}:
		return NewFields(p, Fields(message))
	case Fields:
		return NewFields(p, message)
	case nil:
		return NewLineMessage(p)
	default:
		return NewFormattedMessage(p, "%+v", message)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/PieterD/kafka"
	"github.com/PieterD/kafka/internal/killchan"
	"github.com/Shopify/sarama"
)

var (
	fPeers     = flag.String("peers", os.Getenv("ZOOKEEPER_PEERS"), "List of Zookeeper peer addresses (Defaults to ZOOKEEPER_PEERS env)")
	fPartition = flag.Int("partition", -1, "Partition to receive from")
	fTopic     = flag.String("topic", "", "Topic to receive from")
	fOffset    = flag.String("offset", "newest", "newest, oldest, manual")
	fManual    = flag.Int64("manual", -1, "Set offset when -offset=manual")
	fVerbose   = flag.Bool("verbose", false, "Print message details")

	logger = log.New(os.Stderr, "consumer", log.LstdFlags)
)

func flagbad(f string, i ...interface{}) {
	fmt.Fprintf(os.Stderr, f, i...)
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *fPeers == "" {
		flagbad("-peers is empty\n")
	}
	if *fPartition == -1 {
		flagbad("-partition is empty\n")
	}
	if *fTopic == "" {
		flagbad("-topic is empty\n")
	}
	var offset int64
	switch *fOffset {
	case "newest":
		offset = sarama.OffsetNewest
	case "oldest":
		offset = sarama.OffsetOldest
	case "manual":
		offset = *fManual
		if offset <= -1 {
			flagbad("-manual positive integer expected")
		}
	default:
		flagbad("-offset expects newest, oldest or time\n")
	}

	kfk, err := kafka.New("kafka-consumer", logger, strings.Split(*fPeers, ","))
	if err != nil {
		logger.Panicf("Creating kafka client: %v", err)
	}
	defer kfk.Close()

	err = kfk.Listen(*fTopic, int32(*fPartition), offset)
	if err != nil {
		logger.Panicf("Failed to listen to partition %s:%d: %v", *fTopic, *fPartition, err)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, syscall.SIGINT)
	kc := killchan.New()
	go func() {
		<-sigchan
		kc.Kill()
	}()

	for {
		select {
		case <-kc.Chan():
			logger.Printf("Interrupt\n")
			return
		case message := <-kfk.Incoming():
			if *fVerbose {
				fmt.Printf("%s:%d (offset %d) %s\n", message.Topic, message.Partition, message.Offset, message.Val)
			} else {
				fmt.Printf("%s\n", message.Val)
			}
		}
	}
}
<commit_msg>better error message<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/PieterD/kafka"
	"github.com/PieterD/kafka/internal/killchan"
	"github.com/Shopify/sarama"
)

var (
	fPeers     = flag.String("peers", os.Getenv("ZOOKEEPER_PEERS"), "List of Zookeeper peer addresses (Defaults to ZOOKEEPER_PEERS env)")
	fPartition = flag.Int("partition", -1, "Partition to receive from")
	fTopic     = flag.String("topic", "", "Topic to receive from")
	fOffset    = flag.String("offset", "newest", "newest, oldest, manual")
	fManual    = flag.Int64("manual", -1, "Set offset when -offset=manual")
	fVerbose   = flag.Bool("verbose", false, "Print message details")

	logger = log.New(os.Stderr, "consumer", log.LstdFlags)
)

func flagbad(f string, i ...interface{}) {
	fmt.Fprintf(os.Stderr, f, i...)
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *fPeers == "" {
		flagbad("-peers is empty\n")
	}
	if *fPartition == -1 {
		flagbad("-partition is empty\n")
	}
	if *fTopic == "" {
		flagbad("-topic is empty\n")
	}
	var offset int64
	switch *fOffset {
	case "newest":
		offset = sarama.OffsetNewest
	case "oldest":
		offset = sarama.OffsetOldest
	case "manual":
		offset = *fManual
		if offset <= -1 {
			flagbad("-manual positive integer expected")
		}
	default:
		flagbad("-offset expects newest, oldest or time\n")
	}

	kfk, err := kafka.New("kafka-consumer", logger, strings.Split(*fPeers, ","))
	if err != nil {
		logger.Panicf("Failed to create kafka client: %v", err)
	}
	defer kfk.Close()

	err = kfk.Listen(*fTopic, int32(*fPartition), offset)
	if err != nil {
		logger.Panicf("Failed to listen to partition %s:%d: %v", *fTopic, *fPartition, err)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, syscall.SIGINT)
	kc := killchan.New()
	go func() {
		<-sigchan
		kc.Kill()
	}()

	for {
		select {
		case <-kc.Chan():
			logger.Printf("Interrupt\n")
			return
		case message := <-kfk.Incoming():
			if *fVerbose {
				fmt.Printf("%s:%d (offset %d) %s\n", message.Topic, message.Partition, message.Offset, message.Val)
			} else {
				fmt.Printf("%s\n", message.Val)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
)

// Build some graphs, and then print some code that can be used to generate the
// graphs in SageMath.
func main() {
	// Diamond Graph
	gDiamond := NewLowBlockTimeGraph()
	d1 := gDiamond.CreateNode(gDiamond.GenesisNode())
	d2 := gDiamond.CreateNode(gDiamond.GenesisNode())
	_ = gDiamond.CreateNode(d1, d2)
	fmt.Printf("\n# Diamond Graph\n%s", gDiamond.SageGen())

	// Pentagon Graph
	gPentagon := NewLowBlockTimeGraph()
	p1 := gPentagon.CreateNode(gPentagon.GenesisNode())
	p2 := gPentagon.CreateNode(gPentagon.GenesisNode())
	p3 := gPentagon.CreateNode(p1)
	_ = gPentagon.CreateNode(p2, p3)
	fmt.Printf("\n# Pentagon Graph\n%s", gPentagon.SageGen())

	// Double Diamond Graph
	gDDiamond := NewLowBlockTimeGraph()
	dd1 := gDDiamond.CreateNode(gDDiamond.GenesisNode())
	dd2 := gDDiamond.CreateNode(gDDiamond.GenesisNode())
	dd3 := gDDiamond.CreateNode(dd1, dd2)
	dd4 := gDDiamond.CreateNode(dd2)
	_ = gDDiamond.CreateNode(dd3, dd4)
	fmt.Printf("\n# Double Diamond Graph\n%s", gDDiamond.SageGen())

	// Nested Diamond Graph
	ndg := NewLowBlockTimeGraph()
	nd1 := ndg.CreateNode(ndg.GenesisNode())
	nd2 := ndg.CreateNode(ndg.GenesisNode())
	nd3 := ndg.CreateNode(nd1)
	nd4 := ndg.CreateNode(nd1, nd2)
	fmt.Printf("\n# Nested Diamond Graph M1\n%s", ndg.SageGen())
	nd5 := ndg.CreateNode(nd2)
	nd6 := ndg.CreateNode(nd3, nd4)
	fmt.Printf("\n# Nested Diamond Graph M2\n%s", ndg.SageGen())
	nd7 := ndg.CreateNode(nd4, nd5)
	fmt.Printf("\n# Nested Diamond Graph M3\n%s", ndg.SageGen())
	nd8 := ndg.CreateNode(nd6, nd7)
	fmt.Printf("\n# Nested Diamond Graph M4\n%s", ndg.SageGen())
	nd9 := ndg.CreateNode(nd8)
	nd10 := ndg.CreateNode(nd9)
	_ = ndg.CreateNode(nd10)
	fmt.Printf("\n# Nested Diamond Graph\n%s", ndg.SageGen())

	// Ongoing Graph.
	og := NewLowBlockTimeGraph()
	o1 := og.CreateNode(og.GenesisNode())
	o2 := og.CreateNode(og.GenesisNode())
	o3 := og.CreateNode(og.GenesisNode())
	o4 := og.CreateNode(og.GenesisNode())
	o5 := og.CreateNode(og.GenesisNode())
	o11 := og.CreateNode(o1, o2)
	o12 := og.CreateNode(o1, o2, o3)
	o13 := og.CreateNode(o2, o3, o4)
	o14 := og.CreateNode(o3, o4, o5)
	o15 := og.CreateNode(o4, o5, o5)
	o21 := og.CreateNode(o11, o12)
	o22 := og.CreateNode(o11, o12, o13)
	o23 := og.CreateNode(o12, o13, o14)
	o24 := og.CreateNode(o13, o14, o15)
	o25 := og.CreateNode(o14, o15, o15)
	o31 := og.CreateNode(o21, o22)
	o32 := og.CreateNode(o21, o22, o23)
	o33 := og.CreateNode(o22, o23, o24)
	o34 := og.CreateNode(o23, o24, o25)
	o35 := og.CreateNode(o24, o25, o25)
	o41 := og.CreateNode(o31, o32)
	o42 := og.CreateNode(o31, o32, o33)
	o43 := og.CreateNode(o32, o33, o34)
	o44 := og.CreateNode(o33, o34, o35)
	o45 := og.CreateNode(o34, o35, o35)
	_ = og.CreateNode(o41, o42, o43, o44, o45)
	fmt.Printf("\n# Ongoing Graph\n%s", og.SageGen())

	// Impossibility Proof Graph
	ip := NewLowBlockTimeGraph()
	ip1 := ip.CreateNode(ip.GenesisNode())
	ip2 := ip.CreateNode(ip.GenesisNode())
	ip3 := ip.CreateNode(ip1)
	ip4 := ip.CreateNode(ip2)
	ip5 := ip.CreateNode(ip3)
	ip6 := ip.CreateNode(ip4, ip5)
	ip7 := ip.CreateNode(ip4)
	ip8 := ip.CreateNode(ip7)
	ip9 := ip.CreateNode(ip8)
	ip10 := ip.CreateNode(ip9)
	_ = ip.CreateNode(ip10, ip6)
	fmt.Printf("\n# Impossibility Proof Graph\n%s", ip.SageGen())

	// Abstain Graph
	a := NewLowBlockTimeGraph()
	a1 := a.CreateNode(a.GenesisNode())
	a2 := a.CreateNode(a1)
	a3 := a.CreateNode(a2)
	a4 := a.CreateNode(a3)
	a5 := a.CreateNode(a4)
	a6 := a.CreateNode(a.GenesisNode())
	a7 := a.CreateNode(a6)
	a8 := a.CreateNode(a7)
	a9 := a.CreateNode(a8)
	a10 := a.CreateNode(a9)
	a11 := a.CreateNode(a10)
	a12 := a.CreateNode(a11)
	a13 := a.CreateNode(a12)
	a14 := a.CreateNode(a13)
	a15 := a.CreateNode(a14)
	a16 := a.CreateNode(a15)
	a17 := a.CreateNode(a16)
	_ = a.CreateNode(a17, a5)
	fmt.Printf("\n# Abstain Graph\n%s", a.SageGen())

	// Leech Graph
	l := NewLowBlockTimeGraph()
	l1 := l.CreateNode(l.GenesisNode())
	l2 := l.CreateNode(l1)
	l3 := l.CreateNode(l2)
	l4 := l.CreateNode(l3)
	l5 := l.CreateNode(l4)
	l6 := l.CreateNode(l5)
	l7 := l.CreateNode(l6)
	l8 := l.CreateNode(l7)
	l9 := l.CreateNode(l8)
	la := l.CreateNode(l9)
	lb := l.CreateNode(la)
	lc := l.CreateNode(lb)
	ld := l.CreateNode(l.GenesisNode())
	le := l.CreateNode(ld, l2)
	lf := l.CreateNode(le, l3)
	lg := l.CreateNode(lf, l5)
	lh := l.CreateNode(lg, l7)
	li := l.CreateNode(lh, l8)
	lj := l.CreateNode(li, l9)
	lk := l.CreateNode(lj, lb)
	_ = l.CreateNode(lk, lc)
	fmt.Printf("\n# Leech Graph\n%s", l.SageGen())

	// Low Latency Adversary Graph
	lla := NewLowBlockTimeGraph()
	lla1 := lla.CreateNode(lla.GenesisNode())
	lla2 := lla.CreateNode(lla.GenesisNode())
	lla3 := lla.CreateNode(lla.GenesisNode())
	lla4 := lla.CreateNode(lla.GenesisNode())
	lla5 := lla.CreateNode(lla.GenesisNode())
	lla6 := lla.CreateNode(lla5)
	lla7 := lla.CreateNode(lla1, lla2, lla3, lla4, lla6)
	lla8 := lla.CreateNode(lla6)
	lla9 := lla.CreateNode(lla8)
	_ = lla.CreateNode(lla7, lla9)
	fmt.Printf("\n# Low Latenct Adversary Graph\n%s", lla.SageGen())
}
<commit_msg>remove partial graph orderings<commit_after>package main

import (
	"fmt"
)

// Build some graphs, and then print some code that can be used to generate the
// graphs in SageMath.
func main() {
	// Diamond Graph
	gDiamond := NewLowBlockTimeGraph()
	d1 := gDiamond.CreateNode(gDiamond.GenesisNode())
	d2 := gDiamond.CreateNode(gDiamond.GenesisNode())
	_ = gDiamond.CreateNode(d1, d2)
	fmt.Printf("\n# Diamond Graph\n%s", gDiamond.SageGen())

	// Pentagon Graph
	gPentagon := NewLowBlockTimeGraph()
	p1 := gPentagon.CreateNode(gPentagon.GenesisNode())
	p2 := gPentagon.CreateNode(gPentagon.GenesisNode())
	p3 := gPentagon.CreateNode(p1)
	_ = gPentagon.CreateNode(p2, p3)
	fmt.Printf("\n# Pentagon Graph\n%s", gPentagon.SageGen())

	// Double Diamond Graph
	gDDiamond := NewLowBlockTimeGraph()
	dd1 := gDDiamond.CreateNode(gDDiamond.GenesisNode())
	dd2 := gDDiamond.CreateNode(gDDiamond.GenesisNode())
	dd3 := gDDiamond.CreateNode(dd1, dd2)
	dd4 := gDDiamond.CreateNode(dd2)
	_ = gDDiamond.CreateNode(dd3, dd4)
	fmt.Printf("\n# Double Diamond Graph\n%s", gDDiamond.SageGen())

	// Nested Diamond Graph
	ndg := NewLowBlockTimeGraph()
	nd1 := ndg.CreateNode(ndg.GenesisNode())
	nd2 := ndg.CreateNode(ndg.GenesisNode())
	nd3 := ndg.CreateNode(nd1)
	nd4 := ndg.CreateNode(nd1, nd2)
	nd5 := ndg.CreateNode(nd2)
	nd6 := ndg.CreateNode(nd3, nd4)
	nd7 := ndg.CreateNode(nd4, nd5)
	nd8 := ndg.CreateNode(nd6, nd7)
	nd9 := ndg.CreateNode(nd8)
	nd10 := ndg.CreateNode(nd9)
	_ = ndg.CreateNode(nd10)
	fmt.Printf("\n# Nested Diamond Graph\n%s", ndg.SageGen())

	// Ongoing Graph.
	og := NewLowBlockTimeGraph()
	o1 := og.CreateNode(og.GenesisNode())
	o2 := og.CreateNode(og.GenesisNode())
	o3 := og.CreateNode(og.GenesisNode())
	o4 := og.CreateNode(og.GenesisNode())
	o5 := og.CreateNode(og.GenesisNode())
	o11 := og.CreateNode(o1, o2)
	o12 := og.CreateNode(o1, o2, o3)
	o13 := og.CreateNode(o2, o3, o4)
	o14 := og.CreateNode(o3, o4, o5)
	o15 := og.CreateNode(o4, o5, o5)
	o21 := og.CreateNode(o11, o12)
	o22 := og.CreateNode(o11, o12, o13)
	o23 := og.CreateNode(o12, o13, o14)
	o24 := og.CreateNode(o13, o14, o15)
	o25 := og.CreateNode(o14, o15, o15)
	o31 := og.CreateNode(o21, o22)
	o32 := og.CreateNode(o21, o22, o23)
	o33 := og.CreateNode(o22, o23, o24)
	o34 := og.CreateNode(o23, o24, o25)
	o35 := og.CreateNode(o24, o25, o25)
	o41 := og.CreateNode(o31, o32)
	o42 := og.CreateNode(o31, o32, o33)
	o43 := og.CreateNode(o32, o33, o34)
	o44 := og.CreateNode(o33, o34, o35)
	o45 := og.CreateNode(o34, o35, o35)
	_ = og.CreateNode(o41, o42, o43, o44, o45)
	fmt.Printf("\n# Ongoing Graph\n%s", og.SageGen())

	// Impossibility Proof Graph
	ip := NewLowBlockTimeGraph()
	ip1 := ip.CreateNode(ip.GenesisNode())
	ip2 := ip.CreateNode(ip.GenesisNode())
	ip3 := ip.CreateNode(ip1)
	ip4 := ip.CreateNode(ip2)
	ip5 := ip.CreateNode(ip3)
	ip6 := ip.CreateNode(ip4, ip5)
	ip7 := ip.CreateNode(ip4)
	ip8 := ip.CreateNode(ip7)
	ip9 := ip.CreateNode(ip8)
	ip10 := ip.CreateNode(ip9)
	_ = ip.CreateNode(ip10, ip6)
	fmt.Printf("\n# Impossibility Proof Graph\n%s", ip.SageGen())

	// Abstain Graph
	a := NewLowBlockTimeGraph()
	a1 := a.CreateNode(a.GenesisNode())
	a2 := a.CreateNode(a1)
	a3 := a.CreateNode(a2)
	a4 := a.CreateNode(a3)
	a5 := a.CreateNode(a4)
	a6 := a.CreateNode(a.GenesisNode())
	a7 := a.CreateNode(a6)
	a8 := a.CreateNode(a7)
	a9 := a.CreateNode(a8)
	a10 := a.CreateNode(a9)
	a11 := a.CreateNode(a10)
	a12 := a.CreateNode(a11)
	a13 := a.CreateNode(a12)
	a14 := a.CreateNode(a13)
	a15 := a.CreateNode(a14)
	a16 := a.CreateNode(a15)
	a17 := a.CreateNode(a16)
	_ = a.CreateNode(a17, a5)
	fmt.Printf("\n# Abstain Graph\n%s", a.SageGen())

	// Leech Graph
	l := NewLowBlockTimeGraph()
	l1 := l.CreateNode(l.GenesisNode())
	l2 := l.CreateNode(l1)
	l3 := l.CreateNode(l2)
	l4 := l.CreateNode(l3)
	l5 := l.CreateNode(l4)
	l6 := l.CreateNode(l5)
	l7 := l.CreateNode(l6)
	l8 := l.CreateNode(l7)
	l9 := l.CreateNode(l8)
	la := l.CreateNode(l9)
	lb := l.CreateNode(la)
	lc := l.CreateNode(lb)
	ld := l.CreateNode(l.GenesisNode())
	le := l.CreateNode(ld, l2)
	lf := l.CreateNode(le, l3)
	lg := l.CreateNode(lf, l5)
	lh := l.CreateNode(lg, l7)
	li := l.CreateNode(lh, l8)
	lj := l.CreateNode(li, l9)
	lk := l.CreateNode(lj, lb)
	_ = l.CreateNode(lk, lc)
	fmt.Printf("\n# Leech Graph\n%s", l.SageGen())

	// Low Latency Adversary Graph
	lla := NewLowBlockTimeGraph()
	lla1 := lla.CreateNode(lla.GenesisNode())
	lla2 := lla.CreateNode(lla.GenesisNode())
	lla3 := lla.CreateNode(lla.GenesisNode())
	lla4 := lla.CreateNode(lla.GenesisNode())
	lla5 := lla.CreateNode(lla.GenesisNode())
	lla6 := lla.CreateNode(lla5)
	lla7 := lla.CreateNode(lla1, lla2, lla3, lla4, lla6)
	lla8 := lla.CreateNode(lla6)
	lla9 := lla.CreateNode(lla8)
	_ = lla.CreateNode(lla7, lla9)
	fmt.Printf("\n# Low Latenct Adversary Graph\n%s", lla.SageGen())
}
<|endoftext|>


<commit_before>package main_test

//go:generate v23 test generate .

import (
	"os"

	"v.io/core/veyron/lib/testutil/v23tests"
	_ "v.io/core/veyron/profiles/static"
)

func V23TestAgentd(t *v23tests.T) {
	vrunBin := t.BuildGoPkg("v.io/core/veyron/tools/vrun")
	pingpongBin := t.BuildGoPkg("v.io/core/veyron/security/agent/pingpong")
	agentdBin := t.BuildGoPkg("v.io/core/veyron/security/agent/agentd")
	helperBin := t.BuildGoPkg("v.io/core/veyron/tools/vrun/test")
	principalBin := t.BuildGoPkg("v.io/core/veyron/tools/principal")

	creds := t.TempDir()
	agentdBin.WithEnv("VEYRON_CREDENTIALS="+creds).Start("--no_passphrase",
		"--additional_principals="+creds,
		helperBin.Path(),
		vrunBin.Path(),
		pingpongBin.Path(),
		principalBin.Path()).WaitOrDie(os.Stdout, os.Stderr)
}
<commit_msg>TBR veyron/tools/vrun/test: remove old generated file<commit_after><|endoftext|>


<commit_before>package bootkube

import (
	"fmt"

	"github.com/coreos-inc/pluton"
	"github.com/coreos-inc/pluton/spawn"

	"github.com/coreos/mantle/kola/cluster"
	"github.com/coreos/mantle/platform"
)

// Run various destruction tests
func bootkubeDestruction(c cluster.TestCluster) error {
	bc, err := spawn.MakeBootkubeCluster(c)
	if err != nil {
		return err
	}

	if err := masterRestart(bc); err != nil {
		return fmt.Errorf("masterRestart: %s", err)
	}

	// TODO: add more destructive tests. Also test that workloads started
	// before any destruction test are still working after the destruction
	// test rather then creating a new workload after the destructive test.

	return nil
}

// Restart master node and run nginxCheck
func masterRestart(c *pluton.Cluster) error {
	if err := platform.Reboot(c.Masters[0]); err != nil {
		return err
	}

	if err := c.NodeCheck(20); err != nil {
		return fmt.Errorf("nodeCheck: %s", err)
	}

	if err := nginxCheck(c); err != nil {
		return fmt.Errorf("nginxCheck: %s", err)
	}
	return nil
}
<commit_msg>tests/bootkube: bump destruction nodecheck timeout<commit_after>package bootkube

import (
	"fmt"

	"github.com/coreos-inc/pluton"
	"github.com/coreos-inc/pluton/spawn"

	"github.com/coreos/mantle/kola/cluster"
	"github.com/coreos/mantle/platform"
)

// Run various destruction tests
func bootkubeDestruction(c cluster.TestCluster) error {
	bc, err := spawn.MakeBootkubeCluster(c)
	if err != nil {
		return err
	}

	if err := masterRestart(bc); err != nil {
		return fmt.Errorf("masterRestart: %s", err)
	}

	// TODO: add more destructive tests. Also test that workloads started
	// before any destruction test are still working after the destruction
	// test rather then creating a new workload after the destructive test.

	return nil
}

// Restart master node and run nginxCheck
func masterRestart(c *pluton.Cluster) error {
	if err := platform.Reboot(c.Masters[0]); err != nil {
		return err
	}

	if err := c.NodeCheck(25); err != nil {
		return fmt.Errorf("nodeCheck: %s", err)
	}

	if err := nginxCheck(c); err != nil {
		return fmt.Errorf("nginxCheck: %s", err)
	}
	return nil
}
<|endoftext|>