<commit_before>// Package auth manages auth tokens for command line usage.
package auth

import (
	"bufio"
	"context"
	"golang.org/x/crypto/ssh/terminal"
	"io"
	"log"
	"os"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	gstatus "google.golang.org/grpc/status"

	"pixur.org/pixur/api"
	"pixur.org/pixur/be/status"
	"pixur.org/pixur/tools/config"
)

const (
	authPwtHeaderKey = "pixur-auth-token"
	pixPwtHeaderKey  = "pixur-pix-token"
	httpHeaderKey    = "pixur-http-header-bin"
)

// cc is optional
func Auth(ctx context.Context, cc **grpc.ClientConn) (context.Context, status.S) {
	conf, sts := getAndUpdateConfig(ctx, cc)
	if sts != nil {
		return nil, sts
	}

	md := make(metadata.MD, 2)
	if conf.AuthToken != "" {
		md.Set(authPwtHeaderKey, conf.AuthToken)
	}
	if conf.PixToken != "" {
		md.Set(pixPwtHeaderKey, conf.PixToken)
	}
	return metadata.NewOutgoingContext(ctx, md), nil
}

func getAndUpdateConfig(ctx context.Context, cc **grpc.ClientConn) (
	_ *config.Config, stscap status.S) {
	oldconf, sts := config.GetConfig()
	if sts != nil {
		if sts.Code() == codes.NotFound {
			oldconf = &config.Config{}
		} else {
			return nil, sts
		}
	}

	newconf := proto.Clone(oldconf).(*config.Config)
	if sts := updateConfig(ctx, cc, newconf); sts != nil {
		return nil, sts
	}
	if !proto.Equal(oldconf, newconf) {
		if sts := config.SetConfig(newconf); sts != nil {
			status.ReplaceOrSuppress(&stscap, sts)
		}
	}

	return newconf, nil
}

// shouldRefresh determines if a token refresh is needed.  notAfter indicates when the token is
// expired and must be non-nil.  softNotAfter indicates that the token is still valid, but should be
// refreshed.  softNotAfter may be nil.  The return value useToken indicates that the auth token
// can be reusued.  useCreds indicates if a ident/secret is needed.  At most one of useToken or
// useCreds can be true.
func shouldRefresh(now func() time.Time, notAfter, softNotAfter *tspb.Timestamp) (
	useToken, useCreds bool, _ status.S) {
	if notAfter == nil {
		return false, false, status.Internal(nil, "missing not after")
	}
	na, err := ptypes.Timestamp(notAfter)
	if err != nil {
		return false, false, status.Internal(err, "bad not after")
	}
	ts := now()
	if ts.After(na) {
		return false, true, nil
	}
	if softNotAfter != nil {
		sna, err := ptypes.Timestamp(softNotAfter)
		if err != nil {
			return false, false, status.Internal(err, "bad soft not after")
		}
		if ts.After(sna) {
			return true, false, nil
		}
	}
	return false, false, nil
}

func updateConfig(ctx context.Context, cc **grpc.ClientConn, conf *config.Config) (stscap status.S) {
	var (
		useToken bool
		useCreds bool = true
	)
	if conf.AuthPayload != nil {
		var sts status.S
		useToken, useCreds, sts =
			shouldRefresh(time.Now, conf.AuthPayload.NotAfter, conf.AuthPayload.SoftNotAfter)
		if sts != nil {
			return sts
		}
		if !useToken && !useCreds {
			return nil
		}
	}

	var conn *grpc.ClientConn
	if cc == nil || *cc == nil {
		if conf.PixurTarget == "" {
			target, sts := getTargetFromTerm(os.Stderr, os.Stdin)
			if sts != nil {
				return sts
			}
			conf.PixurTarget = target
		}
		var err error
		conn, err = grpc.DialContext(ctx, conf.PixurTarget, grpc.WithInsecure())
		if err != nil {
			return status.From(err)
		}
		if cc != nil {
			*cc = conn
		} else {
			defer func() {

				if err := conn.Close(); err != nil {
					status.ReplaceOrSuppress(&stscap, status.From(err))

				}
			}()
		}

	}

	client := api.NewPixurServiceClient(conn)

	if useToken && conf.AuthToken != "" {
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			PreviousAuthToken: conf.AuthToken,
		})
		if err != nil {
			if gstatus.Code(err) == codes.Unauthenticated {
				log.Println("Previous Auth token denied, falling back to creds", err)
				useCreds = true
			} else {
				return status.From(err)
			}
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	if useCreds {
		ident, secret, sts := getIdentSecretFromTerm(os.Stderr, os.Stdin)
		if sts != nil {
			return sts
		}
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			Ident:  ident,
			Secret: secret,
		})
		if err != nil {
			return status.From(err)
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	panic("unreachable")
}

func getTargetFromTerm(w io.Writer, r io.Reader) (string, status.S) {
	if _, err := w.Write([]byte("gRPC Target of Pixur Backend (e.g. dns:///localhost:8079): \n")); err != nil {
		return "", status.From(err)
	}
	br := bufio.NewReader(r)
	target, err := br.ReadString('\n')
	if err != nil {
		return "", status.From(err)
	}

	return target[:len(target)-1], nil
}

func getIdentSecretFromTerm(w io.Writer, r *os.File) (string, string, status.S) {
	if _, err := w.Write([]byte("User Ident (e.g. admin): \n")); err != nil {
		return "", "", status.From(err)
	}
	br := bufio.NewReader(r)
	ident, err := br.ReadString('\n')
	if err != nil {
		return "", "", status.From(err)
	}

	if _, err := w.Write([]byte("Secret (e.g. 12345): \n")); err != nil {
		return "", "", status.From(err)
	}
	secret, err := terminal.ReadPassword(int(r.Fd()))
	if err != nil {
		return "", "", status.From(err)
	}

	return ident[:len(ident)-1], string(secret), nil
}
<commit_msg>tools: unconditionally get a conn for auth<commit_after>// Package auth manages auth tokens for command line usage.
package auth

import (
	"bufio"
	"context"
	"golang.org/x/crypto/ssh/terminal"
	"io"
	"log"
	"os"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	tspb "github.com/golang/protobuf/ptypes/timestamp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	gstatus "google.golang.org/grpc/status"

	"pixur.org/pixur/api"
	"pixur.org/pixur/be/status"
	"pixur.org/pixur/tools/config"
)

const (
	authPwtHeaderKey = "pixur-auth-token"
	pixPwtHeaderKey  = "pixur-pix-token"
	httpHeaderKey    = "pixur-http-header-bin"
)

// cc is optional
func Auth(ctx context.Context, cc **grpc.ClientConn) (context.Context, status.S) {
	conf, sts := getAndUpdateConfig(ctx, cc)
	if sts != nil {
		return nil, sts
	}

	md := make(metadata.MD, 2)
	if conf.AuthToken != "" {
		md.Set(authPwtHeaderKey, conf.AuthToken)
	}
	if conf.PixToken != "" {
		md.Set(pixPwtHeaderKey, conf.PixToken)
	}
	return metadata.NewOutgoingContext(ctx, md), nil
}

func getAndUpdateConfig(ctx context.Context, cc **grpc.ClientConn) (
	_ *config.Config, stscap status.S) {
	oldconf, sts := config.GetConfig()
	if sts != nil {
		if sts.Code() == codes.NotFound {
			oldconf = &config.Config{}
		} else {
			return nil, sts
		}
	}

	newconf := proto.Clone(oldconf).(*config.Config)
	if sts := updateConfig(ctx, cc, newconf); sts != nil {
		return nil, sts
	}
	if !proto.Equal(oldconf, newconf) {
		if sts := config.SetConfig(newconf); sts != nil {
			status.ReplaceOrSuppress(&stscap, sts)
		}
	}

	return newconf, nil
}

// shouldRefresh determines if a token refresh is needed.  notAfter indicates when the token is
// expired and must be non-nil.  softNotAfter indicates that the token is still valid, but should be
// refreshed.  softNotAfter may be nil.  The return value useToken indicates that the auth token
// can be reusued.  useCreds indicates if a ident/secret is needed.  At most one of useToken or
// useCreds can be true.
func shouldRefresh(now func() time.Time, notAfter, softNotAfter *tspb.Timestamp) (
	useToken, useCreds bool, _ status.S) {
	if notAfter == nil {
		return false, false, status.Internal(nil, "missing not after")
	}
	na, err := ptypes.Timestamp(notAfter)
	if err != nil {
		return false, false, status.Internal(err, "bad not after")
	}
	ts := now()
	if ts.After(na) {
		return false, true, nil
	}
	if softNotAfter != nil {
		sna, err := ptypes.Timestamp(softNotAfter)
		if err != nil {
			return false, false, status.Internal(err, "bad soft not after")
		}
		if ts.After(sna) {
			return true, false, nil
		}
	}
	return false, false, nil
}

func updateConfig(ctx context.Context, cc **grpc.ClientConn, conf *config.Config) (stscap status.S) {
	var conn *grpc.ClientConn
	if cc == nil || *cc == nil {
		if conf.PixurTarget == "" {
			target, sts := getTargetFromTerm(os.Stderr, os.Stdin)
			if sts != nil {
				return sts
			}
			conf.PixurTarget = target
		}
		var err error
		conn, err = grpc.DialContext(ctx, conf.PixurTarget, grpc.WithInsecure())
		if err != nil {
			return status.From(err)
		}
		if cc != nil {
			*cc = conn
		} else {
			defer func() {
				if err := conn.Close(); err != nil {
					status.ReplaceOrSuppress(&stscap, status.From(err))
				}
			}()
		}
	}
	var (
		useToken bool
		useCreds bool = true
	)
	if conf.AuthPayload != nil {
		var sts status.S
		useToken, useCreds, sts =
			shouldRefresh(time.Now, conf.AuthPayload.NotAfter, conf.AuthPayload.SoftNotAfter)
		if sts != nil {
			return sts
		}
		if !useToken && !useCreds {
			return nil
		}
	}

	client := api.NewPixurServiceClient(conn)

	if useToken && conf.AuthToken != "" {
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			PreviousAuthToken: conf.AuthToken,
		})
		if err != nil {
			if gstatus.Code(err) == codes.Unauthenticated {
				log.Println("Previous Auth token denied, falling back to creds", err)
				useCreds = true
			} else {
				return status.From(err)
			}
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	if useCreds {
		ident, secret, sts := getIdentSecretFromTerm(os.Stderr, os.Stdin)
		if sts != nil {
			return sts
		}
		res, err := client.GetRefreshToken(ctx, &api.GetRefreshTokenRequest{
			Ident:  ident,
			Secret: secret,
		})
		if err != nil {
			return status.From(err)
		} else {
			conf.AuthToken, conf.PixToken = res.AuthToken, res.PixToken
			conf.AuthPayload, conf.PixPayload = res.AuthPayload, res.PixPayload
			return nil
		}
	}

	panic("unreachable")
}

func getTargetFromTerm(w io.Writer, r io.Reader) (string, status.S) {
	if _, err := w.Write([]byte("gRPC Target of Pixur Backend (e.g. dns:///localhost:8079): \n")); err != nil {
		return "", status.From(err)
	}
	br := bufio.NewReader(r)
	target, err := br.ReadString('\n')
	if err != nil {
		return "", status.From(err)
	}

	return target[:len(target)-1], nil
}

func getIdentSecretFromTerm(w io.Writer, r *os.File) (string, string, status.S) {
	if _, err := w.Write([]byte("User Ident (e.g. admin): \n")); err != nil {
		return "", "", status.From(err)
	}
	br := bufio.NewReader(r)
	ident, err := br.ReadString('\n')
	if err != nil {
		return "", "", status.From(err)
	}

	if _, err := w.Write([]byte("Secret (e.g. 12345): \n")); err != nil {
		return "", "", status.From(err)
	}
	secret, err := terminal.ReadPassword(int(r.Fd()))
	if err != nil {
		return "", "", status.From(err)
	}

	return ident[:len(ident)-1], string(secret), nil
}
<|endoftext|>


<commit_before>package main_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"

	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
)

var _ = Describe("Pushing an app", func() {
	It("stops the time for pushing an app", func() {
		for index := 0; index < loopCount; index++ {
			startTime := time.Now()

			appName := generator.PrefixedRandomName("APP")
			Expect(
				cf.Cf("push", appName, "-p", "assets/dora").Wait(cfPushTimeout)).
				To(Exit(0))

			statsdClient.Timing(asSparseMetric("cf-push"), time.Since(startTime).Seconds()*1000)
		}
	})
})

func asSparseMetric(metricName string) string {
	return metricsPrefix + metricName + ".sparse-avg"
}
<commit_msg>Delete app after pushing and measuring push time<commit_after>package main_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"

	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
)

var _ = Describe("Pushing an app", func() {
	It("stops the time for pushing an app", func() {
		for index := 0; index < loopCount; index++ {
			startTime := time.Now()

			appName := generator.PrefixedRandomName("APP")
			Expect(
				cf.Cf("push", appName, "-p", "assets/dora").Wait(cfPushTimeout)).
				To(Exit(0))

			statsdClient.Timing(asSparseMetric("cf-push"), time.Since(startTime).Seconds()*1000)

			Expect(
				cf.Cf("delete", appName, "-f").Wait(cfPushTimeout)).
				To(Exit(0))
		}
	})
})

func asSparseMetric(metricName string) string {
	return metricsPrefix + metricName + ".sparse-avg"
}
<|endoftext|>


<commit_before>// UNREVIEWED

// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package noder

import "internal/pkgbits"

type codeStmt int

func (c codeStmt) Marker() pkgbits.SyncMarker { return pkgbits.SyncStmt1 }
func (c codeStmt) Value() int                 { return int(c) }

const (
	stmtEnd codeStmt = iota
	stmtLabel
	stmtBlock
	stmtExpr
	stmtSend
	stmtAssign
	stmtAssignOp
	stmtIncDec
	stmtBranch
	stmtCall
	stmtReturn
	stmtIf
	stmtFor
	stmtSwitch
	stmtSelect

	// TODO(mdempsky): Remove after we don't care about toolstash -cmp.
	stmtTypeDeclHack
)

type codeExpr int

func (c codeExpr) Marker() pkgbits.SyncMarker { return pkgbits.SyncExpr }
func (c codeExpr) Value() int                 { return int(c) }

// TODO(mdempsky): Split expr into addr, for lvalues.
const (
	exprNone codeExpr = iota
	exprConst
	exprType  // type expression
	exprLocal // local variable
	exprName  // global variable or function
	exprBlank
	exprCompLit
	exprFuncLit
	exprSelector
	exprIndex
	exprSlice
	exprAssert
	exprUnaryOp
	exprBinaryOp
	exprCall
	exprConvert
)

type codeDecl int

func (c codeDecl) Marker() pkgbits.SyncMarker { return pkgbits.SyncDecl }
func (c codeDecl) Value() int                 { return int(c) }

const (
	declEnd codeDecl = iota
	declFunc
	declMethod
	declVar
	declOther
)
<commit_msg>cmd/compile: remove unified IR stmtTypeDeclHack<commit_after>// UNREVIEWED

// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package noder

import "internal/pkgbits"

type codeStmt int

func (c codeStmt) Marker() pkgbits.SyncMarker { return pkgbits.SyncStmt1 }
func (c codeStmt) Value() int                 { return int(c) }

const (
	stmtEnd codeStmt = iota
	stmtLabel
	stmtBlock
	stmtExpr
	stmtSend
	stmtAssign
	stmtAssignOp
	stmtIncDec
	stmtBranch
	stmtCall
	stmtReturn
	stmtIf
	stmtFor
	stmtSwitch
	stmtSelect
)

type codeExpr int

func (c codeExpr) Marker() pkgbits.SyncMarker { return pkgbits.SyncExpr }
func (c codeExpr) Value() int                 { return int(c) }

// TODO(mdempsky): Split expr into addr, for lvalues.
const (
	exprNone codeExpr = iota
	exprConst
	exprType  // type expression
	exprLocal // local variable
	exprName  // global variable or function
	exprBlank
	exprCompLit
	exprFuncLit
	exprSelector
	exprIndex
	exprSlice
	exprAssert
	exprUnaryOp
	exprBinaryOp
	exprCall
	exprConvert
)

type codeDecl int

func (c codeDecl) Marker() pkgbits.SyncMarker { return pkgbits.SyncDecl }
func (c codeDecl) Value() int                 { return int(c) }

const (
	declEnd codeDecl = iota
	declFunc
	declMethod
	declVar
	declOther
)
<|endoftext|>


<commit_before>package gogl

import "sync"

// VertexSet uses maps to express a value-less (empty struct), indexed
// unordered list. See
// https://groups.google.com/forum/#!searchin/golang-nuts/map/golang-nuts/H2cXpwisEUE/1X2FV-rODfIJ
type VertexSet map[Vertex]struct{}
type adjacencyList map[Vertex]VertexSet

// Helper to not have to write struct{} everywhere.
var keyExists = struct{}{}

type al struct {
	list adjacencyList
}

type AdjacencyList struct {
	al
	size uint
	mu   sync.RWMutex
}

// Composite literal to create a new AdjacencyList.
func NewAdjacencyList() *AdjacencyList {
	// Cannot assign to promoted fields in a composite literals.
	list := &AdjacencyList{}
	list.list = make(map[Vertex]VertexSet)
	return list
}

func NewAdjacencyListFromEdgeSet(set []Edge) *AdjacencyList {
	g := NewAdjacencyList()

	for _, edge := range set {
		g.addEdge(edge)
	}

	return g
}

func (g *AdjacencyList) EachVertex(f func(vertex Vertex)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for v := range g.list {
		f(v)
	}
}

func (g *AdjacencyList) EachEdge(f func(edge Edge)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for source, adjacent := range g.list {
		for target, _ := range adjacent {
			f(&BaseEdge{u: source, v: target})
		}
	}
}

func (g *AdjacencyList) EachAdjacent(vertex Vertex, f func(target Vertex)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if _, exists := g.list[vertex]; exists {
		for adjacent, _ := range g.list[vertex] {
			f(adjacent)
		}
	}
}

func (g *AdjacencyList) HasVertex(vertex Vertex) (exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	exists = g.hasVertex(vertex)
	return
}

func (g *AdjacencyList) OutDegree(vertex Vertex) (degree uint, exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if exists = g.hasVertex(vertex); exists {
		degree = uint(len(g.list[vertex]))
	}
	return
}

// Getting InDegree is inefficient for directed adjacency lists
func (g *AdjacencyList) InDegree(vertex Vertex) (degree uint, exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if exists = g.hasVertex(vertex); exists {

		f := func(v Vertex) {
			if v == vertex {
				degree++
			}
		}

		// This results in a double read-lock. Should be fine.
		for e := range g.list {
			g.EachAdjacent(e, f)
		}
	}

	return
}

func (g *AdjacencyList) hasVertex(vertex Vertex) (exists bool) {
	_, exists = g.list[vertex]
	return
}

func (g *AdjacencyList) Order() (length uint) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	length = uint(len(g.list))

	g.mu.RUnlock()
	return
}

func (g *AdjacencyList) Size() uint {
	return g.size
}

func (g *AdjacencyList) Density() float64 {
	g.mu.RLock()
	defer g.mu.RUnlock()

	order := float64(g.Order())
	return (2 * float64(g.Size())) / (order * (order - 1))
}

func (g *AdjacencyList) AddVertex(vertex Vertex) (success bool) {
	g.mu.Lock()
	defer g.mu.Unlock()

	success = g.addVertex(vertex)
	return
}

func (g *AdjacencyList) addVertex(vertex Vertex) (success bool) {
	if exists := g.hasVertex(vertex); !exists {
		// TODO experiment with different lengths...possibly by analyzing existing density?
		g.list[vertex] = make(VertexSet, 10)
		success = true
	}

	return
}

func (g *AdjacencyList) RemoveVertex(vertex Vertex) (success bool) {
	g.mu.Lock()
	defer g.mu.Unlock()

	if g.hasVertex(vertex) {
		// TODO Is the expensive search good to do here and now...
		// while read-locked?
		delete(g.list, vertex)

		// TODO consider chunking the list and parallelizing into goroutines
		for _, adjacent := range g.list {
			if _, has := adjacent[vertex]; has {
				delete(adjacent, vertex)
				g.size--
			}
		}

		success = true
	}
	return
}

func (g *AdjacencyList) AddEdge(edge Edge) bool {
	g.mu.Lock()
	defer g.mu.Unlock()

	return g.addEdge(edge)
}

func (g *AdjacencyList) addEdge(edge Edge) (exists bool) {
	g.addVertex(edge.Source())
	g.addVertex(edge.Target())

	if _, exists = g.list[edge.Source()][edge.Target()]; !exists {
		g.list[edge.Source()][edge.Target()] = keyExists
	}
	return !exists
}

func (g *AdjacencyList) RemoveEdge(edge Edge) {
	g.mu.Lock()
	defer g.mu.Unlock()

	delete(g.list[edge.Source()], edge.Target())
}
<commit_msg>Shorten Order() implementation a tad.<commit_after>package gogl

import "sync"

// VertexSet uses maps to express a value-less (empty struct), indexed
// unordered list. See
// https://groups.google.com/forum/#!searchin/golang-nuts/map/golang-nuts/H2cXpwisEUE/1X2FV-rODfIJ
type VertexSet map[Vertex]struct{}
type adjacencyList map[Vertex]VertexSet

// Helper to not have to write struct{} everywhere.
var keyExists = struct{}{}

type al struct {
	list adjacencyList
}

type AdjacencyList struct {
	al
	size uint
	mu   sync.RWMutex
}

// Composite literal to create a new AdjacencyList.
func NewAdjacencyList() *AdjacencyList {
	// Cannot assign to promoted fields in a composite literals.
	list := &AdjacencyList{}
	list.list = make(map[Vertex]VertexSet)
	return list
}

func NewAdjacencyListFromEdgeSet(set []Edge) *AdjacencyList {
	g := NewAdjacencyList()

	for _, edge := range set {
		g.addEdge(edge)
	}

	return g
}

func (g *AdjacencyList) EachVertex(f func(vertex Vertex)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for v := range g.list {
		f(v)
	}
}

func (g *AdjacencyList) EachEdge(f func(edge Edge)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for source, adjacent := range g.list {
		for target, _ := range adjacent {
			f(&BaseEdge{u: source, v: target})
		}
	}
}

func (g *AdjacencyList) EachAdjacent(vertex Vertex, f func(target Vertex)) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if _, exists := g.list[vertex]; exists {
		for adjacent, _ := range g.list[vertex] {
			f(adjacent)
		}
	}
}

func (g *AdjacencyList) HasVertex(vertex Vertex) (exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	exists = g.hasVertex(vertex)
	return
}

func (g *AdjacencyList) OutDegree(vertex Vertex) (degree uint, exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if exists = g.hasVertex(vertex); exists {
		degree = uint(len(g.list[vertex]))
	}
	return
}

// Getting InDegree is inefficient for directed adjacency lists
func (g *AdjacencyList) InDegree(vertex Vertex) (degree uint, exists bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	if exists = g.hasVertex(vertex); exists {

		f := func(v Vertex) {
			if v == vertex {
				degree++
			}
		}

		// This results in a double read-lock. Should be fine.
		for e := range g.list {
			g.EachAdjacent(e, f)
		}
	}

	return
}

func (g *AdjacencyList) hasVertex(vertex Vertex) (exists bool) {
	_, exists = g.list[vertex]
	return
}

func (g *AdjacencyList) Order() uint {
	g.mu.RLock()
	defer g.mu.RUnlock()

	return uint(len(g.list))
}

func (g *AdjacencyList) Size() uint {
	return g.size
}

func (g *AdjacencyList) Density() float64 {
	g.mu.RLock()
	defer g.mu.RUnlock()

	order := float64(g.Order())
	return (2 * float64(g.Size())) / (order * (order - 1))
}

func (g *AdjacencyList) AddVertex(vertex Vertex) (success bool) {
	g.mu.Lock()
	defer g.mu.Unlock()

	success = g.addVertex(vertex)
	return
}

func (g *AdjacencyList) addVertex(vertex Vertex) (success bool) {
	if exists := g.hasVertex(vertex); !exists {
		// TODO experiment with different lengths...possibly by analyzing existing density?
		g.list[vertex] = make(VertexSet, 10)
		success = true
	}

	return
}

func (g *AdjacencyList) RemoveVertex(vertex Vertex) (success bool) {
	g.mu.Lock()
	defer g.mu.Unlock()

	if g.hasVertex(vertex) {
		// TODO Is the expensive search good to do here and now...
		// while read-locked?
		delete(g.list, vertex)

		// TODO consider chunking the list and parallelizing into goroutines
		for _, adjacent := range g.list {
			if _, has := adjacent[vertex]; has {
				delete(adjacent, vertex)
				g.size--
			}
		}

		success = true
	}
	return
}

func (g *AdjacencyList) AddEdge(edge Edge) bool {
	g.mu.Lock()
	defer g.mu.Unlock()

	return g.addEdge(edge)
}

func (g *AdjacencyList) addEdge(edge Edge) (exists bool) {
	g.addVertex(edge.Source())
	g.addVertex(edge.Target())

	if _, exists = g.list[edge.Source()][edge.Target()]; !exists {
		g.list[edge.Source()][edge.Target()] = keyExists
	}
	return !exists
}

func (g *AdjacencyList) RemoveEdge(edge Edge) {
	g.mu.Lock()
	defer g.mu.Unlock()

	delete(g.list[edge.Source()], edge.Target())
}
<|endoftext|>


<commit_before>package server

import (
	"database/sql"
	"servant/conf"
	"net/http"
	"encoding/json"
)

type DatabaseServer struct {
	*Session
}

type sqlResult []map[string]string

func NewDatabaseServer(sess *Session) Handler {
	return DatabaseServer{
		Session:sess,
	}
}

func (self DatabaseServer) findDatabaseQueryConfig() (*conf.Database, *conf.Query) {
	dbConf, ok := self.config.Databases[self.group]
	if !ok {
		return nil, nil
	}
	qConf, ok := dbConf.Queries[self.item]
	if !ok {
		return dbConf, nil
	}
	return dbConf, qConf
}

func (self DatabaseServer) serve() {
	method := self.req.Method
	if method != "GET" {
		self.ErrorEnd(http.StatusMethodNotAllowed, "not allow method: %s", method)
		return
	}
	dbConf, queryConf := self.findDatabaseQueryConfig()
	if dbConf == nil {
		self.ErrorEnd(http.StatusNotFound, "database not found")
		return
	}
	if queryConf == nil {
		self.ErrorEnd(http.StatusNotFound, "query not found")
		return
	}

	db, err := sql.Open(dbConf.Driver, dbConf.Dsn)
	if err != nil {
		self.ErrorEnd(http.StatusInternalServerError, "driver init failed")
		return
	}
	defer db.Close()
	data := make([]sqlResult, 0, 1)
	for _, sql := range(queryConf.Sqls) {
		sql, params := replaceSqlParams(sql, self.req.URL.Query())

		result, err := dbQuery(db, sql, params)
		if err != nil {
			self.ErrorEnd(http.StatusInternalServerError, "query %s failed: %s", sql, err)
			return
		}
		data = append(data, result)
	}
	buf, err := json.Marshal(data)
	if err != nil {
		self.ErrorEnd(http.StatusInternalServerError, "json marshal failed: %s", err)
		return
	}
	self.resp.Write(buf)
}

func replaceSqlParams(inSql string, query map[string][]string) (outSql string, params []interface{}){
	params = make([]interface{}, 0, 4)
	return paramRe.ReplaceAllStringFunc(inSql, func(s string) string {
		v, ok := query[s[2:len(s) - 1]]
		if ok {
			params = append(params, v[0])
 		} else {
			params = append(params, "")
		}
		return "?"
	}), params
}

func dbQuery(db *sql.DB, sql string, params []interface{}) (sqlResult, error) {
	rows, err := db.Query(sql, params...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return rowsToResult(rows)
}

func rowsToResult(rows *sql.Rows) (sqlResult, error) {
	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	ret := make([]map[string]string, 0, 1)
	row := make([]interface{}, len(columns))
	for i, _ := range(row) {
		row[i] = ""
	}
	for rows.Next() {
		err = rows.Scan(row...)
		if err != nil {
			return ret, err
		}
		mapRow := make(map[string]string)
		for i, column := range(columns) {
			mapRow[column] = *(row[i]).(*string)
		}
		ret = append(ret, mapRow)
	}
	return ret, nil
}<commit_msg>add end log<commit_after>package server

import (
	"database/sql"
	"servant/conf"
	"net/http"
	"encoding/json"
)

type DatabaseServer struct {
	*Session
}

type sqlResult []map[string]string

func NewDatabaseServer(sess *Session) Handler {
	return DatabaseServer{
		Session:sess,
	}
}

func (self DatabaseServer) findDatabaseQueryConfig() (*conf.Database, *conf.Query) {
	dbConf, ok := self.config.Databases[self.group]
	if !ok {
		return nil, nil
	}
	qConf, ok := dbConf.Queries[self.item]
	if !ok {
		return dbConf, nil
	}
	return dbConf, qConf
}

func (self DatabaseServer) serve() {
	method := self.req.Method
	if method != "GET" {
		self.ErrorEnd(http.StatusMethodNotAllowed, "not allow method: %s", method)
		return
	}
	dbConf, queryConf := self.findDatabaseQueryConfig()
	if dbConf == nil {
		self.ErrorEnd(http.StatusNotFound, "database not found")
		return
	}
	if queryConf == nil {
		self.ErrorEnd(http.StatusNotFound, "query not found")
		return
	}

	db, err := sql.Open(dbConf.Driver, dbConf.Dsn)
	if err != nil {
		self.ErrorEnd(http.StatusInternalServerError, "driver init failed")
		return
	}
	defer db.Close()
	data := make([]sqlResult, 0, 1)
	for _, sql := range(queryConf.Sqls) {
		sql, params := replaceSqlParams(sql, self.req.URL.Query())

		result, err := dbQuery(db, sql, params)
		if err != nil {
			self.ErrorEnd(http.StatusInternalServerError, "query %s failed: %s", sql, err)
			return
		}
		data = append(data, result)
	}
	buf, err := json.Marshal(data)
	if err != nil {
		self.ErrorEnd(http.StatusInternalServerError, "json marshal failed: %s", err)
		return
	}
	self.resp.Write(buf)
	self.GoodEnd("execution done")
}

func replaceSqlParams(inSql string, query map[string][]string) (outSql string, params []interface{}){
	params = make([]interface{}, 0, 4)
	return paramRe.ReplaceAllStringFunc(inSql, func(s string) string {
		v, ok := query[s[2:len(s) - 1]]
		if ok {
			params = append(params, v[0])
 		} else {
			params = append(params, "")
		}
		return "?"
	}), params
}

func dbQuery(db *sql.DB, sql string, params []interface{}) (sqlResult, error) {
	rows, err := db.Query(sql, params...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return rowsToResult(rows)
}

func rowsToResult(rows *sql.Rows) (sqlResult, error) {
	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	ret := make([]map[string]string, 0, 1)
	row := make([]interface{}, len(columns))
	for i, _ := range(row) {
		row[i] = ""
	}
	for rows.Next() {
		err = rows.Scan(row...)
		if err != nil {
			return ret, err
		}
		mapRow := make(map[string]string)
		for i, column := range(columns) {
			mapRow[column] = *(row[i]).(*string)
		}
		ret = append(ret, mapRow)
	}
	return ret, nil
}<|endoftext|>


<commit_before>// +build !linux !amd64

package archive

import "syscall"

func getLastAccess(stat *syscall.Stat_t) syscall.Timespec {
	return syscall.Timespec{}
}

func getLastModification(stat *syscall.Stat_t) syscall.Timespec {
	return syscall.Timespec{}
}

func LUtimesNano(path string, ts []syscall.Timespec) error {
	return ErrNotImplemented
}

func UtimesNano(path string, ts []syscall.Timespec) error {
	return ErrNotImplemented
}
<commit_msg>Remove archive/stat_unsupported.go because it is not used<commit_after><|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	"github.com/jacobsa/comeback/repr"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestDirectoryRestorer(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func returnEntries(entries []*fs.DirectoryEntry) oglemock.Action {
	data, err := repr.Marshal(entries)
	AssertEq(nil, err)

	return oglemock.Return(data, nil)
}

func makeStrPtr(s string) *string {
	return &s
}

type DirectoryRestorerTest struct {
	blobStore    mock_blob.MockStore
	fileSystem   mock_fs.MockFileSystem
	fileRestorer mock_backup.MockFileRestorer
	wrapped      mock_backup.MockDirectoryRestorer

	dirRestorer backup.DirectoryRestorer

	score    blob.Score
	basePath string
	relPath  string

	err error
}

func init() { RegisterTestSuite(&DirectoryRestorerTest{}) }

func (t *DirectoryRestorerTest) SetUp(i *TestInfo) {
	var err error

	// Create dependencies.
	t.blobStore = mock_blob.NewMockStore(i.MockController, "blobStore")
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.fileRestorer = mock_backup.NewMockFileRestorer(i.MockController, "fileRestorer")
	t.wrapped = mock_backup.NewMockDirectoryRestorer(i.MockController, "wrapped")

	// Create restorer.
	t.dirRestorer, err = backup.NewNonRecursiveDirectoryRestorer(
		t.blobStore,
		t.fileSystem,
		t.fileRestorer,
		t.wrapped,
	)

	AssertEq(nil, err)
}

func (t *DirectoryRestorerTest) call() {
	t.err = t.dirRestorer.RestoreDirectory(t.score, t.basePath, t.relPath)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *DirectoryRestorerTest) CallsBlobStore() {
	t.score = []byte("taco")

	// Blob store
	ExpectCall(t.blobStore, "Load")(DeepEquals(t.score)).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) BlobStoreReturnsError() {
	// Blob store
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Load")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) BlobStoreReturnsJunk() {
	// Blob store
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(oglemock.Return([]byte("taco"), nil))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("invalid")))
	ExpectThat(t.err, Error(HasSubstr("data")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) NoEntries() {
	// Blob store
	entries := []*fs.DirectoryEntry{}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectEq(nil, t.err)
}

func (t *DirectoryRestorerTest) FileEntry_CallsLinkForHardLink() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name:           "taco",
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr("burrito/enchilada"),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(
		"/foo/burrito/enchilada",
		"/foo/bar/baz/taco",
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) FileEntry_LinkReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr(""),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("CreateHardLink")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) FileEntry_LinkSucceeds() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr(""),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(Any(), Any()).
		WillOnce(oglemock.Return(nil))

	// Call
	t.call()

	ExpectEq(nil, t.err)
}

func (t *DirectoryRestorerTest) FileEntry_CallsRestoreFile() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name:           "taco",
			Type:           fs.TypeFile,
			Permissions: 0712,
			Scores: []blob.Score{
				blob.ComputeScore([]byte("burrito")),
				blob.ComputeScore([]byte("enchilada")),
			},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File restorer
	ExpectCall(t.fileRestorer, "RestoreFile")(
		DeepEquals(entries[0].Scores),
		"/foo/bar/baz/taco",
		0712,
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) FileEntry_RestoreFileReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File restorer
	ExpectCall(t.fileRestorer, "RestoreFile")(Any(), Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("RestoreFile")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) DirEntry_ZeroScores() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("directory")))
	ExpectThat(t.err, Error(HasSubstr("entry")))
	ExpectThat(t.err, Error(HasSubstr("exactly one")))
	ExpectThat(t.err, Error(HasSubstr("score")))
}

func (t *DirectoryRestorerTest) DirEntry_TwoScores() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{
				blob.ComputeScore([]byte("a")),
				blob.ComputeScore([]byte("b")),
			},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("directory")))
	ExpectThat(t.err, Error(HasSubstr("entry")))
	ExpectThat(t.err, Error(HasSubstr("exactly one")))
	ExpectThat(t.err, Error(HasSubstr("score")))
}

func (t *DirectoryRestorerTest) DirEntry_CallsMkdir() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name: "taco",
			Type:           fs.TypeDirectory,
			Permissions: 0712,
			Scores: []blob.Score{blob.ComputeScore([]byte(""))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")("/foo/bar/baz/taco", 0712).
		WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) DirEntry_MkdirReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{blob.ComputeScore([]byte(""))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")(Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Mkdir")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) DirEntry_CallsWrapped() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name: "taco",
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{blob.ComputeScore([]byte("burrito"))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")(Any(), Any()).
		WillOnce(oglemock.Return(nil))

	// Wrapped
	ExpectCall(t.wrapped, "RestoreDirectory")(
		DeepEquals(blob.ComputeScore([]byte("burrito"))),
		"/foo",
		"bar/baz/taco",
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) DirEntry_WrappedReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) SymlinkEntry_CallsSymlink() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) SymlinkEntry_SymlinkReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) PipeEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) PipeEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) BlockDevEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) BlockDevEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CharDevEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CharDevEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CallsChown() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) ChownReturnsErrorForOneEntry() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CallsSetModTime() {
	ExpectEq("TODO", "")
	// NOTE: Not for devices (see restore.go)
}

func (t *DirectoryRestorerTest) SetModTimeReturnsErrorForOneEntry() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) EverythingSucceeds() {
	ExpectEq("TODO", "")
}
<commit_msg>DirEntry_WrappedReturnsError<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	"github.com/jacobsa/comeback/repr"
	. "github.com/jacobsa/oglematchers"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestDirectoryRestorer(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func returnEntries(entries []*fs.DirectoryEntry) oglemock.Action {
	data, err := repr.Marshal(entries)
	AssertEq(nil, err)

	return oglemock.Return(data, nil)
}

func makeStrPtr(s string) *string {
	return &s
}

type DirectoryRestorerTest struct {
	blobStore    mock_blob.MockStore
	fileSystem   mock_fs.MockFileSystem
	fileRestorer mock_backup.MockFileRestorer
	wrapped      mock_backup.MockDirectoryRestorer

	dirRestorer backup.DirectoryRestorer

	score    blob.Score
	basePath string
	relPath  string

	err error
}

func init() { RegisterTestSuite(&DirectoryRestorerTest{}) }

func (t *DirectoryRestorerTest) SetUp(i *TestInfo) {
	var err error

	// Create dependencies.
	t.blobStore = mock_blob.NewMockStore(i.MockController, "blobStore")
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.fileRestorer = mock_backup.NewMockFileRestorer(i.MockController, "fileRestorer")
	t.wrapped = mock_backup.NewMockDirectoryRestorer(i.MockController, "wrapped")

	// Create restorer.
	t.dirRestorer, err = backup.NewNonRecursiveDirectoryRestorer(
		t.blobStore,
		t.fileSystem,
		t.fileRestorer,
		t.wrapped,
	)

	AssertEq(nil, err)
}

func (t *DirectoryRestorerTest) call() {
	t.err = t.dirRestorer.RestoreDirectory(t.score, t.basePath, t.relPath)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *DirectoryRestorerTest) CallsBlobStore() {
	t.score = []byte("taco")

	// Blob store
	ExpectCall(t.blobStore, "Load")(DeepEquals(t.score)).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) BlobStoreReturnsError() {
	// Blob store
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Load")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) BlobStoreReturnsJunk() {
	// Blob store
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(oglemock.Return([]byte("taco"), nil))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("invalid")))
	ExpectThat(t.err, Error(HasSubstr("data")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) NoEntries() {
	// Blob store
	entries := []*fs.DirectoryEntry{}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectEq(nil, t.err)
}

func (t *DirectoryRestorerTest) FileEntry_CallsLinkForHardLink() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name:           "taco",
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr("burrito/enchilada"),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(
		"/foo/burrito/enchilada",
		"/foo/bar/baz/taco",
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) FileEntry_LinkReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr(""),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("CreateHardLink")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) FileEntry_LinkSucceeds() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
			HardLinkTarget: makeStrPtr(""),
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "CreateHardLink")(Any(), Any()).
		WillOnce(oglemock.Return(nil))

	// Call
	t.call()

	ExpectEq(nil, t.err)
}

func (t *DirectoryRestorerTest) FileEntry_CallsRestoreFile() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name:           "taco",
			Type:           fs.TypeFile,
			Permissions: 0712,
			Scores: []blob.Score{
				blob.ComputeScore([]byte("burrito")),
				blob.ComputeScore([]byte("enchilada")),
			},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File restorer
	ExpectCall(t.fileRestorer, "RestoreFile")(
		DeepEquals(entries[0].Scores),
		"/foo/bar/baz/taco",
		0712,
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) FileEntry_RestoreFileReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeFile,
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File restorer
	ExpectCall(t.fileRestorer, "RestoreFile")(Any(), Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("RestoreFile")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) DirEntry_ZeroScores() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("directory")))
	ExpectThat(t.err, Error(HasSubstr("entry")))
	ExpectThat(t.err, Error(HasSubstr("exactly one")))
	ExpectThat(t.err, Error(HasSubstr("score")))
}

func (t *DirectoryRestorerTest) DirEntry_TwoScores() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{
				blob.ComputeScore([]byte("a")),
				blob.ComputeScore([]byte("b")),
			},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("directory")))
	ExpectThat(t.err, Error(HasSubstr("entry")))
	ExpectThat(t.err, Error(HasSubstr("exactly one")))
	ExpectThat(t.err, Error(HasSubstr("score")))
}

func (t *DirectoryRestorerTest) DirEntry_CallsMkdir() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name: "taco",
			Type:           fs.TypeDirectory,
			Permissions: 0712,
			Scores: []blob.Score{blob.ComputeScore([]byte(""))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")("/foo/bar/baz/taco", 0712).
		WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) DirEntry_MkdirReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{blob.ComputeScore([]byte(""))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")(Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(HasSubstr("Mkdir")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectoryRestorerTest) DirEntry_CallsWrapped() {
	t.basePath = "/foo"
	t.relPath = "bar/baz"

	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Name: "taco",
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{blob.ComputeScore([]byte("burrito"))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")(Any(), Any()).
		WillOnce(oglemock.Return(nil))

	// Wrapped
	ExpectCall(t.wrapped, "RestoreDirectory")(
		DeepEquals(blob.ComputeScore([]byte("burrito"))),
		"/foo",
		"bar/baz/taco",
	).WillOnce(oglemock.Return(errors.New("")))

	// Call
	t.call()
}

func (t *DirectoryRestorerTest) DirEntry_WrappedReturnsError() {
	// Blob store
	entries := []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type:           fs.TypeDirectory,
			Scores: []blob.Score{blob.ComputeScore([]byte(""))},
		},
	}

	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(returnEntries(entries))

	// File system
	ExpectCall(t.fileSystem, "Mkdir")(Any(), Any()).
		WillOnce(oglemock.Return(nil))

	// Wrapped
	ExpectCall(t.wrapped, "RestoreDirectory")(Any(), Any(), Any()).
		WillOnce(oglemock.Return(errors.New("taco")))

	// Call
	t.call()

	ExpectThat(t.err, Error(Equals("taco")))
}

func (t *DirectoryRestorerTest) SymlinkEntry_CallsSymlink() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) SymlinkEntry_SymlinkReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) PipeEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) PipeEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) BlockDevEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) BlockDevEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CharDevEntry_CallsCreate() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CharDevEntry_CreateReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CallsChown() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) ChownReturnsErrorForOneEntry() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) CallsSetModTime() {
	ExpectEq("TODO", "")
	// NOTE: Not for devices (see restore.go)
}

func (t *DirectoryRestorerTest) SetModTimeReturnsErrorForOneEntry() {
	ExpectEq("TODO", "")
}

func (t *DirectoryRestorerTest) EverythingSucceeds() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * 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 stackdriver

import (
	"bytes"
	"fmt"
	"path"
	"sort"
	"sync"
	"time"

	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/pkg/errors"
	labelpb "google.golang.org/genproto/googleapis/api/label"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

type Metric struct {
	Name      string
	Value     float64
	Labels    map[string]string
	EventTime time.Time
	Unit      string // TODO Should this be "1" if it's empty?
}

func (m *Metric) Hash() string {
	var b bytes.Buffer
	b.Write([]byte(m.Name))

	// Extract keys to a slice and sort it
	keys := make([]string, len(m.Labels), len(m.Labels))
	for k, _ := range m.Labels {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		b.Write([]byte(k))
		b.Write([]byte(m.Labels[k]))
	}
	return b.String()
}

type MetricAdapter interface {
	PostMetrics([]Metric) error
}

type Heartbeater interface {
	Start()
	Increment(string)
	Stop()
}

type metricAdapter struct {
	projectID             string
	client                MetricClient
	descriptors           map[string]struct{}
	createDescriptorMutex *sync.Mutex
	heartbeater           Heartbeater
}

func NewMetricAdapter(projectID string, client MetricClient, heartbeater Heartbeater) (MetricAdapter, error) {
	ma := &metricAdapter{
		projectID:             projectID,
		client:                client,
		createDescriptorMutex: &sync.Mutex{},
		descriptors:           map[string]struct{}{},
		heartbeater:           heartbeater,
	}

	err := ma.fetchMetricDescriptorNames()
	return ma, err
}

func (ma *metricAdapter) PostMetrics(metrics []Metric) error {
	projectName := path.Join("projects", ma.projectID)
	var timeSerieses []*monitoringpb.TimeSeries

	for _, metric := range metrics {
		ma.heartbeater.Increment("metrics.count")

		err := ma.ensureMetricDescriptor(metric)
		if err != nil {
			return err
		}

		metricType := path.Join("custom.googleapis.com", metric.Name)
		timeSeries := monitoringpb.TimeSeries{
			Metric: &metricpb.Metric{
				Type:   metricType,
				Labels: metric.Labels,
			},
			Points: points(metric.Value, metric.EventTime),
		}
		timeSerieses = append(timeSerieses, &timeSeries)
	}

	request := &monitoringpb.CreateTimeSeriesRequest{
		Name:       projectName,
		TimeSeries: timeSerieses,
	}

	ma.heartbeater.Increment("metrics.requests")
	err := ma.client.Post(request)
	if err != nil {
		ma.heartbeater.Increment("metrics.errors")
	}
	err = errors.Wrapf(err, "Request: %+v", request)
	return err
}

func (ma *metricAdapter) CreateMetricDescriptor(metric Metric) error {
	projectName := path.Join("projects", ma.projectID)
	metricType := path.Join("custom.googleapis.com", metric.Name)
	metricName := path.Join(projectName, "metricDescriptors", metricType)

	var labelDescriptors []*labelpb.LabelDescriptor
	for key := range metric.Labels {
		labelDescriptors = append(labelDescriptors, &labelpb.LabelDescriptor{
			Key:       key,
			ValueType: labelpb.LabelDescriptor_STRING,
		})
	}

	req := &monitoringpb.CreateMetricDescriptorRequest{
		Name: projectName,
		MetricDescriptor: &metricpb.MetricDescriptor{
			Name:        metricName,
			Type:        metricType,
			Labels:      labelDescriptors,
			MetricKind:  metricpb.MetricDescriptor_GAUGE,
			ValueType:   metricpb.MetricDescriptor_DOUBLE,
			Unit:        metric.Unit,
			Description: "stackdriver-nozzle created custom metric.",
			DisplayName: metric.Name, // TODO
		},
	}

	return ma.client.CreateMetricDescriptor(req)
}

func (ma *metricAdapter) fetchMetricDescriptorNames() error {
	req := &monitoringpb.ListMetricDescriptorsRequest{
		Name:   fmt.Sprintf("projects/%s", ma.projectID),
		Filter: "metric.type = starts_with(\"custom.googleapis.com/\")",
	}

	descriptors, err := ma.client.ListMetricDescriptors(req)
	if err != nil {
		return err
	}

	for _, descriptor := range descriptors {
		ma.descriptors[descriptor.Name] = struct{}{}
	}
	return nil
}

func (ma *metricAdapter) ensureMetricDescriptor(metric Metric) error {
	if metric.Unit == "" {
		return nil
	}

	ma.createDescriptorMutex.Lock()
	defer ma.createDescriptorMutex.Unlock()

	if _, ok := ma.descriptors[metric.Name]; ok {
		return nil
	}

	err := ma.CreateMetricDescriptor(metric)
	if err != nil {
		return err
	}
	ma.descriptors[metric.Name] = struct{}{}
	return nil
}

func points(value float64, eventTime time.Time) []*monitoringpb.Point {
	timeStamp := timestamp.Timestamp{
		Seconds: eventTime.Unix(),
		Nanos:   int32(eventTime.Nanosecond()),
	}
	point := &monitoringpb.Point{
		Interval: &monitoringpb.TimeInterval{
			EndTime:   &timeStamp,
			StartTime: &timeStamp,
		},
		Value: &monitoringpb.TypedValue{
			Value: &monitoringpb.TypedValue_DoubleValue{
				DoubleValue: value,
			},
		},
	}
	return []*monitoringpb.Point{point}
}
<commit_msg>stackdriver: Don't send an empty metrics slice<commit_after>/*
 * Copyright 2017 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * 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 stackdriver

import (
	"bytes"
	"fmt"
	"path"
	"sort"
	"sync"
	"time"

	"github.com/golang/protobuf/ptypes/timestamp"
	"github.com/pkg/errors"
	labelpb "google.golang.org/genproto/googleapis/api/label"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

type Metric struct {
	Name      string
	Value     float64
	Labels    map[string]string
	EventTime time.Time
	Unit      string // TODO Should this be "1" if it's empty?
}

func (m *Metric) Hash() string {
	var b bytes.Buffer
	b.Write([]byte(m.Name))

	// Extract keys to a slice and sort it
	keys := make([]string, len(m.Labels), len(m.Labels))
	for k, _ := range m.Labels {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		b.Write([]byte(k))
		b.Write([]byte(m.Labels[k]))
	}
	return b.String()
}

type MetricAdapter interface {
	PostMetrics([]Metric) error
}

type Heartbeater interface {
	Start()
	Increment(string)
	Stop()
}

type metricAdapter struct {
	projectID             string
	client                MetricClient
	descriptors           map[string]struct{}
	createDescriptorMutex *sync.Mutex
	heartbeater           Heartbeater
}

func NewMetricAdapter(projectID string, client MetricClient, heartbeater Heartbeater) (MetricAdapter, error) {
	ma := &metricAdapter{
		projectID:             projectID,
		client:                client,
		createDescriptorMutex: &sync.Mutex{},
		descriptors:           map[string]struct{}{},
		heartbeater:           heartbeater,
	}

	err := ma.fetchMetricDescriptorNames()
	return ma, err
}

func (ma *metricAdapter) PostMetrics(metrics []Metric) error {
	if len(metrics) == 0 {
		return nil
	}

	projectName := path.Join("projects", ma.projectID)
	var timeSerieses []*monitoringpb.TimeSeries

	for _, metric := range metrics {
		ma.heartbeater.Increment("metrics.count")

		err := ma.ensureMetricDescriptor(metric)
		if err != nil {
			return err
		}

		metricType := path.Join("custom.googleapis.com", metric.Name)
		timeSeries := monitoringpb.TimeSeries{
			Metric: &metricpb.Metric{
				Type:   metricType,
				Labels: metric.Labels,
			},
			Points: points(metric.Value, metric.EventTime),
		}
		timeSerieses = append(timeSerieses, &timeSeries)
	}

	request := &monitoringpb.CreateTimeSeriesRequest{
		Name:       projectName,
		TimeSeries: timeSerieses,
	}

	ma.heartbeater.Increment("metrics.requests")
	err := ma.client.Post(request)
	if err != nil {
		ma.heartbeater.Increment("metrics.errors")
	}
	err = errors.Wrapf(err, "Request: %+v", request)
	return err
}

func (ma *metricAdapter) CreateMetricDescriptor(metric Metric) error {
	projectName := path.Join("projects", ma.projectID)
	metricType := path.Join("custom.googleapis.com", metric.Name)
	metricName := path.Join(projectName, "metricDescriptors", metricType)

	var labelDescriptors []*labelpb.LabelDescriptor
	for key := range metric.Labels {
		labelDescriptors = append(labelDescriptors, &labelpb.LabelDescriptor{
			Key:       key,
			ValueType: labelpb.LabelDescriptor_STRING,
		})
	}

	req := &monitoringpb.CreateMetricDescriptorRequest{
		Name: projectName,
		MetricDescriptor: &metricpb.MetricDescriptor{
			Name:        metricName,
			Type:        metricType,
			Labels:      labelDescriptors,
			MetricKind:  metricpb.MetricDescriptor_GAUGE,
			ValueType:   metricpb.MetricDescriptor_DOUBLE,
			Unit:        metric.Unit,
			Description: "stackdriver-nozzle created custom metric.",
			DisplayName: metric.Name, // TODO
		},
	}

	return ma.client.CreateMetricDescriptor(req)
}

func (ma *metricAdapter) fetchMetricDescriptorNames() error {
	req := &monitoringpb.ListMetricDescriptorsRequest{
		Name:   fmt.Sprintf("projects/%s", ma.projectID),
		Filter: "metric.type = starts_with(\"custom.googleapis.com/\")",
	}

	descriptors, err := ma.client.ListMetricDescriptors(req)
	if err != nil {
		return err
	}

	for _, descriptor := range descriptors {
		ma.descriptors[descriptor.Name] = struct{}{}
	}
	return nil
}

func (ma *metricAdapter) ensureMetricDescriptor(metric Metric) error {
	if metric.Unit == "" {
		return nil
	}

	ma.createDescriptorMutex.Lock()
	defer ma.createDescriptorMutex.Unlock()

	if _, ok := ma.descriptors[metric.Name]; ok {
		return nil
	}

	err := ma.CreateMetricDescriptor(metric)
	if err != nil {
		return err
	}
	ma.descriptors[metric.Name] = struct{}{}
	return nil
}

func points(value float64, eventTime time.Time) []*monitoringpb.Point {
	timeStamp := timestamp.Timestamp{
		Seconds: eventTime.Unix(),
		Nanos:   int32(eventTime.Nanosecond()),
	}
	point := &monitoringpb.Point{
		Interval: &monitoringpb.TimeInterval{
			EndTime:   &timeStamp,
			StartTime: &timeStamp,
		},
		Value: &monitoringpb.TypedValue{
			Value: &monitoringpb.TypedValue_DoubleValue{
				DoubleValue: value,
			},
		},
	}
	return []*monitoringpb.Point{point}
}
<|endoftext|>


<commit_before>package s3

import (
	"net/http"
	"strings"

	"github.com/pachyderm/pachyderm/src/server/pkg/s3server"
)

// The S3 storage class that all PFS content will be reported to be stored in
const storageClass = "STANDARD"

// The S3 user associated with all PFS content
var defaultUser = s3server.User{ID: "00000000000000000000000000000000", DisplayName: "pachyderm"}

func bucketArgs(r *http.Request, bucket string) (string, string, *s3server.S3Error) {
	parts := strings.SplitN(bucket, ".", 2)
	if len(parts) != 2 || len(parts[0]) == 0 || len(parts[1]) == 0 {
		return "", "", s3server.InvalidBucketNameError(r)
	}
	return parts[0], parts[1], nil
}
<commit_msg>Fixed reversed ordering of bucket/branch names<commit_after>package s3

import (
	"net/http"
	"strings"

	"github.com/pachyderm/pachyderm/src/server/pkg/s3server"
)

// The S3 storage class that all PFS content will be reported to be stored in
const storageClass = "STANDARD"

// The S3 user associated with all PFS content
var defaultUser = s3server.User{ID: "00000000000000000000000000000000", DisplayName: "pachyderm"}

func bucketArgs(r *http.Request, bucket string) (string, string, *s3server.S3Error) {
	parts := strings.SplitN(bucket, ".", 2)
	if len(parts) != 2 || len(parts[0]) == 0 || len(parts[1]) == 0 {
		return "", "", s3server.InvalidBucketNameError(r)
	}
	return parts[1], parts[0], nil
}
<|endoftext|>


<commit_before>package sql

import (
	"bufio"
	"fmt"
	"io"
	"strings"
)

type pgDumpReader struct {
	Header []byte
	Footer []byte
	rd     *bufio.Reader
}

func NewPGDumpReader(r *bufio.Reader) *pgDumpReader {
	return &pgDumpReader{
		rd: r,
	}
}

// ReadRows parses the pgdump file and populates the header and the footer
// It returns EOF when done, and at that time both the Header and Footer will
// be populated. Both header and footer are required. If either are missing, an
// error is returned
func (r *pgDumpReader) ReadRows(count int64) (rowsDump []byte, rowsRead int64, err error) {
	endLine := "\\.\n" // Trailing '\.' denotes the end of the row inserts
	if len(r.Header) == 0 {
		done := false
		for !done {
			b, err := r.rd.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					return nil, 0, fmt.Errorf("invalid header - missing row inserts")
				}
				return nil, 0, err
			}
			if strings.HasPrefix(string(b), "COPY") {
				done = true
			}
			r.Header = append(r.Header, b...)
		}
		fmt.Printf("read off the header (%v) w err(%v)\n", string(r.Header), err)
	}

	for rowsRead = 0; rowsRead < count; rowsRead++ {
		row, _err := r.rd.ReadBytes('\n')
		err = _err
		if string(row) == endLine {
			r.Footer = append(r.Footer, row...)
			err = r.readFooter()
			if count == 1 {
				// In this case, when we see and endline, we don't want to return any content
				return nil, 0, io.EOF
			}
			break
		}
		if err == io.EOF && len(r.Footer) == 0 {
			return nil, 0, fmt.Errorf("invalid pgdump - missing footer\n")
		}
		rowsDump = append(rowsDump, row...)
	}
	return rowsDump, rowsRead, err
}

func (r *pgDumpReader) readFooter() error {
	for true {
		b, err := r.rd.ReadBytes('\n')
		r.Footer = append(r.Footer, b...)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Remove debug statement and fix error string<commit_after>package sql

import (
	"bufio"
	"fmt"
	"io"
	"strings"
)

type pgDumpReader struct {
	Header []byte
	Footer []byte
	rd     *bufio.Reader
}

func NewPGDumpReader(r *bufio.Reader) *pgDumpReader {
	return &pgDumpReader{
		rd: r,
	}
}

// ReadRows parses the pgdump file and populates the header and the footer
// It returns EOF when done, and at that time both the Header and Footer will
// be populated. Both header and footer are required. If either are missing, an
// error is returned
func (r *pgDumpReader) ReadRows(count int64) (rowsDump []byte, rowsRead int64, err error) {
	endLine := "\\.\n" // Trailing '\.' denotes the end of the row inserts
	if len(r.Header) == 0 {
		done := false
		for !done {
			b, err := r.rd.ReadBytes('\n')
			if err != nil {
				if err == io.EOF {
					return nil, 0, fmt.Errorf("invalid header - missing row inserts")
				}
				return nil, 0, err
			}
			if strings.HasPrefix(string(b), "COPY") {
				done = true
			}
			r.Header = append(r.Header, b...)
		}
	}

	for rowsRead = 0; rowsRead < count; rowsRead++ {
		row, _err := r.rd.ReadBytes('\n')
		err = _err
		if string(row) == endLine {
			r.Footer = append(r.Footer, row...)
			err = r.readFooter()
			if count == 1 {
				// In this case, when we see and endline, we don't want to return any content
				return nil, 0, io.EOF
			}
			break
		}
		if err == io.EOF && len(r.Footer) == 0 {
			return nil, 0, fmt.Errorf("invalid pgdump - missing footer")
		}
		rowsDump = append(rowsDump, row...)
	}
	return rowsDump, rowsRead, err
}

func (r *pgDumpReader) readFooter() error {
	for true {
		b, err := r.rd.ReadBytes('\n')
		r.Footer = append(r.Footer, b...)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package instruments

import (
	"testing"
	"testing/quick"
)

var quantileTests = []struct {
	values   []int64
	quantile float64
	value    int64
}{
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.75,
		value:    43,
	},
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.25,
		value:    15,
	},
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.5,
		value:    40,
	},
	{
		values:   []int64{},
		quantile: 0.95,
		value:    0,
	},
}

func TestQuantile(t *testing.T) {
	for i, qt := range quantileTests {
		q := Quantile(qt.values, qt.quantile)
		if q != qt.value {
			t.Errorf("%d: wrong value returned for quantile %g, got %d wants %d", i, qt.quantile, q, qt.value)
		}
	}
}

func TestMean(t *testing.T) {
	m := Mean([]int64{1, 2, 3, 4})
	if m != 2.5 {
		t.Error("wrong calculated mean")
	}
}

var minMaxTests = []struct {
	values []int64
	min    int64
	max    int64
}{
	{
		values: []int64{},
		min:    0,
		max:    0,
	},
	{
		values: []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		min:    6,
		max:    49,
	},
	{
		values: []int64{15, 6, 36, 39, 7, 40, 43, 43, 47, 49},
		min:    6,
		max:    49,
	},
}

func TestMinMax(t *testing.T) {
	for i, mm := range minMaxTests {
		min := Min(mm.values)
		if min != mm.min {
			t.Errorf("%d: expected min value %d got %d", i, mm.min, min)
		}
		max := Max(mm.values)
		if max != mm.max {
			t.Errorf("%d: expected max value %d got %d", i, mm.max, max)
		}
	}
}

func TestMin(t *testing.T) {
	err := quick.Check(func(values []int64) bool {
		min := Min(values)
		for _, v := range values {
			if v < min {
				return false
			}
		}
		return true
	}, nil)
	if err != nil {
		t.Error(err)
	}
}

func TestMax(t *testing.T) {
	err := quick.Check(func(values []int64) bool {
		max := Max(values)
		for _, v := range values {
			if v > max {
				return false
			}
		}
		return true
	}, nil)
	if err != nil {
		t.Error(err)
	}
}
<commit_msg>make min/max test more readable<commit_after>package instruments

import (
	"testing"
	"testing/quick"
)

var quantileTests = []struct {
	values   []int64
	quantile float64
	value    int64
}{
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.75,
		value:    43,
	},
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.25,
		value:    15,
	},
	{
		values:   []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		quantile: 0.5,
		value:    40,
	},
	{
		values:   []int64{},
		quantile: 0.95,
		value:    0,
	},
}

func TestQuantile(t *testing.T) {
	for i, qt := range quantileTests {
		q := Quantile(qt.values, qt.quantile)
		if q != qt.value {
			t.Errorf("%d: wrong value returned for quantile %g, got %d wants %d", i, qt.quantile, q, qt.value)
		}
	}
}

func TestMean(t *testing.T) {
	m := Mean([]int64{1, 2, 3, 4})
	if m != 2.5 {
		t.Error("wrong calculated mean")
	}
}

var minMaxTests = []struct {
	values []int64
	min    int64
	max    int64
}{
	{
		values: []int64{},
		min:    0,
		max:    0,
	},
	{
		values: []int64{6, 7, 15, 36, 39, 40, 41, 43, 43, 47, 49},
		min:    6,
		max:    49,
	},
	{
		values: []int64{15, 6, 36, 39, 7, 40, 43, 43, 47, 49},
		min:    6,
		max:    49,
	},
}

func TestMinMax(t *testing.T) {
	for i, mm := range minMaxTests {
		min := Min(mm.values)
		if min != mm.min {
			t.Errorf("%d: expected min value %d got %d", i, mm.min, min)
		}
		max := Max(mm.values)
		if max != mm.max {
			t.Errorf("%d: expected max value %d got %d", i, mm.max, max)
		}
	}
}

func TestMin(t *testing.T) {
	min := func(values []int64) bool {
		min := Min(values)
		for _, v := range values {
			if v < min {
				return false
			}
		}
		return true
	}
	if err := quick.Check(min, nil); err != nil {
		t.Error(err)
	}
}

func TestMax(t *testing.T) {
	max := func(values []int64) bool {
		max := Max(values)
		for _, v := range values {
			if v > max {
				return false
			}
		}
		return true
	}
	if err := quick.Check(max, nil); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code,string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code,string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3,string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name,string,omitempty"`   // Name of the country
	Region         string  `json:"region,string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code,integer,omitempty"`    // Metro code within the country
	City           string  `json:"city,string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code,integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code,string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude,float"`                  // Latitude
	Longitude      float64 `json:"longitude,float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	for _, ip := range ips {
		if ip == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		ip, _ = web100.NormalizeIPv6(ip)
		reqData = append(reqData, RequestData{ip, 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for index, ip := range ips {
		data, ok := annotationData[ip+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[index] = *data.Geo

	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<commit_msg>use separate slice for normalized<commit_after>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code,string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code,string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3,string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name,string,omitempty"`   // Name of the country
	Region         string  `json:"region,string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code,integer,omitempty"`    // Metro code within the country
	City           string  `json:"city,string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code,integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code,string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude,float"`                  // Latitude
	Longitude      float64 `json:"longitude,float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	normalized := make([]string, len(ips))
	for i := range ips {
		if ips[i] == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		var err error
		normalized[i], err = web100.NormalizeIPv6(ips[i])
		if err != nil {
			log.Println(err)
		}
		reqData = append(reqData, RequestData{normalized[i], 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for i := range normalized {
		data, ok := annotationData[normalized[i]+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[i] = *data.Geo
	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<|endoftext|>