<commit_before>package starbound

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"errors"
	"io"
)

var (
	ErrDidNotReachLeaf  = errors.New("starbound: did not reach a leaf node")
	ErrInvalidHeader    = errors.New("starbound: invalid header")
	ErrInvalidKeyLength = errors.New("starbound: invalid key length")
	ErrInvalidSBON      = errors.New("starbound: invalid SBON data")
	ErrKeyNotFound      = errors.New("starbound: key not found")
)

const (
	WorldDatabaseName = "World4"
)

type Tile struct {
	ForegroundMaterial    int16
	ForegroundHueShift    uint8
	ForegroundVariant     uint8
	ForegroundMod         int16
	ForegroundModHueShift uint8
	BackgroundMaterial    int16
	BackgroundHueShift    uint8
	BackgroundVariant     uint8
	BackgroundMod         int16
	BackgroundModHueShift uint8
	Liquid                uint8
	LiquidLevel           float32
	LiquidPressure        float32
	LiquidInfinite        uint8 // bool
	Collision             uint8
	DungeonID             uint16
	Biome1, Biome2        uint8
	Indestructible        uint8 // bool
}

// NewWorld creates and initializes a new World using r as the data source.
func NewWorld(r io.ReaderAt) (w *World, err error) {
	db, err := NewBTreeDB5(r)
	if err != nil {
		return
	}
	if db.Name != WorldDatabaseName || db.KeySize != 5 {
		return nil, ErrInvalidHeader
	}
	return &World{db}, nil
}

// A World is a representation of a Starbound world, enabling read access to
// individual regions in the world as well as its metadata.
type World struct {
	*BTreeDB5
}

func (w *World) Get(layer, x, y int) (data []byte, err error) {
	src, err := w.GetReader(layer, x, y)
	if err != nil {
		return
	}
	dst := new(bytes.Buffer)
	_, err = io.Copy(dst, src)
	if err != nil {
		return
	}
	return dst.Bytes(), nil
}

func (w *World) GetReader(layer, x, y int) (r io.Reader, err error) {
	key := []byte{byte(layer), byte(x >> 8), byte(x), byte(y >> 8), byte(y)}
	lr, err := w.BTreeDB5.GetReader(key)
	if err != nil {
		return nil, err
	}
	return zlib.NewReader(lr)
}

func (w *World) GetTiles(x, y int) (t []Tile, err error) {
	r, err := w.GetReader(1, x, y)
	if err != nil {
		return
	}
	// Ignore the first three bytes.
	// TODO: Do something with these bytes?
	discard := make([]byte, 3)
	_, err = io.ReadFull(r, discard)
	if err != nil {
		return
	}
	t = make([]Tile, 1024) // 32x32 tiles in a region
	err = binary.Read(r, binary.BigEndian, t)
	return
}

// Reads a 32-bit integer from the provided buffer and offset.
func getInt(data []byte, n int) int {
	return int(data[n])<<24 | int(data[n+1])<<16 | int(data[n+2])<<8 | int(data[n+3])
}

type logger interface {
	Fatalf(format string, args ...interface{})
}
<commit_msg>Remove another constant<commit_after>package starbound

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"errors"
	"io"
)

var (
	ErrDidNotReachLeaf  = errors.New("starbound: did not reach a leaf node")
	ErrInvalidHeader    = errors.New("starbound: invalid header")
	ErrInvalidKeyLength = errors.New("starbound: invalid key length")
	ErrInvalidSBON      = errors.New("starbound: invalid SBON data")
	ErrKeyNotFound      = errors.New("starbound: key not found")
)

type Tile struct {
	ForegroundMaterial    int16
	ForegroundHueShift    uint8
	ForegroundVariant     uint8
	ForegroundMod         int16
	ForegroundModHueShift uint8
	BackgroundMaterial    int16
	BackgroundHueShift    uint8
	BackgroundVariant     uint8
	BackgroundMod         int16
	BackgroundModHueShift uint8
	Liquid                uint8
	LiquidLevel           float32
	LiquidPressure        float32
	LiquidInfinite        uint8 // bool
	Collision             uint8
	DungeonID             uint16
	Biome1, Biome2        uint8
	Indestructible        uint8 // bool
}

// NewWorld creates and initializes a new World using r as the data source.
func NewWorld(r io.ReaderAt) (w *World, err error) {
	db, err := NewBTreeDB5(r)
	if err != nil {
		return
	}
	if db.Name != "World4" || db.KeySize != 5 {
		return nil, ErrInvalidHeader
	}
	return &World{db}, nil
}

// A World is a representation of a Starbound world, enabling read access to
// individual regions in the world as well as its metadata.
type World struct {
	*BTreeDB5
}

func (w *World) Get(layer, x, y int) (data []byte, err error) {
	src, err := w.GetReader(layer, x, y)
	if err != nil {
		return
	}
	dst := new(bytes.Buffer)
	_, err = io.Copy(dst, src)
	if err != nil {
		return
	}
	return dst.Bytes(), nil
}

func (w *World) GetReader(layer, x, y int) (r io.Reader, err error) {
	key := []byte{byte(layer), byte(x >> 8), byte(x), byte(y >> 8), byte(y)}
	lr, err := w.BTreeDB5.GetReader(key)
	if err != nil {
		return nil, err
	}
	return zlib.NewReader(lr)
}

func (w *World) GetTiles(x, y int) (t []Tile, err error) {
	r, err := w.GetReader(1, x, y)
	if err != nil {
		return
	}
	// Ignore the first three bytes.
	// TODO: Do something with these bytes?
	discard := make([]byte, 3)
	_, err = io.ReadFull(r, discard)
	if err != nil {
		return
	}
	t = make([]Tile, 1024) // 32x32 tiles in a region
	err = binary.Read(r, binary.BigEndian, t)
	return
}

// Reads a 32-bit integer from the provided buffer and offset.
func getInt(data []byte, n int) int {
	return int(data[n])<<24 | int(data[n+1])<<16 | int(data[n+2])<<8 | int(data[n+3])
}

type logger interface {
	Fatalf(format string, args ...interface{})
}
<|endoftext|>


<commit_before>package docker

import (
	"github.com/jfrogdev/jfrog-cli-go/jfrog-cli/artifactory/utils"
	"io"
	"os/exec"
	"path"
	"strings"
)

func New(imageTag string) Image {
	return &image{tag: imageTag}
}

// Docker image
type Image interface {
	Push() error
	Id() (string, error)
	ParentId() (string, error)
	Tag() string
	Path() string
	Name() string
}

// Internal implementation of docker image
type image struct {
	tag string
}

// Push docker image
func (image *image) Push() error {
	cmd := &pushCmd{image: image}
	return utils.RunCmd(cmd)
}

// Get docker image tag
func (image *image) Tag() string {
	return image.tag
}

// Get docker image ID
func (image *image) Id() (string, error) {
	cmd := &getImageIdCmd{image: image}
	content, err := utils.RunCmdOutput(cmd)
	return strings.Trim(string(content), "\n"), err
}

// Get docker parent image ID
func (image *image) ParentId() (string, error) {
	cmd := &getParentId{image: image}
	content, err := utils.RunCmdOutput(cmd)
	return strings.Trim(string(content), "\n"), err
}

// Get docker image relative path in Artifactory
func (image *image) Path() string {
	indexOfFirstSlash := strings.Index(image.tag, "/")
	indexOfLastColon := strings.LastIndex(image.tag, ":")

	if indexOfLastColon < 0 || indexOfLastColon < indexOfFirstSlash {
		return path.Join(image.tag[indexOfFirstSlash:], "latest")
	}
	return path.Join(image.tag[indexOfFirstSlash:indexOfLastColon], image.tag[indexOfLastColon+1:])
}

// Get docker image name
func (image *image) Name() string {
	indexOfFirstSlash := strings.Index(image.tag, "/")
	indexOfLastColon := strings.LastIndex(image.tag, ":")

	if indexOfLastColon < 0 || indexOfLastColon < indexOfFirstSlash {
		return image.tag[indexOfFirstSlash+1:] + ":latest"
	}
	return image.tag[indexOfFirstSlash+1:]
}

// Image push command
type pushCmd struct {
	image *image
}

func (pushCmd *pushCmd) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "push")
	cmd = append(cmd, pushCmd.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (pushCmd *pushCmd) GetEnv() map[string]string {
	return map[string]string{}
}

func (pushCmd *pushCmd) GetStdWriter() io.WriteCloser {
	return nil
}
func (pushCmd *pushCmd) GetErrWriter() io.WriteCloser {
	return nil
}

// Image get image id command
type getImageIdCmd struct {
	image *image
}

func (getImageId *getImageIdCmd) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "images")
	cmd = append(cmd, "--format", "{{.ID}}")
	cmd = append(cmd, "--no-trunc")
	cmd = append(cmd, getImageId.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (getImageId *getImageIdCmd) GetEnv() map[string]string {
	return map[string]string{}
}

func (getImageId *getImageIdCmd) GetStdWriter() io.WriteCloser {
	return nil
}
func (getImageId *getImageIdCmd) GetErrWriter() io.WriteCloser {
	return nil
}

// Image get parent image id command
type getParentId struct {
	image *image
}

func (getImageId *getParentId) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "inspect")
	cmd = append(cmd, "--format", "{{.Parent}}")
	cmd = append(cmd, getImageId.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (getImageId *getParentId) GetEnv() map[string]string {
	return map[string]string{}
}

func (getImageId *getParentId) GetStdWriter() io.WriteCloser {
	return nil
}
func (getImageId *getParentId) GetErrWriter() io.WriteCloser {
	return nil
}
<commit_msg>Docker upload fails #160<commit_after>package docker

import (
	"github.com/jfrogdev/jfrog-cli-go/jfrog-cli/artifactory/utils"
	"io"
	"os/exec"
	"path"
	"strings"
)

func New(imageTag string) Image {
	return &image{tag: imageTag}
}

// Docker image
type Image interface {
	Push() error
	Id() (string, error)
	ParentId() (string, error)
	Tag() string
	Path() string
	Name() string
}

// Internal implementation of docker image
type image struct {
	tag string
}

// Push docker image
func (image *image) Push() error {
	cmd := &pushCmd{image: image}
	return utils.RunCmd(cmd)
}

// Get docker image tag
func (image *image) Tag() string {
	return image.tag
}

// Get docker image ID
func (image *image) Id() (string, error) {
	cmd := &getImageIdCmd{image: image}
	content, err := utils.RunCmdOutput(cmd)
	return strings.Trim(string(content), "\n"), err
}

// Get docker parent image ID
func (image *image) ParentId() (string, error) {
	cmd := &getParentId{image: image}
	content, err := utils.RunCmdOutput(cmd)
	return strings.Trim(string(content), "\n"), err
}

// Get docker image relative path in Artifactory
func (image *image) Path() string {
	indexOfLastSlash := strings.LastIndex(image.tag, "/")
	indexOfLastColon := strings.LastIndex(image.tag, ":")

	if indexOfLastColon < 0 || indexOfLastColon < indexOfLastSlash {
		return path.Join(image.tag[indexOfLastSlash:], "latest")
	}
	return path.Join(image.tag[indexOfLastSlash:indexOfLastColon], image.tag[indexOfLastColon+1:])
}

// Get docker image name
func (image *image) Name() string {
	indexOfLastSlash := strings.LastIndex(image.tag, "/")
	indexOfLastColon := strings.LastIndex(image.tag, ":")

	if indexOfLastColon < 0 || indexOfLastColon < indexOfLastSlash {
		return image.tag[indexOfLastSlash+1:] + ":latest"
	}
	return image.tag[indexOfLastSlash+1:]
}

// Image push command
type pushCmd struct {
	image *image
}

func (pushCmd *pushCmd) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "push")
	cmd = append(cmd, pushCmd.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (pushCmd *pushCmd) GetEnv() map[string]string {
	return map[string]string{}
}

func (pushCmd *pushCmd) GetStdWriter() io.WriteCloser {
	return nil
}
func (pushCmd *pushCmd) GetErrWriter() io.WriteCloser {
	return nil
}

// Image get image id command
type getImageIdCmd struct {
	image *image
}

func (getImageId *getImageIdCmd) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "images")
	cmd = append(cmd, "--format", "{{.ID}}")
	cmd = append(cmd, "--no-trunc")
	cmd = append(cmd, getImageId.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (getImageId *getImageIdCmd) GetEnv() map[string]string {
	return map[string]string{}
}

func (getImageId *getImageIdCmd) GetStdWriter() io.WriteCloser {
	return nil
}
func (getImageId *getImageIdCmd) GetErrWriter() io.WriteCloser {
	return nil
}

// Image get parent image id command
type getParentId struct {
	image *image
}

func (getImageId *getParentId) GetCmd() *exec.Cmd {
	var cmd []string
	cmd = append(cmd, "docker")
	cmd = append(cmd, "inspect")
	cmd = append(cmd, "--format", "{{.Parent}}")
	cmd = append(cmd, getImageId.image.tag)
	return exec.Command(cmd[0], cmd[1:]...)
}

func (getImageId *getParentId) GetEnv() map[string]string {
	return map[string]string{}
}

func (getImageId *getParentId) GetStdWriter() io.WriteCloser {
	return nil
}
func (getImageId *getParentId) GetErrWriter() io.WriteCloser {
	return nil
}
<|endoftext|>


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

package cli

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"strings"

	ipb "kythe.io/kythe/proto/identifier_proto"
)

type identCommand struct {
	corpora, languages string
}

func (identCommand) Name() string     { return "identifier" }
func (identCommand) Synopsis() string { return "list tickets associated with a given identifier" }
func (identCommand) Usage() string    { return "<identifier>" }
func (c *identCommand) SetFlags(flag *flag.FlagSet) {
	flag.StringVar(&c.corpora, "corpora", "", "Comma-separated list of corpora with which to restrict matches")
	flag.StringVar(&c.languages, "languages", "", "Comma-separated list of languages with which to restrict matches")
}
func (c identCommand) Run(ctx context.Context, flag *flag.FlagSet, api API) error {
	if flag.NArg() == 0 {
		return errors.New("identifier missing")
	} else if flag.NArg() > 1 {
		return fmt.Errorf("only 1 identifier may be given; found: %v", flag.Args())
	}

	req := &ipb.FindRequest{
		Identifier: flag.Arg(0),
	}
	if c.corpora != "" {
		req.Corpus = strings.Split(c.corpora, ",")
	}
	if c.languages != "" {
		req.Languages = strings.Split(c.languages, ",")
	}

	LogRequest(req)
	reply, err := api.IdentifierService.Find(ctx, req)
	if err != nil {
		return err
	}

	return c.displayMatches(reply)
}

func (c identCommand) displayMatches(reply *ipb.FindReply) error {
	if DisplayJSON {
		return PrintJSONMessage(reply)
	}

	for _, m := range reply.Matches {
		kind := m.NodeKind
		if m.NodeSubkind != "" {
			kind += "/" + m.NodeSubkind
		}
		fmt.Sprintf("%s [kind: %s]", m.Ticket, kind)
	}
	return nil
}
<commit_msg>Fix Kythe IdentifierService CLI printing<commit_after>/*
 * Copyright 2017 Google Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cli

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"strings"

	ipb "kythe.io/kythe/proto/identifier_proto"
)

type identCommand struct {
	corpora, languages string
}

func (identCommand) Name() string     { return "identifier" }
func (identCommand) Synopsis() string { return "list tickets associated with a given identifier" }
func (identCommand) Usage() string    { return "<identifier>" }
func (c *identCommand) SetFlags(flag *flag.FlagSet) {
	flag.StringVar(&c.corpora, "corpora", "", "Comma-separated list of corpora with which to restrict matches")
	flag.StringVar(&c.languages, "languages", "", "Comma-separated list of languages with which to restrict matches")
}
func (c identCommand) Run(ctx context.Context, flag *flag.FlagSet, api API) error {
	if flag.NArg() == 0 {
		return errors.New("identifier missing")
	} else if flag.NArg() > 1 {
		return fmt.Errorf("only 1 identifier may be given; found: %v", flag.Args())
	}

	req := &ipb.FindRequest{
		Identifier: flag.Arg(0),
	}
	if c.corpora != "" {
		req.Corpus = strings.Split(c.corpora, ",")
	}
	if c.languages != "" {
		req.Languages = strings.Split(c.languages, ",")
	}

	LogRequest(req)
	reply, err := api.IdentifierService.Find(ctx, req)
	if err != nil {
		return err
	}

	return c.displayMatches(reply)
}

func (c identCommand) displayMatches(reply *ipb.FindReply) error {
	if DisplayJSON {
		return PrintJSONMessage(reply)
	}

	for _, m := range reply.Matches {
		kind := m.NodeKind
		if m.NodeSubkind != "" {
			kind += "/" + m.NodeSubkind
		}
		fmt.Printf("%s [kind: %s]\n", m.Ticket, kind)
	}
	return nil
}
<|endoftext|>


<commit_before>package registry

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/docker/docker/pkg/httputils"
)

const DockerDigestHeader = "Docker-Content-Digest"

func getV2Builder(e *Endpoint) *v2.URLBuilder {
	if e.URLBuilder == nil {
		e.URLBuilder = v2.NewURLBuilder(e.URL)
	}
	return e.URLBuilder
}

func (r *Session) V2RegistryEndpoint(index *IndexInfo) (ep *Endpoint, err error) {
	// TODO check if should use Mirror
	if index.Official {
		ep, err = newEndpoint(REGISTRYSERVER, true, nil)
		if err != nil {
			return
		}
		err = validateEndpoint(ep)
		if err != nil {
			return
		}
	} else if r.indexEndpoint.String() == index.GetAuthConfigKey() {
		ep = r.indexEndpoint
	} else {
		ep, err = NewEndpoint(index, nil)
		if err != nil {
			return
		}
	}

	ep.URLBuilder = v2.NewURLBuilder(ep.URL)
	return
}

// GetV2Authorization gets the authorization needed to the given image
// If readonly access is requested, then the authorization may
// only be used for Get operations.
func (r *Session) GetV2Authorization(ep *Endpoint, imageName string, readOnly bool) (auth *RequestAuthorization, err error) {
	scopes := []string{"pull"}
	if !readOnly {
		scopes = append(scopes, "push")
	}

	logrus.Debugf("Getting authorization for %s %s", imageName, scopes)
	return NewRequestAuthorization(r.GetAuthConfig(true), ep, "repository", imageName, scopes), nil
}

//
// 1) Check if TarSum of each layer exists /v2/
//  1.a) if 200, continue
//  1.b) if 300, then push the
//  1.c) if anything else, err
// 2) PUT the created/signed manifest
//
func (r *Session) GetV2ImageManifest(ep *Endpoint, imageName, tagName string, auth *RequestAuthorization) ([]byte, string, error) {
	routeURL, err := getV2Builder(ep).BuildManifestURL(imageName, tagName)
	if err != nil {
		return nil, "", err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return nil, "", err
	}
	if err := auth.Authorize(req); err != nil {
		return nil, "", err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return nil, "", err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return nil, "", errLoginRequired
		} else if res.StatusCode == 404 {
			return nil, "", ErrDoesNotExist
		}
		return nil, "", httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to fetch for %s:%s", res.StatusCode, imageName, tagName), res)
	}

	manifestBytes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, "", fmt.Errorf("Error while reading the http response: %s", err)
	}

	return manifestBytes, res.Header.Get(DockerDigestHeader), nil
}

// - Succeeded to head image blob (already exists)
// - Failed with no error (continue to Push the Blob)
// - Failed with error
func (r *Session) HeadV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, auth *RequestAuthorization) (bool, error) {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return false, err
	}

	method := "HEAD"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return false, err
	}
	if err := auth.Authorize(req); err != nil {
		return false, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return false, err
	}
	res.Body.Close() // close early, since we're not needing a body on this call .. yet?
	switch {
	case res.StatusCode >= 200 && res.StatusCode < 400:
		// return something indicating no push needed
		return true, nil
	case res.StatusCode == 401:
		return false, errLoginRequired
	case res.StatusCode == 404:
		// return something indicating blob push needed
		return false, nil
	}

	return false, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying head request for %s - %s", res.StatusCode, imageName, dgst), res)
}

func (r *Session) GetV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, blobWrtr io.Writer, auth *RequestAuthorization) error {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return err
	}
	if err := auth.Authorize(req); err != nil {
		return err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return errLoginRequired
		}
		return httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to pull %s blob", res.StatusCode, imageName), res)
	}

	_, err = io.Copy(blobWrtr, res.Body)
	return err
}

func (r *Session) GetV2ImageBlobReader(ep *Endpoint, imageName string, dgst digest.Digest, auth *RequestAuthorization) (io.ReadCloser, int64, error) {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return nil, 0, err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return nil, 0, err
	}
	if err := auth.Authorize(req); err != nil {
		return nil, 0, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return nil, 0, errLoginRequired
		}
		return nil, 0, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to pull %s blob - %s", res.StatusCode, imageName, dgst), res)
	}
	lenStr := res.Header.Get("Content-Length")
	l, err := strconv.ParseInt(lenStr, 10, 64)
	if err != nil {
		return nil, 0, err
	}

	return res.Body, l, err
}

// Push the image to the server for storage.
// 'layer' is an uncompressed reader of the blob to be pushed.
// The server will generate it's own checksum calculation.
func (r *Session) PutV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, blobRdr io.Reader, auth *RequestAuthorization) error {
	location, err := r.initiateBlobUpload(ep, imageName, auth)
	if err != nil {
		return err
	}

	method := "PUT"
	logrus.Debugf("[registry] Calling %q %s", method, location)
	req, err := http.NewRequest(method, location, ioutil.NopCloser(blobRdr))
	if err != nil {
		return err
	}
	queryParams := req.URL.Query()
	queryParams.Add("digest", dgst.String())
	req.URL.RawQuery = queryParams.Encode()
	if err := auth.Authorize(req); err != nil {
		return err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 201 {
		if res.StatusCode == 401 {
			return errLoginRequired
		}
		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return err
		}
		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to push %s blob - %s", res.StatusCode, imageName, dgst), res)
	}

	return nil
}

// initiateBlobUpload gets the blob upload location for the given image name.
func (r *Session) initiateBlobUpload(ep *Endpoint, imageName string, auth *RequestAuthorization) (location string, err error) {
	routeURL, err := getV2Builder(ep).BuildBlobUploadURL(imageName)
	if err != nil {
		return "", err
	}

	logrus.Debugf("[registry] Calling %q %s", "POST", routeURL)
	req, err := http.NewRequest("POST", routeURL, nil)
	if err != nil {
		return "", err
	}

	if err := auth.Authorize(req); err != nil {
		return "", err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return "", err
	}

	if res.StatusCode != http.StatusAccepted {
		if res.StatusCode == http.StatusUnauthorized {
			return "", errLoginRequired
		}
		if res.StatusCode == http.StatusNotFound {
			return "", ErrDoesNotExist
		}

		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return "", err
		}

		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return "", httputils.NewHTTPRequestError(fmt.Sprintf("Server error: unexpected %d response status trying to initiate upload of %s", res.StatusCode, imageName), res)
	}

	if location = res.Header.Get("Location"); location == "" {
		return "", fmt.Errorf("registry did not return a Location header for resumable blob upload for image %s", imageName)
	}

	return
}

// Finally Push the (signed) manifest of the blobs we've just pushed
func (r *Session) PutV2ImageManifest(ep *Endpoint, imageName, tagName string, signedManifest, rawManifest []byte, auth *RequestAuthorization) (digest.Digest, error) {
	routeURL, err := getV2Builder(ep).BuildManifestURL(imageName, tagName)
	if err != nil {
		return "", err
	}

	method := "PUT"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, bytes.NewReader(signedManifest))
	if err != nil {
		return "", err
	}
	if err := auth.Authorize(req); err != nil {
		return "", err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	// All 2xx and 3xx responses can be accepted for a put.
	if res.StatusCode >= 400 {
		if res.StatusCode == 401 {
			return "", errLoginRequired
		}
		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return "", err
		}
		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return "", httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to push %s:%s manifest", res.StatusCode, imageName, tagName), res)
	}

	hdrDigest, err := digest.ParseDigest(res.Header.Get(DockerDigestHeader))
	if err != nil {
		return "", fmt.Errorf("invalid manifest digest from registry: %s", err)
	}

	dgstVerifier, err := digest.NewDigestVerifier(hdrDigest)
	if err != nil {
		return "", fmt.Errorf("invalid manifest digest from registry: %s", err)
	}

	dgstVerifier.Write(rawManifest)

	if !dgstVerifier.Verified() {
		computedDigest, _ := digest.FromBytes(rawManifest)
		return "", fmt.Errorf("unable to verify manifest digest: registry has %q, computed %q", hdrDigest, computedDigest)
	}

	return hdrDigest, nil
}

type remoteTags struct {
	Name string   `json:"name"`
	Tags []string `json:"tags"`
}

// Given a repository name, returns a json array of string tags
func (r *Session) GetV2RemoteTags(ep *Endpoint, imageName string, auth *RequestAuthorization) ([]string, error) {
	routeURL, err := getV2Builder(ep).BuildTagsURL(imageName)
	if err != nil {
		return nil, err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return nil, err
	}
	if err := auth.Authorize(req); err != nil {
		return nil, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return nil, errLoginRequired
		} else if res.StatusCode == 404 {
			return nil, ErrDoesNotExist
		}
		return nil, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to fetch for %s", res.StatusCode, imageName), res)
	}

	var remote remoteTags
	if err := json.NewDecoder(res.Body).Decode(&remote); err != nil {
		return nil, fmt.Errorf("Error while decoding the http response: %s", err)
	}
	return remote.Tags, nil
}
<commit_msg>Properly verify manifests and layer digests on pull<commit_after>package registry

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/digest"
	"github.com/docker/distribution/registry/api/v2"
	"github.com/docker/docker/pkg/httputils"
)

const DockerDigestHeader = "Docker-Content-Digest"

func getV2Builder(e *Endpoint) *v2.URLBuilder {
	if e.URLBuilder == nil {
		e.URLBuilder = v2.NewURLBuilder(e.URL)
	}
	return e.URLBuilder
}

func (r *Session) V2RegistryEndpoint(index *IndexInfo) (ep *Endpoint, err error) {
	// TODO check if should use Mirror
	if index.Official {
		ep, err = newEndpoint(REGISTRYSERVER, true, nil)
		if err != nil {
			return
		}
		err = validateEndpoint(ep)
		if err != nil {
			return
		}
	} else if r.indexEndpoint.String() == index.GetAuthConfigKey() {
		ep = r.indexEndpoint
	} else {
		ep, err = NewEndpoint(index, nil)
		if err != nil {
			return
		}
	}

	ep.URLBuilder = v2.NewURLBuilder(ep.URL)
	return
}

// GetV2Authorization gets the authorization needed to the given image
// If readonly access is requested, then the authorization may
// only be used for Get operations.
func (r *Session) GetV2Authorization(ep *Endpoint, imageName string, readOnly bool) (auth *RequestAuthorization, err error) {
	scopes := []string{"pull"}
	if !readOnly {
		scopes = append(scopes, "push")
	}

	logrus.Debugf("Getting authorization for %s %s", imageName, scopes)
	return NewRequestAuthorization(r.GetAuthConfig(true), ep, "repository", imageName, scopes), nil
}

//
// 1) Check if TarSum of each layer exists /v2/
//  1.a) if 200, continue
//  1.b) if 300, then push the
//  1.c) if anything else, err
// 2) PUT the created/signed manifest
//

// GetV2ImageManifest simply fetches the bytes of a manifest and the remote
// digest, if available in the request. Note that the application shouldn't
// rely on the untrusted remoteDigest, and should also verify against a
// locally provided digest, if applicable.
func (r *Session) GetV2ImageManifest(ep *Endpoint, imageName, tagName string, auth *RequestAuthorization) (remoteDigest digest.Digest, p []byte, err error) {
	routeURL, err := getV2Builder(ep).BuildManifestURL(imageName, tagName)
	if err != nil {
		return "", nil, err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return "", nil, err
	}

	if err := auth.Authorize(req); err != nil {
		return "", nil, err
	}

	res, err := r.client.Do(req)
	if err != nil {
		return "", nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return "", nil, errLoginRequired
		} else if res.StatusCode == 404 {
			return "", nil, ErrDoesNotExist
		}
		return "", nil, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to fetch for %s:%s", res.StatusCode, imageName, tagName), res)
	}

	p, err = ioutil.ReadAll(res.Body)
	if err != nil {
		return "", nil, fmt.Errorf("Error while reading the http response: %s", err)
	}

	dgstHdr := res.Header.Get(DockerDigestHeader)
	if dgstHdr != "" {
		remoteDigest, err = digest.ParseDigest(dgstHdr)
		if err != nil {
			// NOTE(stevvooe): Including the remote digest is optional. We
			// don't need to verify against it, but it is good practice.
			remoteDigest = ""
			logrus.Debugf("error parsing remote digest when fetching %v: %v", routeURL, err)
		}
	}

	return
}

// - Succeeded to head image blob (already exists)
// - Failed with no error (continue to Push the Blob)
// - Failed with error
func (r *Session) HeadV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, auth *RequestAuthorization) (bool, error) {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return false, err
	}

	method := "HEAD"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return false, err
	}
	if err := auth.Authorize(req); err != nil {
		return false, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return false, err
	}
	res.Body.Close() // close early, since we're not needing a body on this call .. yet?
	switch {
	case res.StatusCode >= 200 && res.StatusCode < 400:
		// return something indicating no push needed
		return true, nil
	case res.StatusCode == 401:
		return false, errLoginRequired
	case res.StatusCode == 404:
		// return something indicating blob push needed
		return false, nil
	}

	return false, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying head request for %s - %s", res.StatusCode, imageName, dgst), res)
}

func (r *Session) GetV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, blobWrtr io.Writer, auth *RequestAuthorization) error {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return err
	}
	if err := auth.Authorize(req); err != nil {
		return err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return errLoginRequired
		}
		return httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to pull %s blob", res.StatusCode, imageName), res)
	}

	_, err = io.Copy(blobWrtr, res.Body)
	return err
}

func (r *Session) GetV2ImageBlobReader(ep *Endpoint, imageName string, dgst digest.Digest, auth *RequestAuthorization) (io.ReadCloser, int64, error) {
	routeURL, err := getV2Builder(ep).BuildBlobURL(imageName, dgst)
	if err != nil {
		return nil, 0, err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return nil, 0, err
	}
	if err := auth.Authorize(req); err != nil {
		return nil, 0, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return nil, 0, errLoginRequired
		}
		return nil, 0, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to pull %s blob - %s", res.StatusCode, imageName, dgst), res)
	}
	lenStr := res.Header.Get("Content-Length")
	l, err := strconv.ParseInt(lenStr, 10, 64)
	if err != nil {
		return nil, 0, err
	}

	return res.Body, l, err
}

// Push the image to the server for storage.
// 'layer' is an uncompressed reader of the blob to be pushed.
// The server will generate it's own checksum calculation.
func (r *Session) PutV2ImageBlob(ep *Endpoint, imageName string, dgst digest.Digest, blobRdr io.Reader, auth *RequestAuthorization) error {
	location, err := r.initiateBlobUpload(ep, imageName, auth)
	if err != nil {
		return err
	}

	method := "PUT"
	logrus.Debugf("[registry] Calling %q %s", method, location)
	req, err := http.NewRequest(method, location, ioutil.NopCloser(blobRdr))
	if err != nil {
		return err
	}
	queryParams := req.URL.Query()
	queryParams.Add("digest", dgst.String())
	req.URL.RawQuery = queryParams.Encode()
	if err := auth.Authorize(req); err != nil {
		return err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != 201 {
		if res.StatusCode == 401 {
			return errLoginRequired
		}
		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return err
		}
		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to push %s blob - %s", res.StatusCode, imageName, dgst), res)
	}

	return nil
}

// initiateBlobUpload gets the blob upload location for the given image name.
func (r *Session) initiateBlobUpload(ep *Endpoint, imageName string, auth *RequestAuthorization) (location string, err error) {
	routeURL, err := getV2Builder(ep).BuildBlobUploadURL(imageName)
	if err != nil {
		return "", err
	}

	logrus.Debugf("[registry] Calling %q %s", "POST", routeURL)
	req, err := http.NewRequest("POST", routeURL, nil)
	if err != nil {
		return "", err
	}

	if err := auth.Authorize(req); err != nil {
		return "", err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return "", err
	}

	if res.StatusCode != http.StatusAccepted {
		if res.StatusCode == http.StatusUnauthorized {
			return "", errLoginRequired
		}
		if res.StatusCode == http.StatusNotFound {
			return "", ErrDoesNotExist
		}

		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return "", err
		}

		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return "", httputils.NewHTTPRequestError(fmt.Sprintf("Server error: unexpected %d response status trying to initiate upload of %s", res.StatusCode, imageName), res)
	}

	if location = res.Header.Get("Location"); location == "" {
		return "", fmt.Errorf("registry did not return a Location header for resumable blob upload for image %s", imageName)
	}

	return
}

// Finally Push the (signed) manifest of the blobs we've just pushed
func (r *Session) PutV2ImageManifest(ep *Endpoint, imageName, tagName string, signedManifest, rawManifest []byte, auth *RequestAuthorization) (digest.Digest, error) {
	routeURL, err := getV2Builder(ep).BuildManifestURL(imageName, tagName)
	if err != nil {
		return "", err
	}

	method := "PUT"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)
	req, err := http.NewRequest(method, routeURL, bytes.NewReader(signedManifest))
	if err != nil {
		return "", err
	}
	if err := auth.Authorize(req); err != nil {
		return "", err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()

	// All 2xx and 3xx responses can be accepted for a put.
	if res.StatusCode >= 400 {
		if res.StatusCode == 401 {
			return "", errLoginRequired
		}
		errBody, err := ioutil.ReadAll(res.Body)
		if err != nil {
			return "", err
		}
		logrus.Debugf("Unexpected response from server: %q %#v", errBody, res.Header)
		return "", httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to push %s:%s manifest", res.StatusCode, imageName, tagName), res)
	}

	hdrDigest, err := digest.ParseDigest(res.Header.Get(DockerDigestHeader))
	if err != nil {
		return "", fmt.Errorf("invalid manifest digest from registry: %s", err)
	}

	dgstVerifier, err := digest.NewDigestVerifier(hdrDigest)
	if err != nil {
		return "", fmt.Errorf("invalid manifest digest from registry: %s", err)
	}

	dgstVerifier.Write(rawManifest)

	if !dgstVerifier.Verified() {
		computedDigest, _ := digest.FromBytes(rawManifest)
		return "", fmt.Errorf("unable to verify manifest digest: registry has %q, computed %q", hdrDigest, computedDigest)
	}

	return hdrDigest, nil
}

type remoteTags struct {
	Name string   `json:"name"`
	Tags []string `json:"tags"`
}

// Given a repository name, returns a json array of string tags
func (r *Session) GetV2RemoteTags(ep *Endpoint, imageName string, auth *RequestAuthorization) ([]string, error) {
	routeURL, err := getV2Builder(ep).BuildTagsURL(imageName)
	if err != nil {
		return nil, err
	}

	method := "GET"
	logrus.Debugf("[registry] Calling %q %s", method, routeURL)

	req, err := http.NewRequest(method, routeURL, nil)
	if err != nil {
		return nil, err
	}
	if err := auth.Authorize(req); err != nil {
		return nil, err
	}
	res, err := r.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		if res.StatusCode == 401 {
			return nil, errLoginRequired
		} else if res.StatusCode == 404 {
			return nil, ErrDoesNotExist
		}
		return nil, httputils.NewHTTPRequestError(fmt.Sprintf("Server error: %d trying to fetch for %s", res.StatusCode, imageName), res)
	}

	var remote remoteTags
	if err := json.NewDecoder(res.Body).Decode(&remote); err != nil {
		return nil, fmt.Errorf("Error while decoding the http response: %s", err)
	}
	return remote.Tags, nil
}
<|endoftext|>


<commit_before>package main

import (
	"math"
	"math/rand"
)

// Distribution provides an interface to model a statistical distribution.
type Distribution interface {
	Advance()
	Get() float64 // should be idempotent
}

// NormalDistribution models a normal distribution.
type NormalDistribution struct {
	Mean   float64
	StdDev float64

	value float64
}

func ND(mean, stddev float64) *NormalDistribution {
	return &NormalDistribution{Mean: mean, StdDev: stddev}
}

// Advance advances this distribution. Since a normal distribution is
// stateless, this is just overwrites the internal cache value.
func (d *NormalDistribution) Advance() {
	d.value = rand.NormFloat64()*d.StdDev + d.Mean
}

// Get returns the last computed value for this distribution.
func (d *NormalDistribution) Get() float64 {
	return d.value
}

// UniformDistribution models a uniform distribution.
type UniformDistribution struct {
	Low  float64
	High float64

	value float64
}

func UD(low, high float64) *UniformDistribution {
	return &UniformDistribution{Low: low, High: high}
}

// Advance advances this distribution. Since a uniform distribution is
// stateless, this is just overwrites the internal cache value.
func (d *UniformDistribution) Advance() {
	x := rand.Float64() // uniform
	x *= d.High - d.Low
	x += d.Low
	d.value = x
}

// Get computes and returns the next value in the distribution.
func (d *UniformDistribution) Get() float64 {
	return d.value
}

// RandomWalkDistribution is a stateful random walk. Initialize it with an
// underlying distribution, which is used to compute the new step value.
type RandomWalkDistribution struct {
	Step Distribution

	State float64 // optional
}

func WD(step Distribution, state float64) *RandomWalkDistribution {
	return &RandomWalkDistribution{Step: step, State: state}
}

// Advance computes the next value of this distribution and stores it.
func (d *RandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += d.Step.Get()
}

// Get returns the last computed value for this distribution.
func (d *RandomWalkDistribution) Get() float64 {
	return d.State
}

// ClampedRandomWalkDistribution is a stateful random walk, with minimum and
// maximum bounds. Initialize it with a Min, Max, and an underlying
// distribution, which is used to compute the new step value.
type ClampedRandomWalkDistribution struct {
	Step Distribution
	Min  float64
	Max  float64

	State float64 // optional
}

func CWD(step Distribution, min, max, state float64) *ClampedRandomWalkDistribution {
	return &ClampedRandomWalkDistribution{
		Step: step,
		Min: min,
		Max: max,

		State: state,
	}
}

// Advance computes the next value of this distribution and stores it.
func (d *ClampedRandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += d.Step.Get()
	if d.State > d.Max {
		d.State = d.Max
	}
	if d.State < d.Min {
		d.State = d.Min
	}
}

// Get returns the last computed value for this distribution.
func (d *ClampedRandomWalkDistribution) Get() float64 {
	return d.State
}

// MonotonicRandomWalkDistribution is a stateful random walk that only
// increases. Initialize it with a Start and an underlying distribution,
// which is used to compute the new step value. The sign of any value of the
// u.d. is always made positive.
type MonotonicRandomWalkDistribution struct {
	Step  Distribution
	State float64
}

// Advance computes the next value of this distribution and stores it.
func (d *MonotonicRandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += math.Abs(d.Step.Get())
}

func (d *MonotonicRandomWalkDistribution) Get() float64 {
	return d.State
}

func MWD(step Distribution, state float64) *MonotonicRandomWalkDistribution {
	return &MonotonicRandomWalkDistribution{Step: step, State: state}
}

type ConstantDistribution struct {
	State float64
}

func (d *ConstantDistribution) Advacne() {
}

func (d *ConstantDistribution) Get() float64 {
	return d.State
}
<commit_msg>Fixed typo in ConstantDistribution.Advance method name<commit_after>package main

import (
	"math"
	"math/rand"
)

// Distribution provides an interface to model a statistical distribution.
type Distribution interface {
	Advance()
	Get() float64 // should be idempotent
}

// NormalDistribution models a normal distribution.
type NormalDistribution struct {
	Mean   float64
	StdDev float64

	value float64
}

func ND(mean, stddev float64) *NormalDistribution {
	return &NormalDistribution{Mean: mean, StdDev: stddev}
}

// Advance advances this distribution. Since a normal distribution is
// stateless, this is just overwrites the internal cache value.
func (d *NormalDistribution) Advance() {
	d.value = rand.NormFloat64()*d.StdDev + d.Mean
}

// Get returns the last computed value for this distribution.
func (d *NormalDistribution) Get() float64 {
	return d.value
}

// UniformDistribution models a uniform distribution.
type UniformDistribution struct {
	Low  float64
	High float64

	value float64
}

func UD(low, high float64) *UniformDistribution {
	return &UniformDistribution{Low: low, High: high}
}

// Advance advances this distribution. Since a uniform distribution is
// stateless, this is just overwrites the internal cache value.
func (d *UniformDistribution) Advance() {
	x := rand.Float64() // uniform
	x *= d.High - d.Low
	x += d.Low
	d.value = x
}

// Get computes and returns the next value in the distribution.
func (d *UniformDistribution) Get() float64 {
	return d.value
}

// RandomWalkDistribution is a stateful random walk. Initialize it with an
// underlying distribution, which is used to compute the new step value.
type RandomWalkDistribution struct {
	Step Distribution

	State float64 // optional
}

func WD(step Distribution, state float64) *RandomWalkDistribution {
	return &RandomWalkDistribution{Step: step, State: state}
}

// Advance computes the next value of this distribution and stores it.
func (d *RandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += d.Step.Get()
}

// Get returns the last computed value for this distribution.
func (d *RandomWalkDistribution) Get() float64 {
	return d.State
}

// ClampedRandomWalkDistribution is a stateful random walk, with minimum and
// maximum bounds. Initialize it with a Min, Max, and an underlying
// distribution, which is used to compute the new step value.
type ClampedRandomWalkDistribution struct {
	Step Distribution
	Min  float64
	Max  float64

	State float64 // optional
}

func CWD(step Distribution, min, max, state float64) *ClampedRandomWalkDistribution {
	return &ClampedRandomWalkDistribution{
		Step: step,
		Min: min,
		Max: max,

		State: state,
	}
}

// Advance computes the next value of this distribution and stores it.
func (d *ClampedRandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += d.Step.Get()
	if d.State > d.Max {
		d.State = d.Max
	}
	if d.State < d.Min {
		d.State = d.Min
	}
}

// Get returns the last computed value for this distribution.
func (d *ClampedRandomWalkDistribution) Get() float64 {
	return d.State
}

// MonotonicRandomWalkDistribution is a stateful random walk that only
// increases. Initialize it with a Start and an underlying distribution,
// which is used to compute the new step value. The sign of any value of the
// u.d. is always made positive.
type MonotonicRandomWalkDistribution struct {
	Step  Distribution
	State float64
}

// Advance computes the next value of this distribution and stores it.
func (d *MonotonicRandomWalkDistribution) Advance() {
	d.Step.Advance()
	d.State += math.Abs(d.Step.Get())
}

func (d *MonotonicRandomWalkDistribution) Get() float64 {
	return d.State
}

func MWD(step Distribution, state float64) *MonotonicRandomWalkDistribution {
	return &MonotonicRandomWalkDistribution{Step: step, State: state}
}

type ConstantDistribution struct {
	State float64
}

func (d *ConstantDistribution) Advance() {
}

func (d *ConstantDistribution) Get() float64 {
	return d.State
}
<|endoftext|>


<commit_before>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	log "github.com/funkygao/log4go"
)

type subServer struct {
	*webServer

	idleConnsWg   sync.WaitGroup      // wait for all inflight http connections done
	idleConns     map[string]net.Conn // in keep-alive state http connections
	closedConnCh  chan string         // channel of remote addr
	idleConnsLock sync.Mutex

	auditor log.Logger

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics        *subMetrics
	throttleSubStatus *ratelimiter.LeakyBuckets
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:         newWebServer("sub", httpAddr, httpsAddr, maxClients, gw),
		closedConnCh:      make(chan string, 1<<10),
		idleConns:         make(map[string]net.Conn, 10000), // TODO
		wsReadLimit:       8 << 10,
		wsPongWait:        time.Minute,
		throttleSubStatus: ratelimiter.NewLeakyBuckets(60, time.Minute),
		ackShutdown:       0,
		ackCh:             make(chan ackOffsets, 100),
		ackedOffsets:      make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c.RemoteAddr().String())
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c.RemoteAddr().String()] = c
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		remoteAddr := c.RemoteAddr().String()
		if Options.EnableClientStats {
			this.gw.clientStates.UnregisterSubClient(remoteAddr)
		}

		this.closedConnCh <- remoteAddr
		this.idleConnsWg.Done()
	}
}

func (this *subServer) waitExit(server *http.Server, listener net.Listener, exit <-chan struct{}) {
	<-exit

	// HTTP response will have "Connection: close"
	server.SetKeepAlivesEnabled(false)

	// avoid new connections
	if err := listener.Close(); err != nil {
		log.Error(err.Error())
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for _, c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	log.Trace("%s waiting for all connected client close...", this.name)
	if waitTimeout(&this.idleConnsWg, Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()

	this.gw.wg.Done()
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("commit offset {C:%s T:%s G:%s P:%d O:%d}", cluster, topic, group, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group,
						strconv.Itoa(partition), offset); err != nil {
						log.Error("commitOffsets: %v", err)
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<commit_msg>enhancement: reduce memory<commit_after>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	log "github.com/funkygao/log4go"
)

type subServer struct {
	*webServer

	idleConnsWg   sync.WaitGroup // wait for all inflight http connections done
	closedConnCh  chan string    // channel of remote addr
	idleConns     map[net.Conn]struct{}
	idleConnsLock sync.Mutex

	auditor log.Logger

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics        *subMetrics
	throttleSubStatus *ratelimiter.LeakyBuckets
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:         newWebServer("sub", httpAddr, httpsAddr, maxClients, gw),
		closedConnCh:      make(chan string, 1<<10),
		idleConns:         make(map[net.Conn]struct{}, 200),
		wsReadLimit:       8 << 10,
		wsPongWait:        time.Minute,
		throttleSubStatus: ratelimiter.NewLeakyBuckets(60, time.Minute),
		ackShutdown:       0,
		ackCh:             make(chan ackOffsets, 100),
		ackedOffsets:      make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c] = struct{}{}
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		remoteAddr := c.RemoteAddr().String()
		if Options.EnableClientStats {
			this.gw.clientStates.UnregisterSubClient(remoteAddr)
		}

		this.closedConnCh <- remoteAddr
		this.idleConnsWg.Done()
	}
}

func (this *subServer) waitExit(server *http.Server, listener net.Listener, exit <-chan struct{}) {
	<-exit

	// HTTP response will have "Connection: close"
	server.SetKeepAlivesEnabled(false)

	// avoid new connections
	if err := listener.Close(); err != nil {
		log.Error(err.Error())
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	log.Trace("%s waiting for all connected client close...", this.name)
	if waitTimeout(&this.idleConnsWg, Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()

	this.gw.wg.Done()
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("commit offset {C:%s T:%s G:%s P:%d O:%d}", cluster, topic, group, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group,
						strconv.Itoa(partition), offset); err != nil {
						log.Error("commitOffsets: %v", err)
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<|endoftext|>


<commit_before>// Package collect provides functions for sending data to OpenTSDB.
//
// The "collect" namespace is used (i.e., <metric_root>.collect) to collect
// program and queue metrics.
package collect

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
)

var (
	// Freq is how often metrics are sent to OpenTSDB.
	Freq = time.Second * 15

	// MaxQueueLen is the maximum size of the queue, above which incoming data will
	// be discarded. Defaults to about 150MB.
	MaxQueueLen = 200000

	// BatchSize is the maximum length of data points sent at once to OpenTSDB.
	BatchSize = 50

	// Debug enables debug logging.
	Debug = false

	// Dropped is the number of dropped data points due to a full queue.
	dropped int64

	// Sent is the number of sent data points.
	sent int64

	tchan               chan *opentsdb.DataPoint
	tsdbURL             string
	osHostname          string
	metricRoot          string
	queue               opentsdb.MultiDataPoint
	qlock, mlock, slock sync.Mutex   // Locks for queues, maps, stats.
	counters                         = make(map[string]*addMetric)
	sets                             = make(map[string]*setMetric)
	client              *http.Client = &http.Client{
		Transport: &timeoutTransport{Transport: new(http.Transport)},
		Timeout:   time.Minute,
	}
)

type timeoutTransport struct {
	*http.Transport
	Timeout time.Time
}

func (t *timeoutTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	if time.Now().After(t.Timeout) {
		t.Transport.CloseIdleConnections()
		t.Timeout = time.Now().Add(time.Minute * 5)
	}
	return t.Transport.RoundTrip(r)
}

// InitChan is similar to Init, but uses the given channel instead of creating a
// new one.
func InitChan(tsdbhost, metric_root string, ch chan *opentsdb.DataPoint) error {
	if tchan != nil {
		return fmt.Errorf("cannot init twice")
	}
	if err := checkClean(metric_root, "metric root"); err != nil {
		return err
	}
	if tsdbhost == "" {
		return fmt.Errorf("must specify non-empty tsdb host")
	}
	u := url.URL{
		Scheme: "http",
		Path:   "/api/put",
	}
	if !strings.Contains(tsdbhost, ":") {
		tsdbhost += ":4242"
	}
	u.Host = tsdbhost
	tsdbURL = u.String()
	metricRoot = metric_root + "."
	tchan = ch
	go func() {
		for dp := range tchan {
			qlock.Lock()
			for {
				if len(queue) > MaxQueueLen {
					slock.Lock()
					dropped++
					slock.Unlock()
					break
				}
				queue = append(queue, dp)
				select {
				case dp = <-tchan:
					continue
				default:
				}
				break
			}
			qlock.Unlock()
		}
	}()
	go send()

	go collect()
	Set("collect.dropped", nil, func() (i float64) {
		slock.Lock()
		i = float64(dropped)
		slock.Unlock()
		return
	})
	Set("collect.sent", nil, func() (i float64) {
		slock.Lock()
		i = float64(sent)
		slock.Unlock()
		return
	})
	Set("collect.alloc", nil, func() float64 {
		var ms runtime.MemStats
		runtime.ReadMemStats(&ms)
		return float64(ms.Alloc)
	})
	Set("collect.goroutines", nil, func() float64 {
		return float64(runtime.NumGoroutine())
	})
	return nil
}

// Init sets up the channels and the queue for sending data to OpenTSDB. It also
// sets up the basename for all metrics.
func Init(tsdbhost, metric_root string) error {
	return InitChan(tsdbhost, metric_root, make(chan *opentsdb.DataPoint))
}

func setHostName() error {
	h, err := os.Hostname()
	if err != nil {
		return err
	}
	osHostname = strings.ToLower(strings.SplitN(h, ".", 2)[0])
	if err := checkClean(osHostname, "host tag"); err != nil {
		return err
	}
	return nil
}

type setMetric struct {
	metric string
	ts     opentsdb.TagSet
	f      func() float64
}

func Set(metric string, ts opentsdb.TagSet, f func() float64) error {
	if err := check(metric, &ts); err != nil {
		return err
	}
	tss := metric + ts.String()
	mlock.Lock()
	sets[tss] = &setMetric{metric, ts.Copy(), f}
	mlock.Unlock()
	return nil
}

type addMetric struct {
	metric string
	ts     opentsdb.TagSet
	value  float64
}

// Add takes a metric and increments a counter for that metric. The metric name
// is appended to the basename specified in the Init function.
func Add(metric string, ts opentsdb.TagSet, inc float64) error {
	if err := check(metric, &ts); err != nil {
		return err
	}
	tss := metric + ts.String()
	mlock.Lock()
	if counters[tss] == nil {
		counters[tss] = &addMetric{
			metric: metric,
			ts:     ts.Copy(),
		}
	}
	counters[tss].value += inc
	mlock.Unlock()
	return nil
}

func check(metric string, ts *opentsdb.TagSet) error {
	if err := checkClean(metric, "metric"); err != nil {
		return err
	}
	for k, v := range *ts {
		if err := checkClean(k, "tagk"); err != nil {
			return err
		}
		if err := checkClean(v, "tagv"); err != nil {
			return err
		}
	}
	if osHostname == "" {
		if err := setHostName(); err != nil {
			return err
		}
	}
	if *ts == nil {
		*ts = make(opentsdb.TagSet)
	}
	if (*ts)["host"] == "" {
		(*ts)["host"] = osHostname
	}
	return nil
}

func checkClean(s, t string) error {
	if sc, err := opentsdb.Clean(s); s != sc || err != nil {
		if err != nil {
			return err
		}
		return fmt.Errorf("%s %s may only contain a to z, A to Z, 0 to 9, -, _, ., / or Unicode letters and may not be empty", t, s)
	}
	return nil
}

func collect() {
	for {
		mlock.Lock()
		now := time.Now().Unix()
		for _, c := range counters {
			dp := &opentsdb.DataPoint{
				Metric:    metricRoot + c.metric,
				Timestamp: now,
				Value:     c.value,
				Tags:      c.ts,
			}
			tchan <- dp
		}
		for _, s := range sets {
			dp := &opentsdb.DataPoint{
				Metric:    metricRoot + s.metric,
				Timestamp: now,
				Value:     s.f(),
				Tags:      s.ts,
			}
			tchan <- dp
		}
		mlock.Unlock()
		time.Sleep(Freq)
	}
}
<commit_msg>cmd/scollector: Put puts a metric on directly on the datapoint channel.<commit_after>// Package collect provides functions for sending data to OpenTSDB.
//
// The "collect" namespace is used (i.e., <metric_root>.collect) to collect
// program and queue metrics.
package collect

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/StackExchange/scollector/opentsdb"
)

var (
	// Freq is how often metrics are sent to OpenTSDB.
	Freq = time.Second * 15

	// MaxQueueLen is the maximum size of the queue, above which incoming data will
	// be discarded. Defaults to about 150MB.
	MaxQueueLen = 200000

	// BatchSize is the maximum length of data points sent at once to OpenTSDB.
	BatchSize = 50

	// Debug enables debug logging.
	Debug = false

	// Dropped is the number of dropped data points due to a full queue.
	dropped int64

	// Sent is the number of sent data points.
	sent int64

	tchan               chan *opentsdb.DataPoint
	tsdbURL             string
	osHostname          string
	metricRoot          string
	queue               opentsdb.MultiDataPoint
	qlock, mlock, slock sync.Mutex   // Locks for queues, maps, stats.
	counters                         = make(map[string]*addMetric)
	sets                             = make(map[string]*setMetric)
	client              *http.Client = &http.Client{
		Transport: &timeoutTransport{Transport: new(http.Transport)},
		Timeout:   time.Minute,
	}
)

type timeoutTransport struct {
	*http.Transport
	Timeout time.Time
}

func (t *timeoutTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	if time.Now().After(t.Timeout) {
		t.Transport.CloseIdleConnections()
		t.Timeout = time.Now().Add(time.Minute * 5)
	}
	return t.Transport.RoundTrip(r)
}

// InitChan is similar to Init, but uses the given channel instead of creating a
// new one.
func InitChan(tsdbhost, metric_root string, ch chan *opentsdb.DataPoint) error {
	if tchan != nil {
		return fmt.Errorf("cannot init twice")
	}
	if err := checkClean(metric_root, "metric root"); err != nil {
		return err
	}
	if tsdbhost == "" {
		return fmt.Errorf("must specify non-empty tsdb host")
	}
	u := url.URL{
		Scheme: "http",
		Path:   "/api/put",
	}
	if !strings.Contains(tsdbhost, ":") {
		tsdbhost += ":4242"
	}
	u.Host = tsdbhost
	tsdbURL = u.String()
	metricRoot = metric_root + "."
	tchan = ch
	go func() {
		for dp := range tchan {
			qlock.Lock()
			for {
				if len(queue) > MaxQueueLen {
					slock.Lock()
					dropped++
					slock.Unlock()
					break
				}
				queue = append(queue, dp)
				select {
				case dp = <-tchan:
					continue
				default:
				}
				break
			}
			qlock.Unlock()
		}
	}()
	go send()

	go collect()
	Set("collect.dropped", nil, func() (i float64) {
		slock.Lock()
		i = float64(dropped)
		slock.Unlock()
		return
	})
	Set("collect.sent", nil, func() (i float64) {
		slock.Lock()
		i = float64(sent)
		slock.Unlock()
		return
	})
	Set("collect.alloc", nil, func() float64 {
		var ms runtime.MemStats
		runtime.ReadMemStats(&ms)
		return float64(ms.Alloc)
	})
	Set("collect.goroutines", nil, func() float64 {
		return float64(runtime.NumGoroutine())
	})
	return nil
}

// Init sets up the channels and the queue for sending data to OpenTSDB. It also
// sets up the basename for all metrics.
func Init(tsdbhost, metric_root string) error {
	return InitChan(tsdbhost, metric_root, make(chan *opentsdb.DataPoint))
}

func setHostName() error {
	h, err := os.Hostname()
	if err != nil {
		return err
	}
	osHostname = strings.ToLower(strings.SplitN(h, ".", 2)[0])
	if err := checkClean(osHostname, "host tag"); err != nil {
		return err
	}
	return nil
}

type setMetric struct {
	metric string
	ts     opentsdb.TagSet
	f      func() float64
}

func Set(metric string, ts opentsdb.TagSet, f func() float64) error {
	if err := check(metric, &ts); err != nil {
		return err
	}
	tss := metric + ts.String()
	mlock.Lock()
	sets[tss] = &setMetric{metric, ts.Copy(), f}
	mlock.Unlock()
	return nil
}

type addMetric struct {
	metric string
	ts     opentsdb.TagSet
	value  float64
}

// Add takes a metric and increments a counter for that metric. The metric name
// is appended to the basename specified in the Init function.
func Add(metric string, ts opentsdb.TagSet, inc float64) error {
	if err := check(metric, &ts); err != nil {
		return err
	}
	tss := metric + ts.String()
	mlock.Lock()
	if counters[tss] == nil {
		counters[tss] = &addMetric{
			metric: metric,
			ts:     ts.Copy(),
		}
	}
	counters[tss].value += inc
	mlock.Unlock()
	return nil
}

// Put puts a metric on directly on the datapoint channel. This is useful for capturing "events" that have a gauge value
func Put(metric string, ts opentsdb.TagSet, timestamp int64, v float64) error {
	if err := check(metric, &ts); err != nil {
		return err
	}
	dp := &opentsdb.DataPoint{
		Metric:    metricRoot + metric,
		Timestamp: timestamp,
		Value:     v,
		Tags:      ts,
	}
	tchan <- dp
	return nil
}

func check(metric string, ts *opentsdb.TagSet) error {
	if err := checkClean(metric, "metric"); err != nil {
		return err
	}
	for k, v := range *ts {
		if err := checkClean(k, "tagk"); err != nil {
			return err
		}
		if err := checkClean(v, "tagv"); err != nil {
			return err
		}
	}
	if osHostname == "" {
		if err := setHostName(); err != nil {
			return err
		}
	}
	if *ts == nil {
		*ts = make(opentsdb.TagSet)
	}
	if (*ts)["host"] == "" {
		(*ts)["host"] = osHostname
	}
	return nil
}

func checkClean(s, t string) error {
	if sc, err := opentsdb.Clean(s); s != sc || err != nil {
		if err != nil {
			return err
		}
		return fmt.Errorf("%s %s may only contain a to z, A to Z, 0 to 9, -, _, ., / or Unicode letters and may not be empty", t, s)
	}
	return nil
}

func collect() {
	for {
		mlock.Lock()
		now := time.Now().Unix()
		for _, c := range counters {
			dp := &opentsdb.DataPoint{
				Metric:    metricRoot + c.metric,
				Timestamp: now,
				Value:     c.value,
				Tags:      c.ts,
			}
			tchan <- dp
		}
		for _, s := range sets {
			dp := &opentsdb.DataPoint{
				Metric:    metricRoot + s.metric,
				Timestamp: now,
				Value:     s.f(),
				Tags:      s.ts,
			}
			tchan <- dp
		}
		mlock.Unlock()
		time.Sleep(Freq)
	}
}
<|endoftext|>


<commit_before>package internal

import (
	"log"
	"os"

	"github.com/v2ray/v2ray-core/common/platform"
)

type LogWriter interface {
	Log(LogEntry)
}

type NoOpLogWriter struct {
}

func (this *NoOpLogWriter) Log(entry LogEntry) {
	entry.Release()
}

type StdOutLogWriter struct {
	logger *log.Logger
}

func NewStdOutLogWriter() LogWriter {
	return &StdOutLogWriter{
		logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
	}
}

func (this *StdOutLogWriter) Log(log LogEntry) {
	this.logger.Print(log.String() + platform.LineSeparator())
	log.Release()
}

type FileLogWriter struct {
	queue  chan LogEntry
	logger *log.Logger
	file   *os.File
}

func (this *FileLogWriter) Log(log LogEntry) {
	select {
	case this.queue <- log:
	default:
		log.Release()
		// We don't expect this to happen, but don't want to block main thread as well.
	}
}

func (this *FileLogWriter) run() {
	for {
		entry, open := <-this.queue
		if !open {
			break
		}
		this.logger.Print(entry.String() + platform.LineSeparator())
		entry.Release()
		entry = nil
	}
}

func (this *FileLogWriter) Close() {
	this.file.Close()
}

func NewFileLogWriter(path string) (*FileLogWriter, error) {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}
	logger := &FileLogWriter{
		queue:  make(chan LogEntry, 16),
		logger: log.New(file, "", log.Ldate|log.Ltime),
		file:   file,
	}
	go logger.run()
	return logger, nil
}
<commit_msg>realtime logger<commit_after>package internal

import (
	"log"
	"os"

	"github.com/v2ray/v2ray-core/common/platform"
)

type LogWriter interface {
	Log(LogEntry)
}

type NoOpLogWriter struct {
}

func (this *NoOpLogWriter) Log(entry LogEntry) {
	entry.Release()
}

type StdOutLogWriter struct {
	logger *log.Logger
}

func NewStdOutLogWriter() LogWriter {
	return &StdOutLogWriter{
		logger: log.New(os.Stdout, "", log.Ldate|log.Ltime),
	}
}

func (this *StdOutLogWriter) Log(log LogEntry) {
	this.logger.Print(log.String() + platform.LineSeparator())
	log.Release()
}

type FileLogWriter struct {
	queue  chan string
	logger *log.Logger
	file   *os.File
}

func (this *FileLogWriter) Log(log LogEntry) {
	select {
	case this.queue <- log.String():
	default:
		// We don't expect this to happen, but don't want to block main thread as well.
	}
	log.Release()
}

func (this *FileLogWriter) run() {
	for {
		entry, open := <-this.queue
		if !open {
			break
		}
		this.logger.Print(entry + platform.LineSeparator())
	}
}

func (this *FileLogWriter) Close() {
	this.file.Close()
}

func NewFileLogWriter(path string) (*FileLogWriter, error) {
	file, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return nil, err
	}
	logger := &FileLogWriter{
		queue:  make(chan string, 16),
		logger: log.New(file, "", log.Ldate|log.Ltime),
		file:   file,
	}
	go logger.run()
	return logger, nil
}
<|endoftext|>


<commit_before>package store

import (
	"container/heap"
	"container/vector"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// Special values for a revision.
const (
	Missing = int64(-iota)
	Clobber
	Dir
	nop
)

// TODO revisit this when package regexp is more complete (e.g. do Unicode)
const charPat = `[a-zA-Z0-9.\-]`

var pathRe = mustBuildRe(charPat)

var Any = MustCompileGlob("/**")

var ErrTooLate = os.NewError("too late")

var (
	ErrBadMutation = os.NewError("bad mutation")
	ErrRevMismatch = os.NewError("rev mismatch")
	ErrBadPath     = os.NewError("bad path")
)

func mustBuildRe(p string) *regexp.Regexp {
	return regexp.MustCompile(`^/$|^(/` + p + `+)+$`)
}

// Applies mutations sent on Ops in sequence according to field Seqn. Any
// errors that occur will be written to ErrorPath. Duplicate operations at a
// given position are sliently ignored.
type Store struct {
	Ops     chan<- Op
	Seqns   <-chan int64
	Waiting <-chan int
	watchCh chan *watch
	watches []*watch
	todo    *vector.Vector
	state   *state
	head    int64
	log     map[int64]Event
	cleanCh chan int64
	flush   chan bool
}

// Represents an operation to apply to the store at position Seqn.
//
// If Mut is Nop, no change will be made, but an event will still be sent.
type Op struct {
	Seqn int64
	Mut  string
}

// Satisfies vector.LessInterface.
func (x Op) Less(y interface{}) bool {
	return x.Seqn < y.(Op).Seqn
}

type state struct {
	ver  int64
	root node
}

type watch struct {
	glob *Glob
	rev  int64
	c    chan<- Event
}

// Creates a new, empty data store. Mutations will be applied in order,
// starting at number 1 (number 0 can be thought of as the creation of the
// store).
func New() *Store {
	ops := make(chan Op)
	seqns := make(chan int64)
	watches := make(chan int)

	st := &Store{
		Ops:     ops,
		Seqns:   seqns,
		Waiting: watches,
		watchCh: make(chan *watch),
		todo:    new(vector.Vector),
		watches: []*watch{},
		state:   &state{0, emptyDir},
		log:     map[int64]Event{},
		cleanCh: make(chan int64),
		flush:   make(chan bool),
	}

	go st.process(ops, seqns, watches)
	return st
}

func split(path string) []string {
	if path == "/" {
		return []string{}
	}
	return strings.Split(path[1:], "/", -1)
}

func join(parts []string) string {
	return "/" + strings.Join(parts, "/")
}

func checkPath(k string) os.Error {
	if !pathRe.MatchString(k) {
		return ErrBadPath
	}
	return nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will set
// the contents of the file at `path` to `body` iff `rev` is greater than
// of equal to the file's revision at the time of application, with
// one exception: if `rev` is Clobber, the file will be set unconditionally.
func EncodeSet(path, body string, rev int64) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return strconv.Itoa64(rev) + ":" + path + "=" + body, nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will cause
// the file at `path` to be deleted iff `rev` is greater than
// of equal to the file's revision at the time of application, with
// one exception: if `rev` is Clobber, the file will be deleted
// unconditionally.
func EncodeDel(path string, rev int64) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return strconv.Itoa64(rev) + ":" + path, nil
}

// MustEncodeSet is like EncodeSet but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeSet(path, body string, rev int64) (mutation string) {
	m, err := EncodeSet(path, body, rev)
	if err != nil {
		panic(err)
	}
	return m
}

// MustEncodeDel is like EncodeDel but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeDel(path string, rev int64) (mutation string) {
	m, err := EncodeDel(path, rev)
	if err != nil {
		panic(err)
	}
	return m
}

func decode(mutation string) (path, v string, rev int64, keep bool, err os.Error) {
	cm := strings.Split(mutation, ":", 2)

	if len(cm) != 2 {
		err = ErrBadMutation
		return
	}

	rev, err = strconv.Atoi64(cm[0])
	if err != nil {
		return
	}

	kv := strings.Split(cm[1], "=", 2)

	if err = checkPath(kv[0]); err != nil {
		return
	}

	switch len(kv) {
	case 1:
		return kv[0], "", rev, false, nil
	case 2:
		return kv[0], kv[1], rev, true, nil
	}
	panic("unreachable")
}

func (st *Store) notify(e Event, ws []*watch) (nws []*watch) {
	for _, w := range ws {
		if e.Seqn >= w.rev && w.glob.Match(e.Path) {
			w.c <- e
		} else {
			nws = append(nws, w)
		}
	}

	return nws
}

func (st *Store) closeWatches() {
	for _, w := range st.watches {
		close(w.c)
	}
}

func (st *Store) process(ops <-chan Op, seqns chan<- int64, watches chan<- int) {
	defer st.closeWatches()

	for {
		var flush bool
		ver, values := st.state.ver, st.state.root

		// Take any incoming requests and queue them up.
		select {
		case a := <-ops:
			if closed(ops) {
				return
			}

			if a.Seqn > ver {
				heap.Push(st.todo, a)
			}
		case w := <-st.watchCh:
			n, ws := w.rev, []*watch{w}
			for ; len(ws) > 0 && n < st.head; n++ {
				ws = []*watch{}
			}
			for ; len(ws) > 0 && n <= ver; n++ {
				ws = st.notify(st.log[n], ws)
			}

			st.watches = append(st.watches, ws...)
		case seqn := <-st.cleanCh:
			for ; st.head <= seqn; st.head++ {
				st.log[st.head] = Event{}, false
			}
		case seqns <- ver:
			// nothing to do here
		case watches <- len(st.watches):
			// nothing to do here
		case flush = <-st.flush:
			// nothing
		}

		var ev Event
		// If we have any mutations that can be applied, do them.
		for st.todo.Len() > 0 {
			t := st.todo.At(0).(Op)
			if flush && ver < t.Seqn {
				ver = t.Seqn - 1
			}
			if t.Seqn > ver+1 {
				break
			}

			heap.Pop(st.todo)
			if t.Seqn < ver+1 {
				continue
			}

			values, ev = values.apply(t.Seqn, t.Mut)
			st.state = &state{ev.Seqn, values}
			ver = ev.Seqn
			if !flush {
				st.log[ev.Seqn] = ev
				st.watches = st.notify(ev, st.watches)
			}
		}

		// A flush just gets one final event.
		if flush {
			st.log[ev.Seqn] = ev
			st.watches = st.notify(ev, st.watches)
			st.head = ver + 1
		}
	}
}

// Returns a point-in-time snapshot of the contents of the store.
func (st *Store) Snap() (ver int64, g Getter) {
	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	p := st.state

	return p.ver, p.root
}

// Gets the value stored at `path`, if any.
//
// If no value is stored at `path`, `rev` will be `Missing` and `value` will be
// nil.
//
// if `path` is a directory, `rev` will be `Dir` and `value` will be a list of
// entries.
//
// Otherwise, `rev` is the revision and `value[0]` is the body.
func (st *Store) Get(path string) (value []string, rev int64) {
	_, g := st.Snap()
	return g.Get(path)
}

func (st *Store) Stat(path string) (int32, int64) {
	_, g := st.Snap()
	return g.Stat(path)
}

// Apply all operations in the internal queue, even if there are gaps in the
// sequence (gaps will be treated as no-ops). This is only useful for
// bootstrapping a store from a point-in-time snapshot of another store.
func (st *Store) Flush() {
	st.flush <- true
}

// Returns a chan that will receive a single event representing the
// first change made to any file matching glob on or after rev.
//
// If rev is less than any value passed to st.Clean, WaitGlob will return
// ErrTooLate.
func (st *Store) Wait(glob *Glob, rev int64) (<-chan Event, os.Error) {
	if rev < 1 {
		return nil, ErrTooLate
	}

	ch := make(chan Event, 1)
	wt := &watch{
		glob: glob,
		rev:  rev,
		c:    ch,
	}
	st.watchCh <- wt

	if rev < st.head {
		return nil, ErrTooLate
	}
	return ch, nil
}

func (st *Store) Clean(seqn int64) {
	st.cleanCh <- seqn
}
<commit_msg>fix inaccurate comment<commit_after>package store

import (
	"container/heap"
	"container/vector"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// Special values for a revision.
const (
	Missing = int64(-iota)
	Clobber
	Dir
	nop
)

// TODO revisit this when package regexp is more complete (e.g. do Unicode)
const charPat = `[a-zA-Z0-9.\-]`

var pathRe = mustBuildRe(charPat)

var Any = MustCompileGlob("/**")

var ErrTooLate = os.NewError("too late")

var (
	ErrBadMutation = os.NewError("bad mutation")
	ErrRevMismatch = os.NewError("rev mismatch")
	ErrBadPath     = os.NewError("bad path")
)

func mustBuildRe(p string) *regexp.Regexp {
	return regexp.MustCompile(`^/$|^(/` + p + `+)+$`)
}

// Applies mutations sent on Ops in sequence according to field Seqn. Any
// errors that occur will be written to ErrorPath. Duplicate operations at a
// given position are sliently ignored.
type Store struct {
	Ops     chan<- Op
	Seqns   <-chan int64
	Waiting <-chan int
	watchCh chan *watch
	watches []*watch
	todo    *vector.Vector
	state   *state
	head    int64
	log     map[int64]Event
	cleanCh chan int64
	flush   chan bool
}

// Represents an operation to apply to the store at position Seqn.
//
// If Mut is Nop, no change will be made, but an event will still be sent.
type Op struct {
	Seqn int64
	Mut  string
}

// Satisfies vector.LessInterface.
func (x Op) Less(y interface{}) bool {
	return x.Seqn < y.(Op).Seqn
}

type state struct {
	ver  int64
	root node
}

type watch struct {
	glob *Glob
	rev  int64
	c    chan<- Event
}

// Creates a new, empty data store. Mutations will be applied in order,
// starting at number 1 (number 0 can be thought of as the creation of the
// store).
func New() *Store {
	ops := make(chan Op)
	seqns := make(chan int64)
	watches := make(chan int)

	st := &Store{
		Ops:     ops,
		Seqns:   seqns,
		Waiting: watches,
		watchCh: make(chan *watch),
		todo:    new(vector.Vector),
		watches: []*watch{},
		state:   &state{0, emptyDir},
		log:     map[int64]Event{},
		cleanCh: make(chan int64),
		flush:   make(chan bool),
	}

	go st.process(ops, seqns, watches)
	return st
}

func split(path string) []string {
	if path == "/" {
		return []string{}
	}
	return strings.Split(path[1:], "/", -1)
}

func join(parts []string) string {
	return "/" + strings.Join(parts, "/")
}

func checkPath(k string) os.Error {
	if !pathRe.MatchString(k) {
		return ErrBadPath
	}
	return nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will set
// the contents of the file at `path` to `body` iff `rev` is greater than
// of equal to the file's revision at the time of application, with
// one exception: if `rev` is Clobber, the file will be set unconditionally.
func EncodeSet(path, body string, rev int64) (mutation string, err os.Error) {
	if err = checkPath(path); err != nil {
		return
	}
	return strconv.Itoa64(rev) + ":" + path + "=" + body, nil
}

// Returns a mutation that can be applied to a `Store`. The mutation will cause
// the file at `path` to be deleted iff `rev` is greater than
// of equal to the file's revision at the time of application, with
// one exception: if `rev` is Clobber, the file will be deleted
// unconditionally.
func EncodeDel(path string, rev int64) (mutation string, err os.Error) {
	if err := checkPath(path); err != nil {
		return
	}
	return strconv.Itoa64(rev) + ":" + path, nil
}

// MustEncodeSet is like EncodeSet but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeSet(path, body string, rev int64) (mutation string) {
	m, err := EncodeSet(path, body, rev)
	if err != nil {
		panic(err)
	}
	return m
}

// MustEncodeDel is like EncodeDel but panics if the mutation cannot be
// encoded. It simplifies safe initialization of global variables holding
// mutations.
func MustEncodeDel(path string, rev int64) (mutation string) {
	m, err := EncodeDel(path, rev)
	if err != nil {
		panic(err)
	}
	return m
}

func decode(mutation string) (path, v string, rev int64, keep bool, err os.Error) {
	cm := strings.Split(mutation, ":", 2)

	if len(cm) != 2 {
		err = ErrBadMutation
		return
	}

	rev, err = strconv.Atoi64(cm[0])
	if err != nil {
		return
	}

	kv := strings.Split(cm[1], "=", 2)

	if err = checkPath(kv[0]); err != nil {
		return
	}

	switch len(kv) {
	case 1:
		return kv[0], "", rev, false, nil
	case 2:
		return kv[0], kv[1], rev, true, nil
	}
	panic("unreachable")
}

func (st *Store) notify(e Event, ws []*watch) (nws []*watch) {
	for _, w := range ws {
		if e.Seqn >= w.rev && w.glob.Match(e.Path) {
			w.c <- e
		} else {
			nws = append(nws, w)
		}
	}

	return nws
}

func (st *Store) closeWatches() {
	for _, w := range st.watches {
		close(w.c)
	}
}

func (st *Store) process(ops <-chan Op, seqns chan<- int64, watches chan<- int) {
	defer st.closeWatches()

	for {
		var flush bool
		ver, values := st.state.ver, st.state.root

		// Take any incoming requests and queue them up.
		select {
		case a := <-ops:
			if closed(ops) {
				return
			}

			if a.Seqn > ver {
				heap.Push(st.todo, a)
			}
		case w := <-st.watchCh:
			n, ws := w.rev, []*watch{w}
			for ; len(ws) > 0 && n < st.head; n++ {
				ws = []*watch{}
			}
			for ; len(ws) > 0 && n <= ver; n++ {
				ws = st.notify(st.log[n], ws)
			}

			st.watches = append(st.watches, ws...)
		case seqn := <-st.cleanCh:
			for ; st.head <= seqn; st.head++ {
				st.log[st.head] = Event{}, false
			}
		case seqns <- ver:
			// nothing to do here
		case watches <- len(st.watches):
			// nothing to do here
		case flush = <-st.flush:
			// nothing
		}

		var ev Event
		// If we have any mutations that can be applied, do them.
		for st.todo.Len() > 0 {
			t := st.todo.At(0).(Op)
			if flush && ver < t.Seqn {
				ver = t.Seqn - 1
			}
			if t.Seqn > ver+1 {
				break
			}

			heap.Pop(st.todo)
			if t.Seqn < ver+1 {
				continue
			}

			values, ev = values.apply(t.Seqn, t.Mut)
			st.state = &state{ev.Seqn, values}
			ver = ev.Seqn
			if !flush {
				st.log[ev.Seqn] = ev
				st.watches = st.notify(ev, st.watches)
			}
		}

		// A flush just gets one final event.
		if flush {
			st.log[ev.Seqn] = ev
			st.watches = st.notify(ev, st.watches)
			st.head = ver + 1
		}
	}
}

// Returns a point-in-time snapshot of the contents of the store.
func (st *Store) Snap() (ver int64, g Getter) {
	// WARNING: Be sure to read the pointer value of st.state only once. If you
	// need multiple accesses, copy the pointer first.
	p := st.state

	return p.ver, p.root
}

// Gets the value stored at `path`, if any.
//
// If no value is stored at `path`, `rev` will be `Missing` and `value` will be
// nil.
//
// if `path` is a directory, `rev` will be `Dir` and `value` will be a list of
// entries.
//
// Otherwise, `rev` is the revision and `value[0]` is the body.
func (st *Store) Get(path string) (value []string, rev int64) {
	_, g := st.Snap()
	return g.Get(path)
}

func (st *Store) Stat(path string) (int32, int64) {
	_, g := st.Snap()
	return g.Stat(path)
}

// Apply all operations in the internal queue, even if there are gaps in the
// sequence (gaps will be treated as no-ops). This is only useful for
// bootstrapping a store from a point-in-time snapshot of another store.
func (st *Store) Flush() {
	st.flush <- true
}

// Returns a chan that will receive a single event representing the
// first change made to any file matching glob on or after rev.
//
// If rev is less than any value passed to st.Clean, Wait will return
// ErrTooLate.
func (st *Store) Wait(glob *Glob, rev int64) (<-chan Event, os.Error) {
	if rev < 1 {
		return nil, ErrTooLate
	}

	ch := make(chan Event, 1)
	wt := &watch{
		glob: glob,
		rev:  rev,
		c:    ch,
	}
	st.watchCh <- wt

	if rev < st.head {
		return nil, ErrTooLate
	}
	return ch, nil
}

func (st *Store) Clean(seqn int64) {
	st.cleanCh <- seqn
}
<|endoftext|>


<commit_before>package youtubejoinplugin

import (
	"encoding/json"
	"errors"
	"log"
	"sync"

	"google.golang.org/api/youtube/v3"

	"github.com/iopred/bruxism"
)

// YouTubeJoinPlugin is a plugin that monitors channels, and when they go live, will join the service
type YouTubeJoinPlugin struct {
	sync.RWMutex
	ytLiveChannel *bruxism.YTLiveChannel
	liveVideoChan chan *youtube.Video
	Channels      map[string]bool
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *YouTubeJoinPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	return nil
}

// Message handler.
func (p *YouTubeJoinPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	if service.IsBotOwner(message) || service.IsChannelOwner(message) && bruxism.MatchesCommand(service, "leave", message) {
		p.Unmonitor(message.Channel())
	}
}

// Load will load plugin state from a byte array.
func (p *YouTubeJoinPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if service.Name() != bruxism.YouTubeServiceName {
		panic("YouTubeJoin plugin only supports YouTube.")
	}

	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}

	for channel, _ := range p.Channels {
		p.ytLiveChannel.MonitorAll(channel, p.liveVideoChan)
	}

	go p.Run(bot, service)

	return nil
}

// Run will poll YouTube for channels going live and send messages.
func (p *YouTubeJoinPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	p.RLock()
	lvc := p.liveVideoChan
	p.RUnlock()
	for {
		v := <-lvc
		p.RLock()
		if p.Channels[v.Snippet.ChannelId] && v.LiveStreamingDetails != nil && v.LiveStreamingDetails.ActiveLiveChatId != "" {
			service.(*bruxism.YouTube).JoinVideo(v)
		}
		p.RUnlock()
	}
}

func (p *YouTubeJoinPlugin) Monitor(channel string) error {
	p.Lock()
	defer p.Unlock()
	if p.Channels[channel] {
		return errors.New("already monitoring that channel")
	}

	p.Channels[channel] = true

	return p.ytLiveChannel.MonitorAll(channel, p.liveVideoChan)
}

func (p *YouTubeJoinPlugin) Unmonitor(channel string) error {
	p.Lock()
	defer p.Unlock()
	if !p.Channels[channel] {
		return errors.New("not monitoring that channel")
	}

	delete(p.Channels, channel)

	return p.ytLiveChannel.Unmonitor(channel, p.liveVideoChan)
}

// Save will save plugin state to a byte array.
func (p *YouTubeJoinPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *YouTubeJoinPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return nil
}

// Name returns the name of the plugin.
func (p *YouTubeJoinPlugin) Name() string {
	return "YouTubeJoin"
}

// New will create a new YouTubeLive plugin.
func New(ytLiveChannel *bruxism.YTLiveChannel) *YouTubeJoinPlugin {
	return &YouTubeJoinPlugin{
		ytLiveChannel: ytLiveChannel,
		liveVideoChan: make(chan *youtube.Video),
		Channels:      map[string]bool{},
	}
}
<commit_msg>Update text for debugging.<commit_after>package youtubejoinplugin

import (
	"encoding/json"
	"errors"
	"log"
	"sync"

	"google.golang.org/api/youtube/v3"

	"github.com/iopred/bruxism"
)

// YouTubeJoinPlugin is a plugin that monitors channels, and when they go live, will join the service
type YouTubeJoinPlugin struct {
	sync.RWMutex
	ytLiveChannel *bruxism.YTLiveChannel
	liveVideoChan chan *youtube.Video
	Channels      map[string]bool
}

// Help returns a list of help strings that are printed when the user requests them.
func (p *YouTubeJoinPlugin) Help(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message, detailed bool) []string {
	return nil
}

// Message handler.
func (p *YouTubeJoinPlugin) Message(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) {
	if service.IsBotOwner(message) || service.IsChannelOwner(message) && bruxism.MatchesCommand(service, "leave", message) {
		p.Unmonitor(message.Channel())
	}
}

// Load will load plugin state from a byte array.
func (p *YouTubeJoinPlugin) Load(bot *bruxism.Bot, service bruxism.Service, data []byte) error {
	if service.Name() != bruxism.YouTubeServiceName {
		panic("YouTubeJoin plugin only supports YouTube.")
	}

	if data != nil {
		if err := json.Unmarshal(data, p); err != nil {
			log.Println("Error loading data", err)
		}
	}

	for channel, _ := range p.Channels {
		p.ytLiveChannel.MonitorAll(channel, p.liveVideoChan)
	}

	go p.Run(bot, service)

	return nil
}

// Run will poll YouTube for channels going live and send messages.
func (p *YouTubeJoinPlugin) Run(bot *bruxism.Bot, service bruxism.Service) {
	p.RLock()
	lvc := p.liveVideoChan
	p.RUnlock()
	for {
		v := <-lvc
		p.RLock()
		if p.Channels[v.Snippet.ChannelId] && v.LiveStreamingDetails != nil && v.LiveStreamingDetails.ActiveLiveChatId != "" {
			service.(*bruxism.YouTube).JoinVideo(v)
		}
		p.RUnlock()
	}
}

func (p *YouTubeJoinPlugin) Monitor(channel string) error {
	p.Lock()
	defer p.Unlock()
	if p.Channels[channel] {
		return errors.New("already tracking that channel")
	}

	p.Channels[channel] = true

	return p.ytLiveChannel.MonitorAll(channel, p.liveVideoChan)
}

func (p *YouTubeJoinPlugin) Unmonitor(channel string) error {
	p.Lock()
	defer p.Unlock()
	if !p.Channels[channel] {
		return errors.New("not tracking that channel")
	}

	delete(p.Channels, channel)

	return p.ytLiveChannel.Unmonitor(channel, p.liveVideoChan)
}

// Save will save plugin state to a byte array.
func (p *YouTubeJoinPlugin) Save() ([]byte, error) {
	return json.Marshal(p)
}

// Stats will return the stats for a plugin.
func (p *YouTubeJoinPlugin) Stats(bot *bruxism.Bot, service bruxism.Service, message bruxism.Message) []string {
	return nil
}

// Name returns the name of the plugin.
func (p *YouTubeJoinPlugin) Name() string {
	return "YouTubeJoin"
}

// New will create a new YouTubeLive plugin.
func New(ytLiveChannel *bruxism.YTLiveChannel) *YouTubeJoinPlugin {
	return &YouTubeJoinPlugin{
		ytLiveChannel: ytLiveChannel,
		liveVideoChan: make(chan *youtube.Video),
		Channels:      map[string]bool{},
	}
}
<|endoftext|>