<commit_before>package main

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"github.com/ericchiang/k8s"
	"github.com/ericchiang/k8s/api/resource"
	apiv1 "github.com/ericchiang/k8s/api/v1"
	batchv1 "github.com/ericchiang/k8s/apis/batch/v1"
	metav1 "github.com/ericchiang/k8s/apis/meta/v1"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/zerolog/log"
)

// CiBuilderClient is the interface for running kubernetes commands specific to this application
type CiBuilderClient interface {
	CreateCiBuilderJob(CiBuilderParams) (*batchv1.Job, error)
	RemoveCiBuilderJob(string) error
}

// CiBuilderParams contains the parameters required to create a ci builder job
type CiBuilderParams struct {
	RepoFullName         string
	RepoURL              string
	RepoBranch           string
	RepoRevision         string
	EnvironmentVariables map[string]string
}

type ciBuilderClientImpl struct {
	KubeClient *k8s.Client
}

// newCiBuilderClient return a estafette ci builder client
func newCiBuilderClient() (ciBuilderClient CiBuilderClient, err error) {

	kubeClient, err := k8s.NewInClusterClient()
	if err != nil {
		log.Error().Err(err).Msg("Creating k8s client failed")
		return
	}

	ciBuilderClient = &ciBuilderClientImpl{
		KubeClient: kubeClient,
	}

	return
}

// CreateCiBuilderJob creates an estafette-ci-builder job in Kubernetes to run the estafette build
func (cbc *ciBuilderClientImpl) CreateCiBuilderJob(ciBuilderParams CiBuilderParams) (job *batchv1.Job, err error) {

	// create job name of max 63 chars
	re := regexp.MustCompile("[^a-zA-Z0-9]+")
	repoName := re.ReplaceAllString(ciBuilderParams.RepoFullName, "-")
	if len(repoName) > 50 {
		repoName = repoName[:50]
	}
	jobName := strings.ToLower(fmt.Sprintf("build-%v-%v", repoName, ciBuilderParams.RepoRevision[:6]))

	// create envvars for job
	estafetteGitNameName := "ESTAFETTE_GIT_NAME"
	estafetteGitNameValue := ciBuilderParams.RepoFullName
	estafetteGitURLName := "ESTAFETTE_GIT_URL"
	estafetteGitURLValue := ciBuilderParams.RepoURL
	estafetteGitBranchName := "ESTAFETTE_GIT_BRANCH"
	estafetteGitBranchValue := ciBuilderParams.RepoBranch
	estafetteGitRevisionName := "ESTAFETTE_GIT_REVISION"
	estafetteGitRevisionValue := ciBuilderParams.RepoRevision
	estafetteBuildJobNameName := "ESTAFETTE_BUILD_JOB_NAME"
	estafetteBuildJobNameValue := jobName
	estafetteCiServerBaseURLName := "ESTAFETTE_CI_SERVER_BASE_URL"
	estafetteCiServerBaseURLValue := *estafetteCiServerBaseURL

	// temporarily pass build version equal to revision from the outside until estafette supports versioning
	estafetteBuildVersionName := "ESTAFETTE_BUILD_VERSION"
	estafetteBuildVersionValue := ciBuilderParams.RepoRevision
	estafetteBuildVersionPatchName := "ESTAFETTE_BUILD_VERSION_PATCH"
	estafetteBuildVersionPatchValue := "1"

	environmentVariables := []*apiv1.EnvVar{
		&apiv1.EnvVar{
			Name:  &estafetteGitNameName,
			Value: &estafetteGitNameValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitURLName,
			Value: &estafetteGitURLValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitBranchName,
			Value: &estafetteGitBranchValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitRevisionName,
			Value: &estafetteGitRevisionValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildVersionName,
			Value: &estafetteBuildVersionValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildVersionPatchName,
			Value: &estafetteBuildVersionPatchValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildJobNameName,
			Value: &estafetteBuildJobNameValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteCiServerBaseURLName,
			Value: &estafetteCiServerBaseURLValue,
		},
	}

	for key, value := range ciBuilderParams.EnvironmentVariables {
		environmentVariables = append(environmentVariables, &apiv1.EnvVar{
			Name:  &key,
			Value: &value,
		})
	}

	// define resource request and limit values to fit reasonably well inside a n1-highmem-4 machine
	cpuRequest := "1.0"
	cpuLimit := "3.0"
	memoryRequest := "2.0Gi"
	memoryLimit := "20.0Gi"

	// other job config
	containerName := "estafette-ci-builder"
	image := fmt.Sprintf("estafette/estafette-ci-builder:%v", *estafetteCiBuilderVersion)
	restartPolicy := "Never"
	privileged := true

	job = &batchv1.Job{
		Metadata: &metav1.ObjectMeta{
			Name:      &jobName,
			Namespace: &cbc.KubeClient.Namespace,
			Labels: map[string]string{
				"createdBy": "estafette",
			},
		},
		Spec: &batchv1.JobSpec{
			Template: &apiv1.PodTemplateSpec{
				Metadata: &metav1.ObjectMeta{
					Labels: map[string]string{
						"createdBy": "estafette",
					},
				},
				Spec: &apiv1.PodSpec{
					Containers: []*apiv1.Container{
						&apiv1.Container{
							Name:  &containerName,
							Image: &image,
							Env:   environmentVariables,
							SecurityContext: &apiv1.SecurityContext{
								Privileged: &privileged,
							},
							Resources: &apiv1.ResourceRequirements{
								Requests: map[string]*resource.Quantity{
									"cpu":    &resource.Quantity{String_: &cpuRequest},
									"memory": &resource.Quantity{String_: &memoryRequest},
								},
								Limits: map[string]*resource.Quantity{
									"cpu":    &resource.Quantity{String_: &cpuLimit},
									"memory": &resource.Quantity{String_: &memoryLimit},
								},
							},
						},
					},
					RestartPolicy: &restartPolicy,
				},
			},
		},
	}

	// track call via prometheus
	outgoingAPIRequestTotal.With(prometheus.Labels{"target": "kubernetes"}).Inc()

	job, err = cbc.KubeClient.BatchV1().CreateJob(context.Background(), job)

	return
}

// RemoveCiBuilderJob waits for a job to finish and then removes it
func (cbc *ciBuilderClientImpl) RemoveCiBuilderJob(jobName string) (err error) {

	// watch for job updates
	watcher, err := cbc.KubeClient.BatchV1().WatchJobs(context.Background(), cbc.KubeClient.Namespace)
	if err != nil {
		log.Error().Err(err).
			Str("jobName", jobName).
			Msgf("WatchJobs call for job %v failed", jobName)
		return
	}

	// wait for job to succeed
	for {
		event, job, err := watcher.Next()
		if err != nil {
			log.Error().Err(err)
			return err
		}

		if *event.Type == k8s.EventModified && *job.Metadata.Name == jobName && *job.Status.Succeeded == 1 {
			break
		}
	}

	// delete job
	err = cbc.KubeClient.BatchV1().DeleteJob(context.Background(), jobName, cbc.KubeClient.Namespace)
	if err != nil {
		log.Error().Err(err).
			Str("jobName", jobName).
			Msgf("Deleting job %v failed", jobName)
		return
	}

	return
}
<commit_msg>pass ESTAFETTE_GCR_PROJECT for now<commit_after>package main

import (
	"context"
	"fmt"
	"regexp"
	"strings"

	"github.com/ericchiang/k8s"
	"github.com/ericchiang/k8s/api/resource"
	apiv1 "github.com/ericchiang/k8s/api/v1"
	batchv1 "github.com/ericchiang/k8s/apis/batch/v1"
	metav1 "github.com/ericchiang/k8s/apis/meta/v1"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/rs/zerolog/log"
)

// CiBuilderClient is the interface for running kubernetes commands specific to this application
type CiBuilderClient interface {
	CreateCiBuilderJob(CiBuilderParams) (*batchv1.Job, error)
	RemoveCiBuilderJob(string) error
}

// CiBuilderParams contains the parameters required to create a ci builder job
type CiBuilderParams struct {
	RepoFullName         string
	RepoURL              string
	RepoBranch           string
	RepoRevision         string
	EnvironmentVariables map[string]string
}

type ciBuilderClientImpl struct {
	KubeClient *k8s.Client
}

// newCiBuilderClient return a estafette ci builder client
func newCiBuilderClient() (ciBuilderClient CiBuilderClient, err error) {

	kubeClient, err := k8s.NewInClusterClient()
	if err != nil {
		log.Error().Err(err).Msg("Creating k8s client failed")
		return
	}

	ciBuilderClient = &ciBuilderClientImpl{
		KubeClient: kubeClient,
	}

	return
}

// CreateCiBuilderJob creates an estafette-ci-builder job in Kubernetes to run the estafette build
func (cbc *ciBuilderClientImpl) CreateCiBuilderJob(ciBuilderParams CiBuilderParams) (job *batchv1.Job, err error) {

	// create job name of max 63 chars
	re := regexp.MustCompile("[^a-zA-Z0-9]+")
	repoName := re.ReplaceAllString(ciBuilderParams.RepoFullName, "-")
	if len(repoName) > 50 {
		repoName = repoName[:50]
	}
	jobName := strings.ToLower(fmt.Sprintf("build-%v-%v", repoName, ciBuilderParams.RepoRevision[:6]))

	// create envvars for job
	estafetteGitNameName := "ESTAFETTE_GIT_NAME"
	estafetteGitNameValue := ciBuilderParams.RepoFullName
	estafetteGitURLName := "ESTAFETTE_GIT_URL"
	estafetteGitURLValue := ciBuilderParams.RepoURL
	estafetteGitBranchName := "ESTAFETTE_GIT_BRANCH"
	estafetteGitBranchValue := ciBuilderParams.RepoBranch
	estafetteGitRevisionName := "ESTAFETTE_GIT_REVISION"
	estafetteGitRevisionValue := ciBuilderParams.RepoRevision
	estafetteBuildJobNameName := "ESTAFETTE_BUILD_JOB_NAME"
	estafetteBuildJobNameValue := jobName
	estafetteCiServerBaseURLName := "ESTAFETTE_CI_SERVER_BASE_URL"
	estafetteCiServerBaseURLValue := *estafetteCiServerBaseURL

	// temporarily pass build version equal to revision from the outside until estafette supports versioning
	estafetteBuildVersionName := "ESTAFETTE_BUILD_VERSION"
	estafetteBuildVersionValue := ciBuilderParams.RepoRevision
	estafetteBuildVersionPatchName := "ESTAFETTE_BUILD_VERSION_PATCH"
	estafetteBuildVersionPatchValue := "1"
	estafetteGcrProjectName := "ESTAFETTE_GCR_PROJECT"
	estafetteGcrProjectValue := "travix-com"

	environmentVariables := []*apiv1.EnvVar{
		&apiv1.EnvVar{
			Name:  &estafetteGitNameName,
			Value: &estafetteGitNameValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitURLName,
			Value: &estafetteGitURLValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitBranchName,
			Value: &estafetteGitBranchValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGitRevisionName,
			Value: &estafetteGitRevisionValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildVersionName,
			Value: &estafetteBuildVersionValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildVersionPatchName,
			Value: &estafetteBuildVersionPatchValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteBuildJobNameName,
			Value: &estafetteBuildJobNameValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteCiServerBaseURLName,
			Value: &estafetteCiServerBaseURLValue,
		},
		&apiv1.EnvVar{
			Name:  &estafetteGcrProjectName,
			Value: &estafetteGcrProjectValue,
		},
	}

	for key, value := range ciBuilderParams.EnvironmentVariables {
		environmentVariables = append(environmentVariables, &apiv1.EnvVar{
			Name:  &key,
			Value: &value,
		})
	}

	// define resource request and limit values to fit reasonably well inside a n1-highmem-4 machine
	cpuRequest := "1.0"
	cpuLimit := "3.0"
	memoryRequest := "2.0Gi"
	memoryLimit := "20.0Gi"

	// other job config
	containerName := "estafette-ci-builder"
	image := fmt.Sprintf("estafette/estafette-ci-builder:%v", *estafetteCiBuilderVersion)
	restartPolicy := "Never"
	privileged := true

	job = &batchv1.Job{
		Metadata: &metav1.ObjectMeta{
			Name:      &jobName,
			Namespace: &cbc.KubeClient.Namespace,
			Labels: map[string]string{
				"createdBy": "estafette",
			},
		},
		Spec: &batchv1.JobSpec{
			Template: &apiv1.PodTemplateSpec{
				Metadata: &metav1.ObjectMeta{
					Labels: map[string]string{
						"createdBy": "estafette",
					},
				},
				Spec: &apiv1.PodSpec{
					Containers: []*apiv1.Container{
						&apiv1.Container{
							Name:  &containerName,
							Image: &image,
							Env:   environmentVariables,
							SecurityContext: &apiv1.SecurityContext{
								Privileged: &privileged,
							},
							Resources: &apiv1.ResourceRequirements{
								Requests: map[string]*resource.Quantity{
									"cpu":    &resource.Quantity{String_: &cpuRequest},
									"memory": &resource.Quantity{String_: &memoryRequest},
								},
								Limits: map[string]*resource.Quantity{
									"cpu":    &resource.Quantity{String_: &cpuLimit},
									"memory": &resource.Quantity{String_: &memoryLimit},
								},
							},
						},
					},
					RestartPolicy: &restartPolicy,
				},
			},
		},
	}

	// track call via prometheus
	outgoingAPIRequestTotal.With(prometheus.Labels{"target": "kubernetes"}).Inc()

	job, err = cbc.KubeClient.BatchV1().CreateJob(context.Background(), job)

	return
}

// RemoveCiBuilderJob waits for a job to finish and then removes it
func (cbc *ciBuilderClientImpl) RemoveCiBuilderJob(jobName string) (err error) {

	// watch for job updates
	watcher, err := cbc.KubeClient.BatchV1().WatchJobs(context.Background(), cbc.KubeClient.Namespace)
	if err != nil {
		log.Error().Err(err).
			Str("jobName", jobName).
			Msgf("WatchJobs call for job %v failed", jobName)
		return
	}

	// wait for job to succeed
	for {
		event, job, err := watcher.Next()
		if err != nil {
			log.Error().Err(err)
			return err
		}

		if *event.Type == k8s.EventModified && *job.Metadata.Name == jobName && *job.Status.Succeeded == 1 {
			break
		}
	}

	// delete job
	err = cbc.KubeClient.BatchV1().DeleteJob(context.Background(), jobName, cbc.KubeClient.Namespace)
	if err != nil {
		log.Error().Err(err).
			Str("jobName", jobName).
			Msgf("Deleting job %v failed", jobName)
		return
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2017 The etcd 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 etcdserver

import (
	"fmt"
	"os"
	"time"

	"go.etcd.io/etcd/etcdserver/api/snap"
	"go.etcd.io/etcd/etcdserver/cindex"
	"go.etcd.io/etcd/lease"
	"go.etcd.io/etcd/mvcc"
	"go.etcd.io/etcd/mvcc/backend"
	"go.etcd.io/etcd/raft/raftpb"

	"go.uber.org/zap"
)

func newBackend(cfg ServerConfig) backend.Backend {
	bcfg := backend.DefaultBackendConfig()
	bcfg.Path = cfg.backendPath()
	if cfg.BackendBatchLimit != 0 {
		bcfg.BatchLimit = cfg.BackendBatchLimit
		if cfg.Logger != nil {
			cfg.Logger.Info("setting backend batch limit", zap.Int("batch limit", cfg.BackendBatchLimit))
		}
	}
	if cfg.BackendBatchInterval != 0 {
		bcfg.BatchInterval = cfg.BackendBatchInterval
		if cfg.Logger != nil {
			cfg.Logger.Info("setting backend batch interval", zap.Duration("batch interval", cfg.BackendBatchInterval))
		}
	}
	bcfg.BackendFreelistType = cfg.BackendFreelistType
	bcfg.Logger = cfg.Logger
	if cfg.QuotaBackendBytes > 0 && cfg.QuotaBackendBytes != DefaultQuotaBytes {
		// permit 10% excess over quota for disarm
		bcfg.MmapSize = uint64(cfg.QuotaBackendBytes + cfg.QuotaBackendBytes/10)
	}
	return backend.New(bcfg)
}

// openSnapshotBackend renames a snapshot db to the current etcd db and opens it.
func openSnapshotBackend(cfg ServerConfig, ss *snap.Snapshotter, snapshot raftpb.Snapshot) (backend.Backend, error) {
	snapPath, err := ss.DBFilePath(snapshot.Metadata.Index)
	if err != nil {
		return nil, fmt.Errorf("failed to find database snapshot file (%v)", err)
	}
	if err := os.Rename(snapPath, cfg.backendPath()); err != nil {
		return nil, fmt.Errorf("failed to rename database snapshot file (%v)", err)
	}
	return openBackend(cfg), nil
}

// openBackend returns a backend using the current etcd db.
func openBackend(cfg ServerConfig) backend.Backend {
	fn := cfg.backendPath()

	now, beOpened := time.Now(), make(chan backend.Backend)
	go func() {
		beOpened <- newBackend(cfg)
	}()

	select {
	case be := <-beOpened:
		cfg.Logger.Info("opened backend db", zap.String("path", fn), zap.Duration("took", time.Since(now)))
		return be

	case <-time.After(10 * time.Second):
		cfg.Logger.Info(
			"db file is flocked by another process, or taking too long",
			zap.String("path", fn),
			zap.Duration("took", time.Since(now)),
		)
	}

	return <-beOpened
}

// recoverBackendSnapshot recovers the DB from a snapshot in case etcd crashes
// before updating the backend db after persisting raft snapshot to disk,
// violating the invariant snapshot.Metadata.Index < db.consistentIndex. In this
// case, replace the db with the snapshot db sent by the leader.
func recoverSnapshotBackend(cfg ServerConfig, oldbe backend.Backend, snapshot raftpb.Snapshot) (backend.Backend, error) {
	ci := cindex.NewConsistentIndex(oldbe.BatchTx())
	kv := mvcc.New(cfg.Logger, oldbe, &lease.FakeLessor{}, ci, mvcc.StoreConfig{CompactionBatchLimit: cfg.CompactionBatchLimit})
	defer kv.Close()
	if snapshot.Metadata.Index <= kv.ConsistentIndex() {
		return oldbe, nil
	}
	oldbe.Close()
	return openSnapshotBackend(cfg, snap.New(cfg.Logger, cfg.SnapDir()), snapshot)
}
<commit_msg>etcdserver: remove redundant storage restore operation<commit_after>// Copyright 2017 The etcd 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 etcdserver

import (
	"fmt"
	"os"
	"time"

	"go.etcd.io/etcd/etcdserver/api/snap"
	"go.etcd.io/etcd/etcdserver/cindex"
	"go.etcd.io/etcd/mvcc/backend"
	"go.etcd.io/etcd/raft/raftpb"

	"go.uber.org/zap"
)

func newBackend(cfg ServerConfig) backend.Backend {
	bcfg := backend.DefaultBackendConfig()
	bcfg.Path = cfg.backendPath()
	if cfg.BackendBatchLimit != 0 {
		bcfg.BatchLimit = cfg.BackendBatchLimit
		if cfg.Logger != nil {
			cfg.Logger.Info("setting backend batch limit", zap.Int("batch limit", cfg.BackendBatchLimit))
		}
	}
	if cfg.BackendBatchInterval != 0 {
		bcfg.BatchInterval = cfg.BackendBatchInterval
		if cfg.Logger != nil {
			cfg.Logger.Info("setting backend batch interval", zap.Duration("batch interval", cfg.BackendBatchInterval))
		}
	}
	bcfg.BackendFreelistType = cfg.BackendFreelistType
	bcfg.Logger = cfg.Logger
	if cfg.QuotaBackendBytes > 0 && cfg.QuotaBackendBytes != DefaultQuotaBytes {
		// permit 10% excess over quota for disarm
		bcfg.MmapSize = uint64(cfg.QuotaBackendBytes + cfg.QuotaBackendBytes/10)
	}
	return backend.New(bcfg)
}

// openSnapshotBackend renames a snapshot db to the current etcd db and opens it.
func openSnapshotBackend(cfg ServerConfig, ss *snap.Snapshotter, snapshot raftpb.Snapshot) (backend.Backend, error) {
	snapPath, err := ss.DBFilePath(snapshot.Metadata.Index)
	if err != nil {
		return nil, fmt.Errorf("failed to find database snapshot file (%v)", err)
	}
	if err := os.Rename(snapPath, cfg.backendPath()); err != nil {
		return nil, fmt.Errorf("failed to rename database snapshot file (%v)", err)
	}
	return openBackend(cfg), nil
}

// openBackend returns a backend using the current etcd db.
func openBackend(cfg ServerConfig) backend.Backend {
	fn := cfg.backendPath()

	now, beOpened := time.Now(), make(chan backend.Backend)
	go func() {
		beOpened <- newBackend(cfg)
	}()

	select {
	case be := <-beOpened:
		cfg.Logger.Info("opened backend db", zap.String("path", fn), zap.Duration("took", time.Since(now)))
		return be

	case <-time.After(10 * time.Second):
		cfg.Logger.Info(
			"db file is flocked by another process, or taking too long",
			zap.String("path", fn),
			zap.Duration("took", time.Since(now)),
		)
	}

	return <-beOpened
}

// recoverBackendSnapshot recovers the DB from a snapshot in case etcd crashes
// before updating the backend db after persisting raft snapshot to disk,
// violating the invariant snapshot.Metadata.Index < db.consistentIndex. In this
// case, replace the db with the snapshot db sent by the leader.
func recoverSnapshotBackend(cfg ServerConfig, oldbe backend.Backend, snapshot raftpb.Snapshot) (backend.Backend, error) {
	ci := cindex.NewConsistentIndex(oldbe.BatchTx())
	if snapshot.Metadata.Index <= ci.ConsistentIndex() {
		return oldbe, nil
	}
	oldbe.Close()
	return openSnapshotBackend(cfg, snap.New(cfg.Logger, cfg.SnapDir()), snapshot)
}
<|endoftext|>


<commit_before>package cmap

import (
	"bytes"
	"encoding/json"
	"io"
	"sync"
	"sync/atomic"

	"github.com/OneOfOne/lfchan"
)

// IgnoreValue can be returned from the func called to NewFromJSON to ignore setting the value
var IgnoreValue = &struct{ bool }{true}

// DefaultShardCount is the default number of shards to use when New() or NewFromJSON() are called.
const DefaultShardCount = 1 << 4 // 16

// ForeachFunc is a function that gets passed to Foreach, returns true to break early
type ForEachFunc func(key string, val interface{}) (BreakEarly bool)

type lockedMap struct {
	m map[string]interface{}
	l sync.RWMutex
}

func (ms *lockedMap) Set(key string, v interface{}) {
	ms.l.Lock()
	ms.m[key] = v
	ms.l.Unlock()
}

func (ms *lockedMap) Get(key string) interface{} {
	ms.l.RLock()
	v := ms.m[key]
	ms.l.RUnlock()
	return v
}

func (ms *lockedMap) Has(key string) bool {
	ms.l.RLock()
	_, ok := ms.m[key]
	ms.l.RUnlock()
	return ok
}

func (ms *lockedMap) Delete(key string) {
	ms.l.Lock()
	delete(ms.m, key)
	ms.l.Unlock()
}

func (ms *lockedMap) DeleteAndGet(key string) interface{} {
	ms.l.Lock()
	v := ms.m[key]
	delete(ms.m, key)
	ms.l.Unlock()
	return v
}

func (ms *lockedMap) Len() int {
	ms.l.RLock()
	ln := len(ms.m)
	ms.l.RUnlock()
	return ln
}

func (ms *lockedMap) ForEach(fn ForEachFunc) {
	ms.l.RLock()
	for k, v := range ms.m {
		if fn(k, v) {
			break
		}
	}
	ms.l.RUnlock()
}

func (ms *lockedMap) iter(ch KeyValueChan, wg *sync.WaitGroup) {
	var kv KeyValue
	ms.l.RLock()
	for k, v := range ms.m {
		kv.Key, kv.Value = k, v
		if !ch.ch.Send(&kv, true) {
			break
		}
	}
	ms.l.RUnlock()
	wg.Done()
}

// CMap is a sharded thread-safe concurrent map.
type CMap struct {
	shards []lockedMap
	l      uint64
}

// New is an alias for NewSize(DefaultShardCount)
func New() CMap { return NewSize(DefaultShardCount) }

// NewSize returns a CMap with the specific shardSize, note that for performance reasons,
// shardCount must be a power of 2
func NewSize(shardCount int) CMap {
	// must be a power of 2
	if shardCount == 0 {
		shardCount = DefaultShardCount
	} else if shardCount&(shardCount-1) != 0 {
		panic("shardCount must be a power of 2")
	}
	cm := CMap{
		shards: make([]lockedMap, shardCount),
		l:      uint64(shardCount) - 1,
	}
	for i := range cm.shards {
		cm.shards[i].m = make(map[string]interface{}, shardCount/2)
	}
	return cm
}

// NewFromJSON is an alias for NewSizeFromJSON(DefaultShardCount, r, fn)
func NewFromJSON(r io.Reader, fn func(v interface{}) interface{}) (CMap, error) {
	return NewSizeFromJSON(DefaultShardCount, r, fn)
}

// NewFromJSON returns a cmap constructed from json, fn will return the "proper" value, for example:
// json by default reads all numbers as float64, so fn(v) where v is supposed to be an int should look like:
// 	func(v interface{}) interface{} { n, _ := v.(json.Number).Int64(); return int(n) }
// note that by default all numbers will be json.Number
func NewSizeFromJSON(shardCount int, r io.Reader, fn func(v interface{}) interface{}) (CMap, error) {
	//TODO use json.RawMessage
	cm := NewSize(shardCount)
	dec := json.NewDecoder(r)
	dec.UseNumber()
	var key string
	for dec.More() {
		t, err := dec.Token()
		if err == io.EOF {
			break
		}
		if err != nil {
			return cm, err
		}
		if t, ok := t.(string); ok && key == "" {
			key = t
			continue
		}

		if key != "" {
			if v := fn(t); v != IgnoreValue {
				cm.shard(key).m[key] = v // no need to use locks for this
			}
			key = ""
		}
	}

	return cm, nil
}

func (cm CMap) shard(key string) *lockedMap {
	h := FNV64aString(key)
	return &cm.shards[h&cm.l]
}

func (cm CMap) Set(key string, val interface{})     { cm.shard(key).Set(key, val) }
func (cm CMap) Get(key string) interface{}          { return cm.shard(key).Get(key) }
func (cm CMap) Has(key string) bool                 { return cm.shard(key).Has(key) }
func (cm CMap) Delete(key string)                   { cm.shard(key).Delete(key) }
func (cm CMap) DeleteAndGet(key string) interface{} { return cm.shard(key).DeleteAndGet(key) }

func (cm CMap) Foreach(fn ForEachFunc) {
	for i := range cm.shards {
		cm.shards[i].ForEach(fn)
	}
}

func (cm CMap) ForEachParallel(fn ForEachFunc) {
	var (
		wg   sync.WaitGroup
		exit uint32
	)
	wg.Add(len(cm.shards))
	for i := range cm.shards {
		go func(i int) {
			cm.shards[i].ForEach(func(k string, v interface{}) bool {
				if atomic.LoadUint32(&exit) == 1 {
					return true
				}
				b := fn(k, v)
				if b {
					atomic.StoreUint32(&exit, 1)
				}
				return b
			})
			wg.Done()
		}(i)
	}
	wg.Wait()
}

type KeyValue struct {
	Key   string
	Value interface{}
}

// Iter is an alias for IterBuffered(1)
func (cm CMap) Iter() KeyValueChan { return cm.IterBuffered(1) }

// IterBuffered returns a buffered channel sz, to return an unbuffered channel you can pass 1
// calling breakLoop will close the channel and consume any remaining values in it.
// note that calling breakLoop() will show as a race on the race detector but it's more or less a "safe" race,
// and it is the only clean way to break out of a channel.
func (cm CMap) IterBuffered(sz int) KeyValueChan {
	ch := KeyValueChan{lfchan.NewSize(sz)}
	go func() {
		var wg sync.WaitGroup
		wg.Add(len(cm.shards))
		for i := range cm.shards {
			go cm.shards[i].iter(ch, &wg)
		}
		wg.Wait()
		ch.ch.Close()
	}()
	return ch
}

func (cm CMap) Len() int {
	ln := 0
	for i := range cm.shards {
		ln += cm.shards[i].Len()
	}
	return ln
}

func (cm CMap) MarshalJSON() ([]byte, error) {
	var buf bytes.Buffer
	buf.WriteByte('{')
	enc := json.NewEncoder(&buf)
	cm.Foreach(func(k string, v interface{}) bool {
		buf.WriteString(`"` + k + `":`)
		enc.Encode(v)
		buf.Bytes()[buf.Len()-1] = ','
		return false
	})
	if buf.Bytes()[buf.Len()-1] == ',' {
		buf.Truncate(buf.Len() - 1)
	}
	buf.WriteByte('}')
	return buf.Bytes(), nil
}

type KeyValueChan struct {
	ch *lfchan.Chan
}

func (kvch KeyValueChan) Recv() *KeyValue {
	v, ok := kvch.ch.Recv(true)
	if !ok {
		return nil
	}
	return v.(*KeyValue)
}

func (kvch KeyValueChan) Break() {
	kvch.ch.Close()
	for {
		if _, ok := kvch.ch.Recv(true); !ok {
			return
		}
	}
}
<commit_msg>update to use non-pointer chans<commit_after>package cmap

import (
	"bytes"
	"encoding/json"
	"io"
	"sync"
	"sync/atomic"

	"github.com/OneOfOne/lfchan"
)

// IgnoreValue can be returned from the func called to NewFromJSON to ignore setting the value
var IgnoreValue = &struct{ bool }{true}

// DefaultShardCount is the default number of shards to use when New() or NewFromJSON() are called.
const DefaultShardCount = 1 << 4 // 16

// ForeachFunc is a function that gets passed to Foreach, returns true to break early
type ForEachFunc func(key string, val interface{}) (BreakEarly bool)

type lockedMap struct {
	m map[string]interface{}
	l sync.RWMutex
}

func (ms *lockedMap) Set(key string, v interface{}) {
	ms.l.Lock()
	ms.m[key] = v
	ms.l.Unlock()
}

func (ms *lockedMap) Get(key string) interface{} {
	ms.l.RLock()
	v := ms.m[key]
	ms.l.RUnlock()
	return v
}

func (ms *lockedMap) Has(key string) bool {
	ms.l.RLock()
	_, ok := ms.m[key]
	ms.l.RUnlock()
	return ok
}

func (ms *lockedMap) Delete(key string) {
	ms.l.Lock()
	delete(ms.m, key)
	ms.l.Unlock()
}

func (ms *lockedMap) DeleteAndGet(key string) interface{} {
	ms.l.Lock()
	v := ms.m[key]
	delete(ms.m, key)
	ms.l.Unlock()
	return v
}

func (ms *lockedMap) Len() int {
	ms.l.RLock()
	ln := len(ms.m)
	ms.l.RUnlock()
	return ln
}

func (ms *lockedMap) ForEach(fn ForEachFunc) {
	ms.l.RLock()
	for k, v := range ms.m {
		if fn(k, v) {
			break
		}
	}
	ms.l.RUnlock()
}

func (ms *lockedMap) iter(ch KeyValueChan, wg *sync.WaitGroup) {
	var kv KeyValue
	ms.l.RLock()
	for k, v := range ms.m {
		kv.Key, kv.Value = k, v
		if !ch.v.Send(&kv, true) {
			break
		}
	}
	ms.l.RUnlock()
	wg.Done()
}

// CMap is a sharded thread-safe concurrent map.
type CMap struct {
	shards []lockedMap
	l      uint64
}

// New is an alias for NewSize(DefaultShardCount)
func New() CMap { return NewSize(DefaultShardCount) }

// NewSize returns a CMap with the specific shardSize, note that for performance reasons,
// shardCount must be a power of 2
func NewSize(shardCount int) CMap {
	// must be a power of 2
	if shardCount == 0 {
		shardCount = DefaultShardCount
	} else if shardCount&(shardCount-1) != 0 {
		panic("shardCount must be a power of 2")
	}
	cm := CMap{
		shards: make([]lockedMap, shardCount),
		l:      uint64(shardCount) - 1,
	}
	for i := range cm.shards {
		cm.shards[i].m = make(map[string]interface{}, shardCount/2)
	}
	return cm
}

// NewFromJSON is an alias for NewSizeFromJSON(DefaultShardCount, r, fn)
func NewFromJSON(r io.Reader, fn func(v interface{}) interface{}) (CMap, error) {
	return NewSizeFromJSON(DefaultShardCount, r, fn)
}

// NewFromJSON returns a cmap constructed from json, fn will return the "proper" value, for example:
// json by default reads all numbers as float64, so fn(v) where v is supposed to be an int should look like:
// 	func(v interface{}) interface{} { n, _ := v.(json.Number).Int64(); return int(n) }
// note that by default all numbers will be json.Number
func NewSizeFromJSON(shardCount int, r io.Reader, fn func(v interface{}) interface{}) (CMap, error) {
	//TODO use json.RawMessage
	cm := NewSize(shardCount)
	dec := json.NewDecoder(r)
	dec.UseNumber()
	var key string
	for dec.More() {
		t, err := dec.Token()
		if err == io.EOF {
			break
		}
		if err != nil {
			return cm, err
		}
		if t, ok := t.(string); ok && key == "" {
			key = t
			continue
		}

		if key != "" {
			if v := fn(t); v != IgnoreValue {
				cm.shard(key).m[key] = v // no need to use locks for this
			}
			key = ""
		}
	}

	return cm, nil
}

func (cm CMap) shard(key string) *lockedMap {
	h := FNV64aString(key)
	return &cm.shards[h&cm.l]
}

func (cm CMap) Set(key string, val interface{})     { cm.shard(key).Set(key, val) }
func (cm CMap) Get(key string) interface{}          { return cm.shard(key).Get(key) }
func (cm CMap) Has(key string) bool                 { return cm.shard(key).Has(key) }
func (cm CMap) Delete(key string)                   { cm.shard(key).Delete(key) }
func (cm CMap) DeleteAndGet(key string) interface{} { return cm.shard(key).DeleteAndGet(key) }

func (cm CMap) Foreach(fn ForEachFunc) {
	for i := range cm.shards {
		cm.shards[i].ForEach(fn)
	}
}

func (cm CMap) ForEachParallel(fn ForEachFunc) {
	var (
		wg   sync.WaitGroup
		exit uint32
	)
	wg.Add(len(cm.shards))
	for i := range cm.shards {
		go func(i int) {
			cm.shards[i].ForEach(func(k string, v interface{}) bool {
				if atomic.LoadUint32(&exit) == 1 {
					return true
				}
				b := fn(k, v)
				if b {
					atomic.StoreUint32(&exit, 1)
				}
				return b
			})
			wg.Done()
		}(i)
	}
	wg.Wait()
}

type KeyValue struct {
	Key   string
	Value interface{}
}

// Iter is an alias for IterBuffered(1)
func (cm CMap) Iter() KeyValueChan { return cm.IterBuffered(1) }

// IterBuffered returns a buffered channel sz, to return an unbuffered channel you can pass 1
// calling breakLoop will close the channel and consume any remaining values in it.
// note that calling breakLoop() will show as a race on the race detector but it's more or less a "safe" race,
// and it is the only clean way to break out of a channel.
func (cm CMap) IterBuffered(sz int) KeyValueChan {
	ch := KeyValueChan{lfchan.NewSize(sz)}
	go func() {
		var wg sync.WaitGroup
		wg.Add(len(cm.shards))
		for i := range cm.shards {
			go cm.shards[i].iter(ch, &wg)
		}
		wg.Wait()
		ch.Break()
	}()
	return ch
}

func (cm CMap) Len() int {
	ln := 0
	for i := range cm.shards {
		ln += cm.shards[i].Len()
	}
	return ln
}

func (cm CMap) MarshalJSON() ([]byte, error) {
	var buf bytes.Buffer
	buf.WriteByte('{')
	enc := json.NewEncoder(&buf)
	cm.Foreach(func(k string, v interface{}) bool {
		buf.WriteString(`"` + k + `":`)
		enc.Encode(v)
		buf.Bytes()[buf.Len()-1] = ','
		return false
	})
	if buf.Bytes()[buf.Len()-1] == ',' {
		buf.Truncate(buf.Len() - 1)
	}
	buf.WriteByte('}')
	return buf.Bytes(), nil
}

type KeyValueChan struct {
	v lfchan.Chan
}

func (ch KeyValueChan) Recv() *KeyValue {
	v, ok := ch.v.Recv(true)
	if !ok {
		return nil
	}
	return v.(*KeyValue)
}

func (ch KeyValueChan) Break() {
	ch.v.Close()
	for {
		if _, ok := ch.v.Recv(true); !ok {
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/kr/pretty"
	"github.com/pkg/errors"
)

type StrategyCommon struct {
	System
	Logger
	ConfigGetter
	Downloader
	Runner
}

func (sc *StrategyCommon) Templater(version string, archMap map[string]string, system System) Templater {
	archKey := fmt.Sprintf("%s_%s", system.OS(), system.Arch())
	sc.Debugf("Arch key: %s", archKey)
	value := archMap[archKey]
	return Templater{
		Version:      version,
		OS:           system.OS(),
		Arch:         system.Arch(),
		OSArch:       archKey,
		MappedOSArch: value,
	}
}

type Strategy interface {
	Run([]string) error
}

type DockerData struct {
	Name        string
	Desc        string
	Version     string            `yaml:"version"`
	Image       string            `yaml:"image"`
	MountPwd    bool              `yaml:"mount_pwd"`
	MountPwdAs  string            `yaml:"mount_pwd_as"`
	DockerConn  bool              `yaml:"docker_conn"`
	Interactive bool              `yaml:"interactive"`
	Terminal    string            `yaml:"terminal"`
	PidHost     bool              `yaml:"pid_host"`
	RunAsUser   bool              `yaml:"run_as_user"`
	OSArchMap   map[string]string `yaml:"os_arch_map"`
}

type DockerStrategy struct {
	*StrategyCommon
	Data DockerData
}

type BinaryData struct {
	Name       string
	Desc       string
	Version    string            `yaml:"version"`
	BaseURL    string            `yaml:"base_url"`
	UnpackPath string            `yaml:"unpack_path"`
	OSArchMap  map[string]string `yaml:"os_arch_map"`
}

type BinaryStrategy struct {
	*StrategyCommon
	Data BinaryData
}

func (ds DockerStrategy) Run(extraArgs []string) error {
	// skip if docker not found
	if !ds.CheckCommand("docker", []string{"version"}) {
		ds.Debugf("skipping, docker not available")
		return &SkipError{"docker not available"}
	}

	temp := ds.Templater(ds.Data.Version, ds.Data.OSArchMap, ds.System)
	ds.Debugf("templater: %# v", pretty.Formatter(temp))

	image, err := temp.Template(ds.Data.Image)
	if err != nil {
		return errors.Wrap(err, "unable to template image name")
	}

	// TODO: add flag to force pulling image again
	// err = ds.PullDockerImage(image)
	// if err != nil {
	// 	return errors.Wrap(err, "can't pull image")
	// }

	args := []string{"run"}
	if ds.Data.Interactive {
		args = append(args, "-i")
	}
	// TODO: prompt the user for permission to do more invasive docker binding
	if ds.Data.DockerConn {
		args = append(args, "-v", "/var/run/docker.sock:/var/run/docker.sock")
	}
	if ds.Data.PidHost {
		args = append(args, "--pid", "host")
	}
	if len(ds.Data.MountPwdAs) > 0 {
		wd, _ := os.Getwd()
		args = append(args, "--volume", fmt.Sprintf("%s:%s", wd, ds.Data.MountPwdAs))
	}
	if ds.Data.MountPwd {
		wd, _ := os.Getwd()
		args = append(args, "--volume", fmt.Sprintf("%s:%s", wd, wd))
	}
	if ds.Data.RunAsUser {
		args = append(args, "-u", fmt.Sprintf("%d:%d", ds.UID(), ds.GID()))
	}
	if ds.Data.Terminal != "" {
		// TODO: support 'auto' mode that autodetects if tty is present
		if ds.Data.Terminal == "always" {
			args = append(args, "-t")
		}
	}
	args = append(args, "--rm", image)
	args = append(args, extraArgs...)

	err = ds.RunCommand("docker", args)
	if err != nil {
		return errors.Wrap(err, "can't run image")
	}

	return nil
}

func (bs BinaryStrategy) localHolenPath() (string, error) {
	var holenPath string
	if configDataPath, err := bs.Get("holen.datapath"); err == nil && len(configDataPath) > 0 {
		holenPath = configDataPath
	} else if xdgDataHome := os.Getenv("XDG_DATA_HOME"); len(xdgDataHome) > 0 {
		holenPath = filepath.Join(xdgDataHome, "holen")
	} else {
		var home string
		if home = os.Getenv("HOME"); len(home) == 0 {
			return "", fmt.Errorf("$HOME environment variable not found")
		}
		holenPath = filepath.Join(home, ".local", "share", "holen")
	}
	os.MkdirAll(holenPath, 0755)

	return holenPath, nil
}

func (bs BinaryStrategy) DownloadPath() (string, error) {
	var downloadPath string
	if configDownloadPath, err := bs.Get("binary.download"); err == nil && len(configDownloadPath) > 0 {
		downloadPath = configDownloadPath
	} else {
		holenPath, err := bs.localHolenPath()
		if err != nil {
			return "", errors.Wrap(err, "unable to get holen data path")
		}
		downloadPath = filepath.Join(holenPath, "bin")
	}
	os.MkdirAll(downloadPath, 0755)

	return downloadPath, nil
}

func (bs BinaryStrategy) TempPath() (string, error) {
	var tempPath string

	holenPath, err := bs.localHolenPath()
	if err != nil {
		return "", errors.Wrap(err, "unable to get holen data path")
	}

	tempPath = filepath.Join(holenPath, "tmp")
	os.MkdirAll(tempPath, 0755)

	return tempPath, nil
}

func (bs BinaryStrategy) Run(args []string) error {
	temp := bs.Templater(bs.Data.Version, bs.Data.OSArchMap, bs.System)
	bs.Debugf("templater: %# v", pretty.Formatter(temp))

	dlURL, err := temp.Template(bs.Data.BaseURL)
	if err != nil {
		return errors.Wrap(err, "unable to template url")
	}

	downloadPath, err := bs.DownloadPath()
	if err != nil {
		return errors.Wrap(err, "unable to find download path")
	}
	localPath := filepath.Join(downloadPath, fmt.Sprintf("%s--%s", bs.Data.Name, bs.Data.Version))

	if !bs.FileExists(localPath) {
		if len(bs.Data.UnpackPath) > 0 {
			tempPath, err := bs.TempPath()
			tempdir, err := ioutil.TempDir(tempPath, "holen")
			if err != nil {
				return errors.Wrap(err, "unable to make temporary directory")
			}
			// defer os.RemoveAll(tempdir)

			u, err := url.Parse(dlURL)
			if err != nil {
				return errors.Wrap(err, "unable to parse url")
			}

			fileName := path.Base(u.Path)
			archPath := filepath.Join(tempdir, fileName)
			unpackedPath := filepath.Join(tempdir, "unpacked")

			bs.UserMessage("Downloading %s...\n", dlURL)
			err = bs.DownloadFile(dlURL, archPath)
			if err != nil {
				return errors.Wrap(err, "can't download archive")
			}

			err = bs.UnpackArchive(archPath, unpackedPath)
			if err != nil {
				return errors.Wrap(err, "unable to unpack archive")
			}

			unpackPath, err := temp.Template(bs.Data.UnpackPath)
			if err != nil {
				return errors.Wrap(err, "unable to template unpack_path")
			}
			binPath := filepath.Join(unpackedPath, unpackPath)

			err = os.Rename(binPath, localPath)
			if err != nil {
				return errors.Wrap(err, "unable to move binary into position")
			}
		} else {
			bs.UserMessage("Downloading %s...\n", dlURL)
			err = bs.DownloadFile(dlURL, localPath)
			if err != nil {
				return errors.Wrap(err, "can't download binary")
			}
		}

		err = bs.MakeExecutable(localPath)
		if err != nil {
			return errors.Wrap(err, "unable to make binary executable")
		}
	}

	// TODO: checksum the binary

	err = bs.RunCommand(localPath, args)
	if err != nil {
		return errors.Wrap(err, "can't run binary")
	}

	return nil
}
<commit_msg>uncomment cleanup<commit_after>package main

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

	"github.com/kr/pretty"
	"github.com/pkg/errors"
)

type StrategyCommon struct {
	System
	Logger
	ConfigGetter
	Downloader
	Runner
}

func (sc *StrategyCommon) Templater(version string, archMap map[string]string, system System) Templater {
	archKey := fmt.Sprintf("%s_%s", system.OS(), system.Arch())
	sc.Debugf("Arch key: %s", archKey)
	value := archMap[archKey]
	return Templater{
		Version:      version,
		OS:           system.OS(),
		Arch:         system.Arch(),
		OSArch:       archKey,
		MappedOSArch: value,
	}
}

type Strategy interface {
	Run([]string) error
}

type DockerData struct {
	Name        string
	Desc        string
	Version     string            `yaml:"version"`
	Image       string            `yaml:"image"`
	MountPwd    bool              `yaml:"mount_pwd"`
	MountPwdAs  string            `yaml:"mount_pwd_as"`
	DockerConn  bool              `yaml:"docker_conn"`
	Interactive bool              `yaml:"interactive"`
	Terminal    string            `yaml:"terminal"`
	PidHost     bool              `yaml:"pid_host"`
	RunAsUser   bool              `yaml:"run_as_user"`
	OSArchMap   map[string]string `yaml:"os_arch_map"`
}

type DockerStrategy struct {
	*StrategyCommon
	Data DockerData
}

type BinaryData struct {
	Name       string
	Desc       string
	Version    string            `yaml:"version"`
	BaseURL    string            `yaml:"base_url"`
	UnpackPath string            `yaml:"unpack_path"`
	OSArchMap  map[string]string `yaml:"os_arch_map"`
}

type BinaryStrategy struct {
	*StrategyCommon
	Data BinaryData
}

func (ds DockerStrategy) Run(extraArgs []string) error {
	// skip if docker not found
	if !ds.CheckCommand("docker", []string{"version"}) {
		ds.Debugf("skipping, docker not available")
		return &SkipError{"docker not available"}
	}

	temp := ds.Templater(ds.Data.Version, ds.Data.OSArchMap, ds.System)
	ds.Debugf("templater: %# v", pretty.Formatter(temp))

	image, err := temp.Template(ds.Data.Image)
	if err != nil {
		return errors.Wrap(err, "unable to template image name")
	}

	// TODO: add flag to force pulling image again
	// err = ds.PullDockerImage(image)
	// if err != nil {
	// 	return errors.Wrap(err, "can't pull image")
	// }

	args := []string{"run"}
	if ds.Data.Interactive {
		args = append(args, "-i")
	}
	// TODO: prompt the user for permission to do more invasive docker binding
	if ds.Data.DockerConn {
		args = append(args, "-v", "/var/run/docker.sock:/var/run/docker.sock")
	}
	if ds.Data.PidHost {
		args = append(args, "--pid", "host")
	}
	if len(ds.Data.MountPwdAs) > 0 {
		wd, _ := os.Getwd()
		args = append(args, "--volume", fmt.Sprintf("%s:%s", wd, ds.Data.MountPwdAs))
	}
	if ds.Data.MountPwd {
		wd, _ := os.Getwd()
		args = append(args, "--volume", fmt.Sprintf("%s:%s", wd, wd))
	}
	if ds.Data.RunAsUser {
		args = append(args, "-u", fmt.Sprintf("%d:%d", ds.UID(), ds.GID()))
	}
	if ds.Data.Terminal != "" {
		// TODO: support 'auto' mode that autodetects if tty is present
		if ds.Data.Terminal == "always" {
			args = append(args, "-t")
		}
	}
	args = append(args, "--rm", image)
	args = append(args, extraArgs...)

	err = ds.RunCommand("docker", args)
	if err != nil {
		return errors.Wrap(err, "can't run image")
	}

	return nil
}

func (bs BinaryStrategy) localHolenPath() (string, error) {
	var holenPath string
	if configDataPath, err := bs.Get("holen.datapath"); err == nil && len(configDataPath) > 0 {
		holenPath = configDataPath
	} else if xdgDataHome := os.Getenv("XDG_DATA_HOME"); len(xdgDataHome) > 0 {
		holenPath = filepath.Join(xdgDataHome, "holen")
	} else {
		var home string
		if home = os.Getenv("HOME"); len(home) == 0 {
			return "", fmt.Errorf("$HOME environment variable not found")
		}
		holenPath = filepath.Join(home, ".local", "share", "holen")
	}
	os.MkdirAll(holenPath, 0755)

	return holenPath, nil
}

func (bs BinaryStrategy) DownloadPath() (string, error) {
	var downloadPath string
	if configDownloadPath, err := bs.Get("binary.download"); err == nil && len(configDownloadPath) > 0 {
		downloadPath = configDownloadPath
	} else {
		holenPath, err := bs.localHolenPath()
		if err != nil {
			return "", errors.Wrap(err, "unable to get holen data path")
		}
		downloadPath = filepath.Join(holenPath, "bin")
	}
	os.MkdirAll(downloadPath, 0755)

	return downloadPath, nil
}

func (bs BinaryStrategy) TempPath() (string, error) {
	var tempPath string

	holenPath, err := bs.localHolenPath()
	if err != nil {
		return "", errors.Wrap(err, "unable to get holen data path")
	}

	tempPath = filepath.Join(holenPath, "tmp")
	os.MkdirAll(tempPath, 0755)

	return tempPath, nil
}

func (bs BinaryStrategy) Run(args []string) error {
	temp := bs.Templater(bs.Data.Version, bs.Data.OSArchMap, bs.System)
	bs.Debugf("templater: %# v", pretty.Formatter(temp))

	dlURL, err := temp.Template(bs.Data.BaseURL)
	if err != nil {
		return errors.Wrap(err, "unable to template url")
	}

	downloadPath, err := bs.DownloadPath()
	if err != nil {
		return errors.Wrap(err, "unable to find download path")
	}
	localPath := filepath.Join(downloadPath, fmt.Sprintf("%s--%s", bs.Data.Name, bs.Data.Version))

	if !bs.FileExists(localPath) {
		if len(bs.Data.UnpackPath) > 0 {
			tempPath, err := bs.TempPath()
			tempdir, err := ioutil.TempDir(tempPath, "holen")
			if err != nil {
				return errors.Wrap(err, "unable to make temporary directory")
			}
			defer os.RemoveAll(tempdir)

			u, err := url.Parse(dlURL)
			if err != nil {
				return errors.Wrap(err, "unable to parse url")
			}

			fileName := path.Base(u.Path)
			archPath := filepath.Join(tempdir, fileName)
			unpackedPath := filepath.Join(tempdir, "unpacked")

			bs.UserMessage("Downloading %s...\n", dlURL)
			err = bs.DownloadFile(dlURL, archPath)
			if err != nil {
				return errors.Wrap(err, "can't download archive")
			}

			err = bs.UnpackArchive(archPath, unpackedPath)
			if err != nil {
				return errors.Wrap(err, "unable to unpack archive")
			}

			unpackPath, err := temp.Template(bs.Data.UnpackPath)
			if err != nil {
				return errors.Wrap(err, "unable to template unpack_path")
			}
			binPath := filepath.Join(unpackedPath, unpackPath)

			err = os.Rename(binPath, localPath)
			if err != nil {
				return errors.Wrap(err, "unable to move binary into position")
			}
		} else {
			bs.UserMessage("Downloading %s...\n", dlURL)
			err = bs.DownloadFile(dlURL, localPath)
			if err != nil {
				return errors.Wrap(err, "can't download binary")
			}
		}

		err = bs.MakeExecutable(localPath)
		if err != nil {
			return errors.Wrap(err, "unable to make binary executable")
		}
	}

	// TODO: checksum the binary

	err = bs.RunCommand(localPath, args)
	if err != nil {
		return errors.Wrap(err, "can't run binary")
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cases

import (
	"bytes"
	"fmt"
	"math/rand"

	"go.uber.org/zap"

	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/pd/v4/pkg/codec"
	"github.com/pingcap/pd/v4/server/core"
	"github.com/pingcap/pd/v4/tools/pd-simulator/simulator/info"
	"github.com/pingcap/pd/v4/tools/pd-simulator/simulator/simutil"
)

func newImportData() *Case {
	var simCase Case
	// Initialize the cluster
	for i := 1; i <= 10; i++ {
		simCase.Stores = append(simCase.Stores, &Store{
			ID:        IDAllocator.nextID(),
			Status:    metapb.StoreState_Up,
			Capacity:  1 * TB,
			Available: 900 * GB,
			Version:   "2.1.0",
		})
	}

	storeIDs := rand.Perm(3)
	for i := 0; i < 40; i++ {
		peers := []*metapb.Peer{
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[0] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[1] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[2] + 1)},
		}
		simCase.Regions = append(simCase.Regions, Region{
			ID:     IDAllocator.nextID(),
			Peers:  peers,
			Leader: peers[0],
			Size:   32 * MB,
			Keys:   320000,
		})
	}

	simCase.RegionSplitSize = 64 * MB
	simCase.RegionSplitKeys = 640000
	simCase.TableNumber = 10
	// Events description
	e := &WriteFlowOnSpotDescriptor{}
	table2 := string(codec.EncodeBytes(codec.GenerateTableKey(2)))
	table3 := string(codec.EncodeBytes(codec.GenerateTableKey(3)))
	table5 := string(codec.EncodeBytes(codec.GenerateTableKey(5)))
	e.Step = func(tick int64) map[string]int64 {
		if tick < 100 {
			return map[string]int64{
				table3: 4 * MB,
				table5: 32 * MB,
			}
		}
		return map[string]int64{
			table2: 2 * MB,
			table3: 4 * MB,
			table5: 16 * MB,
		}
	}
	simCase.Events = []EventDescriptor{e}

	// Checker description
	simCase.Checker = func(regions *core.RegionsInfo, stats []info.StoreStats) bool {
		leaderDist := make(map[uint64]int)
		peerDist := make(map[uint64]int)
		leaderTotal := 0
		peerTotal := 0
		res := make([]*core.RegionInfo, 0, 100)
		regions.ScanRangeWithIterator([]byte(table2), func(region *core.RegionInfo) bool {
			if bytes.Compare(region.GetEndKey(), []byte(table3)) < 0 {
				res = append(res, regions.GetRegion(region.GetID()))
				return true
			}
			return false
		})

		for _, r := range res {
			leaderTotal++
			leaderDist[r.GetLeader().GetStoreId()]++
			for _, p := range r.GetPeers() {
				peerDist[p.GetStoreId()]++
				peerTotal++
			}
		}
		if leaderTotal == 0 || peerTotal == 0 {
			return false
		}
		tableLeaderLog := fmt.Sprintf("%d leader:", leaderTotal)
		tablePeerLog := fmt.Sprintf("%d peer: ", peerTotal)
		for storeID := 1; storeID <= 10; storeID++ {
			if leaderCount, ok := leaderDist[uint64(storeID)]; ok {
				tableLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", tableLeaderLog, storeID, float64(leaderCount)/float64(leaderTotal)*100)
			}
		}
		for storeID := 1; storeID <= 10; storeID++ {
			if peerCount, ok := peerDist[uint64(storeID)]; ok {
				tablePeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", tablePeerLog, storeID, float64(peerCount)/float64(peerTotal)*100)
			}
		}
		regionTotal := regions.GetRegionCount()
		totalLeaderLog := fmt.Sprintf("%d leader:", regionTotal)
		totalPeerLog := fmt.Sprintf("%d peer:", regionTotal*3)
		isEnd := true
		for storeID := uint64(1); storeID <= 10; storeID++ {
			regions.GetStoreRegionCount(storeID)
			totalLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalLeaderLog, storeID, float64(regions.GetStoreLeaderCount(storeID))/float64(regionTotal)*100)
			regionProp := float64(regions.GetStoreRegionCount(storeID)) / float64(regionTotal*3) * 100
			if regionProp > 13.8 {
				isEnd = false
			}
			totalPeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalPeerLog, storeID, regionProp)
		}
		simutil.Logger.Info("import data information",
			zap.String("table-leader", tableLeaderLog),
			zap.String("table-peer", tablePeerLog),
			zap.String("total-leader", totalLeaderLog),
			zap.String("total-peer", totalPeerLog))
		return isEnd
	}
	return &simCase
}
<commit_msg>pd-simulator: update import-data case (#2741)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cases

import (
	"bytes"
	"fmt"
	"math/rand"

	"go.uber.org/zap"

	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/pd/v4/pkg/codec"
	"github.com/pingcap/pd/v4/server/core"
	"github.com/pingcap/pd/v4/tools/pd-simulator/simulator/info"
	"github.com/pingcap/pd/v4/tools/pd-simulator/simulator/simutil"
)

func newImportData() *Case {
	var simCase Case
	// Initialize the cluster
	for i := 1; i <= 10; i++ {
		simCase.Stores = append(simCase.Stores, &Store{
			ID:        IDAllocator.nextID(),
			Status:    metapb.StoreState_Up,
			Capacity:  1 * TB,
			Available: 900 * GB,
			Version:   "2.1.0",
		})
	}

	for i := 0; i < getRegionNum(); i++ {
		storeIDs := rand.Perm(10)
		peers := []*metapb.Peer{
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[0] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[1] + 1)},
			{Id: IDAllocator.nextID(), StoreId: uint64(storeIDs[2] + 1)},
		}
		simCase.Regions = append(simCase.Regions, Region{
			ID:     IDAllocator.nextID(),
			Peers:  peers,
			Leader: peers[0],
			Size:   32 * MB,
			Keys:   320000,
		})
	}

	simCase.RegionSplitSize = 64 * MB
	simCase.RegionSplitKeys = 640000
	simCase.TableNumber = 10
	// Events description
	e := &WriteFlowOnSpotDescriptor{}
	table2 := string(codec.EncodeBytes(codec.GenerateTableKey(2)))
	table3 := string(codec.EncodeBytes(codec.GenerateTableKey(3)))
	e.Step = func(tick int64) map[string]int64 {
		if tick > int64(getRegionNum())/10 {
			return nil
		}
		return map[string]int64{
			table2: 32 * MB,
		}
	}
	simCase.Events = []EventDescriptor{e}

	// Checker description
	checkCount := uint64(0)
	simCase.Checker = func(regions *core.RegionsInfo, stats []info.StoreStats) bool {
		leaderDist := make(map[uint64]int)
		peerDist := make(map[uint64]int)
		leaderTotal := 0
		peerTotal := 0
		res := make([]*core.RegionInfo, 0, 100)
		regions.ScanRangeWithIterator([]byte(table2), func(region *core.RegionInfo) bool {
			if bytes.Compare(region.GetEndKey(), []byte(table3)) < 0 {
				res = append(res, regions.GetRegion(region.GetID()))
				return true
			}
			return false
		})

		for _, r := range res {
			leaderTotal++
			leaderDist[r.GetLeader().GetStoreId()]++
			for _, p := range r.GetPeers() {
				peerDist[p.GetStoreId()]++
				peerTotal++
			}
		}
		if leaderTotal == 0 || peerTotal == 0 {
			return false
		}
		tableLeaderLog := fmt.Sprintf("%d leader:", leaderTotal)
		tablePeerLog := fmt.Sprintf("%d peer: ", peerTotal)
		for storeID := 1; storeID <= 10; storeID++ {
			if leaderCount, ok := leaderDist[uint64(storeID)]; ok {
				tableLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", tableLeaderLog, storeID, float64(leaderCount)/float64(leaderTotal)*100)
			}
		}
		for storeID := 1; storeID <= 10; storeID++ {
			if peerCount, ok := peerDist[uint64(storeID)]; ok {
				tablePeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", tablePeerLog, storeID, float64(peerCount)/float64(peerTotal)*100)
			}
		}
		regionTotal := regions.GetRegionCount()
		totalLeaderLog := fmt.Sprintf("%d leader:", regionTotal)
		totalPeerLog := fmt.Sprintf("%d peer:", regionTotal*3)
		isEnd := false
		var regionProps []float64
		for storeID := uint64(1); storeID <= 10; storeID++ {
			regions.GetStoreRegionCount(storeID)
			totalLeaderLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalLeaderLog, storeID, float64(regions.GetStoreLeaderCount(storeID))/float64(regionTotal)*100)
			regionProp := float64(regions.GetStoreRegionCount(storeID)) / float64(regionTotal*3) * 100
			regionProps = append(regionProps, regionProp)
			totalPeerLog = fmt.Sprintf("%s [store %d]:%.2f%%", totalPeerLog, storeID, regionProp)
		}
		simutil.Logger.Info("import data information",
			zap.String("table-leader", tableLeaderLog),
			zap.String("table-peer", tablePeerLog),
			zap.String("total-leader", totalLeaderLog),
			zap.String("total-peer", totalPeerLog))
		checkCount++
		dev := 0.0
		for _, p := range regionProps {
			dev += (p - 10) * (p - 10) / 100
		}
		if dev > 0.005 {
			simutil.Logger.Warn("Not balanced, change scheduler or store limit", zap.Float64("dev score", dev))
		}
		if checkCount > uint64(getRegionNum())/10 {
			isEnd = dev < 0.002
		}
		return isEnd
	}
	return &simCase
}
<|endoftext|>


<commit_before>package natty

import (
	"fmt"
	"log"
	"net"
	"sync"
	"testing"
	"time"
)

const (
	MESSAGE_TEXT = "Hello World"
)

// TestLocal starts up two local Natty instances that communicate with each
// other directly.  Once connected, one Natty sends a UDP packet to the other
// to make sure that the connection works.
func TestLocal(t *testing.T) {
	var offerer *Natty
	var answerer *Natty

	offerer = NewNatty(
		func(msg []byte) {
			// This would be done using a signaling server when talking to a
			// remote Natty
			answerer.Receive(msg)
		},
		nil)

	answerer = NewNatty(
		func(msg []byte) {
			// This would be done using a signaling server when talking to a
			// remote Natty
			offerer.Receive(msg)
		},
		nil)

	var wg sync.WaitGroup
	wg.Add(2)

	// Offerer processing
	go func() {
		defer wg.Done()
		fiveTuple, err := offerer.Offer()
		if err != nil {
			t.Errorf("Offerer had error: %s", err)
			return
		}
		log.Printf("Offerer got 5 tuple: %s", fiveTuple)
		if fiveTuple.Proto != UDP {
			t.Errorf("Protocol was %s instead of udp", fiveTuple.Proto)
			return
		}
		local, remote, err := udpAddresses(fiveTuple)
		if err != nil {
			t.Error("Offerer unable to resolve UDP addresses: %s", err)
			return
		}
		conn, err := net.DialUDP("udp", local, remote)
		if err != nil {
			t.Errorf("Unable to dial UDP: %s", err)
			return
		}
		for i := 0; i < 10; i++ {
			_, err := conn.Write([]byte(MESSAGE_TEXT))
			if err != nil {
				t.Errorf("Offerer unable to write to UDP: %s", err)
				return
			}
		}
	}()

	// Answerer processing
	go func() {
		defer wg.Done()
		fiveTuple, err := answerer.Answer()
		if err != nil {
			t.Errorf("Answerer had error: %s", err)
			return
		}
		if fiveTuple.Proto != UDP {
			t.Errorf("Protocol was %s instead of udp", fiveTuple.Proto)
			return
		}
		log.Printf("Answerer got 5 tuple: %s", fiveTuple)
		local, _, err := udpAddresses(fiveTuple)
		if err != nil {
			t.Errorf("Error in Answerer: %s", err)
			return
		}
		conn, err := net.ListenUDP("udp", local)
		if err != nil {
			t.Errorf("Answerer unable to listen on UDP: %s", err)
			return
		}
		b := make([]byte, 1024)
		for {
			n, addr, err := conn.ReadFrom(b)
			if err != nil {
				t.Errorf("Answerer unable to read from UDP: %s", err)
				return
			}
			if addr.String() != fiveTuple.Remote {
				t.Errorf("UDP package had address %s, expected %s", addr, fiveTuple.Remote)
				return
			}
			msg := string(b[:n])
			if msg != MESSAGE_TEXT {
				log.Printf("Got message '%s', expected '%s'", msg, MESSAGE_TEXT)
			}
			return
		}
	}()

	doneCh := make(chan interface{})
	go func() {
		wg.Wait()
		doneCh <- nil
	}()

	select {
	case <-doneCh:
		return
	case <-time.After(10 * time.Second):
		t.Errorf("Test timed out")
	}
}

func udpAddresses(fiveTuple *FiveTuple) (*net.UDPAddr, *net.UDPAddr, error) {
	if fiveTuple.Proto != UDP {
		return nil, nil, fmt.Errorf("FiveTuple.Proto was not UDP!: %s", fiveTuple.Proto)
	}
	local, err := net.ResolveUDPAddr("udp", fiveTuple.Local)
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to resolve local UDP address %s: %s", fiveTuple.Local)
	}
	remote, err := net.ResolveUDPAddr("udp", fiveTuple.Remote)
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to resolve remote UDP address %s: %s", fiveTuple.Remote)
	}
	return local, remote, nil
}
<commit_msg>In test, added wait condition for offerer to make sure it doesn't try to connect before answerer is listening<commit_after>package natty

import (
	"fmt"
	"log"
	"net"
	"sync"
	"testing"
	"time"
)

const (
	MESSAGE_TEXT = "Hello World"
)

// TestLocal starts up two local Natty instances that communicate with each
// other directly.  Once connected, one Natty sends a UDP packet to the other
// to make sure that the connection works.
func TestLocal(t *testing.T) {
	var offerer *Natty
	var answerer *Natty

	offerer = NewNatty(
		func(msg []byte) {
			// This would be done using a signaling server when talking to a
			// remote Natty
			answerer.Receive(msg)
		},
		nil)

	answerer = NewNatty(
		func(msg []byte) {
			// This would be done using a signaling server when talking to a
			// remote Natty
			offerer.Receive(msg)
		},
		nil)

	var answererReady sync.WaitGroup
	answererReady.Add(1)

	var wg sync.WaitGroup
	wg.Add(2)

	// Offerer processing
	go func() {
		defer wg.Done()
		fiveTuple, err := offerer.Offer()
		if err != nil {
			t.Errorf("Offerer had error: %s", err)
			return
		}
		log.Printf("Offerer got 5 tuple: %s", fiveTuple)
		if fiveTuple.Proto != UDP {
			t.Errorf("Protocol was %s instead of udp", fiveTuple.Proto)
			return
		}
		local, remote, err := udpAddresses(fiveTuple)
		if err != nil {
			t.Error("Offerer unable to resolve UDP addresses: %s", err)
			return
		}
		answererReady.Wait()
		conn, err := net.DialUDP("udp", local, remote)
		if err != nil {
			t.Errorf("Unable to dial UDP: %s", err)
			return
		}
		for i := 0; i < 10; i++ {
			_, err := conn.Write([]byte(MESSAGE_TEXT))
			if err != nil {
				t.Errorf("Offerer unable to write to UDP: %s", err)
				return
			}
		}
	}()

	// Answerer processing
	go func() {
		defer wg.Done()
		fiveTuple, err := answerer.Answer()
		if err != nil {
			t.Errorf("Answerer had error: %s", err)
			return
		}
		if fiveTuple.Proto != UDP {
			t.Errorf("Protocol was %s instead of udp", fiveTuple.Proto)
			return
		}
		log.Printf("Answerer got 5 tuple: %s", fiveTuple)
		local, _, err := udpAddresses(fiveTuple)
		if err != nil {
			t.Errorf("Error in Answerer: %s", err)
			return
		}
		conn, err := net.ListenUDP("udp", local)
		if err != nil {
			t.Errorf("Answerer unable to listen on UDP: %s", err)
			return
		}
		answererReady.Done()
		b := make([]byte, 1024)
		for {
			n, addr, err := conn.ReadFrom(b)
			if err != nil {
				t.Errorf("Answerer unable to read from UDP: %s", err)
				return
			}
			if addr.String() != fiveTuple.Remote {
				t.Errorf("UDP package had address %s, expected %s", addr, fiveTuple.Remote)
				return
			}
			msg := string(b[:n])
			if msg != MESSAGE_TEXT {
				log.Printf("Got message '%s', expected '%s'", msg, MESSAGE_TEXT)
			}
			return
		}
	}()

	doneCh := make(chan interface{})
	go func() {
		wg.Wait()
		doneCh <- nil
	}()

	select {
	case <-doneCh:
		return
	case <-time.After(10 * time.Second):
		t.Errorf("Test timed out")
	}
}

func udpAddresses(fiveTuple *FiveTuple) (*net.UDPAddr, *net.UDPAddr, error) {
	if fiveTuple.Proto != UDP {
		return nil, nil, fmt.Errorf("FiveTuple.Proto was not UDP!: %s", fiveTuple.Proto)
	}
	local, err := net.ResolveUDPAddr("udp", fiveTuple.Local)
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to resolve local UDP address %s: %s", fiveTuple.Local)
	}
	remote, err := net.ResolveUDPAddr("udp", fiveTuple.Remote)
	if err != nil {
		return nil, nil, fmt.Errorf("Unable to resolve remote UDP address %s: %s", fiveTuple.Remote)
	}
	return local, remote, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"syscall"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
)

var logger *log.Logger

// flags
var awsKey string
var awsSecret string
var s3Region string
var start int
var end int
var width int
var bucket string
var prefix string

func main() {
	if bucket == "" {
		logger.Println("Bucket not specified")
		syscall.Exit(1)
	}

	auth := new(aws.Auth)
	auth.AccessKey = awsKey
	auth.SecretKey = awsSecret
	s3c := s3.New(*auth, aws.Regions[s3Region])
	s3bucket := s3c.Bucket(bucket)
	logger.Println(s3bucket)

	// making i with leading zeros with this format
	format := "%0" + strconv.Itoa(width) + "d"
	for i := start; i <= end; i += 1 {
		suffix := fmt.Sprintf(format, i)
		key := prefix + suffix
		logger.Printf("Deleting S3 key: %s", key)
		err := s3bucket.Del(key)
		if err != nil {
			logger.Printf("Got error deleting key: %s", err)
		}
	}
}

func init() {
	logger = log.New(os.Stdout, "", 0)

	flag.StringVar(&awsKey, "awsKey", os.Getenv("AWS_ACCESS_KEY_ID"), "AWS Key. Defaults to env var AWS_ACCESS_KEY_ID")
	flag.StringVar(&awsSecret, "awsSecret", os.Getenv("AWS_SECRET_KEY"), "AWS Secret. Defaults to env var AWS_SECRET_KEY")
	flag.StringVar(&s3Region, "s3Region", "us-east-1", "AWS S3 region")
	flag.IntVar(&start, "start", 0, "Starting number")
	flag.IntVar(&end, "end", 0, "Ending number (inclusive)")
	flag.IntVar(&width, "width", 6, "Key number width (ex. when width = 6, 1 = 000001)")
	flag.StringVar(&bucket, "bucket", "", "Bucket")
	flag.StringVar(&prefix, "prefix", "/", "Key prefix")
	flag.Parse()
}
<commit_msg>better logging. strconv package removed for fmt<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"syscall"

	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
)

var logger *log.Logger

// flags
var awsKey string
var awsSecret string
var s3Region string
var start int
var end int
var width int
var bucket string
var prefix string

func main() {
	if bucket == "" {
		logger.Println("Bucket not specified")
		syscall.Exit(1)
	}

	auth := new(aws.Auth)
	auth.AccessKey = awsKey
	auth.SecretKey = awsSecret
	s3c := s3.New(*auth, aws.Regions[s3Region])
	s3bucket := s3c.Bucket(bucket)

	// making i with leading zeros with this format
	format := fmt.Sprintf("%%0%dd", width)
	for i := start; i <= end; i += 1 {
		suffix := fmt.Sprintf(format, i)
		key := prefix + suffix
		logger.Printf("Deleting S3 key: %s%s", bucket, key)
		err := s3bucket.Del(key)
		if err != nil {
			logger.Printf("Got error deleting key: %s", err)
		}
	}
}

func init() {
	logger = log.New(os.Stdout, "", 0)

	flag.StringVar(&awsKey, "awsKey", os.Getenv("AWS_ACCESS_KEY_ID"), "AWS Key. Defaults to env var AWS_ACCESS_KEY_ID")
	flag.StringVar(&awsSecret, "awsSecret", os.Getenv("AWS_SECRET_KEY"), "AWS Secret. Defaults to env var AWS_SECRET_KEY")
	flag.StringVar(&s3Region, "s3Region", "us-east-1", "AWS S3 region")
	flag.IntVar(&start, "start", 0, "Starting number")
	flag.IntVar(&end, "end", 0, "Ending number (inclusive)")
	flag.IntVar(&width, "width", 6, "Key number width (ex. when width = 6, 1 = 000001)")
	flag.StringVar(&bucket, "bucket", "", "Bucket")
	flag.StringVar(&prefix, "prefix", "/", "Key prefix")
	flag.Parse()
}
<|endoftext|>


<commit_before>package s3resource

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"net/http"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/cheggaaa/pb"
)

//go:generate counterfeiter . S3Client

type S3Client interface {
	BucketFiles(bucketName string, prefixHint string) ([]string, error)
	BucketFileVersions(bucketName string, remotePath string) ([]string, error)

	UploadFile(bucketName string, remotePath string, localPath string, options UploadFileOptions) (string, error)
	DownloadFile(bucketName string, remotePath string, versionID string, localPath string) error

	DeleteFile(bucketName string, remotePath string) error
	DeleteVersionedFile(bucketName string, remotePath string, versionID string) error

	URL(bucketName string, remotePath string, private bool, versionID string) string
}

// 12 retries works out to ~5 mins of total backoff time, though AWS randomizes
// the backoff to some extent so it may be as low as 4 or as high as 8 minutes
const maxRetries = 12

type s3client struct {
	client  *s3.S3
	session *session.Session

	progressOutput io.Writer
}

type UploadFileOptions struct {
	Acl                  string
	ServerSideEncryption string
	KmsKeyId             string
	ContentType          string
}

func NewUploadFileOptions() UploadFileOptions {
	return UploadFileOptions{
		Acl: "private",
	}
}

func NewS3Client(
	progressOutput io.Writer,
	awsConfig *aws.Config,
	useV2Signing bool,
) S3Client {
	sess := session.New(awsConfig)
	client := s3.New(sess, awsConfig)

	if useV2Signing {
		setv2Handlers(client)
	}

	return &s3client{
		client:  client,
		session: sess,

		progressOutput: progressOutput,
	}
}

func NewAwsConfig(
	accessKey string,
	secretKey string,
	sessionToken string,
	regionName string,
	endpoint string,
	disableSSL bool,
	skipSSLVerification bool,
) *aws.Config {
	var creds *credentials.Credentials

	if accessKey == "" && secretKey == "" {
		creds = credentials.AnonymousCredentials
	} else {
		creds = credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)
	}

	if len(regionName) == 0 {
		regionName = "us-east-1"
	}

	var httpClient *http.Client
	if skipSSLVerification {
		httpClient = &http.Client{Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}}
	} else {
		httpClient = http.DefaultClient
	}

	awsConfig := &aws.Config{
		Region:           aws.String(regionName),
		Credentials:      creds,
		S3ForcePathStyle: aws.Bool(true),
		MaxRetries:       aws.Int(maxRetries),
		DisableSSL:       aws.Bool(disableSSL),
		HTTPClient:       httpClient,
	}

	if len(endpoint) != 0 {
		endpoint := fmt.Sprintf("%s", endpoint)
		awsConfig.Endpoint = &endpoint
	}

	return awsConfig
}

func (client *s3client) BucketFiles(bucketName string, prefixHint string) ([]string, error) {
	entries, err := client.getBucketContents(bucketName, prefixHint)

	if err != nil {
		return []string{}, err
	}

	paths := make([]string, 0, len(entries))

	for _, entry := range entries {
		paths = append(paths, *entry.Key)
	}
	return paths, nil
}

func (client *s3client) BucketFileVersions(bucketName string, remotePath string) ([]string, error) {
	isBucketVersioned, err := client.getBucketVersioning(bucketName)
	if err != nil {
		return []string{}, err
	}

	if !isBucketVersioned {
		return []string{}, errors.New("bucket is not versioned")
	}

	bucketFiles, err := client.getVersionedBucketContents(bucketName, remotePath)

	if err != nil {
		return []string{}, err
	}

	versions := make([]string, 0, len(bucketFiles))

	for _, objectVersion := range bucketFiles[remotePath] {
		versions = append(versions, *objectVersion.VersionId)
	}

	return versions, nil
}

func (client *s3client) UploadFile(bucketName string, remotePath string, localPath string, options UploadFileOptions) (string, error) {
	uploader := s3manager.NewUploaderWithClient(client.client)

	if client.isGCSHost() {
		// GCS returns `InvalidArgument` on multipart uploads
		uploader.MaxUploadParts = 1
	}

	stat, err := os.Stat(localPath)
	if err != nil {
		return "", err
	}

	localFile, err := os.Open(localPath)
	if err != nil {
		return "", err
	}

	defer localFile.Close()

	progress := client.newProgressBar(stat.Size())

	progress.Start()
	defer progress.Finish()

	uploadInput := s3manager.UploadInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
		Body:   progressReader{localFile, progress},
		ACL:    aws.String(options.Acl),
	}
	if options.ServerSideEncryption != "" {
		uploadInput.ServerSideEncryption = aws.String(options.ServerSideEncryption)
	}
	if options.KmsKeyId != "" {
		uploadInput.SSEKMSKeyId = aws.String(options.KmsKeyId)
	}
	if options.ContentType != "" {
		uploadInput.ContentType = aws.String(options.ContentType)
	}

	uploadOutput, err := uploader.Upload(&uploadInput)
	if err != nil {
		return "", err
	}

	if uploadOutput.VersionID != nil {
		return *uploadOutput.VersionID, nil
	}

	return "", nil
}

func (client *s3client) DownloadFile(bucketName string, remotePath string, versionID string, localPath string) error {
	headObject := &s3.HeadObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		headObject.VersionId = aws.String(versionID)
	}

	object, err := client.client.HeadObject(headObject)
	if err != nil {
		return err
	}

	progress := client.newProgressBar(*object.ContentLength)

	downloader := s3manager.NewDownloaderWithClient(client.client)

	localFile, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer localFile.Close()

	getObject := &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		getObject.VersionId = aws.String(versionID)
	}

	progress.Start()
	defer progress.Finish()

	_, err = downloader.Download(progressWriterAt{localFile, progress}, getObject)
	if err != nil {
		return err
	}

	return nil
}

func (client *s3client) URL(bucketName string, remotePath string, private bool, versionID string) string {
	getObjectInput := &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		getObjectInput.VersionId = aws.String(versionID)
	}

	awsRequest, _ := client.client.GetObjectRequest(getObjectInput)

	var url string

	if private {
		url, _ = awsRequest.Presign(24 * time.Hour)
	} else {
		awsRequest.Build()
		url = awsRequest.HTTPRequest.URL.String()
	}

	return url
}

func (client *s3client) DeleteVersionedFile(bucketName string, remotePath string, versionID string) error {
	_, err := client.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket:    aws.String(bucketName),
		Key:       aws.String(remotePath),
		VersionId: aws.String(versionID),
	})

	return err
}

func (client *s3client) DeleteFile(bucketName string, remotePath string) error {
	_, err := client.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	})

	return err
}

func (client *s3client) getBucketContents(bucketName string, prefix string) (map[string]*s3.Object, error) {
	bucketContents := map[string]*s3.Object{}
	marker := ""

	for {
		listObjectsResponse, err := client.client.ListObjects(&s3.ListObjectsInput{
			Bucket: aws.String(bucketName),
			Prefix: aws.String(prefix),
			Marker: aws.String(marker),
		})

		if err != nil {
			return bucketContents, err
		}

		lastKey := ""

		for _, key := range listObjectsResponse.Contents {
			bucketContents[*key.Key] = key

			lastKey = *key.Key
		}

		if *listObjectsResponse.IsTruncated {
			prevMarker := marker
			if listObjectsResponse.NextMarker == nil {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				marker = lastKey
			} else {
				marker = *listObjectsResponse.NextMarker
			}
			if marker == prevMarker {
				return nil, errors.New("Unable to list all bucket objects; perhaps this is a CloudFront S3 bucket that needs its `Query String Forwarding and Caching` set to `Forward all, cache based on all`?")
			}
		} else {
			break
		}
	}

	return bucketContents, nil
}

func (client *s3client) getBucketVersioning(bucketName string) (bool, error) {
	params := &s3.GetBucketVersioningInput{
		Bucket: aws.String(bucketName),
	}

	resp, err := client.client.GetBucketVersioning(params)
	if err != nil {
		return false, err
	}

	if resp.Status == nil {
		return false, nil
	}

	return *resp.Status == "Enabled", nil
}

func (client *s3client) getVersionedBucketContents(bucketName string, prefix string) (map[string][]*s3.ObjectVersion, error) {
	versionedBucketContents := map[string][]*s3.ObjectVersion{}
	keyMarker := ""
	versionMarker := ""
	for {

		params := &s3.ListObjectVersionsInput{
			Bucket:    aws.String(bucketName),
			KeyMarker: aws.String(keyMarker),
			Prefix:    aws.String(prefix),
		}

		if versionMarker != "" {
			params.VersionIdMarker = aws.String(versionMarker)
		}

		listObjectVersionsResponse, err := client.client.ListObjectVersions(params)
		if err != nil {
			return versionedBucketContents, err
		}

		lastKey := ""
		lastVersionKey := ""

		for _, objectVersion := range listObjectVersionsResponse.Versions {
			versionedBucketContents[*objectVersion.Key] = append(versionedBucketContents[*objectVersion.Key], objectVersion)

			lastKey = *objectVersion.Key
			lastVersionKey = *objectVersion.VersionId
		}

		if *listObjectVersionsResponse.IsTruncated {
			keyMarker = *listObjectVersionsResponse.NextKeyMarker
			versionMarker = *listObjectVersionsResponse.NextVersionIdMarker
			if keyMarker == "" {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				keyMarker = lastKey
			}

			if versionMarker == "" {
				versionMarker = lastVersionKey
			}
		} else {
			break
		}

	}

	return versionedBucketContents, nil
}

func (client *s3client) newProgressBar(total int64) *pb.ProgressBar {
	progress := pb.New64(total)

	progress.Output = client.progressOutput
	progress.ShowSpeed = true
	progress.Units = pb.U_BYTES
	progress.NotPrint = true

	return progress.SetWidth(80)
}

func (client *s3client) isGCSHost() bool {
	return (client.session.Config.Endpoint != nil && strings.Contains(*client.session.Config.Endpoint, "storage.googleapis.com"))
}
<commit_msg>Automatically adjust part size on upload<commit_after>package s3resource

import (
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"net/http"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/cheggaaa/pb"
)

//go:generate counterfeiter . S3Client

type S3Client interface {
	BucketFiles(bucketName string, prefixHint string) ([]string, error)
	BucketFileVersions(bucketName string, remotePath string) ([]string, error)

	UploadFile(bucketName string, remotePath string, localPath string, options UploadFileOptions) (string, error)
	DownloadFile(bucketName string, remotePath string, versionID string, localPath string) error

	DeleteFile(bucketName string, remotePath string) error
	DeleteVersionedFile(bucketName string, remotePath string, versionID string) error

	URL(bucketName string, remotePath string, private bool, versionID string) string
}

// 12 retries works out to ~5 mins of total backoff time, though AWS randomizes
// the backoff to some extent so it may be as low as 4 or as high as 8 minutes
const maxRetries = 12

type s3client struct {
	client  *s3.S3
	session *session.Session

	progressOutput io.Writer
}

type UploadFileOptions struct {
	Acl                  string
	ServerSideEncryption string
	KmsKeyId             string
	ContentType          string
}

func NewUploadFileOptions() UploadFileOptions {
	return UploadFileOptions{
		Acl: "private",
	}
}

func NewS3Client(
	progressOutput io.Writer,
	awsConfig *aws.Config,
	useV2Signing bool,
) S3Client {
	sess := session.New(awsConfig)
	client := s3.New(sess, awsConfig)

	if useV2Signing {
		setv2Handlers(client)
	}

	return &s3client{
		client:  client,
		session: sess,

		progressOutput: progressOutput,
	}
}

func NewAwsConfig(
	accessKey string,
	secretKey string,
	sessionToken string,
	regionName string,
	endpoint string,
	disableSSL bool,
	skipSSLVerification bool,
) *aws.Config {
	var creds *credentials.Credentials

	if accessKey == "" && secretKey == "" {
		creds = credentials.AnonymousCredentials
	} else {
		creds = credentials.NewStaticCredentials(accessKey, secretKey, sessionToken)
	}

	if len(regionName) == 0 {
		regionName = "us-east-1"
	}

	var httpClient *http.Client
	if skipSSLVerification {
		httpClient = &http.Client{Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}}
	} else {
		httpClient = http.DefaultClient
	}

	awsConfig := &aws.Config{
		Region:           aws.String(regionName),
		Credentials:      creds,
		S3ForcePathStyle: aws.Bool(true),
		MaxRetries:       aws.Int(maxRetries),
		DisableSSL:       aws.Bool(disableSSL),
		HTTPClient:       httpClient,
	}

	if len(endpoint) != 0 {
		endpoint := fmt.Sprintf("%s", endpoint)
		awsConfig.Endpoint = &endpoint
	}

	return awsConfig
}

func (client *s3client) BucketFiles(bucketName string, prefixHint string) ([]string, error) {
	entries, err := client.getBucketContents(bucketName, prefixHint)

	if err != nil {
		return []string{}, err
	}

	paths := make([]string, 0, len(entries))

	for _, entry := range entries {
		paths = append(paths, *entry.Key)
	}
	return paths, nil
}

func (client *s3client) BucketFileVersions(bucketName string, remotePath string) ([]string, error) {
	isBucketVersioned, err := client.getBucketVersioning(bucketName)
	if err != nil {
		return []string{}, err
	}

	if !isBucketVersioned {
		return []string{}, errors.New("bucket is not versioned")
	}

	bucketFiles, err := client.getVersionedBucketContents(bucketName, remotePath)

	if err != nil {
		return []string{}, err
	}

	versions := make([]string, 0, len(bucketFiles))

	for _, objectVersion := range bucketFiles[remotePath] {
		versions = append(versions, *objectVersion.VersionId)
	}

	return versions, nil
}

func (client *s3client) UploadFile(bucketName string, remotePath string, localPath string, options UploadFileOptions) (string, error) {
	uploader := s3manager.NewUploaderWithClient(client.client)

	if client.isGCSHost() {
		// GCS returns `InvalidArgument` on multipart uploads
		uploader.MaxUploadParts = 1
	}

	stat, err := os.Stat(localPath)
	if err != nil {
		return "", err
	}

	localFile, err := os.Open(localPath)
	if err != nil {
		return "", err
	}

	defer localFile.Close()

	// Automatically adjust partsize for larger files.
	fSize := stat.Size()
	if fSize > int64(uploader.MaxUploadParts) * uploader.PartSize {
		partSize := fSize / int64(uploader.MaxUploadParts)
		if fSize % int64(uploader.MaxUploadParts) != 0 {
			partSize++
		}
		uploader.PartSize = partSize
	}

	progress := client.newProgressBar(fSize)

	progress.Start()
	defer progress.Finish()

	uploadInput := s3manager.UploadInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
		Body:   progressReader{localFile, progress},
		ACL:    aws.String(options.Acl),
	}
	if options.ServerSideEncryption != "" {
		uploadInput.ServerSideEncryption = aws.String(options.ServerSideEncryption)
	}
	if options.KmsKeyId != "" {
		uploadInput.SSEKMSKeyId = aws.String(options.KmsKeyId)
	}
	if options.ContentType != "" {
		uploadInput.ContentType = aws.String(options.ContentType)
	}

	uploadOutput, err := uploader.Upload(&uploadInput)
	if err != nil {
		return "", err
	}

	if uploadOutput.VersionID != nil {
		return *uploadOutput.VersionID, nil
	}

	return "", nil
}

func (client *s3client) DownloadFile(bucketName string, remotePath string, versionID string, localPath string) error {
	headObject := &s3.HeadObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		headObject.VersionId = aws.String(versionID)
	}

	object, err := client.client.HeadObject(headObject)
	if err != nil {
		return err
	}

	progress := client.newProgressBar(*object.ContentLength)

	downloader := s3manager.NewDownloaderWithClient(client.client)

	localFile, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer localFile.Close()

	getObject := &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		getObject.VersionId = aws.String(versionID)
	}

	progress.Start()
	defer progress.Finish()

	_, err = downloader.Download(progressWriterAt{localFile, progress}, getObject)
	if err != nil {
		return err
	}

	return nil
}

func (client *s3client) URL(bucketName string, remotePath string, private bool, versionID string) string {
	getObjectInput := &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	}

	if versionID != "" {
		getObjectInput.VersionId = aws.String(versionID)
	}

	awsRequest, _ := client.client.GetObjectRequest(getObjectInput)

	var url string

	if private {
		url, _ = awsRequest.Presign(24 * time.Hour)
	} else {
		awsRequest.Build()
		url = awsRequest.HTTPRequest.URL.String()
	}

	return url
}

func (client *s3client) DeleteVersionedFile(bucketName string, remotePath string, versionID string) error {
	_, err := client.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket:    aws.String(bucketName),
		Key:       aws.String(remotePath),
		VersionId: aws.String(versionID),
	})

	return err
}

func (client *s3client) DeleteFile(bucketName string, remotePath string) error {
	_, err := client.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(remotePath),
	})

	return err
}

func (client *s3client) getBucketContents(bucketName string, prefix string) (map[string]*s3.Object, error) {
	bucketContents := map[string]*s3.Object{}
	marker := ""

	for {
		listObjectsResponse, err := client.client.ListObjects(&s3.ListObjectsInput{
			Bucket: aws.String(bucketName),
			Prefix: aws.String(prefix),
			Marker: aws.String(marker),
		})

		if err != nil {
			return bucketContents, err
		}

		lastKey := ""

		for _, key := range listObjectsResponse.Contents {
			bucketContents[*key.Key] = key

			lastKey = *key.Key
		}

		if *listObjectsResponse.IsTruncated {
			prevMarker := marker
			if listObjectsResponse.NextMarker == nil {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				marker = lastKey
			} else {
				marker = *listObjectsResponse.NextMarker
			}
			if marker == prevMarker {
				return nil, errors.New("Unable to list all bucket objects; perhaps this is a CloudFront S3 bucket that needs its `Query String Forwarding and Caching` set to `Forward all, cache based on all`?")
			}
		} else {
			break
		}
	}

	return bucketContents, nil
}

func (client *s3client) getBucketVersioning(bucketName string) (bool, error) {
	params := &s3.GetBucketVersioningInput{
		Bucket: aws.String(bucketName),
	}

	resp, err := client.client.GetBucketVersioning(params)
	if err != nil {
		return false, err
	}

	if resp.Status == nil {
		return false, nil
	}

	return *resp.Status == "Enabled", nil
}

func (client *s3client) getVersionedBucketContents(bucketName string, prefix string) (map[string][]*s3.ObjectVersion, error) {
	versionedBucketContents := map[string][]*s3.ObjectVersion{}
	keyMarker := ""
	versionMarker := ""
	for {

		params := &s3.ListObjectVersionsInput{
			Bucket:    aws.String(bucketName),
			KeyMarker: aws.String(keyMarker),
			Prefix:    aws.String(prefix),
		}

		if versionMarker != "" {
			params.VersionIdMarker = aws.String(versionMarker)
		}

		listObjectVersionsResponse, err := client.client.ListObjectVersions(params)
		if err != nil {
			return versionedBucketContents, err
		}

		lastKey := ""
		lastVersionKey := ""

		for _, objectVersion := range listObjectVersionsResponse.Versions {
			versionedBucketContents[*objectVersion.Key] = append(versionedBucketContents[*objectVersion.Key], objectVersion)

			lastKey = *objectVersion.Key
			lastVersionKey = *objectVersion.VersionId
		}

		if *listObjectVersionsResponse.IsTruncated {
			keyMarker = *listObjectVersionsResponse.NextKeyMarker
			versionMarker = *listObjectVersionsResponse.NextVersionIdMarker
			if keyMarker == "" {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				keyMarker = lastKey
			}

			if versionMarker == "" {
				versionMarker = lastVersionKey
			}
		} else {
			break
		}

	}

	return versionedBucketContents, nil
}

func (client *s3client) newProgressBar(total int64) *pb.ProgressBar {
	progress := pb.New64(total)

	progress.Output = client.progressOutput
	progress.ShowSpeed = true
	progress.Units = pb.U_BYTES
	progress.NotPrint = true

	return progress.SetWidth(80)
}

func (client *s3client) isGCSHost() bool {
	return (client.session.Config.Endpoint != nil && strings.Contains(*client.session.Config.Endpoint, "storage.googleapis.com"))
}
<|endoftext|>


<commit_before>package rpmpack

import (
	"compress/gzip"
	"io"
	"io/ioutil"
	"reflect"
	"testing"

	"github.com/klauspost/compress/zstd"
	"github.com/ulikunitz/xz"
	"github.com/ulikunitz/xz/lzma"
)

func TestFileOwner(t *testing.T) {
	r, err := NewRPM(RPMMetaData{})
	if err != nil {
		t.Fatalf("NewRPM returned error %v", err)
	}
	group := "testGroup"
	user := "testUser"

	r.AddFile(RPMFile{
		Name:  "/usr/local/hello",
		Body:  []byte("content of the file"),
		Group: group,
		Owner: user,
	})

	if err := r.Write(ioutil.Discard); err != nil {
		t.Errorf("NewRPM returned error %v", err)
	}
	if r.fileowners[0] != user {
		t.Errorf("File owner shoud be %s but is %s", user, r.fileowners[0])
	}
	if r.filegroups[0] != group {
		t.Errorf("File owner shoud be %s but is %s", group, r.filegroups[0])
	}
}

// https://github.com/google/rpmpack/issues/49
func Test100644(t *testing.T) {
	r, err := NewRPM(RPMMetaData{})
	if err != nil {
		t.Fatalf("NewRPM returned error %v", err)
	}
	r.AddFile(RPMFile{
		Name: "/usr/local/hello",
		Body: []byte("content of the file"),
		Mode: 0100644,
	})

	if err := r.Write(ioutil.Discard); err != nil {
		t.Errorf("Write returned error %v", err)
	}
	if r.filemodes[0] != 0100644 {
		t.Errorf("file mode want 0100644, got %o", r.filemodes[0])
	}
	if r.filelinktos[0] != "" {
		t.Errorf("linktos want empty (not a symlink), got %q", r.filelinktos[0])
	}
}

func TestCompression(t *testing.T) {
	testCases := []struct {
		Type           string
		Compressors    []string
		ExpectedWriter io.Writer
	}{
		{
			Type: "gzip",
			Compressors: []string{
				"", "gzip", "gzip:1", "gzip:2", "gzip:3",
				"gzip:4", "gzip:5", "gzip:6", "gzip:7", "gzip:8", "gzip:9",
			},
			ExpectedWriter: &gzip.Writer{},
		},
		{
			Type:           "gzip",
			Compressors:    []string{"gzip:fast", "gzip:10"},
			ExpectedWriter: nil, // gzip requires an integer level from -2 to 9
		},
		{
			Type:           "lzma",
			Compressors:    []string{"lzma"},
			ExpectedWriter: &lzma.Writer{},
		},
		{
			Type:           "lzma",
			Compressors:    []string{"lzma:fast", "lzma:1"},
			ExpectedWriter: nil, // lzma does not support specifying the compression level
		},
		{
			Type:           "xz",
			Compressors:    []string{"xz"},
			ExpectedWriter: &xz.Writer{},
		},
		{
			Type:           "xz",
			Compressors:    []string{"xz:fast", "xz:1"},
			ExpectedWriter: nil, // xz does not support specifying the compression level
		},
		{
			Type: "zstd",
			Compressors: []string{
				"zstd", "zstd:fastest", "zstd:default", "zstd:better",
				"zstd:best", "zstd:BeSt", "zstd:0", "zstd:4", "zstd:8", "zstd:15",
			},
			ExpectedWriter: &zstd.Encoder{},
		},
		{
			Type:           "zstd",
			Compressors:    []string{"xz:worst"},
			ExpectedWriter: nil, // zstd does not support integer compression level
		},
	}

	for _, testCase := range testCases {
		testCase := testCase

		for _, compressor := range testCase.Compressors {
			t.Run(compressor, func(t *testing.T) {
				r, err := NewRPM(RPMMetaData{
					Compressor: compressor,
				})
				if err != nil {
					if testCase.ExpectedWriter == nil {
						return // an error is expected
					}

					t.Fatalf("NewRPM returned error %v", err)
				}

				if testCase.ExpectedWriter == nil {
					t.Fatalf("compressor %q should have produced an error", compressor)
				}

				if r.RPMMetaData.Compressor != testCase.Type {
					t.Fatalf("expected compressor %q, got %q", compressor,
						r.RPMMetaData.Compressor)
				}

				expectedWriterType := reflect.Indirect(reflect.ValueOf(
					testCase.ExpectedWriter)).String()
				actualWriterType := reflect.Indirect(reflect.ValueOf(
					r.compressedPayload)).String()

				if expectedWriterType != actualWriterType {
					t.Fatalf("expected writer to be %T, got %T instead",
						testCase.ExpectedWriter, r.compressedPayload)
				}
			})
		}
	}
}
<commit_msg>Fix comment in compression test<commit_after>package rpmpack

import (
	"compress/gzip"
	"io"
	"io/ioutil"
	"reflect"
	"testing"

	"github.com/klauspost/compress/zstd"
	"github.com/ulikunitz/xz"
	"github.com/ulikunitz/xz/lzma"
)

func TestFileOwner(t *testing.T) {
	r, err := NewRPM(RPMMetaData{})
	if err != nil {
		t.Fatalf("NewRPM returned error %v", err)
	}
	group := "testGroup"
	user := "testUser"

	r.AddFile(RPMFile{
		Name:  "/usr/local/hello",
		Body:  []byte("content of the file"),
		Group: group,
		Owner: user,
	})

	if err := r.Write(ioutil.Discard); err != nil {
		t.Errorf("NewRPM returned error %v", err)
	}
	if r.fileowners[0] != user {
		t.Errorf("File owner shoud be %s but is %s", user, r.fileowners[0])
	}
	if r.filegroups[0] != group {
		t.Errorf("File owner shoud be %s but is %s", group, r.filegroups[0])
	}
}

// https://github.com/google/rpmpack/issues/49
func Test100644(t *testing.T) {
	r, err := NewRPM(RPMMetaData{})
	if err != nil {
		t.Fatalf("NewRPM returned error %v", err)
	}
	r.AddFile(RPMFile{
		Name: "/usr/local/hello",
		Body: []byte("content of the file"),
		Mode: 0100644,
	})

	if err := r.Write(ioutil.Discard); err != nil {
		t.Errorf("Write returned error %v", err)
	}
	if r.filemodes[0] != 0100644 {
		t.Errorf("file mode want 0100644, got %o", r.filemodes[0])
	}
	if r.filelinktos[0] != "" {
		t.Errorf("linktos want empty (not a symlink), got %q", r.filelinktos[0])
	}
}

func TestCompression(t *testing.T) {
	testCases := []struct {
		Type           string
		Compressors    []string
		ExpectedWriter io.Writer
	}{
		{
			Type: "gzip",
			Compressors: []string{
				"", "gzip", "gzip:1", "gzip:2", "gzip:3",
				"gzip:4", "gzip:5", "gzip:6", "gzip:7", "gzip:8", "gzip:9",
			},
			ExpectedWriter: &gzip.Writer{},
		},
		{
			Type:           "gzip",
			Compressors:    []string{"gzip:fast", "gzip:10"},
			ExpectedWriter: nil, // gzip requires an integer level from -2 to 9
		},
		{
			Type:           "lzma",
			Compressors:    []string{"lzma"},
			ExpectedWriter: &lzma.Writer{},
		},
		{
			Type:           "lzma",
			Compressors:    []string{"lzma:fast", "lzma:1"},
			ExpectedWriter: nil, // lzma does not support specifying the compression level
		},
		{
			Type:           "xz",
			Compressors:    []string{"xz"},
			ExpectedWriter: &xz.Writer{},
		},
		{
			Type:           "xz",
			Compressors:    []string{"xz:fast", "xz:1"},
			ExpectedWriter: nil, // xz does not support specifying the compression level
		},
		{
			Type: "zstd",
			Compressors: []string{
				"zstd", "zstd:fastest", "zstd:default", "zstd:better",
				"zstd:best", "zstd:BeSt", "zstd:0", "zstd:4", "zstd:8", "zstd:15",
			},
			ExpectedWriter: &zstd.Encoder{},
		},
		{
			Type:           "zstd",
			Compressors:    []string{"xz:worst"},
			ExpectedWriter: nil, // only integers levels or one of the pre-defined string values
		},
	}

	for _, testCase := range testCases {
		testCase := testCase

		for _, compressor := range testCase.Compressors {
			t.Run(compressor, func(t *testing.T) {
				r, err := NewRPM(RPMMetaData{
					Compressor: compressor,
				})
				if err != nil {
					if testCase.ExpectedWriter == nil {
						return // an error is expected
					}

					t.Fatalf("NewRPM returned error %v", err)
				}

				if testCase.ExpectedWriter == nil {
					t.Fatalf("compressor %q should have produced an error", compressor)
				}

				if r.RPMMetaData.Compressor != testCase.Type {
					t.Fatalf("expected compressor %q, got %q", compressor,
						r.RPMMetaData.Compressor)
				}

				expectedWriterType := reflect.Indirect(reflect.ValueOf(
					testCase.ExpectedWriter)).String()
				actualWriterType := reflect.Indirect(reflect.ValueOf(
					r.compressedPayload)).String()

				if expectedWriterType != actualWriterType {
					t.Fatalf("expected writer to be %T, got %T instead",
						testCase.ExpectedWriter, r.compressedPayload)
				}
			})
		}
	}
}
<|endoftext|>