<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net"

	"golang.org/x/crypto/ssh"

	"github.com/tg123/sshpiper/sshpiperd/auditor"
	"github.com/tg123/sshpiper/sshpiperd/challenger"
	"github.com/tg123/sshpiper/sshpiperd/registry"
	"github.com/tg123/sshpiper/sshpiperd/upstream"
	"log"
)

type piperdConfig struct {
	ListenAddr   string `short:"l" long:"listen" description:"Listening Address" default:"0.0.0.0" env:"SSHPIPERD_LISTENADDR" ini-name:"listen-address"`
	Port         uint   `short:"p" long:"port" description:"Listening Port" default:"2222" env:"SSHPIPERD_PORT" ini-name:"listen-port"`
	PiperKeyFile string `short:"i" long:"server-key" description:"Server key file for SSH Piper" default:"/etc/ssh/ssh_host_rsa_key" env:"SSHPIPERD_SERVER_KEY" ini-name:"server-key"`

	UpstreamDriver   string `short:"u" long:"upstream-driver" description:"Upstream provider driver" default:"workingdir" env:"SSHPIPERD_UPSTREAM_DRIVER" ini-name:"upstream-driver"`
	ChallengerDriver string `short:"c" long:"challenger-driver" description:"Additional challenger name, e.g. pam, empty for no additional challenge" env:"SSHPIPERD_CHALLENGER" ini-name:"challenger-driver"`
	AuditorDriver    string `long:"auditor-driver" description:"Auditor for ssh connections piped by SSH Piper " env:"SSHPIPERD_AUDITOR" ini-name:"auditor-driver"`
}

func getAndInstall(name string, get func(n string) registry.Plugin, install func(plugin registry.Plugin) error, logger *log.Logger) error {
	if name == "" {
		return nil
	}

	p := get(name)

	if p == nil {
		return fmt.Errorf("driver %v not found", name)
	}

	err := p.Init(logger)
	if err != nil {
		return err
	}
	return install(p)
}

func installDrivers(piper *ssh.PiperConfig, config *piperdConfig, logger *log.Logger) (auditor.Provider, error) {

	// install upstreamProvider driver
	if config.UpstreamDriver == "" {
		return nil, fmt.Errorf("must provider upstream driver")
	}

	var bigbro auditor.Provider

	for _, d := range []struct {
		name    string
		get     func(n string) registry.Plugin
		install func(plugin registry.Plugin) error
	}{
		// upstream driver
		{
			config.UpstreamDriver,
			func(n string) registry.Plugin {
				return upstream.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(upstream.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("upstream driver return nil handler")
				}

				piper.FindUpstream = handler
				return nil
			},
		},
		// challenger driver
		{
			config.ChallengerDriver,
			func(n string) registry.Plugin {
				return challenger.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(challenger.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("challenger driver return nil handler")
				}

				piper.AdditionalChallenge = handler
				return nil
			},
		},
		// auditor driver
		{
			config.AuditorDriver,
			func(n string) registry.Plugin {
				return auditor.Get(n)
			},
			func(plugin registry.Plugin) error {
				bigbro = plugin.(auditor.Provider)
				return nil
			},
		},
	} {
		err := getAndInstall(d.name, d.get, d.install, logger)
		if err != nil {
			return nil, err
		}
	}

	return bigbro, nil
}

func startPiper(config *piperdConfig, logger *log.Logger) error {

	logger.Println("sshpiper is about to start")

	piper := &ssh.PiperConfig{}

	bigbro, err := installDrivers(piper, config, logger)

	privateBytes, err := ioutil.ReadFile(config.PiperKeyFile)
	if err != nil {
		return err
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return err
	}

	piper.AddHostKey(private)

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.ListenAddr, config.Port))
	if err != nil {
		return fmt.Errorf("failed to listen for connection: %v", err)
	}
	defer listener.Close()

	logger.Printf("sshpiperd started")

	for {
		c, err := listener.Accept()
		if err != nil {
			logger.Printf("failed to accept connection: %v", err)
			continue
		}

		logger.Printf("connection accepted: %v", c.RemoteAddr())
		go func() {
			p, err := ssh.NewSSHPiperConn(c, piper)

			if err != nil {
				logger.Printf("connection from %v establishing failed reason: %v", c.RemoteAddr(), err)
				return
			}

			if bigbro != nil {
				a, err := bigbro.Create(p.DownstreamConnMeta())
				if err != nil {
					logger.Printf("connection from %v failed to create auditor reason: %v", c.RemoteAddr(), err)
					return
				}
				defer a.Close()

				p.HookUpstreamMsg = a.GetUpstreamHook()
				p.HookDownstreamMsg = a.GetDownstreamHook()
			}

			err = p.Wait()
			logger.Printf("connection from %v closed reason: %v", c.RemoteAddr(), err)
		}()
	}
}
<commit_msg>panic when driver not found<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net"

	"golang.org/x/crypto/ssh"

	"github.com/tg123/sshpiper/sshpiperd/auditor"
	"github.com/tg123/sshpiper/sshpiperd/challenger"
	"github.com/tg123/sshpiper/sshpiperd/registry"
	"github.com/tg123/sshpiper/sshpiperd/upstream"
	"log"
)

type piperdConfig struct {
	ListenAddr   string `short:"l" long:"listen" description:"Listening Address" default:"0.0.0.0" env:"SSHPIPERD_LISTENADDR" ini-name:"listen-address"`
	Port         uint   `short:"p" long:"port" description:"Listening Port" default:"2222" env:"SSHPIPERD_PORT" ini-name:"listen-port"`
	PiperKeyFile string `short:"i" long:"server-key" description:"Server key file for SSH Piper" default:"/etc/ssh/ssh_host_rsa_key" env:"SSHPIPERD_SERVER_KEY" ini-name:"server-key"`

	UpstreamDriver   string `short:"u" long:"upstream-driver" description:"Upstream provider driver" default:"workingdir" env:"SSHPIPERD_UPSTREAM_DRIVER" ini-name:"upstream-driver"`
	ChallengerDriver string `short:"c" long:"challenger-driver" description:"Additional challenger name, e.g. pam, empty for no additional challenge" env:"SSHPIPERD_CHALLENGER" ini-name:"challenger-driver"`
	AuditorDriver    string `long:"auditor-driver" description:"Auditor for ssh connections piped by SSH Piper " env:"SSHPIPERD_AUDITOR" ini-name:"auditor-driver"`
}

func getAndInstall(reg, name string, get func(n string) registry.Plugin, install func(plugin registry.Plugin) error, logger *log.Logger) error {
	if name == "" {
		return nil
	}

	p := get(name)

	if p == nil {
		return fmt.Errorf("%v driver %v not found", reg, name)
	}

	err := p.Init(logger)
	if err != nil {
		return err
	}
	return install(p)
}

func installDrivers(piper *ssh.PiperConfig, config *piperdConfig, logger *log.Logger) (auditor.Provider, error) {

	// install upstreamProvider driver
	if config.UpstreamDriver == "" {
		return nil, fmt.Errorf("must provider upstream driver")
	}

	var bigbro auditor.Provider

	for _, d := range []struct {
		reg     string
		name    string
		get     func(n string) registry.Plugin
		install func(plugin registry.Plugin) error
	}{
		// upstream driver
		{
			"Upstream",
			config.UpstreamDriver,
			func(n string) registry.Plugin {
				return upstream.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(upstream.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("upstream driver return nil handler")
				}

				piper.FindUpstream = handler
				return nil
			},
		},
		// challenger driver
		{
			"Challenger",
			config.ChallengerDriver,
			func(n string) registry.Plugin {
				return challenger.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(challenger.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("challenger driver return nil handler")
				}

				piper.AdditionalChallenge = handler
				return nil
			},
		},
		// auditor driver
		{
			"Auditor",
			config.AuditorDriver,
			func(n string) registry.Plugin {
				return auditor.Get(n)
			},
			func(plugin registry.Plugin) error {
				bigbro = plugin.(auditor.Provider)
				return nil
			},
		},
	} {
		err := getAndInstall(d.reg, d.name, d.get, d.install, logger)
		if err != nil {
			return nil, err
		}
	}

	return bigbro, nil
}

func startPiper(config *piperdConfig, logger *log.Logger) error {

	logger.Println("sshpiper is about to start")

	piper := &ssh.PiperConfig{}

	bigbro, err := installDrivers(piper, config, logger)
	if err != nil {
		return err
	}

	privateBytes, err := ioutil.ReadFile(config.PiperKeyFile)
	if err != nil {
		return err
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return err
	}

	piper.AddHostKey(private)

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.ListenAddr, config.Port))
	if err != nil {
		return fmt.Errorf("failed to listen for connection: %v", err)
	}
	defer listener.Close()

	logger.Printf("sshpiperd started")

	for {
		c, err := listener.Accept()
		if err != nil {
			logger.Printf("failed to accept connection: %v", err)
			continue
		}

		logger.Printf("connection accepted: %v", c.RemoteAddr())
		go func() {
			p, err := ssh.NewSSHPiperConn(c, piper)

			if err != nil {
				logger.Printf("connection from %v establishing failed reason: %v", c.RemoteAddr(), err)
				return
			}

			if bigbro != nil {
				a, err := bigbro.Create(p.DownstreamConnMeta())
				if err != nil {
					logger.Printf("connection from %v failed to create auditor reason: %v", c.RemoteAddr(), err)
					return
				}
				defer a.Close()

				p.HookUpstreamMsg = a.GetUpstreamHook()
				p.HookDownstreamMsg = a.GetDownstreamHook()
			}

			err = p.Wait()
			logger.Printf("connection from %v closed reason: %v", c.RemoteAddr(), err)
		}()
	}
}
<|endoftext|>


<commit_before>package webServer

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"github.com/FidelityInternational/etcd-leader-monitor/bosh"
	"github.com/FidelityInternational/etcd-leader-monitor/etcd"
	"github.com/caarlos0/env"
	"github.com/cloudfoundry-community/gogobosh"
	"net/http"
)

// Controller struct
type Controller struct {
	BoshClient     *gogobosh.Client
	EtcdHTTPClient *http.Client
}

// Config struct
type Config struct {
	CfDeploymentName    string `env:"CF_DEPLOYMENT_NAME" envDefault:"cf-"`
	EtcdJobName         string `env:"ETCD_JOB_NAME" envDefault:"etcd_server"`
	SSLEnabled          bool   `env:"SSL_ENABLED" envDefault:false`
	SkipSSLVerification bool   `env:"SKIP_SSL_VERIFICATION" envDefault:false`
}

// CreateController - returns a populated controller object
func CreateController(boshClient *gogobosh.Client, etcdHTTPClient *http.Client) *Controller {
	return &Controller{
		BoshClient:     boshClient,
		EtcdHTTPClient: etcdHTTPClient,
	}
}

// CheckLeaders - checks if leaders are in a healthy state
func (c *Controller) CheckLeaders(w http.ResponseWriter, r *http.Request) {
	var etcdProtocol = `http`

	fmt.Println("Checking Leaders...")
	fmt.Println("Fetching Bosh deployment...")
	deployments, err := c.BoshClient.GetDeployments()
	if err != nil {
		errorPrint(err, w)
		return
	}
	deployconfig := Config{}
	env.Parse(&deployconfig)
	deployment := bosh.FindDeployment(deployments, fmt.Sprintf("^%s*", deployconfig.CfDeploymentName))
	fmt.Println("Found deployment: ", deployment)
	if deployconfig.SSLEnabled {
		etcdProtocol = "https"
		err = c.LoadCerts(deployconfig, deployment)
		if err != nil {
			errorPrint(err, w)
			return
		}
	}
	fmt.Println("Fetching Etcd IPs from BOSH...")
	boshVMs, err := c.BoshClient.GetDeploymentVMs(deployment)
	if err != nil {
		errorPrint(err, w)
		return
	}
	etcdVMs := bosh.FindVMs(boshVMs, fmt.Sprintf("^%s*", deployconfig.EtcdJobName))
	fmt.Println("Found Etcd VMs")
	httpResponseMessage, err := c.etcdProcess(etcdVMs, etcdProtocol, w)
	if err != nil {
		errorPrint(err, w)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, httpResponseMessage)
}

// LoadCerts - downloads certs from BOSH and configures the EtcdHTTPClient appropriately
func (c *Controller) LoadCerts(deployconfig Config, deployment string) error {
	fmt.Println("Fetching Etcd Certs...")
	boshDeployment, err := c.BoshClient.GetDeployment(deployment)
	if err != nil {
		return err
	}
	etcdCerts := bosh.GetEtcdCerts(boshDeployment.Manifest, fmt.Sprintf("^%s*", deployconfig.EtcdJobName))
	if etcdCerts.ClientKey == "" {
		return fmt.Errorf("Etcd Client Key was blank")
	}
	if etcdCerts.ClientCert == "" {
		return fmt.Errorf("Etcd Client Cert was blank")
	}
	caCert := x509.NewCertPool()
	if !deployconfig.SkipSSLVerification {
		if etcdCerts.CaCert == "" {
			return fmt.Errorf("Etcd CA Cert was blank")
		}
		if !caCert.AppendCertsFromPEM([]byte(etcdCerts.CaCert)) {
			return fmt.Errorf("Could not add CA Cert, CA Cert was likely invalid")
		}
	}
	clientCert, err := tls.X509KeyPair([]byte(etcdCerts.ClientCert), []byte(etcdCerts.ClientKey))
	if err != nil {
		return err
	}
	tlsConfig := &tls.Config{
		RootCAs:            caCert,
		Certificates:       []tls.Certificate{clientCert},
		InsecureSkipVerify: deployconfig.SkipSSLVerification,
	}
	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
	}
	c.EtcdHTTPClient.Transport = tr
	return nil
}

func (c *Controller) etcdProcess(etcdVMs []gogobosh.VM, etcdProtocol string, w http.ResponseWriter) (string, error) {
	var (
		leaderInfo          map[bool]int
		leaderList          map[string]map[bool]int
		leaderCount         int
		httpResponseMessage string
	)

	leaderList = make(map[string]map[bool]int)
	for _, etcdVM := range etcdVMs {
		etcdConfig := &etcd.Config{
			EtcdIP:       etcdVM.IPs[0],
			HTTPClient:   c.EtcdHTTPClient,
			EtcdProtocol: etcdProtocol,
		}
		etcdClient := etcd.NewClient(etcdConfig)
		leader, followers, err := etcdClient.GetLeaderStats()
		if err != nil {
			return "", err
		}
		leaderInfo = make(map[bool]int)
		leaderInfo[leader] = followers
		leaderList[etcdVM.IPs[0]] = leaderInfo
	}
	for _, leaderItem := range leaderList {
		for leader, followers := range leaderItem {
			if leader == true {
				leaderCount++
				if followers != (len(etcdVMs) - 1) {
					httpResponseMessage = `{"healthy": false, "message": "Incorrect number of followers"}`
				}
			}
		}
	}
	if leaderCount > 1 {
		fmt.Println("More than one etcd leader detected, number of leaders: ", leaderCount)
		httpResponseMessage = `{"healthy": false, "message": "Too many leaders"}`
	} else if leaderCount == 0 {
		fmt.Println("Not enough etcd leaders detected, number of leaders: ", leaderCount)
		httpResponseMessage = `{"healthy": false, "message": "Not enough leaders"}`
	} else if httpResponseMessage == "" {
		httpResponseMessage = `{"healthy": true, "message": "Everything is healthy"}`
	}
	return httpResponseMessage, nil
}

func errorPrint(err error, w http.ResponseWriter) {
	if err != nil {
		fmt.Println("An error occured:")
		fmt.Println(err.Error())
		w.WriteHeader(http.StatusInternalServerError)
	}
}
<commit_msg>appease gofmt and fix bad spelling<commit_after>package webServer

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"github.com/FidelityInternational/etcd-leader-monitor/bosh"
	"github.com/FidelityInternational/etcd-leader-monitor/etcd"
	"github.com/caarlos0/env"
	"github.com/cloudfoundry-community/gogobosh"
	"net/http"
)

// Controller struct
type Controller struct {
	BoshClient     *gogobosh.Client
	EtcdHTTPClient *http.Client
}

// Config struct
type Config struct {
	CfDeploymentName    string `env:"CF_DEPLOYMENT_NAME" envDefault:"cf-"`
	EtcdJobName         string `env:"ETCD_JOB_NAME" envDefault:"etcd_server"`
	SSLEnabled          bool   `env:"SSL_ENABLED" envDefault:"false"`
	SkipSSLVerification bool   `env:"SKIP_SSL_VERIFICATION" envDefault:"false"`
}

// CreateController - returns a populated controller object
func CreateController(boshClient *gogobosh.Client, etcdHTTPClient *http.Client) *Controller {
	return &Controller{
		BoshClient:     boshClient,
		EtcdHTTPClient: etcdHTTPClient,
	}
}

// CheckLeaders - checks if leaders are in a healthy state
func (c *Controller) CheckLeaders(w http.ResponseWriter, r *http.Request) {
	var etcdProtocol = `http`

	fmt.Println("Checking Leaders...")
	fmt.Println("Fetching Bosh deployment...")
	deployments, err := c.BoshClient.GetDeployments()
	if err != nil {
		errorPrint(err, w)
		return
	}
	deployconfig := Config{}
	env.Parse(&deployconfig)
	deployment := bosh.FindDeployment(deployments, fmt.Sprintf("^%s*", deployconfig.CfDeploymentName))
	fmt.Println("Found deployment: ", deployment)
	if deployconfig.SSLEnabled {
		etcdProtocol = "https"
		err = c.LoadCerts(deployconfig, deployment)
		if err != nil {
			errorPrint(err, w)
			return
		}
	}
	fmt.Println("Fetching Etcd IPs from BOSH...")
	boshVMs, err := c.BoshClient.GetDeploymentVMs(deployment)
	if err != nil {
		errorPrint(err, w)
		return
	}
	etcdVMs := bosh.FindVMs(boshVMs, fmt.Sprintf("^%s*", deployconfig.EtcdJobName))
	fmt.Println("Found Etcd VMs")
	httpResponseMessage, err := c.etcdProcess(etcdVMs, etcdProtocol, w)
	if err != nil {
		errorPrint(err, w)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, httpResponseMessage)
}

// LoadCerts - downloads certs from BOSH and configures the EtcdHTTPClient appropriately
func (c *Controller) LoadCerts(deployconfig Config, deployment string) error {
	fmt.Println("Fetching Etcd Certs...")
	boshDeployment, err := c.BoshClient.GetDeployment(deployment)
	if err != nil {
		return err
	}
	etcdCerts := bosh.GetEtcdCerts(boshDeployment.Manifest, fmt.Sprintf("^%s*", deployconfig.EtcdJobName))
	if etcdCerts.ClientKey == "" {
		return fmt.Errorf("Etcd Client Key was blank")
	}
	if etcdCerts.ClientCert == "" {
		return fmt.Errorf("Etcd Client Cert was blank")
	}
	caCert := x509.NewCertPool()
	if !deployconfig.SkipSSLVerification {
		if etcdCerts.CaCert == "" {
			return fmt.Errorf("Etcd CA Cert was blank")
		}
		if !caCert.AppendCertsFromPEM([]byte(etcdCerts.CaCert)) {
			return fmt.Errorf("Could not add CA Cert, CA Cert was likely invalid")
		}
	}
	clientCert, err := tls.X509KeyPair([]byte(etcdCerts.ClientCert), []byte(etcdCerts.ClientKey))
	if err != nil {
		return err
	}
	tlsConfig := &tls.Config{
		RootCAs:            caCert,
		Certificates:       []tls.Certificate{clientCert},
		InsecureSkipVerify: deployconfig.SkipSSLVerification,
	}
	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
	}
	c.EtcdHTTPClient.Transport = tr
	return nil
}

func (c *Controller) etcdProcess(etcdVMs []gogobosh.VM, etcdProtocol string, w http.ResponseWriter) (string, error) {
	var (
		leaderInfo          map[bool]int
		leaderList          map[string]map[bool]int
		leaderCount         int
		httpResponseMessage string
	)

	leaderList = make(map[string]map[bool]int)
	for _, etcdVM := range etcdVMs {
		etcdConfig := &etcd.Config{
			EtcdIP:       etcdVM.IPs[0],
			HTTPClient:   c.EtcdHTTPClient,
			EtcdProtocol: etcdProtocol,
		}
		etcdClient := etcd.NewClient(etcdConfig)
		leader, followers, err := etcdClient.GetLeaderStats()
		if err != nil {
			return "", err
		}
		leaderInfo = make(map[bool]int)
		leaderInfo[leader] = followers
		leaderList[etcdVM.IPs[0]] = leaderInfo
	}
	for _, leaderItem := range leaderList {
		for leader, followers := range leaderItem {
			if leader == true {
				leaderCount++
				if followers != (len(etcdVMs) - 1) {
					httpResponseMessage = `{"healthy": false, "message": "Incorrect number of followers"}`
				}
			}
		}
	}
	if leaderCount > 1 {
		fmt.Println("More than one etcd leader detected, number of leaders: ", leaderCount)
		httpResponseMessage = `{"healthy": false, "message": "Too many leaders"}`
	} else if leaderCount == 0 {
		fmt.Println("Not enough etcd leaders detected, number of leaders: ", leaderCount)
		httpResponseMessage = `{"healthy": false, "message": "Not enough leaders"}`
	} else if httpResponseMessage == "" {
		httpResponseMessage = `{"healthy": true, "message": "Everything is healthy"}`
	}
	return httpResponseMessage, nil
}

func errorPrint(err error, w http.ResponseWriter) {
	if err != nil {
		fmt.Println("An error occurred:")
		fmt.Println(err.Error())
		w.WriteHeader(http.StatusInternalServerError)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Linux Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package image

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/opencontainers/image-spec/schema"
	"github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

func layoutValidate(w walker) error {
	var blobsExist, indexExist, layoutExist bool

	if err := w.walk(func(path string, info os.FileInfo, r io.Reader) error {
		if strings.EqualFold(path, "blobs") {
			blobsExist = true
			if !info.IsDir() {
				return fmt.Errorf("blobs is not a directory")
			}

			return nil
		}

		if strings.EqualFold(path, "index.json") {
			indexExist = true
			if info.IsDir() {
				return fmt.Errorf("index.json is a directory")
			}

			var index v1.Index
			buf, err := ioutil.ReadAll(r)
			if err != nil {
				return errors.Wrap(err, "error reading index.json")
			}

			if err := json.Unmarshal(buf, &index); err != nil {
				return errors.Wrap(err, "index.json format mismatch")
			}

			return nil
		}

		if strings.EqualFold(path, "oci-layout") {
			layoutExist = true
			if info.IsDir() {
				return fmt.Errorf("oci-layout is a directory")
			}

			var imageLayout v1.ImageLayout
			buf, err := ioutil.ReadAll(r)
			if err != nil {
				return errors.Wrap(err, "error reading oci-layout")
			}

			if err := schema.ValidatorMediaTypeLayoutHeader.Validate(bytes.NewReader(buf)); err != nil {
				return errors.Wrap(err, "oci-layout: imageLayout validation failed")
			}

			if err := json.Unmarshal(buf, &imageLayout); err != nil {
				return errors.Wrap(err, "oci-layout format mismatch")
			}

			return nil
		}

		return nil
	}); err != nil {
		return err
	}

	if !blobsExist {
		return fmt.Errorf("image layout must contain blobs directory")
	}

	if !indexExist {
		return fmt.Errorf("image layout must contain index.json file")
	}

	if !layoutExist {
		return fmt.Errorf("image layout must contain oci-layout file")
	}

	return nil
}
<commit_msg>image: add validation to index.json<commit_after>// Copyright 2016 The Linux Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package image

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strings"

	"github.com/opencontainers/image-spec/schema"
	"github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

func layoutValidate(w walker) error {
	var blobsExist, indexExist, layoutExist bool

	if err := w.walk(func(path string, info os.FileInfo, r io.Reader) error {
		if strings.EqualFold(path, "blobs") {
			blobsExist = true
			if !info.IsDir() {
				return fmt.Errorf("blobs is not a directory")
			}

			return nil
		}

		if strings.EqualFold(path, "index.json") {
			indexExist = true
			if info.IsDir() {
				return fmt.Errorf("index.json is a directory")
			}

			buf, err := ioutil.ReadAll(r)
			if err != nil {
				return errors.Wrap(err, "error reading index.json")
			}

			if err := schema.ValidatorMediaTypeImageIndex.Validate(bytes.NewReader(buf)); err != nil {
				return errors.Wrap(err, "index.json validation failed")
			}

			return nil
		}

		if strings.EqualFold(path, "oci-layout") {
			layoutExist = true
			if info.IsDir() {
				return fmt.Errorf("oci-layout is a directory")
			}

			var imageLayout v1.ImageLayout
			buf, err := ioutil.ReadAll(r)
			if err != nil {
				return errors.Wrap(err, "error reading oci-layout")
			}

			if err := schema.ValidatorMediaTypeLayoutHeader.Validate(bytes.NewReader(buf)); err != nil {
				return errors.Wrap(err, "oci-layout validation failed")
			}

			if err := json.Unmarshal(buf, &imageLayout); err != nil {
				return errors.Wrap(err, "oci-layout format mismatch")
			}

			return nil
		}

		return nil
	}); err != nil {
		return err
	}

	if !blobsExist {
		return fmt.Errorf("image layout must contain blobs directory")
	}

	if !indexExist {
		return fmt.Errorf("image layout must contain index.json file")
	}

	if !layoutExist {
		return fmt.Errorf("image layout must contain oci-layout file")
	}

	return nil
}
<|endoftext|>


<commit_before><commit_msg>rename parameter<commit_after><|endoftext|>


<commit_before>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package localstore

import (
	"runtime/debug"
	"sync"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/store/localstore/engine"
	"github.com/twinj/uuid"
)

var (
	_ kv.Storage = (*dbStore)(nil)
)

type op int

const (
	opSeek = iota + 1
	opCommit
)

type command struct {
	op    op
	txn   *dbTxn
	args  interface{}
	reply interface{}
	done  chan error
}

type seekReply struct {
	key   []byte
	value []byte
	err   error
}

type commitReply struct {
	err error
}

type seekArgs struct {
	key []byte
}

type commitArgs struct {
}

//scheduler

// Seek searches for the first key in the engine which is >= key in byte order, returns (nil, nil, ErrNotFound)
// if such key is not found.
func (s *dbStore) Seek(key []byte) ([]byte, []byte, error) {
	c := &command{
		op:   opSeek,
		args: &seekArgs{key: key},
		done: make(chan error, 1),
	}

	s.commandCh <- c
	err := <-c.done
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	reply := c.reply.(*seekReply)
	return reply.key, reply.value, nil
}

// Commit writes the changed data in Batch.
func (s *dbStore) CommitTxn(txn *dbTxn) error {
	if len(txn.snapshotVals) == 0 {
		return nil
	}
	c := &command{
		op:   opCommit,
		txn:  txn,
		args: &commitArgs{},
		done: make(chan error, 1),
	}

	s.commandCh <- c
	err := <-c.done
	return errors.Trace(err)
}

func (s *dbStore) seekWorker(seekCh chan *command) {
	for {
		var pending []*command
		select {
		case cmd, ok := <-seekCh:
			if !ok {
				return
			}
			pending = append(pending, cmd)
		L:
			for {
				select {
				case cmd, ok := <-seekCh:
					if !ok {
						break L
					}
					pending = append(pending, cmd)
				default:
					break L
				}
			}
		}

		s.doSeek(pending)
	}
}

func (s *dbStore) scheduler() {
	closed := false
	seekCh := make(chan *command, 1000)
	for i := 0; i < 3; i++ {
		go s.seekWorker(seekCh)
	}

	for {
		select {
		case cmd := <-s.commandCh:
			if closed {
				cmd.done <- ErrDBClosed
				continue
			}
			switch cmd.op {
			case opSeek:
				seekCh <- cmd
			case opCommit:
				s.doCommit(cmd)
			}
		case <-s.closeCh:
			closed = true
			s.wg.Done()
			// notify seek worker to exit
			close(seekCh)
		}
	}
}

func (s *dbStore) tryLock(txn *dbTxn) (err error) {
	// check conflict
	for k := range txn.snapshotVals {
		if _, ok := s.keysLocked[k]; ok {
			return errors.Trace(kv.ErrLockConflict)
		}

		lastVer, ok := s.recentUpdates[k]
		if !ok {
			continue
		}
		// If there's newer version of this key, returns error.
		if lastVer.Cmp(kv.Version{Ver: txn.tid}) > 0 {
			return errors.Trace(kv.ErrConditionNotMatch)
		}
	}

	// record
	for k := range txn.snapshotVals {
		s.keysLocked[k] = txn.tid
	}

	return nil
}

func (s *dbStore) doCommit(cmd *command) {
	txn := cmd.txn
	curVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		log.Fatal(err)
	}
	err = s.tryLock(txn)
	if err != nil {
		cmd.done <- errors.Trace(err)
		return
	}
	// Update commit version.
	txn.version = curVer
	b := s.db.NewBatch()
	txn.WalkBuffer(func(k kv.Key, value []byte) error {
		mvccKey := MvccEncodeVersionKey(kv.Key(k), curVer)
		if len(value) == 0 { // Deleted marker
			b.Put(mvccKey, nil)
		} else {
			b.Put(mvccKey, value)
		}
		return nil
	})
	err = s.writeBatch(b)
	s.unLockKeys(txn)
	cmd.done <- errors.Trace(err)
}

func (s *dbStore) doSeek(seekCmds []*command) {
	keys := make([][]byte, 0, len(seekCmds))
	for _, cmd := range seekCmds {
		keys = append(keys, cmd.args.(*seekArgs).key)
	}

	results := s.db.MultiSeek(keys)

	for i, cmd := range seekCmds {
		reply := &seekReply{}
		var err error
		reply.key, reply.value, err = results[i].Key, results[i].Value, results[i].Err
		cmd.reply = reply
		cmd.done <- errors.Trace(err)
	}
}

func (s *dbStore) NewBatch() engine.Batch {
	return s.db.NewBatch()
}

//end of scheduler

type dbStore struct {
	db engine.DB

	txns          map[uint64]*dbTxn
	keysLocked    map[string]uint64
	recentUpdates map[string]kv.Version
	uuid          string
	path          string
	compactor     *localstoreCompactor
	wg            *sync.WaitGroup

	commandCh chan *command
	closeCh   chan struct{}

	mu     sync.Mutex
	closed bool
}

type storeCache struct {
	mu    sync.Mutex
	cache map[string]*dbStore
}

var (
	globalVersionProvider kv.VersionProvider
	mc                    storeCache

	// ErrDBClosed is the error meaning db is closed and we can use it anymore.
	ErrDBClosed = errors.New("db is closed")
)

func init() {
	mc.cache = make(map[string]*dbStore)
	globalVersionProvider = &LocalVersionProvider{}
}

// Driver implements kv.Driver interface.
type Driver struct {
	// engine.Driver is the engine driver for different local db engine.
	engine.Driver
}

// IsLocalStore checks whether a storage is local or not.
func IsLocalStore(s kv.Storage) bool {
	_, ok := s.(*dbStore)
	return ok
}

// Open opens or creates a storage with specific format for a local engine Driver.
func (d Driver) Open(schema string) (kv.Storage, error) {
	mc.mu.Lock()
	defer mc.mu.Unlock()

	if store, ok := mc.cache[schema]; ok {
		// TODO: check the cache store has the same engine with this Driver.
		log.Info("[kv] cache store", schema)
		return store, nil
	}

	db, err := d.Driver.Open(schema)
	if err != nil {
		return nil, errors.Trace(err)
	}

	log.Info("[kv] New store", schema)
	s := &dbStore{
		txns:          make(map[uint64]*dbTxn),
		keysLocked:    make(map[string]uint64),
		uuid:          uuid.NewV4().String(),
		path:          schema,
		db:            db,
		compactor:     newLocalCompactor(localCompactDefaultPolicy, db),
		commandCh:     make(chan *command, 1000),
		closed:        false,
		closeCh:       make(chan struct{}),
		recentUpdates: make(map[string]kv.Version),
		wg:            &sync.WaitGroup{},
	}
	mc.cache[schema] = s
	s.compactor.Start()
	s.wg.Add(1)
	go s.scheduler()
	return s, nil
}

func (s *dbStore) UUID() string {
	return s.uuid
}

func (s *dbStore) GetSnapshot(ver kv.Version) (kv.Snapshot, error) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil, ErrDBClosed
	}
	s.mu.Unlock()

	currentVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ver.Cmp(currentVer) > 0 {
		ver = currentVer
	}

	return &dbSnapshot{
		store:   s,
		version: ver,
	}, nil
}

func (s *dbStore) CurrentVersion() (kv.Version, error) {
	return globalVersionProvider.CurrentVersion()
}

// Begin transaction
func (s *dbStore) Begin() (kv.Transaction, error) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil, ErrDBClosed
	}
	s.mu.Unlock()

	beginVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		return nil, errors.Trace(err)
	}

	// todo: send to chan
	txn := &dbTxn{
		tid:          beginVer.Ver,
		valid:        true,
		store:        s,
		version:      kv.MinVersion,
		snapshotVals: make(map[string]struct{}),
	}
	log.Debugf("[kv] Begin txn:%d", txn.tid)
	txn.UnionStore = kv.NewUnionStore(newSnapshot(s, beginVer))
	return txn, nil
}

func (s *dbStore) Close() error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return ErrDBClosed
	}

	s.closed = true
	s.mu.Unlock()

	mc.mu.Lock()
	defer mc.mu.Unlock()
	s.compactor.Stop()
	s.closeCh <- struct{}{}
	s.wg.Wait()
	delete(mc.cache, s.path)
	return s.db.Close()
}

func (s *dbStore) writeBatch(b engine.Batch) error {
	if b.Len() == 0 {
		return nil
	}

	if s.closed {
		return errors.Trace(ErrDBClosed)
	}

	err := s.db.Commit(b)
	if err != nil {
		log.Error(err)
		return errors.Trace(err)
	}

	return nil
}

func (s *dbStore) newBatch() engine.Batch {
	return s.db.NewBatch()
}
func (s *dbStore) unLockKeys(txn *dbTxn) error {
	for k := range txn.snapshotVals {
		if tid, ok := s.keysLocked[k]; !ok || tid != txn.tid {
			debug.PrintStack()
			log.Fatalf("should never happend:%v, %v", tid, txn.tid)
		}

		delete(s.keysLocked, k)
		s.recentUpdates[k] = txn.version
	}

	return nil
}
<commit_msg>Add todo<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package localstore

import (
	"runtime/debug"
	"sync"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/store/localstore/engine"
	"github.com/twinj/uuid"
)

var (
	_ kv.Storage = (*dbStore)(nil)
)

type op int

const (
	opSeek = iota + 1
	opCommit
)

type command struct {
	op    op
	txn   *dbTxn
	args  interface{}
	reply interface{}
	done  chan error
}

type seekReply struct {
	key   []byte
	value []byte
	err   error
}

type commitReply struct {
	err error
}

type seekArgs struct {
	key []byte
}

type commitArgs struct {
}

//scheduler

// Seek searches for the first key in the engine which is >= key in byte order, returns (nil, nil, ErrNotFound)
// if such key is not found.
func (s *dbStore) Seek(key []byte) ([]byte, []byte, error) {
	c := &command{
		op:   opSeek,
		args: &seekArgs{key: key},
		done: make(chan error, 1),
	}

	s.commandCh <- c
	err := <-c.done
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	reply := c.reply.(*seekReply)
	return reply.key, reply.value, nil
}

// Commit writes the changed data in Batch.
func (s *dbStore) CommitTxn(txn *dbTxn) error {
	if len(txn.snapshotVals) == 0 {
		return nil
	}
	c := &command{
		op:   opCommit,
		txn:  txn,
		args: &commitArgs{},
		done: make(chan error, 1),
	}

	s.commandCh <- c
	err := <-c.done
	return errors.Trace(err)
}

func (s *dbStore) seekWorker(seekCh chan *command) {
	for {
		var pending []*command
		select {
		case cmd, ok := <-seekCh:
			if !ok {
				return
			}
			pending = append(pending, cmd)
		L:
			for {
				select {
				case cmd, ok := <-seekCh:
					if !ok {
						break L
					}
					pending = append(pending, cmd)
				default:
					break L
				}
			}
		}

		s.doSeek(pending)
	}
}

func (s *dbStore) scheduler() {
	closed := false
	seekCh := make(chan *command, 1000)
	for i := 0; i < 3; i++ {
		go s.seekWorker(seekCh)
	}

	for {
		select {
		case cmd := <-s.commandCh:
			if closed {
				cmd.done <- ErrDBClosed
				continue
			}
			switch cmd.op {
			case opSeek:
				seekCh <- cmd
			case opCommit:
				s.doCommit(cmd)
			}
		case <-s.closeCh:
			closed = true
			s.wg.Done()
			// notify seek worker to exit
			close(seekCh)
		}
	}
}

func (s *dbStore) tryLock(txn *dbTxn) (err error) {
	// check conflict
	for k := range txn.snapshotVals {
		if _, ok := s.keysLocked[k]; ok {
			return errors.Trace(kv.ErrLockConflict)
		}

		lastVer, ok := s.recentUpdates[k]
		if !ok {
			continue
		}
		// If there's newer version of this key, returns error.
		if lastVer.Cmp(kv.Version{Ver: txn.tid}) > 0 {
			return errors.Trace(kv.ErrConditionNotMatch)
		}
	}

	// record
	for k := range txn.snapshotVals {
		s.keysLocked[k] = txn.tid
	}

	return nil
}

func (s *dbStore) doCommit(cmd *command) {
	txn := cmd.txn
	curVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		log.Fatal(err)
	}
	err = s.tryLock(txn)
	if err != nil {
		cmd.done <- errors.Trace(err)
		return
	}
	// Update commit version.
	txn.version = curVer
	b := s.db.NewBatch()
	txn.WalkBuffer(func(k kv.Key, value []byte) error {
		mvccKey := MvccEncodeVersionKey(kv.Key(k), curVer)
		if len(value) == 0 { // Deleted marker
			b.Put(mvccKey, nil)
		} else {
			b.Put(mvccKey, value)
		}
		return nil
	})
	err = s.writeBatch(b)
	s.unLockKeys(txn)
	cmd.done <- errors.Trace(err)
}

func (s *dbStore) doSeek(seekCmds []*command) {
	keys := make([][]byte, 0, len(seekCmds))
	for _, cmd := range seekCmds {
		keys = append(keys, cmd.args.(*seekArgs).key)
	}

	results := s.db.MultiSeek(keys)

	for i, cmd := range seekCmds {
		reply := &seekReply{}
		var err error
		reply.key, reply.value, err = results[i].Key, results[i].Value, results[i].Err
		cmd.reply = reply
		cmd.done <- errors.Trace(err)
	}
}

func (s *dbStore) NewBatch() engine.Batch {
	return s.db.NewBatch()
}

//end of scheduler

type dbStore struct {
	db engine.DB

	txns       map[uint64]*dbTxn
	keysLocked map[string]uint64
	// TODO: clean up recentUpdates
	recentUpdates map[string]kv.Version
	uuid          string
	path          string
	compactor     *localstoreCompactor
	wg            *sync.WaitGroup

	commandCh chan *command
	closeCh   chan struct{}

	mu     sync.Mutex
	closed bool
}

type storeCache struct {
	mu    sync.Mutex
	cache map[string]*dbStore
}

var (
	globalVersionProvider kv.VersionProvider
	mc                    storeCache

	// ErrDBClosed is the error meaning db is closed and we can use it anymore.
	ErrDBClosed = errors.New("db is closed")
)

func init() {
	mc.cache = make(map[string]*dbStore)
	globalVersionProvider = &LocalVersionProvider{}
}

// Driver implements kv.Driver interface.
type Driver struct {
	// engine.Driver is the engine driver for different local db engine.
	engine.Driver
}

// IsLocalStore checks whether a storage is local or not.
func IsLocalStore(s kv.Storage) bool {
	_, ok := s.(*dbStore)
	return ok
}

// Open opens or creates a storage with specific format for a local engine Driver.
func (d Driver) Open(schema string) (kv.Storage, error) {
	mc.mu.Lock()
	defer mc.mu.Unlock()

	if store, ok := mc.cache[schema]; ok {
		// TODO: check the cache store has the same engine with this Driver.
		log.Info("[kv] cache store", schema)
		return store, nil
	}

	db, err := d.Driver.Open(schema)
	if err != nil {
		return nil, errors.Trace(err)
	}

	log.Info("[kv] New store", schema)
	s := &dbStore{
		txns:          make(map[uint64]*dbTxn),
		keysLocked:    make(map[string]uint64),
		uuid:          uuid.NewV4().String(),
		path:          schema,
		db:            db,
		compactor:     newLocalCompactor(localCompactDefaultPolicy, db),
		commandCh:     make(chan *command, 1000),
		closed:        false,
		closeCh:       make(chan struct{}),
		recentUpdates: make(map[string]kv.Version),
		wg:            &sync.WaitGroup{},
	}
	mc.cache[schema] = s
	s.compactor.Start()
	s.wg.Add(1)
	go s.scheduler()
	return s, nil
}

func (s *dbStore) UUID() string {
	return s.uuid
}

func (s *dbStore) GetSnapshot(ver kv.Version) (kv.Snapshot, error) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil, ErrDBClosed
	}
	s.mu.Unlock()

	currentVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ver.Cmp(currentVer) > 0 {
		ver = currentVer
	}

	return &dbSnapshot{
		store:   s,
		version: ver,
	}, nil
}

func (s *dbStore) CurrentVersion() (kv.Version, error) {
	return globalVersionProvider.CurrentVersion()
}

// Begin transaction
func (s *dbStore) Begin() (kv.Transaction, error) {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil, ErrDBClosed
	}
	s.mu.Unlock()

	beginVer, err := globalVersionProvider.CurrentVersion()
	if err != nil {
		return nil, errors.Trace(err)
	}

	// todo: send to chan
	txn := &dbTxn{
		tid:          beginVer.Ver,
		valid:        true,
		store:        s,
		version:      kv.MinVersion,
		snapshotVals: make(map[string]struct{}),
	}
	log.Debugf("[kv] Begin txn:%d", txn.tid)
	txn.UnionStore = kv.NewUnionStore(newSnapshot(s, beginVer))
	return txn, nil
}

func (s *dbStore) Close() error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return ErrDBClosed
	}

	s.closed = true
	s.mu.Unlock()

	mc.mu.Lock()
	defer mc.mu.Unlock()
	s.compactor.Stop()
	s.closeCh <- struct{}{}
	s.wg.Wait()
	delete(mc.cache, s.path)
	return s.db.Close()
}

func (s *dbStore) writeBatch(b engine.Batch) error {
	if b.Len() == 0 {
		return nil
	}

	if s.closed {
		return errors.Trace(ErrDBClosed)
	}

	err := s.db.Commit(b)
	if err != nil {
		log.Error(err)
		return errors.Trace(err)
	}

	return nil
}

func (s *dbStore) newBatch() engine.Batch {
	return s.db.NewBatch()
}
func (s *dbStore) unLockKeys(txn *dbTxn) error {
	for k := range txn.snapshotVals {
		if tid, ok := s.keysLocked[k]; !ok || tid != txn.tid {
			debug.PrintStack()
			log.Fatalf("should never happend:%v, %v", tid, txn.tid)
		}

		delete(s.keysLocked, k)
		s.recentUpdates[k] = txn.version
	}

	return nil
}
<|endoftext|>


<commit_before>package queue

import (
	"fmt"
	"runtime/debug"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type dataToFlush struct {
	startTime time.Time
	stopTime  time.Time
	data      []byte
}

type LogBuffer struct {
	buf           []byte
	idx           []int
	pos           int
	startTime     time.Time
	stopTime      time.Time
	sizeBuf       []byte
	flushInterval time.Duration
	flushFn       func(startTime, stopTime time.Time, buf []byte)
	notifyFn      func()
	isStopping    bool
	flushChan     chan *dataToFlush
	sync.RWMutex
}

func NewLogBuffer(flushInterval time.Duration, flushFn func(startTime, stopTime time.Time, buf []byte), notifyFn func()) *LogBuffer {
	lb := &LogBuffer{
		buf:           make([]byte, 4*1024*1024),
		sizeBuf:       make([]byte, 4),
		flushInterval: flushInterval,
		flushFn:       flushFn,
		notifyFn:      notifyFn,
		flushChan:     make(chan *dataToFlush, 256),
	}
	go lb.loopFlush()
	go lb.loopInterval()
	return lb
}

func (m *LogBuffer) AddToBuffer(ts time.Time, key, data []byte) {

	logEntry := &filer_pb.LogEntry{
		TsNs:             ts.UnixNano(),
		PartitionKeyHash: util.HashToInt32(key),
		Data:             data,
	}

	logEntryData, _ := proto.Marshal(logEntry)

	size := len(logEntryData)

	m.Lock()
	defer func() {
		m.Unlock()
		if m.notifyFn != nil {
			m.notifyFn()
		}
	}()

	if m.pos == 0 {
		m.startTime = ts
	}

	if m.startTime.Add(m.flushInterval).Before(ts) || len(m.buf)-m.pos < size+4 {
		m.flushChan <- m.copyToFlush()
		m.startTime = ts
		if len(m.buf) < size+4 {
			m.buf = make([]byte, 2*size+4)
		}
	}
	m.stopTime = ts

	m.idx = append(m.idx, m.pos)
	util.Uint32toBytes(m.sizeBuf, uint32(size))
	copy(m.buf[m.pos:m.pos+4], m.sizeBuf)
	copy(m.buf[m.pos+4:m.pos+4+size], logEntryData)
	m.pos += size + 4
}

func (m *LogBuffer) Shutdown() {
	if m.isStopping {
		return
	}
	m.isStopping = true
	m.Lock()
	toFlush := m.copyToFlush()
	m.Unlock()
	m.flushChan <- toFlush
	close(m.flushChan)
}

func (m *LogBuffer) loopFlush() {
	for d := range m.flushChan {
		if d != nil {
			m.flushFn(d.startTime, d.stopTime, d.data)
		}
	}
}

func (m *LogBuffer) loopInterval() {
	for !m.isStopping {
		m.Lock()
		toFlush := m.copyToFlush()
		m.Unlock()
		m.flushChan <- toFlush
		time.Sleep(m.flushInterval)
	}
}

func (m *LogBuffer) copyToFlush() *dataToFlush {

	if m.flushFn != nil && m.pos > 0 {
		fmt.Printf("flush buffer %d pos %d empty space %d\n", len(m.buf), m.pos, len(m.buf)-m.pos)
		debug.PrintStack()
		d := &dataToFlush{
			startTime: m.startTime,
			stopTime:  m.stopTime,
			data:      copiedBytes(m.buf[:m.pos]),
		}
		m.pos = 0
		m.idx = m.idx[:0]
		return d
	}
	return nil
}

func (m *LogBuffer) ReadFromBuffer(lastReadTime time.Time) (ts time.Time, bufferCopy []byte) {
	m.RLock()
	defer m.RUnlock()

	// fmt.Printf("read from buffer: %v\n", lastReadTime)

	if lastReadTime.Equal(m.stopTime) {
		return lastReadTime, nil
	}
	if lastReadTime.After(m.stopTime) {
		// glog.Fatalf("unexpected last read time %v, older than latest %v", lastReadTime, m.stopTime)
		return lastReadTime, nil
	}
	if lastReadTime.Before(m.startTime) {
		return m.stopTime, copiedBytes(m.buf[:m.pos])
	}

	lastTs := lastReadTime.UnixNano()
	l, h := 0, len(m.idx)-1
	/*
		for i, pos := range m.idx {
			ts := readTs(m.buf, pos)
			fmt.Printf("entry %d ts: %v offset:%d\n", i, time.Unix(0,ts), pos)
		}
		fmt.Printf("l=%d, h=%d\n", l, h)
	*/

	for l <= h {
		mid := (l + h) / 2
		pos := m.idx[mid]
		t := readTs(m.buf, m.idx[mid])
		if t <= lastTs {
			l = mid + 1
		} else if lastTs < t {
			var prevT int64
			if mid > 0 {
				prevT = readTs(m.buf, m.idx[mid-1])
			}
			if prevT <= lastTs {
				return time.Unix(0, t), copiedBytes(m.buf[pos:m.pos])
			}
			h = mid - 1
		}
		// fmt.Printf("l=%d, h=%d\n", l, h)
	}

}
func copiedBytes(buf []byte) (copied []byte) {
	copied = make([]byte, len(buf))
	copy(copied, buf)
	return
}

func readTs(buf []byte, pos int) int64 {

	size := util.BytesToUint32(buf[pos : pos+4])
	entryData := buf[pos+4 : pos+4+int(size)]
	logEntry := &filer_pb.LogEntry{}

	err := proto.Unmarshal(entryData, logEntry)
	if err != nil {
		glog.Fatalf("unexpected unmarshal filer_pb.LogEntry: %v", err)
	}
	return logEntry.TsNs

}
<commit_msg>comment out debug messages<commit_after>package queue

import (
	"sync"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type dataToFlush struct {
	startTime time.Time
	stopTime  time.Time
	data      []byte
}

type LogBuffer struct {
	buf           []byte
	idx           []int
	pos           int
	startTime     time.Time
	stopTime      time.Time
	sizeBuf       []byte
	flushInterval time.Duration
	flushFn       func(startTime, stopTime time.Time, buf []byte)
	notifyFn      func()
	isStopping    bool
	flushChan     chan *dataToFlush
	sync.RWMutex
}

func NewLogBuffer(flushInterval time.Duration, flushFn func(startTime, stopTime time.Time, buf []byte), notifyFn func()) *LogBuffer {
	lb := &LogBuffer{
		buf:           make([]byte, 4*1024*1024),
		sizeBuf:       make([]byte, 4),
		flushInterval: flushInterval,
		flushFn:       flushFn,
		notifyFn:      notifyFn,
		flushChan:     make(chan *dataToFlush, 256),
	}
	go lb.loopFlush()
	go lb.loopInterval()
	return lb
}

func (m *LogBuffer) AddToBuffer(ts time.Time, key, data []byte) {

	logEntry := &filer_pb.LogEntry{
		TsNs:             ts.UnixNano(),
		PartitionKeyHash: util.HashToInt32(key),
		Data:             data,
	}

	logEntryData, _ := proto.Marshal(logEntry)

	size := len(logEntryData)

	m.Lock()
	defer func() {
		m.Unlock()
		if m.notifyFn != nil {
			m.notifyFn()
		}
	}()

	if m.pos == 0 {
		m.startTime = ts
	}

	if m.startTime.Add(m.flushInterval).Before(ts) || len(m.buf)-m.pos < size+4 {
		m.flushChan <- m.copyToFlush()
		m.startTime = ts
		if len(m.buf) < size+4 {
			m.buf = make([]byte, 2*size+4)
		}
	}
	m.stopTime = ts

	m.idx = append(m.idx, m.pos)
	util.Uint32toBytes(m.sizeBuf, uint32(size))
	copy(m.buf[m.pos:m.pos+4], m.sizeBuf)
	copy(m.buf[m.pos+4:m.pos+4+size], logEntryData)
	m.pos += size + 4
}

func (m *LogBuffer) Shutdown() {
	if m.isStopping {
		return
	}
	m.isStopping = true
	m.Lock()
	toFlush := m.copyToFlush()
	m.Unlock()
	m.flushChan <- toFlush
	close(m.flushChan)
}

func (m *LogBuffer) loopFlush() {
	for d := range m.flushChan {
		if d != nil {
			m.flushFn(d.startTime, d.stopTime, d.data)
		}
	}
}

func (m *LogBuffer) loopInterval() {
	for !m.isStopping {
		m.Lock()
		toFlush := m.copyToFlush()
		m.Unlock()
		m.flushChan <- toFlush
		time.Sleep(m.flushInterval)
	}
}

func (m *LogBuffer) copyToFlush() *dataToFlush {

	if m.flushFn != nil && m.pos > 0 {
		// fmt.Printf("flush buffer %d pos %d empty space %d\n", len(m.buf), m.pos, len(m.buf)-m.pos)
		d := &dataToFlush{
			startTime: m.startTime,
			stopTime:  m.stopTime,
			data:      copiedBytes(m.buf[:m.pos]),
		}
		m.pos = 0
		m.idx = m.idx[:0]
		return d
	}
	return nil
}

func (m *LogBuffer) ReadFromBuffer(lastReadTime time.Time) (ts time.Time, bufferCopy []byte) {
	m.RLock()
	defer m.RUnlock()

	// fmt.Printf("read from buffer: %v\n", lastReadTime)

	if lastReadTime.Equal(m.stopTime) {
		return lastReadTime, nil
	}
	if lastReadTime.After(m.stopTime) {
		// glog.Fatalf("unexpected last read time %v, older than latest %v", lastReadTime, m.stopTime)
		return lastReadTime, nil
	}
	if lastReadTime.Before(m.startTime) {
		return m.stopTime, copiedBytes(m.buf[:m.pos])
	}

	lastTs := lastReadTime.UnixNano()
	l, h := 0, len(m.idx)-1

	/*
	for i, pos := range m.idx {
		logEntry, ts := readTs(m.buf, pos)
		event := &filer_pb.FullEventNotification{}
		proto.Unmarshal(logEntry.Data, event)
		entry := event.EventNotification.OldEntry
		if entry == nil {
			entry = event.EventNotification.NewEntry
		}
		fmt.Printf("entry %d ts: %v offset:%d dir:%s name:%s\n", i, time.Unix(0, ts), pos, event.Directory, entry.Name)
	}
	fmt.Printf("l=%d, h=%d\n", l, h)
	 */

	for l <= h {
		mid := (l + h) / 2
		pos := m.idx[mid]
		_, t := readTs(m.buf, m.idx[mid])
		if t <= lastTs {
			l = mid + 1
		} else if lastTs < t {
			var prevT int64
			if mid > 0 {
				_, prevT = readTs(m.buf, m.idx[mid-1])
			}
			if prevT <= lastTs {
				// println("found mid = ", mid)
				return time.Unix(0, t), copiedBytes(m.buf[pos:m.pos])
			}
			h = mid - 1
		}
		// fmt.Printf("l=%d, h=%d\n", l, h)
	}

	// FIXME: this could be that the buffer has been flushed already
	// println("not found")
	return lastReadTime, nil

}
func copiedBytes(buf []byte) (copied []byte) {
	copied = make([]byte, len(buf))
	copy(copied, buf)
	return
}

func readTs(buf []byte, pos int) (*filer_pb.LogEntry, int64) {

	size := util.BytesToUint32(buf[pos : pos+4])
	entryData := buf[pos+4 : pos+4+int(size)]
	logEntry := &filer_pb.LogEntry{}

	err := proto.Unmarshal(entryData, logEntry)
	if err != nil {
		glog.Fatalf("unexpected unmarshal filer_pb.LogEntry: %v", err)
	}
	return logEntry, logEntry.TsNs

}
<|endoftext|>


<commit_before>package wdclient

import (
	"errors"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

const (
	maxCursorIndex = 4096
)

type HasLookupFileIdFunction interface {
	GetLookupFileIdFunction() LookupFileIdFunctionType
}

type LookupFileIdFunctionType func(fileId string) (targetUrls []string, err error)

type Location struct {
	Url        string `json:"url,omitempty"`
	PublicUrl  string `json:"publicUrl,omitempty"`
	DataCenter string `json:"dataCenter,omitempty"`
	GrpcPort   int    `json:"grpcPort,omitempty"`
}

func (l Location) ServerAddress() pb.ServerAddress {
	return pb.NewServerAddressWithGrpcPort(l.Url, l.GrpcPort)
}

type vidMap struct {
	sync.RWMutex
	vid2Locations   map[uint32][]Location
	ecVid2Locations map[uint32][]Location
	DataCenter      string
	cursor          int32
}

func newVidMap(dataCenter string) vidMap {
	return vidMap{
		vid2Locations:   make(map[uint32][]Location),
		ecVid2Locations: make(map[uint32][]Location),
		DataCenter:      dataCenter,
		cursor:          -1,
	}
}

func (vc *vidMap) getLocationIndex(length int) (int, error) {
	if length <= 0 {
		return 0, fmt.Errorf("invalid length: %d", length)
	}
	if atomic.LoadInt32(&vc.cursor) == maxCursorIndex {
		atomic.CompareAndSwapInt32(&vc.cursor, maxCursorIndex, -1)
	}
	return int(atomic.AddInt32(&vc.cursor, 1)) % length, nil
}

func (vc *vidMap) LookupVolumeServerUrl(vid string) (serverUrls []string, err error) {
	id, err := strconv.Atoi(vid)
	if err != nil {
		glog.V(1).Infof("Unknown volume id %s", vid)
		return nil, err
	}

	locations, found := vc.GetLocations(uint32(id))
	if !found {
		return nil, fmt.Errorf("volume %d not found", id)
	}
	var sameDcServers, otherDcServers []string
	for _, loc := range locations {
		if vc.DataCenter == "" || loc.DataCenter == "" || vc.DataCenter != loc.DataCenter {
			otherDcServers = append(otherDcServers, loc.Url)
		} else {
			sameDcServers = append(sameDcServers, loc.Url)
		}
	}
	rand.Shuffle(len(sameDcServers), func(i, j int) {
		sameDcServers[i], sameDcServers[j] = sameDcServers[j], sameDcServers[i]
	})
	rand.Shuffle(len(otherDcServers), func(i, j int) {
		otherDcServers[i], otherDcServers[j] = otherDcServers[j], otherDcServers[i]
	})
	serverUrls = append(sameDcServers, otherDcServers...)
	return
}

func (vc *vidMap) GetLookupFileIdFunction() LookupFileIdFunctionType {
	return vc.LookupFileId
}

func (vc *vidMap) LookupFileId(fileId string) (fullUrls []string, err error) {
	parts := strings.Split(fileId, ",")
	if len(parts) != 2 {
		return nil, errors.New("Invalid fileId " + fileId)
	}
	serverUrls, lookupError := vc.LookupVolumeServerUrl(parts[0])
	if lookupError != nil {
		return nil, lookupError
	}
	for _, serverUrl := range serverUrls {
		fullUrls = append(fullUrls, "http://"+serverUrl+"/"+fileId)
	}
	return
}

func (vc *vidMap) GetVidLocations(vid string) (locations []Location, err error) {
	id, err := strconv.Atoi(vid)
	if err != nil {
		glog.V(1).Infof("Unknown volume id %s", vid)
		return nil, fmt.Errorf("Unknown volume id %s", vid)
	}
	foundLocations, found := vc.GetLocations(uint32(id))
	if found {
		return foundLocations, nil
	}
	return nil, fmt.Errorf("volume id %s not found", vid)
}

func (vc *vidMap) GetLocations(vid uint32) (locations []Location, found bool) {
	vc.RLock()
	defer vc.RUnlock()

	glog.V(4).Infof("~ lookup volume id %d: %+v ec:%+v", vid, vc.vid2Locations, vc.ecVid2Locations)

	locations, found = vc.vid2Locations[vid]
	if found && len(locations) > 0 {
		return
	}
	locations, found = vc.ecVid2Locations[vid]
	return
}

func (vc *vidMap) addLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("+ volume id %d: %+v", vid, location)

	locations, found := vc.vid2Locations[vid]
	if !found {
		vc.vid2Locations[vid] = []Location{location}
		return
	}

	for _, loc := range locations {
		if loc.Url == location.Url {
			return
		}
	}

	vc.vid2Locations[vid] = append(locations, location)

}

func (vc *vidMap) addEcLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("+ ec volume id %d: %+v", vid, location)

	locations, found := vc.ecVid2Locations[vid]
	if !found {
		vc.ecVid2Locations[vid] = []Location{location}
		return
	}

	for _, loc := range locations {
		if loc.Url == location.Url {
			return
		}
	}

	vc.ecVid2Locations[vid] = append(locations, location)

}

func (vc *vidMap) deleteLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("- volume id %d: %+v", vid, location)

	locations, found := vc.vid2Locations[vid]
	if !found {
		return
	}

	for i, loc := range locations {
		if loc.Url == location.Url {
			vc.vid2Locations[vid] = append(locations[0:i], locations[i+1:]...)
			break
		}
	}

}

func (vc *vidMap) deleteEcLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("- ec volume id %d: %+v", vid, location)

	locations, found := vc.ecVid2Locations[vid]
	if !found {
		return
	}

	for i, loc := range locations {
		if loc.Url == location.Url {
			vc.ecVid2Locations[vid] = append(locations[0:i], locations[i+1:]...)
			break
		}
	}

}
<commit_msg>fix(wdclient): GetLocations return<commit_after>package wdclient

import (
	"errors"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/chrislusf/seaweedfs/weed/glog"
)

const (
	maxCursorIndex = 4096
)

type HasLookupFileIdFunction interface {
	GetLookupFileIdFunction() LookupFileIdFunctionType
}

type LookupFileIdFunctionType func(fileId string) (targetUrls []string, err error)

type Location struct {
	Url        string `json:"url,omitempty"`
	PublicUrl  string `json:"publicUrl,omitempty"`
	DataCenter string `json:"dataCenter,omitempty"`
	GrpcPort   int    `json:"grpcPort,omitempty"`
}

func (l Location) ServerAddress() pb.ServerAddress {
	return pb.NewServerAddressWithGrpcPort(l.Url, l.GrpcPort)
}

type vidMap struct {
	sync.RWMutex
	vid2Locations   map[uint32][]Location
	ecVid2Locations map[uint32][]Location
	DataCenter      string
	cursor          int32
}

func newVidMap(dataCenter string) vidMap {
	return vidMap{
		vid2Locations:   make(map[uint32][]Location),
		ecVid2Locations: make(map[uint32][]Location),
		DataCenter:      dataCenter,
		cursor:          -1,
	}
}

func (vc *vidMap) getLocationIndex(length int) (int, error) {
	if length <= 0 {
		return 0, fmt.Errorf("invalid length: %d", length)
	}
	if atomic.LoadInt32(&vc.cursor) == maxCursorIndex {
		atomic.CompareAndSwapInt32(&vc.cursor, maxCursorIndex, -1)
	}
	return int(atomic.AddInt32(&vc.cursor, 1)) % length, nil
}

func (vc *vidMap) LookupVolumeServerUrl(vid string) (serverUrls []string, err error) {
	id, err := strconv.Atoi(vid)
	if err != nil {
		glog.V(1).Infof("Unknown volume id %s", vid)
		return nil, err
	}

	locations, found := vc.GetLocations(uint32(id))
	if !found {
		return nil, fmt.Errorf("volume %d not found", id)
	}
	var sameDcServers, otherDcServers []string
	for _, loc := range locations {
		if vc.DataCenter == "" || loc.DataCenter == "" || vc.DataCenter != loc.DataCenter {
			otherDcServers = append(otherDcServers, loc.Url)
		} else {
			sameDcServers = append(sameDcServers, loc.Url)
		}
	}
	rand.Shuffle(len(sameDcServers), func(i, j int) {
		sameDcServers[i], sameDcServers[j] = sameDcServers[j], sameDcServers[i]
	})
	rand.Shuffle(len(otherDcServers), func(i, j int) {
		otherDcServers[i], otherDcServers[j] = otherDcServers[j], otherDcServers[i]
	})
	serverUrls = append(sameDcServers, otherDcServers...)
	return
}

func (vc *vidMap) GetLookupFileIdFunction() LookupFileIdFunctionType {
	return vc.LookupFileId
}

func (vc *vidMap) LookupFileId(fileId string) (fullUrls []string, err error) {
	parts := strings.Split(fileId, ",")
	if len(parts) != 2 {
		return nil, errors.New("Invalid fileId " + fileId)
	}
	serverUrls, lookupError := vc.LookupVolumeServerUrl(parts[0])
	if lookupError != nil {
		return nil, lookupError
	}
	for _, serverUrl := range serverUrls {
		fullUrls = append(fullUrls, "http://"+serverUrl+"/"+fileId)
	}
	return
}

func (vc *vidMap) GetVidLocations(vid string) (locations []Location, err error) {
	id, err := strconv.Atoi(vid)
	if err != nil {
		glog.V(1).Infof("Unknown volume id %s", vid)
		return nil, fmt.Errorf("Unknown volume id %s", vid)
	}
	foundLocations, found := vc.GetLocations(uint32(id))
	if found {
		return foundLocations, nil
	}
	return nil, fmt.Errorf("volume id %s not found", vid)
}

func (vc *vidMap) GetLocations(vid uint32) (locations []Location, found bool) {
	vc.RLock()
	defer vc.RUnlock()

	glog.V(4).Infof("~ lookup volume id %d: %+v ec:%+v", vid, vc.vid2Locations, vc.ecVid2Locations)

	locations, found = vc.vid2Locations[vid]
	if found && len(locations) > 0 {
		return
	}
	locations, found = vc.ecVid2Locations[vid]
	return locations, found && len(locations) > 0
}

func (vc *vidMap) addLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("+ volume id %d: %+v", vid, location)

	locations, found := vc.vid2Locations[vid]
	if !found {
		vc.vid2Locations[vid] = []Location{location}
		return
	}

	for _, loc := range locations {
		if loc.Url == location.Url {
			return
		}
	}

	vc.vid2Locations[vid] = append(locations, location)

}

func (vc *vidMap) addEcLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("+ ec volume id %d: %+v", vid, location)

	locations, found := vc.ecVid2Locations[vid]
	if !found {
		vc.ecVid2Locations[vid] = []Location{location}
		return
	}

	for _, loc := range locations {
		if loc.Url == location.Url {
			return
		}
	}

	vc.ecVid2Locations[vid] = append(locations, location)

}

func (vc *vidMap) deleteLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("- volume id %d: %+v", vid, location)

	locations, found := vc.vid2Locations[vid]
	if !found {
		return
	}

	for i, loc := range locations {
		if loc.Url == location.Url {
			vc.vid2Locations[vid] = append(locations[0:i], locations[i+1:]...)
			break
		}
	}

}

func (vc *vidMap) deleteEcLocation(vid uint32, location Location) {
	vc.Lock()
	defer vc.Unlock()

	glog.V(4).Infof("- ec volume id %d: %+v", vid, location)

	locations, found := vc.ecVid2Locations[vid]
	if !found {
		return
	}

	for i, loc := range locations {
		if loc.Url == location.Url {
			vc.ecVid2Locations[vid] = append(locations[0:i], locations[i+1:]...)
			break
		}
	}

}
<|endoftext|>


<commit_before>package work

import (
	"context"
	"fmt"
	"math"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
)

const (
	etcdHost = "localhost"
	etcdPort = "32379"
)

// generateRandomString is a helper function for getPachClient
func generateRandomString(n int) string {
	rand.Seed(time.Now().UnixNano())
	b := make([]byte, n)
	for i := range b {
		b[i] = byte('a' + rand.Intn(26))
	}
	return string(b)
}

func TestWork(t *testing.T) {
	etcdPrefix := generateRandomString(32)
	workId := "work"
	numShards := 10
	numWorkers := 5
	// Setup etcd client and cleanup previous runs.
	etcdClient, err := getEtcdClient()
	require.NoError(t, err)
	// Setup maps for checking creation, processing, and collection.
	shardsCreated := make(map[string]bool)
	shardsProcessed := make(map[string]bool)
	var processMu sync.Mutex
	shardsCollected := make(map[string]bool)
	// Setup workers.
	failProb := 0.25
	for i := 0; i < numWorkers; i++ {
		go func() {
			for {
				ctx, cancel := context.WithCancel(context.Background())
				w := NewWorker(etcdClient, etcdPrefix, func(_ context.Context, work *Work, shard *Shard) error {
					if rand.Float64() < failProb {
						cancel()
					} else {
						processMu.Lock()
						shardsProcessed[shard.Id] = true
						processMu.Unlock()
					}
					return nil
				})
				if err := w.Run(ctx); err != nil {
					if ctx.Err() == context.Canceled {
						continue
					}
					break
				}
			}
		}()
	}
	// Setup master.
	m := NewMaster(etcdClient, etcdPrefix, func(_ context.Context, shard *Shard) error {
		shardsCollected[shard.Id] = true
		return nil
	})
	// Create work.
	var shards []*Shard
	for i := 0; i < numShards; i++ {
		id := strconv.Itoa(i)
		shards = append(shards, &Shard{Id: id})
		shardsCreated[id] = true
	}
	work := &Work{
		Id:     workId,
		Shards: shards,
	}
	require.NoError(t, m.Run(context.Background(), work))
	require.Equal(t, shardsCreated, shardsProcessed, shardsCollected)
}

func getEtcdClient() (*etcd.Client, error) {
	var etcdClient *etcd.Client
	if err := backoff.Retry(func() error {
		var err error
		etcdClient, err = etcd.New(etcd.Config{
			Endpoints:          []string{net.JoinHostPort(etcdHost, etcdPort)},
			DialOptions:        client.DefaultDialOptions(),
			MaxCallSendMsgSize: math.MaxInt32,
			MaxCallRecvMsgSize: math.MaxInt32,
		})
		if err != nil {
			return fmt.Errorf("failed to initialize etcd client: %v", err)
		}
		return nil
	}, backoff.RetryEvery(time.Second).For(time.Minute)); err != nil {
		return nil, err
	}
	return etcdClient, nil
}
<commit_msg>Lint work package<commit_after>package work

import (
	"context"
	"fmt"
	"math"
	"math/rand"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"

	etcd "github.com/coreos/etcd/clientv3"
	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
)

const (
	etcdHost = "localhost"
	etcdPort = "32379"
)

// generateRandomString is a helper function for getPachClient
func generateRandomString(n int) string {
	rand.Seed(time.Now().UnixNano())
	b := make([]byte, n)
	for i := range b {
		b[i] = byte('a' + rand.Intn(26))
	}
	return string(b)
}

func TestWork(t *testing.T) {
	etcdPrefix := generateRandomString(32)
	workID := "work"
	numShards := 10
	numWorkers := 5
	// Setup etcd client and cleanup previous runs.
	etcdClient, err := getEtcdClient()
	require.NoError(t, err)
	// Setup maps for checking creation, processing, and collection.
	shardsCreated := make(map[string]bool)
	shardsProcessed := make(map[string]bool)
	var processMu sync.Mutex
	shardsCollected := make(map[string]bool)
	// Setup workers.
	failProb := 0.25
	for i := 0; i < numWorkers; i++ {
		go func() {
			for {
				ctx, cancel := context.WithCancel(context.Background())
				w := NewWorker(etcdClient, etcdPrefix, func(_ context.Context, work *Work, shard *Shard) error {
					if rand.Float64() < failProb {
						cancel()
					} else {
						processMu.Lock()
						shardsProcessed[shard.Id] = true
						processMu.Unlock()
					}
					return nil
				})
				if err := w.Run(ctx); err != nil {
					if ctx.Err() == context.Canceled {
						continue
					}
					break
				}
			}
		}()
	}
	// Setup master.
	m := NewMaster(etcdClient, etcdPrefix, func(_ context.Context, shard *Shard) error {
		shardsCollected[shard.Id] = true
		return nil
	})
	// Create work.
	var shards []*Shard
	for i := 0; i < numShards; i++ {
		id := strconv.Itoa(i)
		shards = append(shards, &Shard{Id: id})
		shardsCreated[id] = true
	}
	work := &Work{
		Id:     workID,
		Shards: shards,
	}
	require.NoError(t, m.Run(context.Background(), work))
	require.Equal(t, shardsCreated, shardsProcessed, shardsCollected)
}

func getEtcdClient() (*etcd.Client, error) {
	var etcdClient *etcd.Client
	if err := backoff.Retry(func() error {
		var err error
		etcdClient, err = etcd.New(etcd.Config{
			Endpoints:          []string{net.JoinHostPort(etcdHost, etcdPort)},
			DialOptions:        client.DefaultDialOptions(),
			MaxCallSendMsgSize: math.MaxInt32,
			MaxCallRecvMsgSize: math.MaxInt32,
		})
		if err != nil {
			return fmt.Errorf("failed to initialize etcd client: %v", err)
		}
		return nil
	}, backoff.RetryEvery(time.Second).For(time.Minute)); err != nil {
		return nil, err
	}
	return etcdClient, nil
}
<|endoftext|>


<commit_before>// Package disco handles Terraform's remote service discovery protocol.
//
// This protocol allows mapping from a service hostname, as produced by the
// svchost package, to a set of services supported by that host and the
// endpoint information for each supported service.
package disco

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"time"

	cleanhttp "github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/terraform/svchost"
	"github.com/hashicorp/terraform/svchost/auth"
	"github.com/hashicorp/terraform/terraform"
)

const (
	discoPath        = "/.well-known/terraform.json"
	maxRedirects     = 3               // arbitrary-but-small number to prevent runaway redirect loops
	discoTimeout     = 4 * time.Second // arbitrary-but-small time limit to prevent UI "hangs" during discovery
	maxDiscoDocBytes = 1 * 1024 * 1024 // 1MB - to prevent abusive services from using loads of our memory
)

var userAgent = fmt.Sprintf("Terraform/%s (service discovery)", terraform.VersionString())
var httpTransport = cleanhttp.DefaultPooledTransport() // overridden during tests, to skip TLS verification

// Disco is the main type in this package, which allows discovery on given
// hostnames and caches the results by hostname to avoid repeated requests
// for the same information.
type Disco struct {
	hostCache map[svchost.Hostname]Host
	credsSrc  auth.CredentialsSource
}

func NewDisco() *Disco {
	return &Disco{}
}

// SetCredentialsSource provides a credentials source that will be used to
// add credentials to outgoing discovery requests, where available.
//
// If this method is never called, no outgoing discovery requests will have
// credentials.
func (d *Disco) SetCredentialsSource(src auth.CredentialsSource) {
	d.credsSrc = src
}

// Discover runs the discovery protocol against the given hostname (which must
// already have been validated and prepared with svchost.ForComparison) and
// returns an object describing the services available at that host.
//
// If a given hostname supports no Terraform services at all, a non-nil but
// empty Host object is returned. When giving feedback to the end user about
// such situations, we say e.g. "the host <name> doesn't provide a module
// registry", regardless of whether that is due to that service specifically
// being absent or due to the host not providing Terraform services at all,
// since we don't wish to expose the detail of whole-host discovery to an
// end-user.
func (d *Disco) Discover(host svchost.Hostname) Host {
	if d.hostCache == nil {
		d.hostCache = map[svchost.Hostname]Host{}
	}
	if cache, cached := d.hostCache[host]; cached {
		return cache
	}

	ret := d.discover(host)
	d.hostCache[host] = ret
	return ret
}

// DiscoverServiceURL is a convenience wrapper for discovery on a given
// hostname and then looking up a particular service in the result.
func (d *Disco) DiscoverServiceURL(host svchost.Hostname, serviceID string) *url.URL {
	return d.Discover(host).ServiceURL(serviceID)
}

// discover implements the actual discovery process, with its result cached
// by the public-facing Discover method.
func (d *Disco) discover(host svchost.Hostname) Host {
	discoURL := &url.URL{
		Scheme: "https",
		Host:   string(host),
		Path:   discoPath,
	}
	client := &http.Client{
		Transport: httpTransport,
		Timeout:   discoTimeout,

		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			log.Printf("[DEBUG] Service discovery redirected to %s", req.URL)
			if len(via) > maxRedirects {
				return errors.New("too many redirects") // (this error message will never actually be seen)
			}
			return nil
		},
	}

	var header = http.Header{}
	header.Set("User-Agent", userAgent)

	req := &http.Request{
		Method: "GET",
		URL:    discoURL,
		Header: header,
	}

	if d.credsSrc != nil {
		creds, err := d.credsSrc.ForHost(host)
		if err == nil {
			if creds != nil {
				creds.PrepareRequest(req) // alters req to include credentials
			}
		} else {
			log.Printf("[WARNING] Failed to get credentials for %s: %s (ignoring)", host, err)
		}
	}

	log.Printf("[DEBUG] Service discovery for %s at %s", host, discoURL)

	ret := Host{
		discoURL: discoURL,
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Printf("[WARNING] Failed to request discovery document: %s", err)
		return ret // empty
	}
	if resp.StatusCode != 200 {
		log.Printf("[WARNING] Failed to request discovery document: %s", resp.Status)
		return ret // empty
	}

	// If the client followed any redirects, we will have a new URL to use
	// as our base for relative resolution.
	ret.discoURL = resp.Request.URL

	contentType := resp.Header.Get("Content-Type")
	mediaType, _, err := mime.ParseMediaType(contentType)
	if err != nil {
		log.Printf("[WARNING] Discovery URL has malformed Content-Type %q", contentType)
		return ret // empty
	}
	if mediaType != "application/json" {
		log.Printf("[DEBUG] Discovery URL returned Content-Type %q, rather than application/json", mediaType)
		return ret // empty
	}

	// (this doesn't catch chunked encoding, because ContentLength is -1 in that case...)
	if resp.ContentLength > maxDiscoDocBytes {
		// Size limit here is not a contractual requirement and so we may
		// adjust it over time if we find a different limit is warranted.
		log.Printf("[WARNING] Discovery doc response is too large (got %d bytes; limit %d)", resp.ContentLength, maxDiscoDocBytes)
		return ret // empty
	}

	// If the response is using chunked encoding then we can't predict
	// its size, but we'll at least prevent reading the entire thing into
	// memory.
	lr := io.LimitReader(resp.Body, maxDiscoDocBytes)

	servicesBytes, err := ioutil.ReadAll(lr)
	if err != nil {
		log.Printf("[WARNING] Error reading discovery document body: %s", err)
		return ret // empty
	}

	var services map[string]interface{}
	err = json.Unmarshal(servicesBytes, &services)
	if err != nil {
		log.Printf("[WARNING] Failed to decode discovery document as a JSON object: %s", err)
		return ret // empty
	}

	ret.services = services
	return ret
}

// Forget invalidates any cached record of the given hostname. If the host
// has no cache entry then this is a no-op.
func (d *Disco) Forget(host svchost.Hostname) {
	delete(d.hostCache, host)
}

// ForgetAll is like Forget, but for all of the hostnames that have cache entries.
func (d *Disco) ForgetAll() {
	d.hostCache = nil
}
<commit_msg>use the new version package<commit_after>// Package disco handles Terraform's remote service discovery protocol.
//
// This protocol allows mapping from a service hostname, as produced by the
// svchost package, to a set of services supported by that host and the
// endpoint information for each supported service.
package disco

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime"
	"net/http"
	"net/url"
	"time"

	cleanhttp "github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/terraform/svchost"
	"github.com/hashicorp/terraform/svchost/auth"
	"github.com/hashicorp/terraform/version"
)

const (
	discoPath        = "/.well-known/terraform.json"
	maxRedirects     = 3               // arbitrary-but-small number to prevent runaway redirect loops
	discoTimeout     = 4 * time.Second // arbitrary-but-small time limit to prevent UI "hangs" during discovery
	maxDiscoDocBytes = 1 * 1024 * 1024 // 1MB - to prevent abusive services from using loads of our memory
)

var userAgent = fmt.Sprintf("Terraform/%s (service discovery)", version.String())
var httpTransport = cleanhttp.DefaultPooledTransport() // overridden during tests, to skip TLS verification

// Disco is the main type in this package, which allows discovery on given
// hostnames and caches the results by hostname to avoid repeated requests
// for the same information.
type Disco struct {
	hostCache map[svchost.Hostname]Host
	credsSrc  auth.CredentialsSource
}

func NewDisco() *Disco {
	return &Disco{}
}

// SetCredentialsSource provides a credentials source that will be used to
// add credentials to outgoing discovery requests, where available.
//
// If this method is never called, no outgoing discovery requests will have
// credentials.
func (d *Disco) SetCredentialsSource(src auth.CredentialsSource) {
	d.credsSrc = src
}

// Discover runs the discovery protocol against the given hostname (which must
// already have been validated and prepared with svchost.ForComparison) and
// returns an object describing the services available at that host.
//
// If a given hostname supports no Terraform services at all, a non-nil but
// empty Host object is returned. When giving feedback to the end user about
// such situations, we say e.g. "the host <name> doesn't provide a module
// registry", regardless of whether that is due to that service specifically
// being absent or due to the host not providing Terraform services at all,
// since we don't wish to expose the detail of whole-host discovery to an
// end-user.
func (d *Disco) Discover(host svchost.Hostname) Host {
	if d.hostCache == nil {
		d.hostCache = map[svchost.Hostname]Host{}
	}
	if cache, cached := d.hostCache[host]; cached {
		return cache
	}

	ret := d.discover(host)
	d.hostCache[host] = ret
	return ret
}

// DiscoverServiceURL is a convenience wrapper for discovery on a given
// hostname and then looking up a particular service in the result.
func (d *Disco) DiscoverServiceURL(host svchost.Hostname, serviceID string) *url.URL {
	return d.Discover(host).ServiceURL(serviceID)
}

// discover implements the actual discovery process, with its result cached
// by the public-facing Discover method.
func (d *Disco) discover(host svchost.Hostname) Host {
	discoURL := &url.URL{
		Scheme: "https",
		Host:   string(host),
		Path:   discoPath,
	}
	client := &http.Client{
		Transport: httpTransport,
		Timeout:   discoTimeout,

		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			log.Printf("[DEBUG] Service discovery redirected to %s", req.URL)
			if len(via) > maxRedirects {
				return errors.New("too many redirects") // (this error message will never actually be seen)
			}
			return nil
		},
	}

	var header = http.Header{}
	header.Set("User-Agent", userAgent)

	req := &http.Request{
		Method: "GET",
		URL:    discoURL,
		Header: header,
	}

	if d.credsSrc != nil {
		creds, err := d.credsSrc.ForHost(host)
		if err == nil {
			if creds != nil {
				creds.PrepareRequest(req) // alters req to include credentials
			}
		} else {
			log.Printf("[WARNING] Failed to get credentials for %s: %s (ignoring)", host, err)
		}
	}

	log.Printf("[DEBUG] Service discovery for %s at %s", host, discoURL)

	ret := Host{
		discoURL: discoURL,
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Printf("[WARNING] Failed to request discovery document: %s", err)
		return ret // empty
	}
	if resp.StatusCode != 200 {
		log.Printf("[WARNING] Failed to request discovery document: %s", resp.Status)
		return ret // empty
	}

	// If the client followed any redirects, we will have a new URL to use
	// as our base for relative resolution.
	ret.discoURL = resp.Request.URL

	contentType := resp.Header.Get("Content-Type")
	mediaType, _, err := mime.ParseMediaType(contentType)
	if err != nil {
		log.Printf("[WARNING] Discovery URL has malformed Content-Type %q", contentType)
		return ret // empty
	}
	if mediaType != "application/json" {
		log.Printf("[DEBUG] Discovery URL returned Content-Type %q, rather than application/json", mediaType)
		return ret // empty
	}

	// (this doesn't catch chunked encoding, because ContentLength is -1 in that case...)
	if resp.ContentLength > maxDiscoDocBytes {
		// Size limit here is not a contractual requirement and so we may
		// adjust it over time if we find a different limit is warranted.
		log.Printf("[WARNING] Discovery doc response is too large (got %d bytes; limit %d)", resp.ContentLength, maxDiscoDocBytes)
		return ret // empty
	}

	// If the response is using chunked encoding then we can't predict
	// its size, but we'll at least prevent reading the entire thing into
	// memory.
	lr := io.LimitReader(resp.Body, maxDiscoDocBytes)

	servicesBytes, err := ioutil.ReadAll(lr)
	if err != nil {
		log.Printf("[WARNING] Error reading discovery document body: %s", err)
		return ret // empty
	}

	var services map[string]interface{}
	err = json.Unmarshal(servicesBytes, &services)
	if err != nil {
		log.Printf("[WARNING] Failed to decode discovery document as a JSON object: %s", err)
		return ret // empty
	}

	ret.services = services
	return ret
}

// Forget invalidates any cached record of the given hostname. If the host
// has no cache entry then this is a no-op.
func (d *Disco) Forget(host svchost.Hostname) {
	delete(d.hostCache, host)
}

// ForgetAll is like Forget, but for all of the hostnames that have cache entries.
func (d *Disco) ForgetAll() {
	d.hostCache = nil
}
<|endoftext|>


<commit_before>package main

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

	"net/http/pprof"

	"github.com/gorilla/mux"
)

func setupRouter(app *server) http.Handler {
	r := mux.NewRouter()
	AttachProfiler(r)

	r.Use(handleMiddleware(app.pool))

	r.HandleFunc("/version", app.getVersionHandler).Methods("GET")

	r.HandleFunc("/kdb", app.getKdbHandler).Methods("GET")
	r.HandleFunc("/kdb/{path:.*}", app.getKdbHandler).Methods("GET")
	r.HandleFunc("/kdb/{path:.*}", app.putKdbHandler).Methods("PUT")
	r.HandleFunc("/kdb/{path:.*}", app.deleteKdbHandler).Methods("DELETE")

	r.HandleFunc("/kdbFind/{path:.*}", app.getFindHandler).Methods("GET")

	r.HandleFunc("/kdbMv/{path:.*}", app.postMoveHandler).Methods("POST")

	r.HandleFunc("/kdbMeta/{path:.*}", app.postMetaHandler).Methods("POST")
	r.HandleFunc("/kdbMeta/{path:.*}", app.deleteMetaHandler).Methods("DELETE")

	return r
}

func AttachProfiler(router *mux.Router) {
	router.HandleFunc("/debug/pprof/", pprof.Index)
	router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
	router.HandleFunc("/debug/pprof/profile", pprof.Profile)
	router.HandleFunc("/debug/pprof/symbol", pprof.Symbol)

	// Manually add support for paths linked to by index page at /debug/pprof/
	router.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
	router.Handle("/debug/pprof/heap", pprof.Handler("heap"))
	router.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
	router.Handle("/debug/pprof/block", pprof.Handler("block"))
}

func parseKeyNameFromURL(r *http.Request) string {
	vars := mux.Vars(r)
	keyName := vars["path"]

	if len(keyName) == 0 {
		return "/"
	}

	return keyName
}

func stringBody(r *http.Request) (string, error) {
	value := ""

	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&value); err != nil {
		return "", err
	}

	return value, nil
}

func notFound(w http.ResponseWriter) {
	w.WriteHeader(http.StatusNotFound)
}

func noContent(w http.ResponseWriter) {
	w.WriteHeader(http.StatusNoContent)
}

func badRequest(w http.ResponseWriter) {
	w.WriteHeader(http.StatusBadRequest)
}

func created(w http.ResponseWriter) {
	w.WriteHeader(http.StatusCreated)
}

func writeError(w http.ResponseWriter, err error) {
	badRequest(w)

	writeResponse(w, map[string]string{
		"error": err.Error(),
	})
}

func writeResponse(w http.ResponseWriter, response interface{}) {
	w.Header().Set("Content-Type", "application/json")

	js, _ := json.Marshal(response)

	w.Write(js)
}
<commit_msg>tool: elektrad - remove profiler endpoint<commit_after>package main

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

	"github.com/gorilla/mux"
)

func setupRouter(app *server) http.Handler {
	r := mux.NewRouter()

	r.Use(handleMiddleware(app.pool))

	r.HandleFunc("/version", app.getVersionHandler).Methods("GET")

	r.HandleFunc("/kdb", app.getKdbHandler).Methods("GET")
	r.HandleFunc("/kdb/{path:.*}", app.getKdbHandler).Methods("GET")
	r.HandleFunc("/kdb/{path:.*}", app.putKdbHandler).Methods("PUT")
	r.HandleFunc("/kdb/{path:.*}", app.deleteKdbHandler).Methods("DELETE")

	r.HandleFunc("/kdbFind/{path:.*}", app.getFindHandler).Methods("GET")

	r.HandleFunc("/kdbMv/{path:.*}", app.postMoveHandler).Methods("POST")

	// TODO
	// r.HandleFunc("/kdbCp/{path:.*", app.postCopyHandler).Methods("POST")

	r.HandleFunc("/kdbMeta/{path:.*}", app.postMetaHandler).Methods("POST")
	r.HandleFunc("/kdbMeta/{path:.*}", app.deleteMetaHandler).Methods("DELETE")

	return r
}

func parseKeyNameFromURL(r *http.Request) string {
	vars := mux.Vars(r)
	keyName := vars["path"]

	if len(keyName) == 0 {
		return "/"
	}

	return keyName
}

func stringBody(r *http.Request) (string, error) {
	value := ""

	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&value); err != nil {
		return "", err
	}

	return value, nil
}

func notFound(w http.ResponseWriter) {
	w.WriteHeader(http.StatusNotFound)
}

func noContent(w http.ResponseWriter) {
	w.WriteHeader(http.StatusNoContent)
}

func badRequest(w http.ResponseWriter) {
	w.WriteHeader(http.StatusBadRequest)
}

func created(w http.ResponseWriter) {
	w.WriteHeader(http.StatusCreated)
}

func writeError(w http.ResponseWriter, err error) {
	badRequest(w)

	writeResponse(w, map[string]string{
		"error": err.Error(),
	})
}

func writeResponse(w http.ResponseWriter, response interface{}) {
	w.Header().Set("Content-Type", "application/json")

	js, _ := json.Marshal(response)

	w.Write(js)
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package uniter

import (
	"fmt"

	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/api/watcher"
)

// This module implements a subset of the interface provided by
// state.RelationUnit, as needed by the uniter API.
// Most of this is pretty much a verbatim copy of the code in
// state/relationunit.go, except for a few API-specific changes.

// TODO: Only the calls that do not call the API are implemented,
// without tests for now. Others are added as placeholders,
// the actual implementation will come in a follow-up.

// RelationUnit holds information about a single unit in a relation,
// and allows clients to conveniently access unit-specific
// functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint Endpoint
	scope    string
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() Endpoint {
	return ru.endpoint
}

// PrivateAddress returns the private address of the unit and whether
// it is valid.
//
// NOTE: This differs from state.RelationUnit.PrivateAddress() by
// returning an error instead of a bool, because it needs to make an
// API call.
func (ru *RelationUnit) PrivateAddress() (string, error) {
	return ru.unit.PrivateAddress()
}

// TODO: Translate ErrCannotEnterScope and ErrCannotEnterScopeYet
// errors from EnterScope() into their api/params equivalents with the
// respective codes ( CodeCannotEnterScope and
// CodeCannotEnterScopeYet) and message contents, matching the part of
// the error message after the colon.

// EnterScope ensures that the unit has entered its scope in the relation.
// When the unit has already entered its relation scope, EnterScope will report
// success but make no changes to state.
//
// Otherwise, assuming both the relation and the unit are alive, it will enter
// scope.
//
// If the unit is a principal and the relation has container scope, EnterScope
// will also create the required subordinate unit, if it does not already exist;
// this is because there's no point having a principal in scope if there is no
// corresponding subordinate to join it.
//
// Once a unit has entered a scope, it stays in scope without further
// intervention; the relation will not be able to become Dead until all units
// have departed its scopes.
//
// NOTE: Unlike state.RelatioUnit.EnterScope(), this method does not take
// settings, because uniter only uses this to supply the unit's private
// address, but this is not done at the server-side by the API.
func (ru *RelationUnit) EnterScope() error {
	// TODO: Call Uniter.EnterScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// LeaveScope signals that the unit has left its scope in the relation.
// After the unit has left its relation scope, it is no longer a member
// of the relation; if the relation is dying when its last member unit
// leaves, it is removed immediately. It is not an error to leave a scope
// that the unit is not, or never was, a member of.
func (ru *RelationUnit) LeaveScope() error {
	// TODO: Call Uniter.LeaveScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// Settings returns a Settings which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*Settings, error) {
	// TODO: Call Uniter.ReadSettings(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// ReadSettings returns a map holding the settings of the unit with the
// supplied name within this relation. An error will be returned if the
// relation no longer exists, or if the unit's service is not part of the
// relation, or the settings are invalid; but mere non-existence of the
// unit is not grounds for an error, because the unit settings are
// guaranteed to persist for the lifetime of the relation, regardless
// of the lifetime of the unit.
func (ru *RelationUnit) ReadSettings(uname string) (m map[string]interface{}, err error) {
	// TODO: Convert uname to a remoteUnitTag, then call
	// Uniter.ReadRemoteSettings(), passing ru.relation.Tag(),
	// ru.unit.Tag() (as local unit), and remoteUnitTag as arguments.
	panic("not implemented")
}

// Watch returns a watcher that notifies of changes to conterpart
// units in the relation.
func (ru *RelationUnit) Watch() (watcher.RelationUnitsWatcher, error) {
	var results params.RelationUnitsWatchResults
	args := params.RelationUnits{
		RelationUnits: []params.RelationUnit{
			{Relation: ru.relation.tag, Unit: ru.unit.tag},
		},
	}
	err := ru.st.caller.Call("Uniter", "", "WatchRelationUnits", args, &results)
	if err != nil {
		return nil, err
	}
	if len(results.Results) != 1 {
		return nil, fmt.Errorf("expected one result, got %d", len(results.Results))
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	w := watcher.NewRelationUnitsWatcher(ru.st.caller, result)
	return w, nil
}
<commit_msg>Changes after review<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package uniter

import (
	"fmt"

	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/state/api/watcher"
)

// This module implements a subset of the interface provided by
// state.RelationUnit, as needed by the uniter API.
// Most of this is pretty much a verbatim copy of the code in
// state/relationunit.go, except for a few API-specific changes.

// TODO: Only the calls that do not call the API are implemented,
// without tests for now. Others are added as placeholders,
// the actual implementation will come in a follow-up.

// RelationUnit holds information about a single unit in a relation,
// and allows clients to conveniently access unit-specific
// functionality.
type RelationUnit struct {
	st       *State
	relation *Relation
	unit     *Unit
	endpoint Endpoint
	scope    string
}

// Relation returns the relation associated with the unit.
func (ru *RelationUnit) Relation() *Relation {
	return ru.relation
}

// Endpoint returns the relation endpoint that defines the unit's
// participation in the relation.
func (ru *RelationUnit) Endpoint() Endpoint {
	return ru.endpoint
}

// PrivateAddress returns the private address of the unit and whether
// it is valid.
//
// NOTE: This differs from state.RelationUnit.PrivateAddress() by
// returning an error instead of a bool, because it needs to make an
// API call.
func (ru *RelationUnit) PrivateAddress() (string, error) {
	return ru.unit.PrivateAddress()
}

// TODO: Translate ErrCannotEnterScope and ErrCannotEnterScopeYet
// errors from EnterScope() into their api/params equivalents with the
// respective codes ( CodeCannotEnterScope and
// CodeCannotEnterScopeYet) and message contents, matching the part of
// the error message after the colon.

// EnterScope ensures that the unit has entered its scope in the relation.
// When the unit has already entered its relation scope, EnterScope will report
// success but make no changes to state.
//
// Otherwise, assuming both the relation and the unit are alive, it will enter
// scope.
//
// If the unit is a principal and the relation has container scope, EnterScope
// will also create the required subordinate unit, if it does not already exist;
// this is because there's no point having a principal in scope if there is no
// corresponding subordinate to join it.
//
// Once a unit has entered a scope, it stays in scope without further
// intervention; the relation will not be able to become Dead until all units
// have departed its scopes.
//
// NOTE: Unlike state.RelatioUnit.EnterScope(), this method does not take
// settings, because uniter only uses this to supply the unit's private
// address, but this is not done at the server-side by the API.
func (ru *RelationUnit) EnterScope() error {
	// TODO: Call Uniter.EnterScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// LeaveScope signals that the unit has left its scope in the relation.
// After the unit has left its relation scope, it is no longer a member
// of the relation; if the relation is dying when its last member unit
// leaves, it is removed immediately. It is not an error to leave a scope
// that the unit is not, or never was, a member of.
func (ru *RelationUnit) LeaveScope() error {
	// TODO: Call Uniter.LeaveScope(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// Settings returns a Settings which allows access to the unit's settings
// within the relation.
func (ru *RelationUnit) Settings() (*Settings, error) {
	// TODO: Call Uniter.ReadSettings(), passing ru.relation.Tag() and
	// ru.unit.Tag() as arguments.
	panic("not implemented")
}

// ReadSettings returns a map holding the settings of the unit with the
// supplied name within this relation. An error will be returned if the
// relation no longer exists, or if the unit's service is not part of the
// relation, or the settings are invalid; but mere non-existence of the
// unit is not grounds for an error, because the unit settings are
// guaranteed to persist for the lifetime of the relation, regardless
// of the lifetime of the unit.
func (ru *RelationUnit) ReadSettings(uname string) (m map[string]interface{}, err error) {
	// TODO: Convert uname to a remoteUnitTag, then call
	// Uniter.ReadRemoteSettings(), passing ru.relation.Tag(),
	// ru.unit.Tag() (as local unit), and remoteUnitTag as arguments.
	panic("not implemented")
}

// Watch returns a watcher that notifies of changes to counterpart
// units in the relation.
func (ru *RelationUnit) Watch() (watcher.RelationUnitsWatcher, error) {
	var results params.RelationUnitsWatchResults
	args := params.RelationUnits{
		RelationUnits: []params.RelationUnit{
			{Relation: ru.relation.tag, Unit: ru.unit.tag},
		},
	}
	err := ru.st.caller.Call("Uniter", "", "WatchRelationUnits", args, &results)
	if err != nil {
		return nil, err
	}
	if len(results.Results) != 1 {
		return nil, fmt.Errorf("expected one result, got %d", len(results.Results))
	}
	result := results.Results[0]
	if result.Error != nil {
		return nil, result.Error
	}
	w := watcher.NewRelationUnitsWatcher(ru.st.caller, result)
	return w, nil
}
<|endoftext|>


<commit_before>// Copyright 2015-present, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package objcache

import (
	"testing"

	"github.com/corestoreio/errors"
	"github.com/corestoreio/pkg/util/assert"
)

func TestNewProcessor_NewError(t *testing.T) {

	t.Run("level1 error", func(t *testing.T) {
		p, err := NewService(
			func() (Storager, error) { return nil, errors.NotImplemented.Newf("ups") },
			NewBlackHoleClient(nil),
			nil,
		)
		assert.Nil(t, p)
		assert.True(t, errors.NotImplemented.Match(err), "Error: %s", err)
	})

	t.Run("level2 error", func(t *testing.T) {
		p, err := NewService(
			NewBlackHoleClient(nil),
			func() (Storager, error) { return nil, errors.NotImplemented.Newf("ups") },
			nil,
		)
		assert.Nil(t, p)
		assert.True(t, errors.NotImplemented.Match(err), "Error: %s", err)
	})
}
<commit_msg>storage/objcache: Add test for Marshallers without Codec<commit_after>// Copyright 2015-present, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package objcache

import (
	"context"
	"encoding"
	"testing"

	"github.com/corestoreio/errors"
	"github.com/corestoreio/pkg/util/assert"
)

func TestNewProcessor_NewError(t *testing.T) {

	t.Run("level1 error", func(t *testing.T) {
		p, err := NewService(
			func() (Storager, error) { return nil, errors.NotImplemented.Newf("ups") },
			NewBlackHoleClient(nil),
			nil,
		)
		assert.Nil(t, p)
		assert.True(t, errors.NotImplemented.Match(err), "Error: %s", err)
	})

	t.Run("level2 error", func(t *testing.T) {
		p, err := NewService(
			NewBlackHoleClient(nil),
			func() (Storager, error) { return nil, errors.NotImplemented.Newf("ups") },
			nil,
		)
		assert.Nil(t, p)
		assert.True(t, errors.NotImplemented.Match(err), "Error: %s", err)
	})
}

var (
	_ encoding.TextMarshaler     = (*encodingText)(nil)
	_ encoding.TextUnmarshaler   = (*encodingText)(nil)
	_ encoding.BinaryUnmarshaler = (*encodingBinary)(nil)
	_ encoding.BinaryMarshaler   = (*encodingBinary)(nil)
)

type encodingText string

func (e *encodingText) UnmarshalText(text []byte) error {
	*e = encodingText(text)
	return nil
}
func (e encodingText) MarshalText() (text []byte, err error) { return []byte(e), nil }

type encodingBinary string

func (e *encodingBinary) UnmarshalBinary(text []byte) error {
	*e = encodingBinary(text)
	return nil
}
func (e encodingBinary) MarshalBinary() (text []byte, err error) { return []byte(e), nil }

func TestEncoding_Text_Binary(t *testing.T) {
	t.Parallel()

	// Not using any codec
	p, err := NewService(NewCacheSimpleInmemory, NewCacheSimpleInmemory, &ServiceOptions{Codec: nil})
	assert.NoError(t, err)
	defer assert.NoError(t, p.Close())

	ctx := context.TODO()
	t.Run("Text", func(t *testing.T) {
		obj := encodingText("Hello World 🎉")
		err := p.Set(ctx, "kt", obj, 0)
		assert.NoError(t, err)

		var obj2 encodingText
		err = p.Get(ctx, "kt", &obj2)
		assert.NoError(t, err)
		assert.Exactly(t, obj, obj2)
		assert.NotEmpty(t, obj2)
	})
	t.Run("Binary", func(t *testing.T) {
		obj := encodingBinary("Hello World 🎉")
		err := p.Set(ctx, "kt", obj, 0)
		assert.NoError(t, err)

		var obj2 encodingBinary
		err = p.Get(ctx, "kt", &obj2)
		assert.NoError(t, err)
		assert.Exactly(t, obj, obj2)
		assert.NotEmpty(t, obj2)
	})
}
<|endoftext|>


<commit_before><commit_msg>cp into existing struct<commit_after><|endoftext|>