<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/keys"
	"koding/kites/kloud/koding"
	"log"
	"net/url"
	"os"
	"time"

	"github.com/fatih/structs"
	"github.com/koding/kite"
	kiteconfig "github.com/koding/kite/config"
	"github.com/koding/kite/protocol"
	"github.com/koding/kloud"
	kloudprotocol "github.com/koding/kloud/protocol"
	"github.com/koding/logging"
	"github.com/koding/multiconfig"
)

// Config defines the configuration that Kloud needs to operate.
type Config struct {
	// ---  KLOUD SPECIFIC ---
	IP          string
	Port        int
	Region      string
	Environment string
	Id          string

	// Connect to Koding mongodb
	MongoURL string

	// --- DEVELOPMENT CONFIG ---
	// Show version and exit if enabled
	Version bool

	// Enable debug log mode
	DebugMode bool

	// Enable production mode, operates on production channel
	ProdMode bool

	// Enable test mode, disabled some authentication checks
	TestMode bool

	// Defines the base domain for domain creation
	HostedZone string

	// Defines the default AMI Tag to use for koding provider
	AMITag string

	// --- KLIENT DEVELOPMENT ---
	// KontrolURL to connect and to de deployed with klient
	KontrolURL string

	// Private key to create kite.key
	PrivateKey string

	// Public key to create kite.key
	PublicKey string

	// Contains the users home directory to be added into a image
	TemplateDir string

	// --- KONTROL CONFIGURATION ---
	Public      bool   // Try to register with a public ip
	Proxy       bool   // Try to register behind a koding proxy
	RegisterURL string // Explicitly register with this given url
}

func main() {
	conf := new(Config)

	// Load the config, it's reads environment variables or from flags
	multiconfig.New().MustLoad(conf)

	if conf.Version {
		fmt.Println(kloud.VERSION)
		os.Exit(0)
	}

	fmt.Printf("Kloud loaded with following configuration variables: %+v\n", conf)

	if conf.HostedZone == "" {
		panic("hosted zone is not set. Pass it via -hostedzone or CONFIG_HOSTEDZONE environment variable")
	}

	k := newKite(conf)

	registerURL := k.RegisterURL(!conf.Public)
	if conf.RegisterURL != "" {
		u, err := url.Parse(conf.RegisterURL)
		if err != nil {
			k.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	fmt.Printf("registering with url %+v\n", registerURL)

	if conf.Proxy {
		k.Log.Info("Proxy mode is enabled")
		// Koding proxies in production only
		proxyQuery := &protocol.KontrolQuery{
			Username:    "koding",
			Environment: "production",
			Name:        "proxy",
		}

		k.Log.Info("Seaching proxy: %#v", proxyQuery)
		go k.RegisterToProxy(registerURL, proxyQuery)
	} else {
		if err := k.RegisterForever(registerURL); err != nil {
			k.Log.Fatal(err.Error())
		}
	}

	k.Run()
}

func newKite(conf *Config) *kite.Kite {
	k := kite.New(kloud.NAME, kloud.VERSION)
	k.Config = kiteconfig.MustGet()
	k.Config.Port = conf.Port

	if conf.Region != "" {
		k.Config.Region = conf.Region
	}

	if conf.Environment != "" {
		k.Config.Environment = conf.Environment
	}

	id := uniqueId(k.Config.Port)
	if conf.Id != "" {
		id = conf.Id
	}

	if conf.AMITag != "" {
		k.Log.Warning("Default AMI Tag changed from %s to %s", koding.DefaultCustomAMITag, conf.AMITag)
		koding.DefaultCustomAMITag = conf.AMITag
	}

	modelhelper.Initialize(conf.MongoURL)
	db := modelhelper.Mongo

	kodingProvider := &koding.Provider{
		Kite:         k,
		Log:          newLogger("koding", conf.DebugMode),
		AssigneeName: id,
		Session:      db,
		Test:         conf.TestMode,
		TemplateDir:  conf.TemplateDir,
		HostedZone:   conf.HostedZone,
	}

	go kodingProvider.RunChecker(time.Second * 10)
	go kodingProvider.RunCleaner(time.Minute)

	klientFolder := "klient/development/latest"
	if conf.ProdMode {
		klientFolder = "klient/production/latest"
	}

	privateKey, publicKey := kontrolKeys(conf)

	deployer := &KodingDeploy{
		Kite:              k,
		Log:               newLogger("kloud-deploy", conf.DebugMode),
		KontrolURL:        kontrolURL(conf.KontrolURL),
		KontrolPrivateKey: privateKey,
		KontrolPublicKey:  publicKey,
		Bucket:            newBucket("koding-kites", klientFolder),
		DB:                db,
	}

	kld := kloud.NewKloud()
	kld.Storage = kodingProvider
	kld.Log = newLogger("kloud", conf.DebugMode)

	// be sure it compiles correctly,
	var _ kloudprotocol.Builder = kodingProvider

	err := kld.AddProvider("koding", kodingProvider)
	if err != nil {
		panic(err)
	}

	injectDeploy := func(r *kite.Request) (interface{}, error) {
		d := kloudprotocol.ProviderDeploy{
			KeyName:    keys.DeployKeyName,
			PublicKey:  keys.DeployPublicKey,
			PrivateKey: keys.DeployPrivateKey,
			Username:   r.Username,
		}

		r.Context.Set("deployData", structs.Map(d))
		return true, nil
	}

	k.Handle("build", kld.NewBuild(deployer)).PreHandleFunc(injectDeploy)
	k.HandleFunc("start", kld.Start)
	k.HandleFunc("stop", kld.Stop)
	k.HandleFunc("restart", kld.Restart)
	k.HandleFunc("info", kld.Info)
	k.HandleFunc("destroy", kld.Destroy)
	k.HandleFunc("event", kld.Event)
	k.HandleFunc("domain.set", func(r *kite.Request) (interface{}, error) {
		// let's use the helper function which is doing a lot of things on
		// behalf of us, like document locking, getting the machine document,
		// and so on..
		return kld.ControlFunc(kodingProvider.DomainSet).ServeKite(r)
	})

	return k
}

func uniqueId(port int) string {
	// TODO: add a unique identifier, for letting multiple version of the same
	// worker work on the same hostname.
	hostname, err := os.Hostname()
	if err != nil {
		panic(err) // we should not let it start
	}

	return fmt.Sprintf("%s-%s-%d", kloud.NAME, hostname, port)
}

func newLogger(name string, debug bool) logging.Logger {
	log := logging.NewLogger(name)
	logHandler := logging.NewWriterHandler(os.Stderr)
	logHandler.Colorize = true
	log.SetHandler(logHandler)

	if debug {
		log.SetLevel(logging.DEBUG)
		logHandler.SetLevel(logging.DEBUG)
	}

	return log
}

func kontrolKeys(conf *Config) (string, string) {
	pubKey, err := ioutil.ReadFile(conf.PublicKey)
	if err != nil {
		log.Fatalln(err)
	}
	publicKey := string(pubKey)

	privKey, err := ioutil.ReadFile(conf.PrivateKey)
	if err != nil {
		log.Fatalln(err)
	}
	privateKey := string(privKey)

	return privateKey, publicKey
}

func kontrolURL(ownURL string) string {
	// read kontrolURL from kite.key if it doesn't exist.
	kontrolURL := kiteconfig.MustGet().KontrolURL

	if ownURL != "" {
		u, err := url.Parse(ownURL)
		if err != nil {
			log.Fatalln(err)
		}

		kontrolURL = u.String()
	}

	return kontrolURL
}
<commit_msg>kloud: print klient distribution channel<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/keys"
	"koding/kites/kloud/koding"
	"log"
	"net/url"
	"os"
	"time"

	"github.com/fatih/structs"
	"github.com/koding/kite"
	kiteconfig "github.com/koding/kite/config"
	"github.com/koding/kite/protocol"
	"github.com/koding/kloud"
	kloudprotocol "github.com/koding/kloud/protocol"
	"github.com/koding/logging"
	"github.com/koding/multiconfig"
)

// Config defines the configuration that Kloud needs to operate.
type Config struct {
	// ---  KLOUD SPECIFIC ---
	IP          string
	Port        int
	Region      string
	Environment string
	Id          string

	// Connect to Koding mongodb
	MongoURL string

	// --- DEVELOPMENT CONFIG ---
	// Show version and exit if enabled
	Version bool

	// Enable debug log mode
	DebugMode bool

	// Enable production mode, operates on production channel
	ProdMode bool

	// Enable test mode, disabled some authentication checks
	TestMode bool

	// Defines the base domain for domain creation
	HostedZone string

	// Defines the default AMI Tag to use for koding provider
	AMITag string

	// --- KLIENT DEVELOPMENT ---
	// KontrolURL to connect and to de deployed with klient
	KontrolURL string

	// Private key to create kite.key
	PrivateKey string

	// Public key to create kite.key
	PublicKey string

	// Contains the users home directory to be added into a image
	TemplateDir string

	// --- KONTROL CONFIGURATION ---
	Public      bool   // Try to register with a public ip
	Proxy       bool   // Try to register behind a koding proxy
	RegisterURL string // Explicitly register with this given url
}

func main() {
	conf := new(Config)

	// Load the config, it's reads environment variables or from flags
	multiconfig.New().MustLoad(conf)

	if conf.Version {
		fmt.Println(kloud.VERSION)
		os.Exit(0)
	}

	fmt.Printf("Kloud loaded with following configuration variables: %+v\n", conf)

	if conf.HostedZone == "" {
		panic("hosted zone is not set. Pass it via -hostedzone or CONFIG_HOSTEDZONE environment variable")
	}

	k := newKite(conf)

	registerURL := k.RegisterURL(!conf.Public)
	if conf.RegisterURL != "" {
		u, err := url.Parse(conf.RegisterURL)
		if err != nil {
			k.Log.Fatal("Couldn't parse register url: %s", err)
		}

		registerURL = u
	}

	fmt.Printf("registering with url %+v\n", registerURL)

	if conf.Proxy {
		k.Log.Info("Proxy mode is enabled")
		// Koding proxies in production only
		proxyQuery := &protocol.KontrolQuery{
			Username:    "koding",
			Environment: "production",
			Name:        "proxy",
		}

		k.Log.Info("Seaching proxy: %#v", proxyQuery)
		go k.RegisterToProxy(registerURL, proxyQuery)
	} else {
		if err := k.RegisterForever(registerURL); err != nil {
			k.Log.Fatal(err.Error())
		}
	}

	k.Run()
}

func newKite(conf *Config) *kite.Kite {
	k := kite.New(kloud.NAME, kloud.VERSION)
	k.Config = kiteconfig.MustGet()
	k.Config.Port = conf.Port

	if conf.Region != "" {
		k.Config.Region = conf.Region
	}

	if conf.Environment != "" {
		k.Config.Environment = conf.Environment
	}

	id := uniqueId(k.Config.Port)
	if conf.Id != "" {
		id = conf.Id
	}

	if conf.AMITag != "" {
		k.Log.Warning("Default AMI Tag changed from %s to %s", koding.DefaultCustomAMITag, conf.AMITag)
		koding.DefaultCustomAMITag = conf.AMITag
	}

	klientFolder := "klient/development/latest"
	if conf.ProdMode {
		klientFolder = "klient/production/latest"
	}
	k.Log.Info("Klient distribution channel is: %s", klientFolder)

	modelhelper.Initialize(conf.MongoURL)
	db := modelhelper.Mongo

	kodingProvider := &koding.Provider{
		Kite:         k,
		Log:          newLogger("koding", conf.DebugMode),
		AssigneeName: id,
		Session:      db,
		Test:         conf.TestMode,
		TemplateDir:  conf.TemplateDir,
		HostedZone:   conf.HostedZone,
	}

	go kodingProvider.RunChecker(time.Second * 10)
	go kodingProvider.RunCleaner(time.Minute)

	privateKey, publicKey := kontrolKeys(conf)

	deployer := &KodingDeploy{
		Kite:              k,
		Log:               newLogger("kloud-deploy", conf.DebugMode),
		KontrolURL:        kontrolURL(conf.KontrolURL),
		KontrolPrivateKey: privateKey,
		KontrolPublicKey:  publicKey,
		Bucket:            newBucket("koding-kites", klientFolder),
		DB:                db,
	}

	kld := kloud.NewKloud()
	kld.Storage = kodingProvider
	kld.Log = newLogger("kloud", conf.DebugMode)

	// be sure it compiles correctly,
	var _ kloudprotocol.Builder = kodingProvider

	err := kld.AddProvider("koding", kodingProvider)
	if err != nil {
		panic(err)
	}

	injectDeploy := func(r *kite.Request) (interface{}, error) {
		d := kloudprotocol.ProviderDeploy{
			KeyName:    keys.DeployKeyName,
			PublicKey:  keys.DeployPublicKey,
			PrivateKey: keys.DeployPrivateKey,
			Username:   r.Username,
		}

		r.Context.Set("deployData", structs.Map(d))
		return true, nil
	}

	k.Handle("build", kld.NewBuild(deployer)).PreHandleFunc(injectDeploy)
	k.HandleFunc("start", kld.Start)
	k.HandleFunc("stop", kld.Stop)
	k.HandleFunc("restart", kld.Restart)
	k.HandleFunc("info", kld.Info)
	k.HandleFunc("destroy", kld.Destroy)
	k.HandleFunc("event", kld.Event)
	k.HandleFunc("domain.set", func(r *kite.Request) (interface{}, error) {
		// let's use the helper function which is doing a lot of things on
		// behalf of us, like document locking, getting the machine document,
		// and so on..
		return kld.ControlFunc(kodingProvider.DomainSet).ServeKite(r)
	})

	return k
}

func uniqueId(port int) string {
	// TODO: add a unique identifier, for letting multiple version of the same
	// worker work on the same hostname.
	hostname, err := os.Hostname()
	if err != nil {
		panic(err) // we should not let it start
	}

	return fmt.Sprintf("%s-%s-%d", kloud.NAME, hostname, port)
}

func newLogger(name string, debug bool) logging.Logger {
	log := logging.NewLogger(name)
	logHandler := logging.NewWriterHandler(os.Stderr)
	logHandler.Colorize = true
	log.SetHandler(logHandler)

	if debug {
		log.SetLevel(logging.DEBUG)
		logHandler.SetLevel(logging.DEBUG)
	}

	return log
}

func kontrolKeys(conf *Config) (string, string) {
	pubKey, err := ioutil.ReadFile(conf.PublicKey)
	if err != nil {
		log.Fatalln(err)
	}
	publicKey := string(pubKey)

	privKey, err := ioutil.ReadFile(conf.PrivateKey)
	if err != nil {
		log.Fatalln(err)
	}
	privateKey := string(privKey)

	return privateKey, publicKey
}

func kontrolURL(ownURL string) string {
	// read kontrolURL from kite.key if it doesn't exist.
	kontrolURL := kiteconfig.MustGet().KontrolURL

	if ownURL != "" {
		u, err := url.Parse(ownURL)
		if err != nil {
			log.Fatalln(err)
		}

		kontrolURL = u.String()
	}

	return kontrolURL
}
<|endoftext|>


<commit_before>package amazon

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/awslabs/goformation/v4/cloudformation"
	"github.com/awslabs/goformation/v4/cloudformation/ec2"
	"github.com/awslabs/goformation/v4/cloudformation/iam"

	"github.com/awslabs/goformation/v4/cloudformation"
	"github.com/awslabs/goformation/v4/cloudformation/elasticloadbalancingv2"
	"github.com/compose-spec/compose-go/loader"
	"github.com/compose-spec/compose-go/types"
	"github.com/docker/ecs-plugin/pkg/compose"
	"gotest.tools/assert"
	"gotest.tools/v3/golden"
)

func TestSimpleConvert(t *testing.T) {
	project := load(t, "testdata/input/simple-single-service.yaml")
	result := convertResultAsString(t, project, "TestCluster")
	expected := "simple/simple-cloudformation-conversion.golden"
	golden.Assert(t, result, expected)
}

func TestSimpleWithOverrides(t *testing.T) {
	project := load(t, "testdata/input/simple-single-service.yaml", "testdata/input/simple-single-service-with-overrides.yaml")
	result := convertResultAsString(t, project, "TestCluster")
	expected := "simple/simple-cloudformation-with-overrides-conversion.golden"
	golden.Assert(t, result, expected)
}

func TestRolePolicy(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  foo:
    image: hello_world
    x-aws-pull_credentials: "secret"
`)
	role := template.Resources["FooTaskExecutionRole"].(*iam.Role)
	assert.Check(t, role != nil)
	assert.Check(t, role.ManagedPolicyArns[0] == ECSTaskExecutionPolicy)
	assert.Check(t, role.ManagedPolicyArns[1] == ECRReadOnlyPolicy)
	// We expect an extra policy has been created for x-aws-pull_credentials
	assert.Check(t, len(role.Policies) == 1)
	policy := role.Policies[0].PolicyDocument.(*PolicyDocument)
	expected := []string{"secretsmanager:GetSecretValue", "ssm:GetParameters", "kms:Decrypt"}
	assert.DeepEqual(t, expected, policy.Statement[0].Action)
	assert.DeepEqual(t, []string{"secret"}, policy.Statement[0].Resource)
}

func TestMapNetworksToSecurityGroups(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: hello_world
networks:
  front-tier:
    name: public
  back-tier:
    internal: true
`)
	assert.Check(t, template.Resources["TestPublicNetwork"] != nil)
	assert.Check(t, template.Resources["TestBacktierNetwork"] != nil)
	assert.Check(t, template.Resources["TestBacktierNetworkIngress"] != nil)
	ingress := template.Resources["TestPublicNetworkIngress"].(*ec2.SecurityGroupIngress)
	assert.Check(t, ingress != nil)
	assert.Check(t, ingress.SourceSecurityGroupId == cloudformation.Ref("TestPublicNetwork"))

}

func TestLoadBalancerTypeApplication(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: nginx
    ports:
      - 80:80
`)
	lb := template.Resources["TestLoadBalancer"].(*elasticloadbalancingv2.LoadBalancer)
	assert.Check(t, lb != nil)
	assert.Check(t, lb.Type == elbv2.LoadBalancerTypeEnumApplication)
	assert.Check(t, len(lb.SecurityGroups) > 0)
}

func TestLoadBalancerTypeNetwork(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: nginx
    ports:
      - 80:80
      - 88:88
`)
	lb := template.Resources["TestLoadBalancer"].(*elasticloadbalancingv2.LoadBalancer)
	assert.Check(t, lb != nil)
	assert.Check(t, lb.Type == elbv2.LoadBalancerTypeEnumNetwork)
}

func convertResultAsString(t *testing.T, project *compose.Project, clusterName string) string {
	client, err := NewClient("", clusterName, "")
	assert.NilError(t, err)
	result, err := client.Convert(project)
	assert.NilError(t, err)
	resultAsJSON, err := result.JSON()
	assert.NilError(t, err)
	return fmt.Sprintf("%s\n", string(resultAsJSON))
}

func load(t *testing.T, paths ...string) *compose.Project {
	options := compose.ProjectOptions{
		Name:        t.Name(),
		ConfigPaths: paths,
	}
	project, err := compose.ProjectFromOptions(&options)
	assert.NilError(t, err)
	return project
}

func convertYaml(t *testing.T, yaml string) *cloudformation.Template {
	dict, err := loader.ParseYAML([]byte(yaml))
	assert.NilError(t, err)
	model, err := loader.Load(types.ConfigDetails{
		ConfigFiles: []types.ConfigFile{
			{Config: dict},
		},
	})
	assert.NilError(t, err)
	err = compose.Normalize(model)
	assert.NilError(t, err)
	template, err := client{}.Convert(&compose.Project{
		Config: *model,
		Name:   "test",
	})
	assert.NilError(t, err)
	return template
}
<commit_msg>fix rebase<commit_after>package amazon

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/service/elbv2"
	"github.com/awslabs/goformation/v4/cloudformation"
	"github.com/awslabs/goformation/v4/cloudformation/ec2"
	"github.com/awslabs/goformation/v4/cloudformation/iam"

	"github.com/awslabs/goformation/v4/cloudformation/elasticloadbalancingv2"
	"github.com/compose-spec/compose-go/loader"
	"github.com/compose-spec/compose-go/types"
	"github.com/docker/ecs-plugin/pkg/compose"
	"gotest.tools/assert"
	"gotest.tools/v3/golden"
)

func TestSimpleConvert(t *testing.T) {
	project := load(t, "testdata/input/simple-single-service.yaml")
	result := convertResultAsString(t, project, "TestCluster")
	expected := "simple/simple-cloudformation-conversion.golden"
	golden.Assert(t, result, expected)
}

func TestSimpleWithOverrides(t *testing.T) {
	project := load(t, "testdata/input/simple-single-service.yaml", "testdata/input/simple-single-service-with-overrides.yaml")
	result := convertResultAsString(t, project, "TestCluster")
	expected := "simple/simple-cloudformation-with-overrides-conversion.golden"
	golden.Assert(t, result, expected)
}

func TestRolePolicy(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  foo:
    image: hello_world
    x-aws-pull_credentials: "secret"
`)
	role := template.Resources["FooTaskExecutionRole"].(*iam.Role)
	assert.Check(t, role != nil)
	assert.Check(t, role.ManagedPolicyArns[0] == ECSTaskExecutionPolicy)
	assert.Check(t, role.ManagedPolicyArns[1] == ECRReadOnlyPolicy)
	// We expect an extra policy has been created for x-aws-pull_credentials
	assert.Check(t, len(role.Policies) == 1)
	policy := role.Policies[0].PolicyDocument.(*PolicyDocument)
	expected := []string{"secretsmanager:GetSecretValue", "ssm:GetParameters", "kms:Decrypt"}
	assert.DeepEqual(t, expected, policy.Statement[0].Action)
	assert.DeepEqual(t, []string{"secret"}, policy.Statement[0].Resource)
}

func TestMapNetworksToSecurityGroups(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: hello_world
networks:
  front-tier:
    name: public
  back-tier:
    internal: true
`)
	assert.Check(t, template.Resources["TestPublicNetwork"] != nil)
	assert.Check(t, template.Resources["TestBacktierNetwork"] != nil)
	assert.Check(t, template.Resources["TestBacktierNetworkIngress"] != nil)
	ingress := template.Resources["TestPublicNetworkIngress"].(*ec2.SecurityGroupIngress)
	assert.Check(t, ingress != nil)
	assert.Check(t, ingress.SourceSecurityGroupId == cloudformation.Ref("TestPublicNetwork"))

}

func TestLoadBalancerTypeApplication(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: nginx
    ports:
      - 80:80
`)
	lb := template.Resources["TestLoadBalancer"].(*elasticloadbalancingv2.LoadBalancer)
	assert.Check(t, lb != nil)
	assert.Check(t, lb.Type == elbv2.LoadBalancerTypeEnumApplication)
	assert.Check(t, len(lb.SecurityGroups) > 0)
}

func TestLoadBalancerTypeNetwork(t *testing.T) {
	template := convertYaml(t, `
version: "3"
services:
  test:
    image: nginx
    ports:
      - 80:80
      - 88:88
`)
	lb := template.Resources["TestLoadBalancer"].(*elasticloadbalancingv2.LoadBalancer)
	assert.Check(t, lb != nil)
	assert.Check(t, lb.Type == elbv2.LoadBalancerTypeEnumNetwork)
}

func convertResultAsString(t *testing.T, project *compose.Project, clusterName string) string {
	client, err := NewClient("", clusterName, "")
	assert.NilError(t, err)
	result, err := client.Convert(project)
	assert.NilError(t, err)
	resultAsJSON, err := result.JSON()
	assert.NilError(t, err)
	return fmt.Sprintf("%s\n", string(resultAsJSON))
}

func load(t *testing.T, paths ...string) *compose.Project {
	options := compose.ProjectOptions{
		Name:        t.Name(),
		ConfigPaths: paths,
	}
	project, err := compose.ProjectFromOptions(&options)
	assert.NilError(t, err)
	return project
}

func convertYaml(t *testing.T, yaml string) *cloudformation.Template {
	dict, err := loader.ParseYAML([]byte(yaml))
	assert.NilError(t, err)
	model, err := loader.Load(types.ConfigDetails{
		ConfigFiles: []types.ConfigFile{
			{Config: dict},
		},
	})
	assert.NilError(t, err)
	err = compose.Normalize(model)
	assert.NilError(t, err)
	template, err := client{}.Convert(&compose.Project{
		Config: *model,
		Name:   "test",
	})
	assert.NilError(t, err)
	return template
}
<|endoftext|>


<commit_before>/**********************************************************************************
* Copyright (c) 2009-2017 Misakai Ltd.
* This program is free software: you can redistribute it and/or modify it under the
* terms of the GNU Affero General Public License as published by the  Free Software
* Foundation, either version 3 of the License, or(at your option) any later version.
*
* This program is distributed  in the hope that it  will be useful, but WITHOUT ANY
* WARRANTY;  without even  the implied warranty of MERCHANTABILITY or FITNESS FOR A
* PARTICULAR PURPOSE.  See the GNU Affero General Public License  for  more details.
*
* You should have  received a copy  of the  GNU Affero General Public License along
* with this program. If not, see<http://www.gnu.org/licenses/>.
************************************************************************************/

package cluster

import (
	"bytes"
	"sync"
	"time"

	"github.com/emitter-io/emitter/encoding"
	"github.com/emitter-io/emitter/logging"
	"github.com/emitter-io/emitter/utils"
	"github.com/golang/snappy"
	"github.com/weaveworks/mesh"
)

// PeerSet maintains a map of peers
type peerset struct {
	sync.Mutex
	sender  mesh.Gossip             //  The gossip interface to use for sending.
	peers   *mesh.Peers             // The memberlist discovery mechanism.
	members map[mesh.PeerName]*Peer // The map of members in the peer set.
}

// newPeerSet creates a new peer set for the connection.
func newPeerSet(sender mesh.Gossip, peers *mesh.Peers) *peerset {
	p := &peerset{
		sender:  sender,
		peers:   peers,
		members: make(map[mesh.PeerName]*Peer),
	}

	// Attach the GC callback so we know when a peer is dead.
	peers.OnGC(p.onPeerGC)
	return p
}

// Occurs when a peer is garbage collected.
func (s *peerset) onPeerGC(peer *mesh.Peer) {
	if p := s.Get(peer.Name); p != nil {
		logging.LogTarget("swarm", "peer garbage collected", peer)
		p.Close() // Close the peer on our end

		// TODO: Make sure we remove all subscriptions as well

		// We also need to remove the peer from our set, so next time a new peer can be created.
		s.Lock()
		delete(s.members, p.name)
		s.Unlock()
	}
}

// Get retrieves a peer.
func (s *peerset) Get(name mesh.PeerName) (p *Peer) {
	s.Lock() // TODO: This lock will be contended eventually, need to replace this map by sync.Map
	defer s.Unlock()

	// Get the peer
	if p, ok := s.members[name]; ok {
		return p
	}

	// Create new peer
	p = s.newPeer(name)
	s.members[name] = p
	return p
}

// ------------------------------------------------------------------------------------

// Peer represents a remote peer.
type Peer struct {
	sync.Mutex
	sender  mesh.Gossip   // The gossip interface to use for sending.
	name    mesh.PeerName // The peer name for communicating.
	frame   MessageFrame  // The current message frame.
	closing chan bool     // The closing channel for the peer.
}

// NewPeer creates a new peer for the connection.
func (s *peerset) newPeer(name mesh.PeerName) *Peer {
	peer := &Peer{
		sender:  s.sender,
		name:    name,
		frame:   make(MessageFrame, 0, 64),
		closing: make(chan bool),
	}

	// Spawn the send queue processor
	utils.Repeat(peer.processSendQueue, 5*time.Millisecond, peer.closing)
	return peer
}

// Close termintes the peer and stops everything associated with this peer.
func (p *Peer) Close() {
	close(p.closing)
}

// Send forwards the message to the remote server.
func (p *Peer) Send(ssid []uint32, channel []byte, payload []byte) error {
	p.Lock()
	defer p.Unlock()

	// TODO: Make sure we don't send to a dead peer

	// Send simply appends the message to a frame
	p.frame = append(p.frame, &Message{Ssid: ssid, Channel: channel, Payload: payload})
	return nil
}

// processSendQueue flushes the current frame to the remote server
func (p *Peer) processSendQueue() {
	if len(p.frame) == 0 {
		return // Nothing to send.
	}

	// Compress in-memory. TODO: Optimize the shit out of that, we don't really need to use binc
	buffer := bytes.NewBuffer(nil)
	snappy := snappy.NewBufferedWriter(buffer)
	writer := encoding.NewEncoder(snappy)

	// Encode the current frame
	p.Lock()
	err := writer.Encode(p.frame)
	p.frame = p.frame[:0]
	p.Unlock()

	// Something went wrong during the encoding
	if err != nil {
		logging.LogError("peer", "encoding frame", err)
	}

	// Send the frame directly to the peer.
	if err := snappy.Close(); err == nil {
		p.sender.GossipUnicast(p.name, buffer.Bytes())
	}
}
<commit_msg>fixed peer removal<commit_after>/**********************************************************************************
* Copyright (c) 2009-2017 Misakai Ltd.
* This program is free software: you can redistribute it and/or modify it under the
* terms of the GNU Affero General Public License as published by the  Free Software
* Foundation, either version 3 of the License, or(at your option) any later version.
*
* This program is distributed  in the hope that it  will be useful, but WITHOUT ANY
* WARRANTY;  without even  the implied warranty of MERCHANTABILITY or FITNESS FOR A
* PARTICULAR PURPOSE.  See the GNU Affero General Public License  for  more details.
*
* You should have  received a copy  of the  GNU Affero General Public License along
* with this program. If not, see<http://www.gnu.org/licenses/>.
************************************************************************************/

package cluster

import (
	"bytes"
	"sync"
	"time"

	"github.com/emitter-io/emitter/encoding"
	"github.com/emitter-io/emitter/logging"
	"github.com/emitter-io/emitter/utils"
	"github.com/golang/snappy"
	"github.com/weaveworks/mesh"
)

// PeerSet maintains a map of peers
type peerset struct {
	sender  mesh.Gossip //  The gossip interface to use for sending.
	peers   *mesh.Peers // The memberlist discovery mechanism.
	members *sync.Map   // The map of members in the peer set.
}

// newPeerSet creates a new peer set for the connection.
func newPeerSet(sender mesh.Gossip, peers *mesh.Peers) *peerset {
	p := &peerset{
		sender:  sender,
		peers:   peers,
		members: new(sync.Map),
	}

	// Attach the GC callback so we know when a peer is dead.
	peers.OnGC(p.onPeerGC)
	return p
}

// Occurs when a peer is garbage collected.
func (s *peerset) onPeerGC(p *mesh.Peer) {
	if v, ok := s.members.Load(p.Name); ok {
		peer := v.(*Peer)
		logging.LogTarget("swarm", "peer removed", peer.name)
		peer.Close() // Close the peer on our end

		// TODO: Make sure we remove all subscriptions as well

		// We also need to remove the peer from our set, so next time a new peer can be created.
		s.members.Delete(peer.name)
	}
}

// Get retrieves a peer.
func (s *peerset) Get(name mesh.PeerName) *Peer {
	if p, ok := s.members.Load(name); ok {
		return p.(*Peer)
	}

	// Create new peer and store it
	peer := s.newPeer(name)
	v, ok := s.members.LoadOrStore(name, peer)
	if !ok {
		logging.LogTarget("swarm", "peer created", v)
	}
	return v.(*Peer)
}

// ------------------------------------------------------------------------------------

// Peer represents a remote peer.
type Peer struct {
	sync.Mutex
	sender  mesh.Gossip   // The gossip interface to use for sending.
	name    mesh.PeerName // The peer name for communicating.
	frame   MessageFrame  // The current message frame.
	closing chan bool     // The closing channel for the peer.
}

// NewPeer creates a new peer for the connection.
func (s *peerset) newPeer(name mesh.PeerName) *Peer {
	peer := &Peer{
		sender:  s.sender,
		name:    name,
		frame:   make(MessageFrame, 0, 64),
		closing: make(chan bool),
	}

	// Spawn the send queue processor
	utils.Repeat(peer.processSendQueue, 5*time.Millisecond, peer.closing)
	return peer
}

// Close termintes the peer and stops everything associated with this peer.
func (p *Peer) Close() {
	close(p.closing)
}

// Send forwards the message to the remote server.
func (p *Peer) Send(ssid []uint32, channel []byte, payload []byte) error {
	p.Lock()
	defer p.Unlock()

	// TODO: Make sure we don't send to a dead peer

	// Send simply appends the message to a frame
	p.frame = append(p.frame, &Message{Ssid: ssid, Channel: channel, Payload: payload})
	return nil
}

// processSendQueue flushes the current frame to the remote server
func (p *Peer) processSendQueue() {
	if len(p.frame) == 0 {
		return // Nothing to send.
	}

	// Compress in-memory. TODO: Optimize the shit out of that, we don't really need to use binc
	buffer := bytes.NewBuffer(nil)
	snappy := snappy.NewBufferedWriter(buffer)
	writer := encoding.NewEncoder(snappy)

	// Encode the current frame
	p.Lock()
	err := writer.Encode(p.frame)
	p.frame = p.frame[:0]
	p.Unlock()

	// Something went wrong during the encoding
	if err != nil {
		logging.LogError("peer", "encoding frame", err)
	}

	// Send the frame directly to the peer.
	if err := snappy.Close(); err == nil {
		p.sender.GossipUnicast(p.name, buffer.Bytes())
	}
}
<|endoftext|>


<commit_before>package solr

import (
	"fmt"
	"github.com/samuel/go-zookeeper/zk"
	"time"
)

func (s *solrZkInstance) Listen() error {
	err := s.zookeeper.Connect()
	s.zookeeper.ZKLogger(s.logger)
	s.currentNode = 0
	if err != nil {
		return err
	}
	s.clusterState = ClusterState{}
	var collectionsEvents <-chan zk.Event
	var liveNodeEvents <-chan zk.Event
	connect := func() error {
		collectionsEvents, err = s.initCollectionsListener()
		if err != nil {
			return err
		}
		liveNodeEvents, err = s.initLiveNodesListener()
		if err != nil {
			return err
		}
		return nil
	}
	err = connect()
	if err != nil {
		return err
	}
	//loop forever
	go func() {
		log := s.logger
		sleepTime := s.sleepTimeMS
		logErr := func(err error) {
			log.Error(fmt.Errorf("[go-solr] Error connecting to zk %v sleeping: %d", err, sleepTime))
		}
		for {
			disconnected := false
			select {
			case cEvent := <-collectionsEvents:
				if cEvent.Err != nil {
					log.Debug(fmt.Sprintf("[go-solr] error on cevent %v", cEvent))
					logErr(cEvent.Err)
					sleepTime = backoff(sleepTime)
					continue
				}
				// do something if its not a session or disconnect
				if cEvent.Type == zk.EventNodeDataChanged {
					collections, version, err := s.zookeeper.GetClusterState()
					if err != nil {
						logErr(err)
						sleepTime = backoff(sleepTime)
						continue
					}
					s.setCollections(collections, version)
				}
				disconnected = cEvent.State == zk.StateDisconnected
				sleepTime = s.sleepTimeMS

			case nEvent := <-liveNodeEvents:
				if nEvent.Err != nil {
					logErr(nEvent.Err)
					log.Error(fmt.Errorf("[go-solr] error on nevent %v", nEvent))
					sleepTime = backoff(sleepTime)
					continue
				}
				// do something if its not a session or disconnect
				if nEvent.Type == zk.EventNodeDataChanged || nEvent.Type == zk.EventNodeChildrenChanged {
					liveNodes, err := s.zookeeper.GetLiveNodes()
					if err != nil {
						logErr(err)
						sleepTime = backoff(sleepTime)
						continue
					}
					s.setLiveNodes(liveNodes)
				}
				disconnected = nEvent.State == zk.StateDisconnected
				sleepTime = s.sleepTimeMS
			}
			if !s.zookeeper.IsConnected() || disconnected {
				err = connect()
				if err != nil {
					s.logger.Error(fmt.Errorf("[go-solr] zk connect err %v, sleeping %d", err, sleepTime))
					sleepTime = backoff(sleepTime)
				} else {
					if disconnected {
						sleepTime = backoff(sleepTime)
					} else {
						sleepTime = s.sleepTimeMS
					}
				}
			}
		}
	}()
	s.listening = true
	return nil
}

func backoff(sleepTime int) int {
	time.Sleep(time.Duration(sleepTime) * time.Millisecond)
	return sleepTime * 2
}

func (s *solrZkInstance) initCollectionsListener() (<-chan zk.Event, error) {
	collections, version, collectionsEvents, err := s.zookeeper.GetClusterStateW()
	if err != nil {
		return nil, err
	}
	s.setCollections(collections, version)
	return collectionsEvents, nil
}

func (s *solrZkInstance) initLiveNodesListener() (<-chan zk.Event, error) {
	liveNodes, liveNodeEvents, err := s.zookeeper.GetLiveNodesW()
	if err != nil {
		return nil, err
	}
	s.setLiveNodes(liveNodes)
	return liveNodeEvents, nil
}

// GetClusterState Intentionally return a copy vs a pointer want to be thread safe
func (s *solrZkInstance) GetClusterState() (ClusterState, error) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	return s.clusterState, nil
}

func (s *solrZkInstance) setLiveNodes(nodes []string) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.LiveNodes = nodes
	s.logger.Debug(fmt.Sprintf("go-solr: zk livenodes updated %v ", s.clusterState.LiveNodes))
}

func (s *solrZkInstance) setCollections(collections map[string]Collection, version int) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.Collections = collections
	s.clusterState.Version = version
	s.logger.Debug(fmt.Sprintf("go-solr: zk collections updated %v ", s.clusterState.Collections))
}
<commit_msg>remove logging<commit_after>package solr

import (
	"fmt"
	"github.com/samuel/go-zookeeper/zk"
	"time"
)

func (s *solrZkInstance) Listen() error {
	err := s.zookeeper.Connect()
	s.zookeeper.ZKLogger(s.logger)
	s.currentNode = 0
	if err != nil {
		return err
	}
	s.clusterState = ClusterState{}
	var collectionsEvents <-chan zk.Event
	var liveNodeEvents <-chan zk.Event
	connect := func() error {
		collectionsEvents, err = s.initCollectionsListener()
		if err != nil {
			return err
		}
		liveNodeEvents, err = s.initLiveNodesListener()
		if err != nil {
			return err
		}
		return nil
	}
	err = connect()
	if err != nil {
		return err
	}
	//loop forever
	go func() {
		log := s.logger
		sleepTime := s.sleepTimeMS
		logErr := func(err error) {
			log.Error(fmt.Errorf("[go-solr] Error connecting to zk %v sleeping: %d", err, sleepTime))
		}
		for {
			select {
			case cEvent := <-collectionsEvents:
				if cEvent.Err != nil {
					log.Debug(fmt.Sprintf("[go-solr] error on cevent %v", cEvent))
					logErr(cEvent.Err)
					sleepTime = backoff(sleepTime)
					continue
				}
				// do something if its not a session or disconnect
				if cEvent.Type == zk.EventNodeDataChanged {
					collections, version, err := s.zookeeper.GetClusterState()
					if err != nil {
						logErr(err)
						sleepTime = backoff(sleepTime)
						continue
					}
					s.setCollections(collections, version)
				}
				sleepTime = s.sleepTimeMS

			case nEvent := <-liveNodeEvents:
				if nEvent.Err != nil {
					logErr(nEvent.Err)
					log.Error(fmt.Errorf("[go-solr] error on nevent %v", nEvent))
					sleepTime = backoff(sleepTime)
					continue
				}
				// do something if its not a session or disconnect
				if nEvent.Type == zk.EventNodeDataChanged || nEvent.Type == zk.EventNodeChildrenChanged {
					liveNodes, err := s.zookeeper.GetLiveNodes()
					if err != nil {
						logErr(err)
						sleepTime = backoff(sleepTime)
						continue
					}
					s.setLiveNodes(liveNodes)
				}
				sleepTime = s.sleepTimeMS
			}
			if !s.zookeeper.IsConnected() {
				err = connect()
				if err != nil {
					s.logger.Error(fmt.Errorf("[go-solr] zk connect err %v, sleeping %d", err, sleepTime))
					sleepTime = backoff(sleepTime)
				} else {
					sleepTime = s.sleepTimeMS
				}
			}
		}
	}()
	s.listening = true
	return nil
}

func backoff(sleepTime int) int {
	time.Sleep(time.Duration(sleepTime) * time.Millisecond)
	return sleepTime * 2
}

func (s *solrZkInstance) initCollectionsListener() (<-chan zk.Event, error) {
	collections, version, collectionsEvents, err := s.zookeeper.GetClusterStateW()
	if err != nil {
		return nil, err
	}
	s.setCollections(collections, version)
	return collectionsEvents, nil
}

func (s *solrZkInstance) initLiveNodesListener() (<-chan zk.Event, error) {
	liveNodes, liveNodeEvents, err := s.zookeeper.GetLiveNodesW()
	if err != nil {
		return nil, err
	}
	s.setLiveNodes(liveNodes)
	return liveNodeEvents, nil
}

// GetClusterState Intentionally return a copy vs a pointer want to be thread safe
func (s *solrZkInstance) GetClusterState() (ClusterState, error) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	return s.clusterState, nil
}

func (s *solrZkInstance) setLiveNodes(nodes []string) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.LiveNodes = nodes
	s.logger.Debug(fmt.Sprintf("go-solr: zk livenodes updated %v ", s.clusterState.LiveNodes))
}

func (s *solrZkInstance) setCollections(collections map[string]Collection, version int) {
	s.clusterStateMutex.Lock()
	defer s.clusterStateMutex.Unlock()
	s.clusterState.Collections = collections
	s.clusterState.Version = version
	s.logger.Debug(fmt.Sprintf("go-solr: zk collections updated %v ", s.clusterState.Collections))
}
<|endoftext|>


<commit_before>package pachyderm

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"testing"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/ppsutil"
	"google.golang.org/grpc"
)

func TestJob(t *testing.T) {
	t.Skip()
	dataRepo := uniqueString("TestJob.data")
	pfsClient := getPfsClient(t)
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	commit, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit.Id, "file", 0, strings.NewReader("foo"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit.Id))
	ppsClient := getPpsClient(t)
	job, err := ppsutil.CreateJob(
		ppsClient,
		"",
		[]string{"cp", path.Join("/pfs", dataRepo, "file"), "/pfs/out/file"},
		"",
		1,
		[]*pfs.Commit{commit},
		"",
	)
	require.NoError(t, err)
	listCommitRequest := &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{pps.JobRepo(job)},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err := pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits := listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	var buffer bytes.Buffer
	require.NoError(t, pfsutil.GetFile(pfsClient, pps.JobRepo(job).Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foo", buffer.String())
}

func TestGrep(t *testing.T) {
	t.Skip()
	dataRepo := uniqueString("pachyderm.TestGrep.data")
	pfsClient := getPfsClient(t)
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	commit, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	for i := 0; i < 100; i++ {
		_, err = pfsutil.PutFile(pfsClient, dataRepo, commit.Id, fmt.Sprintf("file%d", i), 0, strings.NewReader("foo\nbar\nfizz\nbuzz\n"))
		require.NoError(t, err)
	}
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit.Id))
	ppsClient := getPpsClient(t)
	_, err = ppsutil.CreateJob(
		ppsClient,
		"",
		[]string{"sh"},
		fmt.Sprintf("grep foo /pfs/%s/* >/pfs/out/foo", dataRepo),
		1,
		[]*pfs.Commit{commit},
		"",
	)
	require.NoError(t, err)
}

func TestPipeline(t *testing.T) {
	pfsClient := getPfsClient(t)
	ppsClient := getPpsClient(t)
	// create repos
	dataRepo := uniqueString("TestPipeline.data")
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	// create pipeline
	pipelineName := uniqueString("pipeline")
	outRepo := pps.PipelineRepo(ppsutil.NewPipeline(pipelineName))
	require.NoError(t, ppsutil.CreatePipeline(
		ppsClient,
		pipelineName,
		"",
		[]string{"cp", path.Join("/pfs", dataRepo, "file"), "/pfs/out/file"},
		"",
		1,
		[]*pfs.Repo{&pfs.Repo{Name: dataRepo}},
	))
	// Do first commit to repo
	log.Printf("Do first commit.")
	commit1, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit1.Id, "file", 0, strings.NewReader("foo"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit1.Id))
	listCommitRequest := &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{outRepo},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err := pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits := listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	var buffer bytes.Buffer
	log.Printf("First commit: %+v", outCommits[0])
	require.NoError(t, pfsutil.GetFile(pfsClient, outRepo.Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foo", buffer.String())
	// Do second commit to repo
	log.Printf("Do second commit.")
	commit2, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit2.Id, "file", 0, strings.NewReader("bar"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit2.Id))
	listCommitRequest = &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{outRepo},
		FromCommit: []*pfs.Commit{outCommits[0].Commit},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err = pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits = listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	buffer = bytes.Buffer{}
	log.Printf("Second commit: %+v", outCommits[0])
	require.NoError(t, pfsutil.GetFile(pfsClient, outRepo.Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foobar", buffer.String())
}

func getPfsClient(t *testing.T) pfs.APIClient {
	pfsdAddr := os.Getenv("PFSD_PORT_650_TCP_ADDR")
	if pfsdAddr == "" {
		t.Error("PFSD_PORT_650_TCP_ADDR not set")
	}
	clientConn, err := grpc.Dial(fmt.Sprintf("%s:650", pfsdAddr), grpc.WithInsecure())
	require.NoError(t, err)
	return pfs.NewAPIClient(clientConn)
}

func getPpsClient(t *testing.T) pps.APIClient {
	ppsdAddr := os.Getenv("PPSD_PORT_651_TCP_ADDR")
	if ppsdAddr == "" {
		t.Error("PPSD_PORT_651_TCP_ADDR not set")
	}
	clientConn, err := grpc.Dial(fmt.Sprintf("%s:651", ppsdAddr), grpc.WithInsecure())
	require.NoError(t, err)
	return pps.NewAPIClient(clientConn)
}

func uniqueString(prefix string) string {
	return prefix + "." + uuid.NewWithoutDashes()[0:12]
}
<commit_msg>Fix pachyderm_test. It passes now!!<commit_after>package pachyderm

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"path"
	"strings"
	"testing"

	"golang.org/x/net/context"

	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pfs/pfsutil"
	"github.com/pachyderm/pachyderm/src/pkg/require"
	"github.com/pachyderm/pachyderm/src/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/pps"
	"github.com/pachyderm/pachyderm/src/pps/ppsutil"
	"google.golang.org/grpc"
)

func TestJob(t *testing.T) {
	t.Skip()
	dataRepo := uniqueString("TestJob.data")
	pfsClient := getPfsClient(t)
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	commit, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit.Id, "file", 0, strings.NewReader("foo"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit.Id))
	ppsClient := getPpsClient(t)
	job, err := ppsutil.CreateJob(
		ppsClient,
		"",
		[]string{"cp", path.Join("/pfs", dataRepo, "file"), "/pfs/out/file"},
		"",
		1,
		[]*pfs.Commit{commit},
		"",
	)
	require.NoError(t, err)
	listCommitRequest := &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{pps.JobRepo(job)},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err := pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits := listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	var buffer bytes.Buffer
	require.NoError(t, pfsutil.GetFile(pfsClient, pps.JobRepo(job).Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foo", buffer.String())
}

func TestGrep(t *testing.T) {
	t.Skip()
	dataRepo := uniqueString("pachyderm.TestGrep.data")
	pfsClient := getPfsClient(t)
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	commit, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	for i := 0; i < 100; i++ {
		_, err = pfsutil.PutFile(pfsClient, dataRepo, commit.Id, fmt.Sprintf("file%d", i), 0, strings.NewReader("foo\nbar\nfizz\nbuzz\n"))
		require.NoError(t, err)
	}
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit.Id))
	ppsClient := getPpsClient(t)
	_, err = ppsutil.CreateJob(
		ppsClient,
		"",
		[]string{"sh"},
		fmt.Sprintf("grep foo /pfs/%s/* >/pfs/out/foo", dataRepo),
		1,
		[]*pfs.Commit{commit},
		"",
	)
	require.NoError(t, err)
}

func TestPipeline(t *testing.T) {
	pfsClient := getPfsClient(t)
	ppsClient := getPpsClient(t)
	// create repos
	dataRepo := uniqueString("TestPipeline.data")
	require.NoError(t, pfsutil.CreateRepo(pfsClient, dataRepo))
	// create pipeline
	pipelineName := uniqueString("pipeline")
	outRepo := pps.PipelineRepo(ppsutil.NewPipeline(pipelineName))
	require.NoError(t, ppsutil.CreatePipeline(
		ppsClient,
		pipelineName,
		"",
		[]string{"cp", path.Join("/pfs", dataRepo, "file"), "/pfs/out/file"},
		"",
		1,
		[]*pfs.Repo{&pfs.Repo{Name: dataRepo}},
	))
	// Do first commit to repo
	log.Printf("Do first commit.")
	commit1, err := pfsutil.StartCommit(pfsClient, dataRepo, "")
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit1.Id, "file", 0, strings.NewReader("foo"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit1.Id))
	listCommitRequest := &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{outRepo},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err := pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits := listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	var buffer bytes.Buffer
	log.Printf("First commit: %+v", outCommits[0])
	require.NoError(t, pfsutil.GetFile(pfsClient, outRepo.Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foo", buffer.String())
	// Do second commit to repo
	log.Printf("Do second commit.")
	commit2, err := pfsutil.StartCommit(pfsClient, dataRepo, commit1.Id)
	require.NoError(t, err)
	_, err = pfsutil.PutFile(pfsClient, dataRepo, commit2.Id, "file", 0, strings.NewReader("bar"))
	require.NoError(t, err)
	require.NoError(t, pfsutil.FinishCommit(pfsClient, dataRepo, commit2.Id))
	listCommitRequest = &pfs.ListCommitRequest{
		Repo:       []*pfs.Repo{outRepo},
		FromCommit: []*pfs.Commit{outCommits[0].Commit},
		CommitType: pfs.CommitType_COMMIT_TYPE_READ,
		Block:      true,
	}
	listCommitResponse, err = pfsClient.ListCommit(
		context.Background(),
		listCommitRequest,
	)
	require.NoError(t, err)
	outCommits = listCommitResponse.CommitInfo
	require.Equal(t, 1, len(outCommits))
	buffer = bytes.Buffer{}
	log.Printf("Second commit: %+v", outCommits[0])
	require.NoError(t, pfsutil.GetFile(pfsClient, outRepo.Name, outCommits[0].Commit.Id, "file", 0, 0, nil, &buffer))
	require.Equal(t, "foobar", buffer.String())
}

func getPfsClient(t *testing.T) pfs.APIClient {
	pfsdAddr := os.Getenv("PFSD_PORT_650_TCP_ADDR")
	if pfsdAddr == "" {
		t.Error("PFSD_PORT_650_TCP_ADDR not set")
	}
	clientConn, err := grpc.Dial(fmt.Sprintf("%s:650", pfsdAddr), grpc.WithInsecure())
	require.NoError(t, err)
	return pfs.NewAPIClient(clientConn)
}

func getPpsClient(t *testing.T) pps.APIClient {
	ppsdAddr := os.Getenv("PPSD_PORT_651_TCP_ADDR")
	if ppsdAddr == "" {
		t.Error("PPSD_PORT_651_TCP_ADDR not set")
	}
	clientConn, err := grpc.Dial(fmt.Sprintf("%s:651", ppsdAddr), grpc.WithInsecure())
	require.NoError(t, err)
	return pps.NewAPIClient(clientConn)
}

func uniqueString(prefix string) string {
	return prefix + "." + uuid.NewWithoutDashes()[0:12]
}
<|endoftext|>


<commit_before>package nsqd

import (
	"bufio"
	"compress/flate"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/bitly/go-nsq"
	"github.com/mreiferson/go-snappystream"
)

const DefaultBufferSize = 16 * 1024

type identifyDataV2 struct {
	ShortId string `json:"short_id"` // TODO: deprecated, remove in 1.0
	LongId  string `json:"long_id"`  // TODO: deprecated, remove in 1.0

	ClientID            string `json:"client_id"`
	Hostname            string `json:"hostname"`
	HeartbeatInterval   int    `json:"heartbeat_interval"`
	OutputBufferSize    int    `json:"output_buffer_size"`
	OutputBufferTimeout int    `json:"output_buffer_timeout"`
	FeatureNegotiation  bool   `json:"feature_negotiation"`
	TLSv1               bool   `json:"tls_v1"`
	Deflate             bool   `json:"deflate"`
	DeflateLevel        int    `json:"deflate_level"`
	Snappy              bool   `json:"snappy"`
	SampleRate          int32  `json:"sample_rate"`
	UserAgent           string `json:"user_agent"`
	MsgTimeout          int    `json:"msg_timeout"`
}

type identifyEvent struct {
	OutputBufferTimeout time.Duration
	HeartbeatInterval   time.Duration
	SampleRate          int32
	MsgTimeout          time.Duration
}

type clientV2 struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	ReadyCount     int64
	LastReadyCount int64
	InFlightCount  int64
	MessageCount   uint64
	FinishCount    uint64
	RequeueCount   uint64

	sync.RWMutex

	ID        int64
	context   *context
	UserAgent string

	// original connection
	net.Conn

	// connections based on negotiated features
	tlsConn     *tls.Conn
	flateWriter *flate.Writer

	// reading/writing interfaces
	Reader *bufio.Reader
	Writer *bufio.Writer

	OutputBufferSize    int
	OutputBufferTimeout time.Duration

	HeartbeatInterval time.Duration

	MsgTimeout time.Duration

	State          int32
	ConnectTime    time.Time
	Channel        *Channel
	ReadyStateChan chan int
	ExitChan       chan int

	ClientID string
	Hostname string

	SampleRate int32

	IdentifyEventChan chan identifyEvent
	SubEventChan      chan *Channel

	TLS     int32
	Snappy  int32
	Deflate int32

	// re-usable buffer for reading the 4-byte lengths off the wire
	lenBuf   [4]byte
	lenSlice []byte
}

func newClientV2(id int64, conn net.Conn, context *context) *clientV2 {
	var identifier string
	if conn != nil {
		identifier, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
	}

	c := &clientV2{
		ID:      id,
		context: context,

		Conn: conn,

		Reader: bufio.NewReaderSize(conn, DefaultBufferSize),
		Writer: bufio.NewWriterSize(conn, DefaultBufferSize),

		OutputBufferSize:    DefaultBufferSize,
		OutputBufferTimeout: 250 * time.Millisecond,

		MsgTimeout: context.nsqd.options.MsgTimeout,

		// ReadyStateChan has a buffer of 1 to guarantee that in the event
		// there is a race the state update is not lost
		ReadyStateChan: make(chan int, 1),
		ExitChan:       make(chan int),
		ConnectTime:    time.Now(),
		State:          nsq.StateInit,

		ClientID: identifier,
		Hostname: identifier,

		SubEventChan:      make(chan *Channel, 1),
		IdentifyEventChan: make(chan identifyEvent, 1),

		// heartbeats are client configurable but default to 30s
		HeartbeatInterval: context.nsqd.options.ClientTimeout / 2,
	}
	c.lenSlice = c.lenBuf[:]
	return c
}

func (c *clientV2) String() string {
	return c.RemoteAddr().String()
}

func (c *clientV2) Identify(data identifyDataV2) error {
	// TODO: for backwards compatibility, remove in 1.0
	hostname := data.Hostname
	if hostname == "" {
		hostname = data.LongId
	}
	// TODO: for backwards compatibility, remove in 1.0
	clientId := data.ClientID
	if clientId == "" {
		clientId = data.ShortId
	}

	c.Lock()
	c.ClientID = clientId
	c.Hostname = hostname
	c.UserAgent = data.UserAgent
	c.Unlock()

	err := c.SetHeartbeatInterval(data.HeartbeatInterval)
	if err != nil {
		return err
	}

	err = c.SetOutputBufferSize(data.OutputBufferSize)
	if err != nil {
		return err
	}

	err = c.SetOutputBufferTimeout(data.OutputBufferTimeout)
	if err != nil {
		return err
	}

	err = c.SetSampleRate(data.SampleRate)
	if err != nil {
		return err
	}

	err = c.SetMsgTimeout(data.MsgTimeout)
	if err != nil {
		return err
	}

	ie := identifyEvent{
		OutputBufferTimeout: c.OutputBufferTimeout,
		HeartbeatInterval:   c.HeartbeatInterval,
		SampleRate:          c.SampleRate,
		MsgTimeout:          c.MsgTimeout,
	}

	// update the client's message pump
	select {
	case c.IdentifyEventChan <- ie:
	default:
	}

	return nil
}

func (c *clientV2) Stats() ClientStats {
	c.RLock()
	// TODO: deprecated, remove in 1.0
	name := c.ClientID

	clientId := c.ClientID
	hostname := c.Hostname
	userAgent := c.UserAgent
	c.RUnlock()
	return ClientStats{
		// TODO: deprecated, remove in 1.0
		Name: name,

		Version:       "V2",
		RemoteAddress: c.RemoteAddr().String(),
		ClientID:      clientId,
		Hostname:      hostname,
		UserAgent:     userAgent,
		State:         atomic.LoadInt32(&c.State),
		ReadyCount:    atomic.LoadInt64(&c.ReadyCount),
		InFlightCount: atomic.LoadInt64(&c.InFlightCount),
		MessageCount:  atomic.LoadUint64(&c.MessageCount),
		FinishCount:   atomic.LoadUint64(&c.FinishCount),
		RequeueCount:  atomic.LoadUint64(&c.RequeueCount),
		ConnectTime:   c.ConnectTime.Unix(),
		SampleRate:    atomic.LoadInt32(&c.SampleRate),
		TLS:           atomic.LoadInt32(&c.TLS) == 1,
		Deflate:       atomic.LoadInt32(&c.Deflate) == 1,
		Snappy:        atomic.LoadInt32(&c.Snappy) == 1,
	}
}

func (c *clientV2) IsReadyForMessages() bool {
	if c.Channel.IsPaused() {
		return false
	}

	readyCount := atomic.LoadInt64(&c.ReadyCount)
	lastReadyCount := atomic.LoadInt64(&c.LastReadyCount)
	inFlightCount := atomic.LoadInt64(&c.InFlightCount)

	if c.context.nsqd.options.Verbose {
		log.Printf("[%s] state rdy: %4d lastrdy: %4d inflt: %4d", c,
			readyCount, lastReadyCount, inFlightCount)
	}

	if inFlightCount >= lastReadyCount || readyCount <= 0 {
		return false
	}

	return true
}

func (c *clientV2) SetReadyCount(count int64) {
	atomic.StoreInt64(&c.ReadyCount, count)
	atomic.StoreInt64(&c.LastReadyCount, count)
	c.tryUpdateReadyState()
}

func (c *clientV2) tryUpdateReadyState() {
	// you can always *try* to write to ReadyStateChan because in the cases
	// where you cannot the message pump loop would have iterated anyway.
	// the atomic integer operations guarantee correctness of the value.
	select {
	case c.ReadyStateChan <- 1:
	default:
	}
}

func (c *clientV2) FinishedMessage() {
	atomic.AddUint64(&c.FinishCount, 1)
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) Empty() {
	atomic.StoreInt64(&c.InFlightCount, 0)
	c.tryUpdateReadyState()
}

func (c *clientV2) SendingMessage() {
	atomic.AddInt64(&c.ReadyCount, -1)
	atomic.AddInt64(&c.InFlightCount, 1)
	atomic.AddUint64(&c.MessageCount, 1)
}

func (c *clientV2) TimedOutMessage() {
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) RequeuedMessage() {
	atomic.AddUint64(&c.RequeueCount, 1)
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) StartClose() {
	// Force the client into ready 0
	c.SetReadyCount(0)
	// mark this client as closing
	atomic.StoreInt32(&c.State, nsq.StateClosing)
}

func (c *clientV2) Pause() {
	c.tryUpdateReadyState()
}

func (c *clientV2) UnPause() {
	c.tryUpdateReadyState()
}

func (c *clientV2) SetHeartbeatInterval(desiredInterval int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case desiredInterval == -1:
		c.HeartbeatInterval = 0
	case desiredInterval == 0:
		// do nothing (use default)
	case desiredInterval >= 1000 &&
		desiredInterval <= int(c.context.nsqd.options.MaxHeartbeatInterval/time.Millisecond):
		c.HeartbeatInterval = time.Duration(desiredInterval) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("heartbeat interval (%d) is invalid", desiredInterval))
	}

	return nil
}

func (c *clientV2) SetOutputBufferSize(desiredSize int) error {
	var size int

	switch {
	case desiredSize == -1:
		// effectively no buffer (every write will go directly to the wrapped net.Conn)
		size = 1
	case desiredSize == 0:
		// do nothing (use default)
	case desiredSize >= 64 && desiredSize <= int(c.context.nsqd.options.MaxOutputBufferSize):
		size = desiredSize
	default:
		return errors.New(fmt.Sprintf("output buffer size (%d) is invalid", desiredSize))
	}

	if size > 0 {
		c.Lock()
		defer c.Unlock()
		c.OutputBufferSize = size
		err := c.Writer.Flush()
		if err != nil {
			return err
		}
		c.Writer = bufio.NewWriterSize(c.Conn, size)
	}

	return nil
}

func (c *clientV2) SetOutputBufferTimeout(desiredTimeout int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case desiredTimeout == -1:
		c.OutputBufferTimeout = 0
	case desiredTimeout == 0:
		// do nothing (use default)
	case desiredTimeout >= 1 &&
		desiredTimeout <= int(c.context.nsqd.options.MaxOutputBufferTimeout/time.Millisecond):
		c.OutputBufferTimeout = time.Duration(desiredTimeout) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("output buffer timeout (%d) is invalid", desiredTimeout))
	}

	return nil
}

func (c *clientV2) SetSampleRate(sampleRate int32) error {
	if sampleRate < 0 || sampleRate > 99 {
		return errors.New(fmt.Sprintf("sample rate (%d) is invalid", sampleRate))
	}
	atomic.StoreInt32(&c.SampleRate, sampleRate)
	return nil
}

func (c *clientV2) SetMsgTimeout(msgTimeout int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case msgTimeout == 0:
		// do nothing (use default)
	case msgTimeout >= 1000 &&
		msgTimeout <= int(c.context.nsqd.options.MaxMsgTimeout/time.Millisecond):
		c.MsgTimeout = time.Duration(msgTimeout) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("msg timeout (%d) is invalid", msgTimeout))
	}

	return nil
}

func (c *clientV2) UpgradeTLS() error {
	c.Lock()
	defer c.Unlock()

	tlsConn := tls.Server(c.Conn, c.context.nsqd.tlsConfig)
	err := tlsConn.Handshake()
	if err != nil {
		return err
	}
	c.tlsConn = tlsConn

	c.Reader = bufio.NewReaderSize(c.tlsConn, DefaultBufferSize)
	c.Writer = bufio.NewWriterSize(c.tlsConn, c.OutputBufferSize)

	atomic.StoreInt32(&c.TLS, 1)

	return nil
}

func (c *clientV2) UpgradeDeflate(level int) error {
	c.Lock()
	defer c.Unlock()

	conn := c.Conn
	if c.tlsConn != nil {
		conn = c.tlsConn
	}

	c.Reader = bufio.NewReaderSize(flate.NewReader(conn), DefaultBufferSize)

	fw, _ := flate.NewWriter(conn, level)
	c.flateWriter = fw
	c.Writer = bufio.NewWriterSize(fw, c.OutputBufferSize)

	atomic.StoreInt32(&c.Deflate, 1)

	return nil
}

func (c *clientV2) UpgradeSnappy() error {
	c.Lock()
	defer c.Unlock()

	conn := c.Conn
	if c.tlsConn != nil {
		conn = c.tlsConn
	}

	c.Reader = bufio.NewReaderSize(snappystream.NewReader(conn, snappystream.SkipVerifyChecksum), DefaultBufferSize)
	c.Writer = bufio.NewWriterSize(snappystream.NewWriter(conn), c.OutputBufferSize)

	atomic.StoreInt32(&c.Snappy, 1)

	return nil
}

func (c *clientV2) Flush() error {
	c.SetWriteDeadline(time.Now().Add(time.Second))

	err := c.Writer.Flush()
	if err != nil {
		return err
	}

	if c.flateWriter != nil {
		return c.flateWriter.Flush()
	}

	return nil
}
<commit_msg>nsqd: set timeout on tls handshake<commit_after>package nsqd

import (
	"bufio"
	"compress/flate"
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/bitly/go-nsq"
	"github.com/mreiferson/go-snappystream"
)

const DefaultBufferSize = 16 * 1024

type identifyDataV2 struct {
	ShortId string `json:"short_id"` // TODO: deprecated, remove in 1.0
	LongId  string `json:"long_id"`  // TODO: deprecated, remove in 1.0

	ClientID            string `json:"client_id"`
	Hostname            string `json:"hostname"`
	HeartbeatInterval   int    `json:"heartbeat_interval"`
	OutputBufferSize    int    `json:"output_buffer_size"`
	OutputBufferTimeout int    `json:"output_buffer_timeout"`
	FeatureNegotiation  bool   `json:"feature_negotiation"`
	TLSv1               bool   `json:"tls_v1"`
	Deflate             bool   `json:"deflate"`
	DeflateLevel        int    `json:"deflate_level"`
	Snappy              bool   `json:"snappy"`
	SampleRate          int32  `json:"sample_rate"`
	UserAgent           string `json:"user_agent"`
	MsgTimeout          int    `json:"msg_timeout"`
}

type identifyEvent struct {
	OutputBufferTimeout time.Duration
	HeartbeatInterval   time.Duration
	SampleRate          int32
	MsgTimeout          time.Duration
}

type clientV2 struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	ReadyCount     int64
	LastReadyCount int64
	InFlightCount  int64
	MessageCount   uint64
	FinishCount    uint64
	RequeueCount   uint64

	sync.RWMutex

	ID        int64
	context   *context
	UserAgent string

	// original connection
	net.Conn

	// connections based on negotiated features
	tlsConn     *tls.Conn
	flateWriter *flate.Writer

	// reading/writing interfaces
	Reader *bufio.Reader
	Writer *bufio.Writer

	OutputBufferSize    int
	OutputBufferTimeout time.Duration

	HeartbeatInterval time.Duration

	MsgTimeout time.Duration

	State          int32
	ConnectTime    time.Time
	Channel        *Channel
	ReadyStateChan chan int
	ExitChan       chan int

	ClientID string
	Hostname string

	SampleRate int32

	IdentifyEventChan chan identifyEvent
	SubEventChan      chan *Channel

	TLS     int32
	Snappy  int32
	Deflate int32

	// re-usable buffer for reading the 4-byte lengths off the wire
	lenBuf   [4]byte
	lenSlice []byte
}

func newClientV2(id int64, conn net.Conn, context *context) *clientV2 {
	var identifier string
	if conn != nil {
		identifier, _, _ = net.SplitHostPort(conn.RemoteAddr().String())
	}

	c := &clientV2{
		ID:      id,
		context: context,

		Conn: conn,

		Reader: bufio.NewReaderSize(conn, DefaultBufferSize),
		Writer: bufio.NewWriterSize(conn, DefaultBufferSize),

		OutputBufferSize:    DefaultBufferSize,
		OutputBufferTimeout: 250 * time.Millisecond,

		MsgTimeout: context.nsqd.options.MsgTimeout,

		// ReadyStateChan has a buffer of 1 to guarantee that in the event
		// there is a race the state update is not lost
		ReadyStateChan: make(chan int, 1),
		ExitChan:       make(chan int),
		ConnectTime:    time.Now(),
		State:          nsq.StateInit,

		ClientID: identifier,
		Hostname: identifier,

		SubEventChan:      make(chan *Channel, 1),
		IdentifyEventChan: make(chan identifyEvent, 1),

		// heartbeats are client configurable but default to 30s
		HeartbeatInterval: context.nsqd.options.ClientTimeout / 2,
	}
	c.lenSlice = c.lenBuf[:]
	return c
}

func (c *clientV2) String() string {
	return c.RemoteAddr().String()
}

func (c *clientV2) Identify(data identifyDataV2) error {
	// TODO: for backwards compatibility, remove in 1.0
	hostname := data.Hostname
	if hostname == "" {
		hostname = data.LongId
	}
	// TODO: for backwards compatibility, remove in 1.0
	clientId := data.ClientID
	if clientId == "" {
		clientId = data.ShortId
	}

	c.Lock()
	c.ClientID = clientId
	c.Hostname = hostname
	c.UserAgent = data.UserAgent
	c.Unlock()

	err := c.SetHeartbeatInterval(data.HeartbeatInterval)
	if err != nil {
		return err
	}

	err = c.SetOutputBufferSize(data.OutputBufferSize)
	if err != nil {
		return err
	}

	err = c.SetOutputBufferTimeout(data.OutputBufferTimeout)
	if err != nil {
		return err
	}

	err = c.SetSampleRate(data.SampleRate)
	if err != nil {
		return err
	}

	err = c.SetMsgTimeout(data.MsgTimeout)
	if err != nil {
		return err
	}

	ie := identifyEvent{
		OutputBufferTimeout: c.OutputBufferTimeout,
		HeartbeatInterval:   c.HeartbeatInterval,
		SampleRate:          c.SampleRate,
		MsgTimeout:          c.MsgTimeout,
	}

	// update the client's message pump
	select {
	case c.IdentifyEventChan <- ie:
	default:
	}

	return nil
}

func (c *clientV2) Stats() ClientStats {
	c.RLock()
	// TODO: deprecated, remove in 1.0
	name := c.ClientID

	clientId := c.ClientID
	hostname := c.Hostname
	userAgent := c.UserAgent
	c.RUnlock()
	return ClientStats{
		// TODO: deprecated, remove in 1.0
		Name: name,

		Version:       "V2",
		RemoteAddress: c.RemoteAddr().String(),
		ClientID:      clientId,
		Hostname:      hostname,
		UserAgent:     userAgent,
		State:         atomic.LoadInt32(&c.State),
		ReadyCount:    atomic.LoadInt64(&c.ReadyCount),
		InFlightCount: atomic.LoadInt64(&c.InFlightCount),
		MessageCount:  atomic.LoadUint64(&c.MessageCount),
		FinishCount:   atomic.LoadUint64(&c.FinishCount),
		RequeueCount:  atomic.LoadUint64(&c.RequeueCount),
		ConnectTime:   c.ConnectTime.Unix(),
		SampleRate:    atomic.LoadInt32(&c.SampleRate),
		TLS:           atomic.LoadInt32(&c.TLS) == 1,
		Deflate:       atomic.LoadInt32(&c.Deflate) == 1,
		Snappy:        atomic.LoadInt32(&c.Snappy) == 1,
	}
}

func (c *clientV2) IsReadyForMessages() bool {
	if c.Channel.IsPaused() {
		return false
	}

	readyCount := atomic.LoadInt64(&c.ReadyCount)
	lastReadyCount := atomic.LoadInt64(&c.LastReadyCount)
	inFlightCount := atomic.LoadInt64(&c.InFlightCount)

	if c.context.nsqd.options.Verbose {
		log.Printf("[%s] state rdy: %4d lastrdy: %4d inflt: %4d", c,
			readyCount, lastReadyCount, inFlightCount)
	}

	if inFlightCount >= lastReadyCount || readyCount <= 0 {
		return false
	}

	return true
}

func (c *clientV2) SetReadyCount(count int64) {
	atomic.StoreInt64(&c.ReadyCount, count)
	atomic.StoreInt64(&c.LastReadyCount, count)
	c.tryUpdateReadyState()
}

func (c *clientV2) tryUpdateReadyState() {
	// you can always *try* to write to ReadyStateChan because in the cases
	// where you cannot the message pump loop would have iterated anyway.
	// the atomic integer operations guarantee correctness of the value.
	select {
	case c.ReadyStateChan <- 1:
	default:
	}
}

func (c *clientV2) FinishedMessage() {
	atomic.AddUint64(&c.FinishCount, 1)
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) Empty() {
	atomic.StoreInt64(&c.InFlightCount, 0)
	c.tryUpdateReadyState()
}

func (c *clientV2) SendingMessage() {
	atomic.AddInt64(&c.ReadyCount, -1)
	atomic.AddInt64(&c.InFlightCount, 1)
	atomic.AddUint64(&c.MessageCount, 1)
}

func (c *clientV2) TimedOutMessage() {
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) RequeuedMessage() {
	atomic.AddUint64(&c.RequeueCount, 1)
	atomic.AddInt64(&c.InFlightCount, -1)
	c.tryUpdateReadyState()
}

func (c *clientV2) StartClose() {
	// Force the client into ready 0
	c.SetReadyCount(0)
	// mark this client as closing
	atomic.StoreInt32(&c.State, nsq.StateClosing)
}

func (c *clientV2) Pause() {
	c.tryUpdateReadyState()
}

func (c *clientV2) UnPause() {
	c.tryUpdateReadyState()
}

func (c *clientV2) SetHeartbeatInterval(desiredInterval int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case desiredInterval == -1:
		c.HeartbeatInterval = 0
	case desiredInterval == 0:
		// do nothing (use default)
	case desiredInterval >= 1000 &&
		desiredInterval <= int(c.context.nsqd.options.MaxHeartbeatInterval/time.Millisecond):
		c.HeartbeatInterval = time.Duration(desiredInterval) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("heartbeat interval (%d) is invalid", desiredInterval))
	}

	return nil
}

func (c *clientV2) SetOutputBufferSize(desiredSize int) error {
	var size int

	switch {
	case desiredSize == -1:
		// effectively no buffer (every write will go directly to the wrapped net.Conn)
		size = 1
	case desiredSize == 0:
		// do nothing (use default)
	case desiredSize >= 64 && desiredSize <= int(c.context.nsqd.options.MaxOutputBufferSize):
		size = desiredSize
	default:
		return errors.New(fmt.Sprintf("output buffer size (%d) is invalid", desiredSize))
	}

	if size > 0 {
		c.Lock()
		defer c.Unlock()
		c.OutputBufferSize = size
		err := c.Writer.Flush()
		if err != nil {
			return err
		}
		c.Writer = bufio.NewWriterSize(c.Conn, size)
	}

	return nil
}

func (c *clientV2) SetOutputBufferTimeout(desiredTimeout int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case desiredTimeout == -1:
		c.OutputBufferTimeout = 0
	case desiredTimeout == 0:
		// do nothing (use default)
	case desiredTimeout >= 1 &&
		desiredTimeout <= int(c.context.nsqd.options.MaxOutputBufferTimeout/time.Millisecond):
		c.OutputBufferTimeout = time.Duration(desiredTimeout) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("output buffer timeout (%d) is invalid", desiredTimeout))
	}

	return nil
}

func (c *clientV2) SetSampleRate(sampleRate int32) error {
	if sampleRate < 0 || sampleRate > 99 {
		return errors.New(fmt.Sprintf("sample rate (%d) is invalid", sampleRate))
	}
	atomic.StoreInt32(&c.SampleRate, sampleRate)
	return nil
}

func (c *clientV2) SetMsgTimeout(msgTimeout int) error {
	c.Lock()
	defer c.Unlock()

	switch {
	case msgTimeout == 0:
		// do nothing (use default)
	case msgTimeout >= 1000 &&
		msgTimeout <= int(c.context.nsqd.options.MaxMsgTimeout/time.Millisecond):
		c.MsgTimeout = time.Duration(msgTimeout) * time.Millisecond
	default:
		return errors.New(fmt.Sprintf("msg timeout (%d) is invalid", msgTimeout))
	}

	return nil
}

func (c *clientV2) UpgradeTLS() error {
	c.Lock()
	defer c.Unlock()

	tlsConn := tls.Server(c.Conn, c.context.nsqd.tlsConfig)
	tlsConn.SetDeadline(time.Now().Add(5 * time.Second))
	err := tlsConn.Handshake()
	if err != nil {
		return err
	}
	c.tlsConn = tlsConn

	c.Reader = bufio.NewReaderSize(c.tlsConn, DefaultBufferSize)
	c.Writer = bufio.NewWriterSize(c.tlsConn, c.OutputBufferSize)

	atomic.StoreInt32(&c.TLS, 1)

	return nil
}

func (c *clientV2) UpgradeDeflate(level int) error {
	c.Lock()
	defer c.Unlock()

	conn := c.Conn
	if c.tlsConn != nil {
		conn = c.tlsConn
	}

	c.Reader = bufio.NewReaderSize(flate.NewReader(conn), DefaultBufferSize)

	fw, _ := flate.NewWriter(conn, level)
	c.flateWriter = fw
	c.Writer = bufio.NewWriterSize(fw, c.OutputBufferSize)

	atomic.StoreInt32(&c.Deflate, 1)

	return nil
}

func (c *clientV2) UpgradeSnappy() error {
	c.Lock()
	defer c.Unlock()

	conn := c.Conn
	if c.tlsConn != nil {
		conn = c.tlsConn
	}

	c.Reader = bufio.NewReaderSize(snappystream.NewReader(conn, snappystream.SkipVerifyChecksum), DefaultBufferSize)
	c.Writer = bufio.NewWriterSize(snappystream.NewWriter(conn), c.OutputBufferSize)

	atomic.StoreInt32(&c.Snappy, 1)

	return nil
}

func (c *clientV2) Flush() error {
	c.SetWriteDeadline(time.Now().Add(time.Second))

	err := c.Writer.Flush()
	if err != nil {
		return err
	}

	if c.flateWriter != nil {
		return c.flateWriter.Flush()
	}

	return nil
}
<|endoftext|>


<commit_before>package client

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/namespaces"
	"github.com/docker/distribution/manifest/schema2"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/util/testutil/httpserver"
	"github.com/moby/buildkit/util/testutil/integration"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

func TestClientIntegration(t *testing.T) {
	integration.Run(t, []integration.Test{
		testCallDiskUsage,
		testBuildMultiMount,
		testBuildHTTPSource,
		testBuildPushAndValidate,
	})
}

func testCallDiskUsage(t *testing.T, sb integration.Sandbox) {
	t.Parallel()
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()
	_, err = c.DiskUsage(context.TODO())
	require.NoError(t, err)
}

func testBuildMultiMount(t *testing.T, sb integration.Sandbox) {
	t.Parallel()
	requiresLinux(t)
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	alpine := llb.Image("docker.io/library/alpine:latest")
	ls := alpine.Run(llb.Shlex("/bin/ls -l"))
	busybox := llb.Image("docker.io/library/busybox:latest")
	cp := ls.Run(llb.Shlex("/bin/cp -a /busybox/etc/passwd baz"))
	cp.AddMount("/busybox", busybox)

	def, err := cp.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.NoError(t, err)
}

func testBuildHTTPSource(t *testing.T, sb integration.Sandbox) {
	t.Parallel()

	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	modTime := time.Now().Add(-24 * time.Hour) // avoid falso positive with current time

	resp := httpserver.Response{
		Etag:         identity.NewID(),
		Content:      []byte("content1"),
		LastModified: &modTime,
	}

	server := httpserver.NewTestServer(map[string]httpserver.Response{
		"/foo": resp,
	})
	defer server.Close()

	// invalid URL first
	st := llb.HTTP(server.URL + "/bar")

	def, err := st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.Error(t, err)
	require.Contains(t, err.Error(), "invalid response status 404")

	// first correct request
	st = llb.HTTP(server.URL + "/foo")

	def, err = st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 1)
	require.Equal(t, server.Stats("/foo").CachedRequests, 0)

	tmpdir, err := ioutil.TempDir("", "buildkit")
	require.NoError(t, err)
	defer os.RemoveAll(tmpdir)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			exporterLocalOutputDir: tmpdir,
		},
	}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 2)
	require.Equal(t, server.Stats("/foo").CachedRequests, 1)

	dt, err := ioutil.ReadFile(filepath.Join(tmpdir, "foo"))
	require.NoError(t, err)
	require.Equal(t, []byte("content1"), dt)

	// test extra options
	st = llb.HTTP(server.URL+"/foo", llb.Filename("bar"), llb.Chmod(0741), llb.Chown(1000, 1000))

	def, err = st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			exporterLocalOutputDir: tmpdir,
		},
	}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 3)
	require.Equal(t, server.Stats("/foo").CachedRequests, 1)

	dt, err = ioutil.ReadFile(filepath.Join(tmpdir, "bar"))
	require.NoError(t, err)
	require.Equal(t, []byte("content1"), dt)

	fi, err := os.Stat(filepath.Join(tmpdir, "bar"))
	require.NoError(t, err)
	require.Equal(t, fi.ModTime().Format(http.TimeFormat), modTime.Format(http.TimeFormat))
	require.Equal(t, int(fi.Mode()&0777), 0741)

	// TODO: check that second request was marked as cached
}

func testBuildPushAndValidate(t *testing.T, sb integration.Sandbox) {
	requiresLinux(t)
	t.Parallel()
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	busybox := llb.Image("busybox:latest")
	st := llb.Scratch()

	run := func(cmd string) {
		st = busybox.Run(llb.Shlex(cmd), llb.Dir("/wd")).AddMount("/wd", st)
	}

	run(`sh -c "mkdir -p foo/sub; echo -n first > foo/sub/bar; chmod 0741 foo;"`)
	run(`sh -c "echo -n second > foo/sub/baz"`)

	def, err := st.Marshal()
	require.NoError(t, err)

	registry, err := sb.NewRegistry()
	if errors.Cause(err) == integration.ErrorRequirements {
		t.Skip(err.Error())
	}
	require.NoError(t, err)

	target := registry + "/buildkit/testpush:latest"

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterImage,
		ExporterAttrs: map[string]string{
			"name": target,
			"push": "true",
		},
	}, nil)
	require.NoError(t, err)

	// test existence of the image with next build
	firstBuild := llb.Image(target)

	def, err = firstBuild.Marshal()
	require.NoError(t, err)

	destDir, err := ioutil.TempDir("", "buildkit")
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			"output": destDir,
		},
	}, nil)
	require.NoError(t, err)

	dt, err := ioutil.ReadFile(filepath.Join(destDir, "foo/sub/bar"))
	require.NoError(t, err)
	require.Equal(t, dt, []byte("first"))

	dt, err = ioutil.ReadFile(filepath.Join(destDir, "foo/sub/baz"))
	require.NoError(t, err)
	require.Equal(t, dt, []byte("second"))

	fi, err := os.Stat(filepath.Join(destDir, "foo"))
	require.NoError(t, err)
	require.Equal(t, 0741, int(fi.Mode()&0777))

	// examine contents of exported tars (requires containerd)
	var cdAddress string
	if cd, ok := sb.(interface {
		ContainerdAddress() string
	}); !ok {
		return
	} else {
		cdAddress = cd.ContainerdAddress()
	}

	// TODO: make public pull helper function so this can be checked for standalone as well

	client, err := containerd.New(cdAddress)
	require.NoError(t, err)
	defer client.Close()

	ctx := namespaces.WithNamespace(context.Background(), "buildkit")

	img, err := client.Pull(ctx, target)
	require.NoError(t, err)

	desc, err := img.Config(ctx)
	require.NoError(t, err)

	dt, err = content.ReadBlob(ctx, img.ContentStore(), desc.Digest)
	require.NoError(t, err)

	var ociimg ocispec.Image
	err = json.Unmarshal(dt, &ociimg)
	require.NoError(t, err)

	require.NotEqual(t, "", ociimg.OS)
	require.NotEqual(t, "", ociimg.Architecture)
	require.NotEqual(t, "", ociimg.Config.WorkingDir)
	require.Equal(t, "layers", ociimg.RootFS.Type)
	require.Equal(t, 2, len(ociimg.RootFS.DiffIDs))
	require.Condition(t, func() bool {
		for _, env := range ociimg.Config.Env {
			if strings.HasPrefix(env, "PATH=") {
				return true
			}
		}
		return false
	})

	dt, err = content.ReadBlob(ctx, img.ContentStore(), img.Target().Digest)
	require.NoError(t, err)

	var mfst schema2.Manifest
	err = json.Unmarshal(dt, &mfst)
	require.NoError(t, err)

	require.Equal(t, schema2.MediaTypeManifest, mfst.MediaType)
	require.Equal(t, 2, len(mfst.Layers))

	dt, err = content.ReadBlob(ctx, img.ContentStore(), mfst.Layers[0].Digest)
	require.NoError(t, err)

	m, err := readTarToMap(dt)
	require.NoError(t, err)

	item, ok := m["foo/"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)
	require.Equal(t, 0741, int(item.header.Mode&0777))

	item, ok = m["foo/sub/"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)

	item, ok = m["foo/sub/bar"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeReg)
	require.Equal(t, []byte("first"), item.data)

	_, ok = m["foo/sub/baz"]
	require.False(t, ok)

	dt, err = content.ReadBlob(ctx, img.ContentStore(), mfst.Layers[1].Digest)
	require.NoError(t, err)

	m, err = readTarToMap(dt)
	require.NoError(t, err)

	item, ok = m["foo/sub/baz"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeReg)
	require.Equal(t, []byte("second"), item.data)

	// TODO: #154 check that the unmodified parents are still in tar
	// item, ok = m["foo/"]
	// require.True(t, ok)
	// require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)
}

func requiresLinux(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skipf("unsupported GOOS: %s", runtime.GOOS)
	}
}

type tarItem struct {
	header *tar.Header
	data   []byte
}

func readTarToMap(dt []byte) (map[string]*tarItem, error) {
	m := map[string]*tarItem{}
	gz, err := gzip.NewReader(bytes.NewBuffer(dt))
	if err != nil {
		return nil, err
	}
	defer gz.Close()
	rdr := tar.NewReader(gz)
	for {
		h, err := rdr.Next()
		if err != nil {
			if err == io.EOF {
				return m, nil
			}
			return nil, err
		}
		var dt []byte
		if h.Typeflag == tar.TypeReg {
			dt, err = ioutil.ReadAll(rdr)
			if err != nil {
				return nil, err
			}
		}
		m[h.Name] = &tarItem{header: h, data: dt}
	}
}
<commit_msg>client: add testcase for exported history array<commit_after>package client

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/namespaces"
	"github.com/docker/distribution/manifest/schema2"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/util/testutil/httpserver"
	"github.com/moby/buildkit/util/testutil/integration"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

func TestClientIntegration(t *testing.T) {
	integration.Run(t, []integration.Test{
		testCallDiskUsage,
		testBuildMultiMount,
		testBuildHTTPSource,
		testBuildPushAndValidate,
	})
}

func testCallDiskUsage(t *testing.T, sb integration.Sandbox) {
	t.Parallel()
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()
	_, err = c.DiskUsage(context.TODO())
	require.NoError(t, err)
}

func testBuildMultiMount(t *testing.T, sb integration.Sandbox) {
	t.Parallel()
	requiresLinux(t)
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	alpine := llb.Image("docker.io/library/alpine:latest")
	ls := alpine.Run(llb.Shlex("/bin/ls -l"))
	busybox := llb.Image("docker.io/library/busybox:latest")
	cp := ls.Run(llb.Shlex("/bin/cp -a /busybox/etc/passwd baz"))
	cp.AddMount("/busybox", busybox)

	def, err := cp.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.NoError(t, err)
}

func testBuildHTTPSource(t *testing.T, sb integration.Sandbox) {
	t.Parallel()

	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	modTime := time.Now().Add(-24 * time.Hour) // avoid falso positive with current time

	resp := httpserver.Response{
		Etag:         identity.NewID(),
		Content:      []byte("content1"),
		LastModified: &modTime,
	}

	server := httpserver.NewTestServer(map[string]httpserver.Response{
		"/foo": resp,
	})
	defer server.Close()

	// invalid URL first
	st := llb.HTTP(server.URL + "/bar")

	def, err := st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.Error(t, err)
	require.Contains(t, err.Error(), "invalid response status 404")

	// first correct request
	st = llb.HTTP(server.URL + "/foo")

	def, err = st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 1)
	require.Equal(t, server.Stats("/foo").CachedRequests, 0)

	tmpdir, err := ioutil.TempDir("", "buildkit")
	require.NoError(t, err)
	defer os.RemoveAll(tmpdir)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			exporterLocalOutputDir: tmpdir,
		},
	}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 2)
	require.Equal(t, server.Stats("/foo").CachedRequests, 1)

	dt, err := ioutil.ReadFile(filepath.Join(tmpdir, "foo"))
	require.NoError(t, err)
	require.Equal(t, []byte("content1"), dt)

	// test extra options
	st = llb.HTTP(server.URL+"/foo", llb.Filename("bar"), llb.Chmod(0741), llb.Chown(1000, 1000))

	def, err = st.Marshal()
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			exporterLocalOutputDir: tmpdir,
		},
	}, nil)
	require.NoError(t, err)

	require.Equal(t, server.Stats("/foo").AllRequests, 3)
	require.Equal(t, server.Stats("/foo").CachedRequests, 1)

	dt, err = ioutil.ReadFile(filepath.Join(tmpdir, "bar"))
	require.NoError(t, err)
	require.Equal(t, []byte("content1"), dt)

	fi, err := os.Stat(filepath.Join(tmpdir, "bar"))
	require.NoError(t, err)
	require.Equal(t, fi.ModTime().Format(http.TimeFormat), modTime.Format(http.TimeFormat))
	require.Equal(t, int(fi.Mode()&0777), 0741)

	// TODO: check that second request was marked as cached
}

func testBuildPushAndValidate(t *testing.T, sb integration.Sandbox) {
	requiresLinux(t)
	t.Parallel()
	c, err := New(sb.Address())
	require.NoError(t, err)
	defer c.Close()

	busybox := llb.Image("busybox:latest")
	st := llb.Scratch()

	run := func(cmd string) {
		st = busybox.Run(llb.Shlex(cmd), llb.Dir("/wd")).AddMount("/wd", st)
	}

	run(`sh -c "mkdir -p foo/sub; echo -n first > foo/sub/bar; chmod 0741 foo;"`)
	run(`true`) // this doesn't create a layer
	run(`sh -c "echo -n second > foo/sub/baz"`)

	def, err := st.Marshal()
	require.NoError(t, err)

	registry, err := sb.NewRegistry()
	if errors.Cause(err) == integration.ErrorRequirements {
		t.Skip(err.Error())
	}
	require.NoError(t, err)

	target := registry + "/buildkit/testpush:latest"

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterImage,
		ExporterAttrs: map[string]string{
			"name": target,
			"push": "true",
		},
	}, nil)
	require.NoError(t, err)

	// test existence of the image with next build
	firstBuild := llb.Image(target)

	def, err = firstBuild.Marshal()
	require.NoError(t, err)

	destDir, err := ioutil.TempDir("", "buildkit")
	require.NoError(t, err)

	err = c.Solve(context.TODO(), def, SolveOpt{
		Exporter: ExporterLocal,
		ExporterAttrs: map[string]string{
			"output": destDir,
		},
	}, nil)
	require.NoError(t, err)

	dt, err := ioutil.ReadFile(filepath.Join(destDir, "foo/sub/bar"))
	require.NoError(t, err)
	require.Equal(t, dt, []byte("first"))

	dt, err = ioutil.ReadFile(filepath.Join(destDir, "foo/sub/baz"))
	require.NoError(t, err)
	require.Equal(t, dt, []byte("second"))

	fi, err := os.Stat(filepath.Join(destDir, "foo"))
	require.NoError(t, err)
	require.Equal(t, 0741, int(fi.Mode()&0777))

	// examine contents of exported tars (requires containerd)
	var cdAddress string
	if cd, ok := sb.(interface {
		ContainerdAddress() string
	}); !ok {
		return
	} else {
		cdAddress = cd.ContainerdAddress()
	}

	// TODO: make public pull helper function so this can be checked for standalone as well

	client, err := containerd.New(cdAddress)
	require.NoError(t, err)
	defer client.Close()

	ctx := namespaces.WithNamespace(context.Background(), "buildkit")

	img, err := client.Pull(ctx, target)
	require.NoError(t, err)

	desc, err := img.Config(ctx)
	require.NoError(t, err)

	dt, err = content.ReadBlob(ctx, img.ContentStore(), desc.Digest)
	require.NoError(t, err)

	var ociimg ocispec.Image
	err = json.Unmarshal(dt, &ociimg)
	require.NoError(t, err)

	require.NotEqual(t, "", ociimg.OS)
	require.NotEqual(t, "", ociimg.Architecture)
	require.NotEqual(t, "", ociimg.Config.WorkingDir)
	require.Equal(t, "layers", ociimg.RootFS.Type)
	require.Equal(t, 2, len(ociimg.RootFS.DiffIDs))
	require.Condition(t, func() bool {
		for _, env := range ociimg.Config.Env {
			if strings.HasPrefix(env, "PATH=") {
				return true
			}
		}
		return false
	})

	require.Equal(t, 3, len(ociimg.History))
	require.Contains(t, ociimg.History[0].CreatedBy, "foo/sub/bar")
	require.Contains(t, ociimg.History[1].CreatedBy, "true")
	require.Contains(t, ociimg.History[2].CreatedBy, "foo/sub/baz")
	require.False(t, ociimg.History[0].EmptyLayer)
	require.True(t, ociimg.History[1].EmptyLayer)
	require.False(t, ociimg.History[2].EmptyLayer)

	dt, err = content.ReadBlob(ctx, img.ContentStore(), img.Target().Digest)
	require.NoError(t, err)

	var mfst schema2.Manifest
	err = json.Unmarshal(dt, &mfst)
	require.NoError(t, err)

	require.Equal(t, schema2.MediaTypeManifest, mfst.MediaType)
	require.Equal(t, 2, len(mfst.Layers))

	dt, err = content.ReadBlob(ctx, img.ContentStore(), mfst.Layers[0].Digest)
	require.NoError(t, err)

	m, err := readTarToMap(dt)
	require.NoError(t, err)

	item, ok := m["foo/"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)
	require.Equal(t, 0741, int(item.header.Mode&0777))

	item, ok = m["foo/sub/"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)

	item, ok = m["foo/sub/bar"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeReg)
	require.Equal(t, []byte("first"), item.data)

	_, ok = m["foo/sub/baz"]
	require.False(t, ok)

	dt, err = content.ReadBlob(ctx, img.ContentStore(), mfst.Layers[1].Digest)
	require.NoError(t, err)

	m, err = readTarToMap(dt)
	require.NoError(t, err)

	item, ok = m["foo/sub/baz"]
	require.True(t, ok)
	require.Equal(t, int32(item.header.Typeflag), tar.TypeReg)
	require.Equal(t, []byte("second"), item.data)

	// TODO: #154 check that the unmodified parents are still in tar
	// item, ok = m["foo/"]
	// require.True(t, ok)
	// require.Equal(t, int32(item.header.Typeflag), tar.TypeDir)
}

func requiresLinux(t *testing.T) {
	if runtime.GOOS != "linux" {
		t.Skipf("unsupported GOOS: %s", runtime.GOOS)
	}
}

type tarItem struct {
	header *tar.Header
	data   []byte
}

func readTarToMap(dt []byte) (map[string]*tarItem, error) {
	m := map[string]*tarItem{}
	gz, err := gzip.NewReader(bytes.NewBuffer(dt))
	if err != nil {
		return nil, err
	}
	defer gz.Close()
	rdr := tar.NewReader(gz)
	for {
		h, err := rdr.Next()
		if err != nil {
			if err == io.EOF {
				return m, nil
			}
			return nil, err
		}
		var dt []byte
		if h.Typeflag == tar.TypeReg {
			dt, err = ioutil.ReadAll(rdr)
			if err != nil {
				return nil, err
			}
		}
		m[h.Name] = &tarItem{header: h, data: dt}
	}
}
<|endoftext|>


<commit_before>package config

import (
	"github.com/micro/cli"
	"github.com/micro/go-micro"
	proto "github.com/micro/go-micro/config/source/service/proto"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/micro/config/db"
	_ "github.com/micro/micro/config/db/cockroach"
	_ "github.com/micro/micro/config/db/etcd"
	_ "github.com/micro/micro/config/db/memory"
	"github.com/micro/micro/config/handler"
)

var (
	Name     = "go.micro.config"
	Database = "memory"
)

func Run(c *cli.Context, srvOpts ...micro.Option) {
	if len(c.GlobalString("server_name")) > 0 {
		Name = c.GlobalString("server_name")
	}

	if len(c.String("watch_topic")) > 0 {
		handler.WatchTopic = c.String("watch_topic")
	}

	if len(c.String("database")) > 0 {
		Database = c.String("database")
	}

	srvOpts = append(srvOpts, micro.Name(Name))

	service := micro.NewService(srvOpts...)
	proto.RegisterConfigHandler(service.Server(), new(handler.Handler))

	_ = service.Server().Subscribe(service.Server().NewSubscriber(handler.WatchTopic, handler.Watcher))

	if err := db.Init(
		db.WithDBName(Database),
		db.WithUrl(c.String("database_url")),
	); err != nil {
		log.Fatalf("micro config init database error: %s", err)
	}

	if err := service.Run(); err != nil {
		log.Fatalf("micro config Run the service error: ", err)
	}
}

func Commands(options ...micro.Option) []cli.Command {
	command := cli.Command{
		Name:  "config",
		Usage: "Run the config server",
		Action: func(c *cli.Context) {
			Run(c, options...)
		},
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "namespace",
				Usage:  "Set the namespace used by the Config Service e.g. go.micro.srv.config",
				EnvVar: "MICRO_CONFIG_NAMESPACE",
			},
			cli.StringFlag{
				Name:   "database_url",
				EnvVar: "MICRO_CONFIG_DATABASE_URL",
				Usage:  "The database URL e.g root:123@(127.0.0.1:3306)/config?charset=utf8&parseTime=true&loc=Asia%2FShanghai",
			},
			cli.StringFlag{
				Name:   "database",
				EnvVar: "MICRO_CONFIG_DATABASE",
				Usage:  "The database e.g mysql(default), postgresql, but now we only support mysql and cockroach(pg).",
			},
			cli.StringFlag{
				Name:   "watch_topic",
				EnvVar: "MICRO_CONFIG_WATCH_TOPIC",
				Usage:  "watch the change event.",
			},
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []cli.Command{command}
}
<commit_msg>update cli to v2<commit_after>package config

import (
	"github.com/micro/cli/v2"
	"github.com/micro/go-micro"
	proto "github.com/micro/go-micro/config/source/service/proto"
	"github.com/micro/go-micro/util/log"
	"github.com/micro/micro/config/db"
	_ "github.com/micro/micro/config/db/cockroach"
	_ "github.com/micro/micro/config/db/etcd"
	_ "github.com/micro/micro/config/db/memory"
	"github.com/micro/micro/config/handler"
)

var (
	Name     = "go.micro.config"
	Database = "memory"
)

func Run(c *cli.Context, srvOpts ...micro.Option) {
	if len(c.String("server_name")) > 0 {
		Name = c.String("server_name")
	}

	if len(c.String("watch_topic")) > 0 {
		handler.WatchTopic = c.String("watch_topic")
	}

	if len(c.String("database")) > 0 {
		Database = c.String("database")
	}

	srvOpts = append(srvOpts, micro.Name(Name))

	service := micro.NewService(srvOpts...)
	proto.RegisterConfigHandler(service.Server(), new(handler.Handler))

	_ = service.Server().Subscribe(service.Server().NewSubscriber(handler.WatchTopic, handler.Watcher))

	if err := db.Init(
		db.WithDBName(Database),
		db.WithUrl(c.String("database_url")),
	); err != nil {
		log.Fatalf("micro config init database error: %s", err)
	}

	if err := service.Run(); err != nil {
		log.Fatalf("micro config Run the service error: ", err)
	}
}

func Commands(options ...micro.Option) []*cli.Command {
	command := &cli.Command{
		Name:  "config",
		Usage: "Run the config server",
		Action: func(ctx *cli.Context) error {
			Run(ctx, options...)
			return nil
		},
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "namespace",
				EnvVars: []string{"MICRO_CONFIG_NAMESPACE"},
				Usage:   "Set the namespace used by the Config Service e.g. go.micro.srv.config",
			},
			&cli.StringFlag{
				Name:    "database_url",
				EnvVars: []string{"MICRO_CONFIG_DATABASE_URL"},
				Usage:   "The database URL e.g root:123@(127.0.0.1:3306)/config?charset=utf8&parseTime=true&loc=Asia%2FShanghai",
			},
			&cli.StringFlag{
				Name:    "database",
				EnvVars: []string{"MICRO_CONFIG_DATABASE"},
				Usage:   "The database e.g mysql(default), postgresql, but now we only support mysql and cockroach(pg).",
			},
			&cli.StringFlag{
				Name:    "watch_topic",
				EnvVars: []string{"MICRO_CONFIG_WATCH_TOPIC"},
				Usage:   "watch the change event.",
			},
		},
	}

	for _, p := range Plugins() {
		if cmds := p.Commands(); len(cmds) > 0 {
			command.Subcommands = append(command.Subcommands, cmds...)
		}

		if flags := p.Flags(); len(flags) > 0 {
			command.Flags = append(command.Flags, flags...)
		}
	}

	return []*cli.Command{command}
}
<|endoftext|>