<commit_before>package server

import (
	"app/server/data"
	"app/server/react"
	. "app/server/utils"

	"github.com/codegangsta/cli"
	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gin-gonic/gin"
)

func Run(args []string) {

	app := cli.NewApp()
	app.Name = "app"
	app.Usage = "lmbd landing server application"

	configFlag := cli.StringFlag{
		Name:   "config, c",
		Value:  "local",
		Usage:  "configuration section name",
		EnvVar: "CONFIG",
	}

	app.Commands = []cli.Command{
		{
			Name:   "run",
			Usage:  "Runs server",
			Action: runServer,
			Flags:  []cli.Flag{configFlag},
		},
	}
	app.Run(args)
}

func runServer(c *cli.Context) {
	kit := NewKit(c, conf)
	kit.Engine = gin.Default()

	kit.Engine.StaticFS("/static", &assetfs.AssetFS{
		Asset:    data.Asset,
		AssetDir: data.AssetDir,
		Prefix:   "static",
	})

	// routes.bind(r)
	react.Bind(kit)
	Must(kit.Engine.Run(":" + kit.Conf.UString("port")))
}
<commit_msg>Favicon redirect<commit_after>package server

import (
	"app/server/data"
	"app/server/react"
	. "app/server/utils"

	"github.com/codegangsta/cli"
	"github.com/elazarl/go-bindata-assetfs"
	"github.com/gin-gonic/gin"
)

func Run(args []string) {

	app := cli.NewApp()
	app.Name = "app"
	app.Usage = "lmbd landing server application"

	configFlag := cli.StringFlag{
		Name:   "config, c",
		Value:  "local",
		Usage:  "configuration section name",
		EnvVar: "CONFIG",
	}

	app.Commands = []cli.Command{
		{
			Name:   "run",
			Usage:  "Runs server",
			Action: runServer,
			Flags:  []cli.Flag{configFlag},
		},
	}
	app.Run(args)
}

func runServer(c *cli.Context) {
	kit := NewKit(c, conf)
	kit.Engine = gin.Default()

	kit.Engine.StaticFS("/static", &assetfs.AssetFS{
		Asset:    data.Asset,
		AssetDir: data.AssetDir,
		Prefix:   "static",
	})

	// Avoid favicon react handling
	kit.Engine.GET("/favicon.ico", func(c *gin.Context) {
		c.Redirect(301, "/static/images/favicon.ico")
	})

	react.Bind(kit)
	Must(kit.Engine.Run(":" + kit.Conf.UString("port")))
}
<|endoftext|>


<commit_before>package ldap

import (
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"

	l "github.com/go-ldap/ldap"
	"github.com/pivotalservices/cf-mgmt/config"
	"github.com/xchapter7x/lo"
)

var (
	attributes = []string{"*"}
)

var (
	userRegexp          = regexp.MustCompile(",[A-Z]+=")
	escapeFilterRegex   = regexp.MustCompile(`([\\\(\)\*\0-\37\177-\377])`)
	unescapeFilterRegex = regexp.MustCompile(`\\([\da-fA-F]{2}|[()\\*])`) // only match \[)*\] or \xx x=a-fA-F
)

const (
	groupFilter                 = "(cn=%s)"
	userFilter                  = "(%s=%s)"
	userFilterWithObjectClass   = "(&(objectclass=%s)(%s=%s))"
	userDNFilter                = "(%s)"
	userDNFilterWithObjectClass = "(&(objectclass=%s)(%s))"
)

func NewManager(ldapConfig *config.LdapConfig) (Manager, error) {
	conn, err := CreateConnection(ldapConfig)
	if err != nil {
		return nil, err
	}
	return &DefaultManager{
		Config:     ldapConfig,
		Connection: conn,
	}, nil
}

func (m *DefaultManager) GetUserDNs(groupName string) ([]string, error) {
	filter := fmt.Sprintf(groupFilter, l.EscapeFilter(groupName))
	var groupEntry *l.Entry
	lo.G.Debug("Searching for group:", filter)
	lo.G.Debug("Using group search base:", m.Config.GroupSearchBase)

	search := l.NewSearchRequest(
		m.Config.GroupSearchBase,
		l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
		filter,
		attributes,
		nil)
	sr, err := m.Connection.Search(search)
	if err != nil {
		lo.G.Error(err)
		return nil, err
	}

	if len(sr.Entries) == 0 {
		lo.G.Errorf("group not found: %s", groupName)
		return []string{}, nil
	}
	if len(sr.Entries) > 1 {
		lo.G.Errorf("multiple groups found for: %s", groupName)
		return []string{}, nil
	}

	groupEntry = sr.Entries[0]
	userDNList := groupEntry.GetAttributeValues(m.Config.GroupAttribute)
	if len(userDNList) == 0 {
		lo.G.Warningf("No users found under group: %s", groupName)
	}

	userMap := make(map[string]string)
	for _, userDN := range userDNList {
		group, cn, err := m.IsGroup(userDN)
		if err != nil {
			return nil, err
		}
		if group {
			if err != nil {
				return nil, err
			}
			nestedUsers, err := m.GetUserDNs(cn)
			if err != nil {
				return nil, err
			}
			for _, nestedUser := range nestedUsers {
				userMap[nestedUser] = nestedUser
			}
		} else {
			userMap[userDN] = userDN
		}
	}
	var userList []string
	for _, userDN := range userMap {
		userList = append(userList, userDN)
	}
	return userList, nil
}

func (m *DefaultManager) getCN(userDN string) (string, error) {
	indexes := userRegexp.FindStringIndex(strings.ToUpper(userDN))
	if len(indexes) == 0 {
		return "", fmt.Errorf("cannot find CN for DN: %s", userDN)
	}
	cn := strings.Replace(userDN[:indexes[0]], "cn=", "", 1)
	cnTemp := UnescapeFilterValue(cn)
	lo.G.Debug("CN unescaped:", cnTemp)

	escapedCN := l.EscapeFilter(strings.Replace(cnTemp, "\\", "", -1))
	lo.G.Debug("CN escaped:", escapedCN)
	return escapedCN, nil
}
func (m *DefaultManager) IsGroup(userDN string) (bool, string, error) {
	if strings.Contains(userDN, m.Config.GroupSearchBase) {
		cn, err := m.getCN(userDN)
		if err != nil {
			return false, "", err
		}
		search := l.NewSearchRequest(
			m.Config.GroupSearchBase,
			l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
			fmt.Sprintf(groupFilter, cn),
			attributes,
			nil)
		sr, err := m.Connection.Search(search)
		if err != nil {
			return false, "", err
		}
		lo.G.Debugf("Found %d entries for userDN %s", len(sr.Entries), userDN)
		return len(sr.Entries) == 1, cn, nil
	} else {
		return false, "", nil
	}
}

func (m *DefaultManager) GetUserByDN(userDN string) (*User, error) {
	lo.G.Debug("User DN:", userDN)
	indexes := userRegexp.FindStringIndex(strings.ToUpper(userDN))
	if len(indexes) == 0 {
		return nil, fmt.Errorf("cannot find CN for user DN: %s", userDN)
	}
	index := indexes[0]
	userCNTemp := UnescapeFilterValue(userDN[:index])
	lo.G.Debug("CN unescaped:", userCNTemp)

	userCN := l.EscapeFilter(strings.Replace(userCNTemp, "\\", "", 1))
	lo.G.Debug("CN escaped:", userCN)
	filter := m.getUserFilterWithCN(userCN)
	return m.searchUser(filter, userDN[index+1:], "")
}

func (m *DefaultManager) GetUserByID(userID string) (*User, error) {
	filter := m.getUserFilter(userID)
	lo.G.Debug("Searching for user:", filter)
	lo.G.Debug("Using user search base:", m.Config.UserSearchBase)
	return m.searchUser(filter, m.Config.UserSearchBase, userID)
}

func (m *DefaultManager) searchUser(filter, searchBase, userID string) (*User, error) {
	lo.G.Debugf("Searching with filter %s", filter)
	search := l.NewSearchRequest(
		searchBase,
		l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
		filter,
		attributes,
		nil)

	sr, err := m.Connection.Search(search)
	if err != nil {
		lo.G.Error(err)
		return nil, err
	}

	if (len(sr.Entries)) == 1 {
		entry := sr.Entries[0]
		user := &User{
			UserDN: entry.DN,
			Email:  entry.GetAttributeValue(m.Config.UserMailAttribute),
		}
		if userID != "" {
			user.UserID = userID
		} else {
			user.UserID = entry.GetAttributeValue(m.Config.UserNameAttribute)
		}
		lo.G.Debugf("Search filter %s returned userDN [%s], email [%s], userID [%s]", filter, user.UserDN, user.Email, user.UserID)
		return user, nil
	}
	lo.G.Errorf("Found %d number of entries for filter %s", len(sr.Entries), filter)
	return nil, nil
}

func UnescapeFilterValue(filter string) string {
	repl := unescapeFilterRegex.ReplaceAllFunc(
		[]byte(filter),
		func(match []byte) []byte {
			// \( \) \\ \*
			if len(match) == 2 {
				return []byte{match[1]}
			}
			// had issues with Decode, TODO fix to use Decode?.
			res, _ := hex.DecodeString(string(match[1:]))
			return res
		},
	)
	return string(repl)
}

func (m *DefaultManager) getUserFilter(userID string) string {
	if m.Config.UserObjectClass == "" {
		return fmt.Sprintf(userFilter, m.Config.UserNameAttribute, userID)
	}
	return fmt.Sprintf(userFilterWithObjectClass, m.Config.UserObjectClass, m.Config.UserNameAttribute, userID)
}

func (m *DefaultManager) getUserFilterWithCN(cn string) string {
	if m.Config.UserObjectClass == "" {
		return fmt.Sprintf(userDNFilter, cn)
	}
	return fmt.Sprintf(userDNFilterWithObjectClass, m.Config.UserObjectClass, cn)
}

func (m *DefaultManager) Close() {
	if m.Connection != nil {
		m.Connection.Close()
	}
}
<commit_msg>adding better logging for result of isGroup ldap search<commit_after>package ldap

import (
	"encoding/hex"
	"fmt"
	"regexp"
	"strings"

	l "github.com/go-ldap/ldap"
	"github.com/pivotalservices/cf-mgmt/config"
	"github.com/xchapter7x/lo"
)

var (
	attributes = []string{"*"}
)

var (
	userRegexp          = regexp.MustCompile(",[A-Z]+=")
	escapeFilterRegex   = regexp.MustCompile(`([\\\(\)\*\0-\37\177-\377])`)
	unescapeFilterRegex = regexp.MustCompile(`\\([\da-fA-F]{2}|[()\\*])`) // only match \[)*\] or \xx x=a-fA-F
)

const (
	groupFilter                 = "(cn=%s)"
	userFilter                  = "(%s=%s)"
	userFilterWithObjectClass   = "(&(objectclass=%s)(%s=%s))"
	userDNFilter                = "(%s)"
	userDNFilterWithObjectClass = "(&(objectclass=%s)(%s))"
)

func NewManager(ldapConfig *config.LdapConfig) (Manager, error) {
	conn, err := CreateConnection(ldapConfig)
	if err != nil {
		return nil, err
	}
	return &DefaultManager{
		Config:     ldapConfig,
		Connection: conn,
	}, nil
}

func (m *DefaultManager) GetUserDNs(groupName string) ([]string, error) {
	filter := fmt.Sprintf(groupFilter, l.EscapeFilter(groupName))
	var groupEntry *l.Entry
	lo.G.Debug("Searching for group:", filter)
	lo.G.Debug("Using group search base:", m.Config.GroupSearchBase)

	search := l.NewSearchRequest(
		m.Config.GroupSearchBase,
		l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
		filter,
		attributes,
		nil)
	sr, err := m.Connection.Search(search)
	if err != nil {
		lo.G.Error(err)
		return nil, err
	}

	if len(sr.Entries) == 0 {
		lo.G.Errorf("group not found: %s", groupName)
		return []string{}, nil
	}
	if len(sr.Entries) > 1 {
		lo.G.Errorf("multiple groups found for: %s", groupName)
		return []string{}, nil
	}

	groupEntry = sr.Entries[0]
	userDNList := groupEntry.GetAttributeValues(m.Config.GroupAttribute)
	if len(userDNList) == 0 {
		lo.G.Warningf("No users found under group: %s", groupName)
	}

	userMap := make(map[string]string)
	for _, userDN := range userDNList {
		group, cn, err := m.IsGroup(userDN)
		if err != nil {
			return nil, err
		}
		if group {
			if err != nil {
				return nil, err
			}
			nestedUsers, err := m.GetUserDNs(cn)
			if err != nil {
				return nil, err
			}
			for _, nestedUser := range nestedUsers {
				userMap[nestedUser] = nestedUser
			}
		} else {
			userMap[userDN] = userDN
		}
	}
	var userList []string
	for _, userDN := range userMap {
		userList = append(userList, userDN)
	}
	return userList, nil
}

func (m *DefaultManager) getCN(userDN string) (string, error) {
	indexes := userRegexp.FindStringIndex(strings.ToUpper(userDN))
	if len(indexes) == 0 {
		return "", fmt.Errorf("cannot find CN for DN: %s", userDN)
	}
	cn := strings.Replace(userDN[:indexes[0]], "cn=", "", 1)
	cnTemp := UnescapeFilterValue(cn)
	lo.G.Debug("CN unescaped:", cnTemp)

	escapedCN := l.EscapeFilter(strings.Replace(cnTemp, "\\", "", -1))
	lo.G.Debug("CN escaped:", escapedCN)
	return escapedCN, nil
}
func (m *DefaultManager) IsGroup(userDN string) (bool, string, error) {
	if strings.Contains(userDN, m.Config.GroupSearchBase) {
		cn, err := m.getCN(userDN)
		if err != nil {
			return false, "", err
		}
		filter := fmt.Sprintf(groupFilter, cn)
		search := l.NewSearchRequest(
			m.Config.GroupSearchBase,
			l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
			filter,
			attributes,
			nil)
		sr, err := m.Connection.Search(search)
		if err != nil {
			return false, "", err
		}
		lo.G.Debugf("Found %d entries for group filter", len(sr.Entries), filter)
		return len(sr.Entries) == 1, cn, nil
	} else {
		return false, "", nil
	}
}

func (m *DefaultManager) GetUserByDN(userDN string) (*User, error) {
	lo.G.Debug("User DN:", userDN)
	indexes := userRegexp.FindStringIndex(strings.ToUpper(userDN))
	if len(indexes) == 0 {
		return nil, fmt.Errorf("cannot find CN for user DN: %s", userDN)
	}
	index := indexes[0]
	userCNTemp := UnescapeFilterValue(userDN[:index])
	lo.G.Debug("CN unescaped:", userCNTemp)

	userCN := l.EscapeFilter(strings.Replace(userCNTemp, "\\", "", 1))
	lo.G.Debug("CN escaped:", userCN)
	filter := m.getUserFilterWithCN(userCN)
	return m.searchUser(filter, userDN[index+1:], "")
}

func (m *DefaultManager) GetUserByID(userID string) (*User, error) {
	filter := m.getUserFilter(userID)
	lo.G.Debug("Searching for user:", filter)
	lo.G.Debug("Using user search base:", m.Config.UserSearchBase)
	return m.searchUser(filter, m.Config.UserSearchBase, userID)
}

func (m *DefaultManager) searchUser(filter, searchBase, userID string) (*User, error) {
	lo.G.Debugf("Searching with filter %s", filter)
	search := l.NewSearchRequest(
		searchBase,
		l.ScopeWholeSubtree, l.NeverDerefAliases, 0, 0, false,
		filter,
		attributes,
		nil)

	sr, err := m.Connection.Search(search)
	if err != nil {
		lo.G.Error(err)
		return nil, err
	}

	if (len(sr.Entries)) == 1 {
		entry := sr.Entries[0]
		user := &User{
			UserDN: entry.DN,
			Email:  entry.GetAttributeValue(m.Config.UserMailAttribute),
		}
		if userID != "" {
			user.UserID = userID
		} else {
			user.UserID = entry.GetAttributeValue(m.Config.UserNameAttribute)
		}
		lo.G.Debugf("Search filter %s returned userDN [%s], email [%s], userID [%s]", filter, user.UserDN, user.Email, user.UserID)
		return user, nil
	}
	lo.G.Errorf("Found %d number of entries for filter %s", len(sr.Entries), filter)
	return nil, nil
}

func UnescapeFilterValue(filter string) string {
	repl := unescapeFilterRegex.ReplaceAllFunc(
		[]byte(filter),
		func(match []byte) []byte {
			// \( \) \\ \*
			if len(match) == 2 {
				return []byte{match[1]}
			}
			// had issues with Decode, TODO fix to use Decode?.
			res, _ := hex.DecodeString(string(match[1:]))
			return res
		},
	)
	return string(repl)
}

func (m *DefaultManager) getUserFilter(userID string) string {
	if m.Config.UserObjectClass == "" {
		return fmt.Sprintf(userFilter, m.Config.UserNameAttribute, userID)
	}
	return fmt.Sprintf(userFilterWithObjectClass, m.Config.UserObjectClass, m.Config.UserNameAttribute, userID)
}

func (m *DefaultManager) getUserFilterWithCN(cn string) string {
	if m.Config.UserObjectClass == "" {
		return fmt.Sprintf(userDNFilter, cn)
	}
	return fmt.Sprintf(userDNFilterWithObjectClass, m.Config.UserObjectClass, cn)
}

func (m *DefaultManager) Close() {
	if m.Connection != nil {
		m.Connection.Close()
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 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 resolve

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/data/id"
	"github.com/google/gapid/gapis/capture"
	"github.com/google/gapid/gapis/database"
	"github.com/google/gapid/gapis/gfxapi"
	"github.com/google/gapid/gapis/service/path"
)

// ResolvedResources contains all of the resolved resources for a
// particular point in the trace.
type ResolvedResources struct {
	resourceMap  gfxapi.ResourceMap
	resources    map[id.ID]gfxapi.Resource
	resourceData map[id.ID]interface{}
}

// Resolve builds a ResolvedResources object for all of the resources
// at the path r.After
func (r *AllResourceDataResolvable) Resolve(ctx context.Context) (interface{}, error) {
	ctx = capture.Put(ctx, r.After.Commands.Capture)
	resources, err := buildResources(ctx, r.After)

	if err != nil {
		return nil, err
	}
	return resources, nil
}

func buildResources(ctx context.Context, p *path.Command) (*ResolvedResources, error) {
	list, err := NCommands(ctx, p.Commands, p.Index+1)
	if err != nil {
		return nil, err
	}
	state := capture.NewState(ctx)
	var currentAtomIndex uint64
	var currentAtomResourceCount int
	idMap := gfxapi.ResourceMap{}

	resources := make(map[id.ID]gfxapi.Resource)

	state.OnResourceCreated = func(r gfxapi.Resource) {
		currentAtomResourceCount++
		i := genResourceID(currentAtomIndex, currentAtomResourceCount)
		idMap[r] = i
		resources[i] = r
	}

	for i, a := range list.Atoms[:p.Index+1] {
		currentAtomResourceCount = 0
		currentAtomIndex = uint64(i)
		a.Mutate(ctx, state, nil /* no builder, just mutate */)
	}

	resourceData := make(map[id.ID]interface{})
	for k, v := range resources {
		res, err := v.ResourceData(ctx, state)
		if err != nil {
			return nil, fmt.Errorf("Could not get data for resource %v", k)
		}
		resourceData[k] = res
	}
	return &ResolvedResources{idMap, resources, resourceData}, nil
}

// ResourceData resolves the data of the specified resource at the specified
// point in the capture.
func ResourceData(ctx context.Context, p *path.ResourceData) (interface{}, error) {
	obj, err := database.Build(ctx, &ResourceDataResolvable{p})
	if err != nil {
		return nil, err
	}
	return obj, nil
}

// Resolve implements the database.Resolver interface.
func (r *ResourceDataResolvable) Resolve(ctx context.Context) (interface{}, error) {
	resources, err := database.Build(ctx, &AllResourceDataResolvable{r.Path.After})
	if err != nil {
		return nil, err
	}
	res, ok := resources.(*ResolvedResources)
	if !ok {
		return nil, fmt.Errorf("Cannot resolve resources at command: %v", r.Path.After)
	}
	id := r.Path.Id.ID()
	if val, ok := res.resourceData[id]; ok {
		return val, nil
	}

	return nil, fmt.Errorf("Cannot find resource with id: %v", id)
}
<commit_msg>Fix resource data bug introduced in commit f7d41ba.<commit_after>// Copyright (C) 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 resolve

import (
	"context"
	"fmt"

	"github.com/google/gapid/core/data/id"
	"github.com/google/gapid/gapis/capture"
	"github.com/google/gapid/gapis/database"
	"github.com/google/gapid/gapis/gfxapi"
	"github.com/google/gapid/gapis/service/path"
)

// ResolvedResources contains all of the resolved resources for a
// particular point in the trace.
type ResolvedResources struct {
	resourceMap  gfxapi.ResourceMap
	resources    map[id.ID]gfxapi.Resource
	resourceData map[id.ID]interface{}
}

// Resolve builds a ResolvedResources object for all of the resources
// at the path r.After
func (r *AllResourceDataResolvable) Resolve(ctx context.Context) (interface{}, error) {
	ctx = capture.Put(ctx, r.After.Commands.Capture)
	resources, err := buildResources(ctx, r.After)

	if err != nil {
		return nil, err
	}
	return resources, nil
}

func buildResources(ctx context.Context, p *path.Command) (*ResolvedResources, error) {
	list, err := NCommands(ctx, p.Commands, p.Index+1)
	if err != nil {
		return nil, err
	}
	state := capture.NewState(ctx)
	var currentAtomIndex uint64
	var currentAtomResourceCount int
	idMap := gfxapi.ResourceMap{}

	resources := make(map[id.ID]gfxapi.Resource)

	state.OnResourceCreated = func(r gfxapi.Resource) {
		currentAtomResourceCount++
		i := genResourceID(currentAtomIndex, currentAtomResourceCount)
		idMap[r] = i
		resources[i] = r
	}

	for i, a := range list.Atoms[:p.Index+1] {
		currentAtomResourceCount = 0
		currentAtomIndex = uint64(i)
		a.Mutate(ctx, state, nil /* no builder, just mutate */)
	}

	resourceData := make(map[id.ID]interface{})
	for k, v := range resources {
		res, err := v.ResourceData(ctx, state)
		if err != nil {
			resourceData[k] = err;
		} else {
			resourceData[k] = res
		}
	}
	return &ResolvedResources{idMap, resources, resourceData}, nil
}

// ResourceData resolves the data of the specified resource at the specified
// point in the capture.
func ResourceData(ctx context.Context, p *path.ResourceData) (interface{}, error) {
	obj, err := database.Build(ctx, &ResourceDataResolvable{p})
	if err != nil {
		return nil, err
	}
	return obj, nil
}

// Resolve implements the database.Resolver interface.
func (r *ResourceDataResolvable) Resolve(ctx context.Context) (interface{}, error) {
	resources, err := database.Build(ctx, &AllResourceDataResolvable{r.Path.After})
	if err != nil {
		return nil, err
	}
	res, ok := resources.(*ResolvedResources)
	if !ok {
		return nil, fmt.Errorf("Cannot resolve resources at command: %v", r.Path.After)
	}
	id := r.Path.Id.ID()
	if val, ok := res.resourceData[id]; ok {
		if err, isErr := val.(error); isErr {
			return nil, err
		}
		return val, nil
	}

	return nil, fmt.Errorf("Cannot find resource with id: %v", id)
}
<|endoftext|>


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

package gcsproxy_test

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"testing"

	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease"
	"github.com/googlecloudplatform/gcsfuse/mutable"
	"github.com/googlecloudplatform/gcsfuse/mutable/mock"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

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

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type SyncTest struct {
	ctx context.Context

	srcObject gcs.Object
	content   mock_mutable.MockContent
	bucket    mock_gcs.MockBucket

	simulatedContents []byte
}

var _ SetUpInterface = &SyncTest{}

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

func (t *SyncTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Set up the source object.
	t.srcObject.Generation = 1234
	t.srcObject.Name = "foo"
	t.srcObject.Size = 17

	// Set up dependencies.
	t.content = mock_mutable.NewMockContent(
		ti.MockController,
		"content")

	t.bucket = mock_gcs.NewMockBucket(
		ti.MockController,
		"bucket")

	// By default, show the content as dirty.
	sr := mutable.StatResult{
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillRepeatedly(Return(sr, nil))

	// Set up fake contents.
	t.simulatedContents = []byte("taco")
	ExpectCall(t.content, "ReadAt")(Any(), Any(), Any()).
		WillRepeatedly(Invoke(t.serveReadAt))

	// And for the released read/write lease.
	leaser := lease.NewFileLeaser("", math.MaxInt32, math.MaxInt32)
	rwl, err := leaser.NewFile()
	AssertEq(nil, err)

	_, err = rwl.Write(t.simulatedContents)
	AssertEq(nil, err)

	ExpectCall(t.content, "Release")().
		WillRepeatedly(Return(rwl))
}

func (t *SyncTest) call() (rl lease.ReadLease, o *gcs.Object, err error) {
	rl, o, err = gcsproxy.Sync(
		t.ctx,
		&t.srcObject,
		t.content,
		t.bucket)

	return
}

func (t *SyncTest) serveReadAt(
	ctx context.Context,
	p []byte,
	offset int64) (n int, err error) {
	// Handle out of range reads.
	if offset > int64(len(t.simulatedContents)) {
		err = io.EOF
		return
	}

	// Copy into the buffer.
	n = copy(p, t.simulatedContents[int(offset):])
	if n < len(p) {
		err = io.EOF
		return
	}

	return
}

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

func (t *SyncTest) StatFails() {
	// Stat
	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(mutable.StatResult{}, errors.New("taco")))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("Stat")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *SyncTest) StatReturnsWackyDirtyThreshold() {
	// Stat
	sr := mutable.StatResult{
		DirtyThreshold: int64(t.srcObject.Size + 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("Stat")))
	ExpectThat(err, Error(HasSubstr("DirtyThreshold")))
	ExpectThat(err, Error(HasSubstr(fmt.Sprint(t.srcObject.Size))))
	ExpectThat(err, Error(HasSubstr(fmt.Sprint(t.srcObject.Size+1))))
}

func (t *SyncTest) StatSaysNotDirty() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size),
		DirtyThreshold: int64(t.srcObject.Size),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// Call
	rl, o, err := t.call()

	AssertEq(nil, err)
	ExpectEq(nil, rl)
	ExpectEq(nil, o)
}

func (t *SyncTest) StatSaysDirty_SameSizeAsSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size),
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *SyncTest) StatSaysDirty_SmallerThanSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size - 1),
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *SyncTest) StatSaysDirty_LargerThanSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size + 1),
		DirtyThreshold: int64(t.srcObject.Size),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *SyncTest) CallsBucket() {
	// CreateObject
	var req *gcs.CreateObjectRequest
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(DoAll(SaveArg(1, &req), Return(nil, errors.New(""))))

	// Call
	t.call()

	AssertNe(nil, req)
	ExpectEq(t.srcObject.Name, req.Name)
	ExpectThat(req.GenerationPrecondition, Pointee(Equals(t.srcObject.Generation)))

	b, err := ioutil.ReadAll(req.Contents)
	AssertEq(nil, err)
	ExpectEq(string(t.simulatedContents), string(b))
}

func (t *SyncTest) BucketFails() {
	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("CreateObject")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *SyncTest) BucketReturnsPreconditionError() {
	// CreateObject
	expected := &gcs.PreconditionError{}
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, expected))

	// Call
	_, _, err := t.call()

	ExpectEq(expected, err)
}

func (t *SyncTest) BucketSucceeds() {
	// CreateObject
	expected := &gcs.Object{}
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(expected, nil))

	// Call
	rl, o, err := t.call()

	AssertEq(nil, err)
	ExpectEq(expected, o)

	_, err = rl.Seek(0, 0)
	AssertEq(nil, err)

	b, err := ioutil.ReadAll(rl)
	AssertEq(nil, err)
	ExpectEq(string(t.simulatedContents), string(b))
}
<commit_msg>Renamed the test.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcsproxy_test

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"testing"

	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease"
	"github.com/googlecloudplatform/gcsfuse/mutable"
	"github.com/googlecloudplatform/gcsfuse/mutable/mock"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

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

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type ObjectSyncerTest struct {
	ctx context.Context

	srcObject gcs.Object
	content   mock_mutable.MockContent
	bucket    mock_gcs.MockBucket

	simulatedContents []byte
}

var _ SetUpInterface = &ObjectSyncerTest{}

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

func (t *ObjectSyncerTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Set up the source object.
	t.srcObject.Generation = 1234
	t.srcObject.Name = "foo"
	t.srcObject.Size = 17

	// Set up dependencies.
	t.content = mock_mutable.NewMockContent(
		ti.MockController,
		"content")

	t.bucket = mock_gcs.NewMockBucket(
		ti.MockController,
		"bucket")

	// By default, show the content as dirty.
	sr := mutable.StatResult{
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillRepeatedly(Return(sr, nil))

	// Set up fake contents.
	t.simulatedContents = []byte("taco")
	ExpectCall(t.content, "ReadAt")(Any(), Any(), Any()).
		WillRepeatedly(Invoke(t.serveReadAt))

	// And for the released read/write lease.
	leaser := lease.NewFileLeaser("", math.MaxInt32, math.MaxInt32)
	rwl, err := leaser.NewFile()
	AssertEq(nil, err)

	_, err = rwl.Write(t.simulatedContents)
	AssertEq(nil, err)

	ExpectCall(t.content, "Release")().
		WillRepeatedly(Return(rwl))
}

func (t *ObjectSyncerTest) call() (rl lease.ReadLease, o *gcs.Object, err error) {
	rl, o, err = gcsproxy.Sync(
		t.ctx,
		&t.srcObject,
		t.content,
		t.bucket)

	return
}

func (t *ObjectSyncerTest) serveReadAt(
	ctx context.Context,
	p []byte,
	offset int64) (n int, err error) {
	// Handle out of range reads.
	if offset > int64(len(t.simulatedContents)) {
		err = io.EOF
		return
	}

	// Copy into the buffer.
	n = copy(p, t.simulatedContents[int(offset):])
	if n < len(p) {
		err = io.EOF
		return
	}

	return
}

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

func (t *ObjectSyncerTest) StatFails() {
	// Stat
	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(mutable.StatResult{}, errors.New("taco")))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("Stat")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *ObjectSyncerTest) StatReturnsWackyDirtyThreshold() {
	// Stat
	sr := mutable.StatResult{
		DirtyThreshold: int64(t.srcObject.Size + 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("Stat")))
	ExpectThat(err, Error(HasSubstr("DirtyThreshold")))
	ExpectThat(err, Error(HasSubstr(fmt.Sprint(t.srcObject.Size))))
	ExpectThat(err, Error(HasSubstr(fmt.Sprint(t.srcObject.Size+1))))
}

func (t *ObjectSyncerTest) StatSaysNotDirty() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size),
		DirtyThreshold: int64(t.srcObject.Size),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// Call
	rl, o, err := t.call()

	AssertEq(nil, err)
	ExpectEq(nil, rl)
	ExpectEq(nil, o)
}

func (t *ObjectSyncerTest) StatSaysDirty_SameSizeAsSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size),
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *ObjectSyncerTest) StatSaysDirty_SmallerThanSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size - 1),
		DirtyThreshold: int64(t.srcObject.Size - 1),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *ObjectSyncerTest) StatSaysDirty_LargerThanSource() {
	// Stat
	sr := mutable.StatResult{
		Size:           int64(t.srcObject.Size + 1),
		DirtyThreshold: int64(t.srcObject.Size),
	}

	ExpectCall(t.content, "Stat")(Any()).
		WillOnce(Return(sr, nil))

	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.call()
}

func (t *ObjectSyncerTest) CallsBucket() {
	// CreateObject
	var req *gcs.CreateObjectRequest
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(DoAll(SaveArg(1, &req), Return(nil, errors.New(""))))

	// Call
	t.call()

	AssertNe(nil, req)
	ExpectEq(t.srcObject.Name, req.Name)
	ExpectThat(req.GenerationPrecondition, Pointee(Equals(t.srcObject.Generation)))

	b, err := ioutil.ReadAll(req.Contents)
	AssertEq(nil, err)
	ExpectEq(string(t.simulatedContents), string(b))
}

func (t *ObjectSyncerTest) BucketFails() {
	// CreateObject
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, _, err := t.call()

	ExpectThat(err, Error(HasSubstr("CreateObject")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *ObjectSyncerTest) BucketReturnsPreconditionError() {
	// CreateObject
	expected := &gcs.PreconditionError{}
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(nil, expected))

	// Call
	_, _, err := t.call()

	ExpectEq(expected, err)
}

func (t *ObjectSyncerTest) BucketSucceeds() {
	// CreateObject
	expected := &gcs.Object{}
	ExpectCall(t.bucket, "CreateObject")(Any(), Any()).
		WillOnce(Return(expected, nil))

	// Call
	rl, o, err := t.call()

	AssertEq(nil, err)
	ExpectEq(expected, o)

	_, err = rl.Seek(0, 0)
	AssertEq(nil, err)

	b, err := ioutil.ReadAll(rl)
	AssertEq(nil, err)
	ExpectEq(string(t.simulatedContents), string(b))
}
<|endoftext|>


<commit_before>package framework

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"strconv"

	"github.com/go-distributed/meritop"
)

/*
The dummy task is designed for regresion test of meritop framework.
This works with tree topology.
The main idea behind the regression test is following:
There will be two kinds of dummyTasks: master and slaves. We will have one master
sits at the top with taskID = 0, and then rest 6 (2^n - 2) tasks forms a tree under
the master. There will be 10 epochs, from 1 to 10, at each epoch, we send out a
vector with all values equal to epochID, and each slave is supposedly return a vector
with all values equals epochID*taskID, the values are reduced back to master, and
master will print out the epochID and aggregated vector. After all 10 epoch, it kills
job.
*/

const (
	NumOfIterations uint64 = uint64(10)
)

// dummyData is used to carry parameter and gradient;
type dummyData struct {
	Value int32
}

// dummyMaster is prototype of parameter server, for now it does not
// carry out optimization yet. But it should be easy to add support when
// this full tests out.
// Note: in theory, since there should be no parent of this, so we should
// add error checing in the right places. We will skip these test for now.
type dummyMaster struct {
	dataChan      chan int32
	finishChan    chan struct{}
	taskStopChan  chan bool
	framework     meritop.Framework
	epoch, taskID uint64
	logger        *log.Logger
	config        map[string]string

	param, gradient *dummyData
	fromChildren    map[uint64]*dummyData
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummyMaster) Init(taskID uint64, framework meritop.Framework) {
	t.taskID = taskID
	t.framework = framework
	// t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	t.logger = log.New(ioutil.Discard, "", log.Ldate|log.Ltime|log.Lshortfile)
}

// Task need to finish up for exit, last chance to save work?
func (t *dummyMaster) Exit() {}

// Ideally, we should also have the following:
func (t *dummyMaster) ParentMetaReady(parentID uint64, meta string) {}
func (t *dummyMaster) ChildMetaReady(childID uint64, meta string) {
	t.logger.Printf("master ChildMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	// Get data from child. When all the data is back, starts the next epoch.
	t.framework.DataRequest(childID, meta)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) SetEpoch(epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.config != nil &&
		t.config["failmaster"] == "yes" &&
		t.config["failepoch"] == strconv.FormatUint(epoch, 10) {
		t.framework.(*framework).stop()
		t.taskStopChan <- true
		return
	}
	t.param = &dummyData{}
	t.gradient = &dummyData{}

	t.epoch = epoch
	t.param.Value = int32(t.epoch)

	// Make sure we have a clean slate.
	t.fromChildren = make(map[uint64]*dummyData)
	t.framework.FlagMetaToChild("ParamReady")
}

// These are payload rpc for application purpose.
func (t *dummyMaster) ServeAsParent(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.param)
	if err != nil {
		t.logger.Fatalf("Master can't encode parameter: %v, error: %v\n", t.param, err)
	}
	return b
}

func (t *dummyMaster) ServeAsChild(fromID uint64, req string) []byte {
	return nil
}

func (t *dummyMaster) ParentDataReady(parentID uint64, req string, resp []byte) {}
func (t *dummyMaster) ChildDataReady(childID uint64, req string, resp []byte) {
	t.logger.Printf("master ChildDataReady, task: %d, epoch: %d, child: %d, ready: %d\n",
		t.taskID, t.epoch, childID, len(t.fromChildren))
	d := new(dummyData)
	json.Unmarshal(resp, d)
	t.fromChildren[childID] = d

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		t.dataChan <- t.gradient.Value
		// TODO(xiaoyunwu) We need to do some test here.

		// In real ML, we modify the gradient first. But here it is noop.
		// Notice that we only
		if t.epoch == NumOfIterations {
			t.framework.ShutdownJob()
			close(t.finishChan)
		} else {
			t.framework.IncEpoch()
		}
	}
}

// dummySlave is an prototype for data shard in machine learning applications.
// It mainly does to things, pass on parameters to its children, and collect
// gradient back then add them together before make it available to its parent.
type dummySlave struct {
	framework     meritop.Framework
	epoch, taskID uint64
	logger        *log.Logger

	param, gradient *dummyData
	fromChildren    map[uint64]*dummyData
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummySlave) Init(taskID uint64, framework meritop.Framework) {
	t.taskID = taskID
	t.framework = framework
	// t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	t.logger = log.New(ioutil.Discard, "", log.Ldate|log.Ltime|log.Lshortfile)
}

// Task need to finish up for exit, last chance to save work?
func (t *dummySlave) Exit() {}

// Ideally, we should also have the following:
func (t *dummySlave) ParentMetaReady(parentID uint64, meta string) {
	t.logger.Printf("slave ParentMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	t.framework.DataRequest(parentID, meta)
}

func (t *dummySlave) ChildMetaReady(childID uint64, meta string) {
	t.logger.Printf("slave ChildMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	t.framework.DataRequest(childID, meta)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummySlave) SetEpoch(epoch uint64) {
	t.logger.Printf("slave SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	t.param = &dummyData{}
	t.gradient = &dummyData{}

	t.epoch = epoch
	// Make sure we have a clean slate.
	t.fromChildren = make(map[uint64]*dummyData)
}

// These are payload rpc for application purpose.
func (t *dummySlave) ServeAsParent(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.param)
	if err != nil {
		t.logger.Fatalf("Slave can't encode parameter: %v, error: %v\n", t.param, err)
	}
	return b
}

func (t *dummySlave) ServeAsChild(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.gradient)
	if err != nil {
		t.logger.Fatalf("Slave can't encode gradient: %v, error: %v\n", t.gradient, err)
	}
	return b
}

func (t *dummySlave) ParentDataReady(parentID uint64, req string, resp []byte) {
	t.logger.Printf("slave ParentDataReady, task: %d, epoch: %d, parent: %d\n", t.taskID, t.epoch, parentID)
	t.param = new(dummyData)
	json.Unmarshal(resp, t.param)

	// We need to carry out local compuation.
	t.gradient.Value = t.param.Value * int32(t.framework.GetTaskID())

	// If this task has children, flag meta so that children can start pull
	// parameter.
	children := t.framework.GetTopology().GetChildren(t.epoch)
	if len(children) != 0 {
		t.framework.FlagMetaToChild("ParamReady")
	} else {
		// On leaf node, we can immediately return by and flag parent
		// that this node is ready.
		t.framework.FlagMetaToParent("GradientReady")
	}
}

func (t *dummySlave) ChildDataReady(childID uint64, req string, resp []byte) {
	t.fromChildren[childID] = new(dummyData)
	json.Unmarshal(resp, t.fromChildren[childID])

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		// In real ML, we add the gradient first.
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		t.framework.FlagMetaToParent("GradientReady")
	}
}

type SimpleTaskBuilder struct {
	GDataChan    chan int32
	FinishChan   chan struct{}
	TaskStopChan chan bool
	Config       map[string]string
}

// Leave it at global level so that we use this to terminate and test.
// cDataChan := make(chan *tDataBundle, 1)

// This method is called once by framework implementation to get the
// right task implementation for the node/task. It requires the taskID
// for current node, and also a global array of tasks.
func (tc SimpleTaskBuilder) GetTask(taskID uint64) meritop.Task {
	if taskID == 0 {
		return &dummyMaster{
			dataChan:     tc.GDataChan,
			finishChan:   tc.FinishChan,
			taskStopChan: tc.TaskStopChan,
			config:       tc.Config,
		}
	}
	return &dummySlave{}
}
<commit_msg>open up regression framework testing logging<commit_after>package framework

import (
	"encoding/json"
	"log"
	"os"
	"strconv"

	"github.com/go-distributed/meritop"
)

/*
The dummy task is designed for regresion test of meritop framework.
This works with tree topology.
The main idea behind the regression test is following:
There will be two kinds of dummyTasks: master and slaves. We will have one master
sits at the top with taskID = 0, and then rest 6 (2^n - 2) tasks forms a tree under
the master. There will be 10 epochs, from 1 to 10, at each epoch, we send out a
vector with all values equal to epochID, and each slave is supposedly return a vector
with all values equals epochID*taskID, the values are reduced back to master, and
master will print out the epochID and aggregated vector. After all 10 epoch, it kills
job.
*/

const (
	NumOfIterations uint64 = uint64(10)
)

// dummyData is used to carry parameter and gradient;
type dummyData struct {
	Value int32
}

// dummyMaster is prototype of parameter server, for now it does not
// carry out optimization yet. But it should be easy to add support when
// this full tests out.
// Note: in theory, since there should be no parent of this, so we should
// add error checing in the right places. We will skip these test for now.
type dummyMaster struct {
	dataChan      chan int32
	finishChan    chan struct{}
	taskStopChan  chan bool
	framework     meritop.Framework
	epoch, taskID uint64
	logger        *log.Logger
	config        map[string]string

	param, gradient *dummyData
	fromChildren    map[uint64]*dummyData
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummyMaster) Init(taskID uint64, framework meritop.Framework) {
	t.taskID = taskID
	t.framework = framework
	t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	// t.logger = log.New(ioutil.Discard, "", log.Ldate|log.Ltime|log.Lshortfile)
}

// Task need to finish up for exit, last chance to save work?
func (t *dummyMaster) Exit() {}

// Ideally, we should also have the following:
func (t *dummyMaster) ParentMetaReady(parentID uint64, meta string) {}
func (t *dummyMaster) ChildMetaReady(childID uint64, meta string) {
	t.logger.Printf("master ChildMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	// Get data from child. When all the data is back, starts the next epoch.
	t.framework.DataRequest(childID, meta)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) SetEpoch(epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.config != nil &&
		t.config["failmaster"] == "yes" &&
		t.config["failepoch"] == strconv.FormatUint(epoch, 10) {
		t.logger.Printf("task %d is doomed to fail at epoch %d\n", t.taskID, epoch)
		t.framework.(*framework).stop()
		t.taskStopChan <- true
		return
	}
	t.param = &dummyData{}
	t.gradient = &dummyData{}

	t.epoch = epoch
	t.param.Value = int32(t.epoch)

	// Make sure we have a clean slate.
	t.fromChildren = make(map[uint64]*dummyData)
	t.framework.FlagMetaToChild("ParamReady")
}

// These are payload rpc for application purpose.
func (t *dummyMaster) ServeAsParent(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.param)
	if err != nil {
		t.logger.Fatalf("Master can't encode parameter: %v, error: %v\n", t.param, err)
	}
	return b
}

func (t *dummyMaster) ServeAsChild(fromID uint64, req string) []byte {
	return nil
}

func (t *dummyMaster) ParentDataReady(parentID uint64, req string, resp []byte) {}
func (t *dummyMaster) ChildDataReady(childID uint64, req string, resp []byte) {
	t.logger.Printf("master ChildDataReady, task: %d, epoch: %d, child: %d, ready: %d\n",
		t.taskID, t.epoch, childID, len(t.fromChildren))
	d := new(dummyData)
	json.Unmarshal(resp, d)
	t.fromChildren[childID] = d

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		t.dataChan <- t.gradient.Value
		// TODO(xiaoyunwu) We need to do some test here.

		// In real ML, we modify the gradient first. But here it is noop.
		// Notice that we only
		if t.epoch == NumOfIterations {
			t.framework.ShutdownJob()
			close(t.finishChan)
		} else {
			t.framework.IncEpoch()
		}
	}
}

// dummySlave is an prototype for data shard in machine learning applications.
// It mainly does to things, pass on parameters to its children, and collect
// gradient back then add them together before make it available to its parent.
type dummySlave struct {
	framework     meritop.Framework
	epoch, taskID uint64
	logger        *log.Logger

	param, gradient *dummyData
	fromChildren    map[uint64]*dummyData
}

// This is useful to bring the task up to speed from scratch or if it recovers.
func (t *dummySlave) Init(taskID uint64, framework meritop.Framework) {
	t.taskID = taskID
	t.framework = framework
	t.logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	// t.logger = log.New(ioutil.Discard, "", log.Ldate|log.Ltime|log.Lshortfile)
}

// Task need to finish up for exit, last chance to save work?
func (t *dummySlave) Exit() {}

// Ideally, we should also have the following:
func (t *dummySlave) ParentMetaReady(parentID uint64, meta string) {
	t.logger.Printf("slave ParentMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	t.framework.DataRequest(parentID, meta)
}

func (t *dummySlave) ChildMetaReady(childID uint64, meta string) {
	t.logger.Printf("slave ChildMetaReady, task: %d, epoch: %d\n", t.taskID, t.epoch)
	t.framework.DataRequest(childID, meta)
}

// This give the task an opportunity to cleanup and regroup.
func (t *dummySlave) SetEpoch(epoch uint64) {
	t.logger.Printf("slave SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	t.param = &dummyData{}
	t.gradient = &dummyData{}

	t.epoch = epoch
	// Make sure we have a clean slate.
	t.fromChildren = make(map[uint64]*dummyData)
}

// These are payload rpc for application purpose.
func (t *dummySlave) ServeAsParent(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.param)
	if err != nil {
		t.logger.Fatalf("Slave can't encode parameter: %v, error: %v\n", t.param, err)
	}
	return b
}

func (t *dummySlave) ServeAsChild(fromID uint64, req string) []byte {
	b, err := json.Marshal(t.gradient)
	if err != nil {
		t.logger.Fatalf("Slave can't encode gradient: %v, error: %v\n", t.gradient, err)
	}
	return b
}

func (t *dummySlave) ParentDataReady(parentID uint64, req string, resp []byte) {
	t.logger.Printf("slave ParentDataReady, task: %d, epoch: %d, parent: %d\n", t.taskID, t.epoch, parentID)
	t.param = new(dummyData)
	json.Unmarshal(resp, t.param)

	// We need to carry out local compuation.
	t.gradient.Value = t.param.Value * int32(t.framework.GetTaskID())

	// If this task has children, flag meta so that children can start pull
	// parameter.
	children := t.framework.GetTopology().GetChildren(t.epoch)
	if len(children) != 0 {
		t.framework.FlagMetaToChild("ParamReady")
	} else {
		// On leaf node, we can immediately return by and flag parent
		// that this node is ready.
		t.framework.FlagMetaToParent("GradientReady")
	}
}

func (t *dummySlave) ChildDataReady(childID uint64, req string, resp []byte) {
	t.fromChildren[childID] = new(dummyData)
	json.Unmarshal(resp, t.fromChildren[childID])

	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		// In real ML, we add the gradient first.
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		t.framework.FlagMetaToParent("GradientReady")
	}
}

type SimpleTaskBuilder struct {
	GDataChan    chan int32
	FinishChan   chan struct{}
	TaskStopChan chan bool
	Config       map[string]string
}

// Leave it at global level so that we use this to terminate and test.
// cDataChan := make(chan *tDataBundle, 1)

// This method is called once by framework implementation to get the
// right task implementation for the node/task. It requires the taskID
// for current node, and also a global array of tasks.
func (tc SimpleTaskBuilder) GetTask(taskID uint64) meritop.Task {
	if taskID == 0 {
		return &dummyMaster{
			dataChan:     tc.GDataChan,
			finishChan:   tc.FinishChan,
			taskStopChan: tc.TaskStopChan,
			config:       tc.Config,
		}
	}
	return &dummySlave{}
}
<|endoftext|>


<commit_before>package effio

// encoding/csv doesn't strip whitespace and does a fair bit of
// work to handle strings & quoting which are totally unnecessary
// overhead for these files so skip it

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"
)

// the input is ints but plotinum uses float64 so might as well
// start there and avoid the type conversions later
type LatRec struct {
	Time float64 `json:"x"` // time offset from beginning of fio run
	Val  float64 `json:"y"` // latency value
	Ddir uint8   `json:"-"` // 0 = read, 1 = write, 2 = trim
	Bsz  uint16  `json:"-"` // block size
}
type LatRecs []*LatRec

// Loads the CSV output by fio into an LatRecs array of LatRec structs.
func LoadCSV(filename string) LatRecs {
	fmt.Printf("Parsing file: '%s' ... ", filename)

	fd, err := os.Open(filename)
	if err != nil {
		fmt.Printf(" Failed.\nCould not open file '%s' for read: %s\n", filename, err)
		return LatRecs{}
	}
	defer fd.Close()

	started := time.Now()
	records := make(LatRecs, 0)

	var tm, perf float64
	var ddir, bsz int
	bfd := bufio.NewReader(fd)
	var lno int = 0
	for {
		line, _, err := bfd.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("\nRead from file '%s' failed: %s", filename, err)
		}
		lno++

		if lno%10000 == 0 {
			fmt.Printf(".")
		}

		// fio always uses ", " instead of "," as far as I can tell
		r := strings.SplitN(string(line), ", ", 4)
		// probably an impartial record at the end of the file
		if len(r) < 4 || r[0] == "" || r[1] == "" {
			continue
		}

		tm, err = strconv.ParseFloat(r[0], 64)
		if err != nil {
			log.Fatalf("\nParsing field 0 failed in file '%s' at line %d: %s", filename, lno, err)
		}
		perf, err = strconv.ParseFloat(r[1], 64)
		if err != nil {
			log.Fatalf("\nParsing field 1 in file '%s' at line %d: %s", filename, lno, err)
		}
		ddir, err = strconv.Atoi(r[2])
		if err != nil {
			log.Fatalf("\nParsing field 2 failed in file '%s' at line %d: %s", filename, lno, err)
		}
		bsz, err = strconv.Atoi(r[3])
		if err != nil {
			log.Fatalf("\nParsing field 3 failed in file '%s' at line %d: %s", filename, lno, err)
		}

		lr := LatRec{tm, perf, uint8(ddir), uint16(bsz)}
		records = append(records, &lr)
	}

	done := time.Now()
	fmt.Printf(" Done.\nRows: %d Elapsed: %s\n", len(records), done.Sub(started).String())

	return records
}

// implement some plotinum interfaces
func (lrs LatRecs) Len() int {
	return len(lrs)
}

func (lrs LatRecs) XY(i int) (float64, float64) {
	return lrs[i].Time, lrs[i].Val
}

func (lrs LatRecs) Value(i int) float64 {
	return lrs[i].Val
}

func (lrs LatRecs) Values(i int) (vals []float64) {
	for _, l := range lrs {
		vals = append(vals, l.Val)
	}
	return
}

func (lrs LatRecs) DumpCSV(fpath string) {
	fd, err := os.OpenFile(fpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Could not open '%s' for write: %s\n", fpath, err)
	}
	defer fd.Close()

	for _, lr := range lrs {
		// TODO: Something isn't right with the sampling below
		// all the samples should always be full
		if lr == nil {
			break
		}
		fmt.Fprintf(fd, "%f,%f,%d\n", lr.Time, lr.Val, lr.Ddir)
	}
}

type LatData struct {
	Min         float64 `json:"min"`
	Max         float64 `json:"max"`
	Count       int     `json:"count"`
	Sum         float64 `json:"sum"`
	Average     float64 `json:"average"`
	Stddev      float64 `json:"stddev"`
	Variance    float64 `json:"variance"`
	P1          float64 `json:"p1"`
	P5          float64 `json:"p5"`
	P10         float64 `json:"p10"`
	P25         float64 `json:"p25"`
	P50         float64 `json:"p50"`
	P75         float64 `json:"p75"`
	P90         float64 `json:"p90"`
	P95         float64 `json:"p95"`
	P99         float64 `json:"p99"`
	BeginTs     float64 `json:"first_ts"`
	EndTs       float64 `json:"last_ts"`
	ElapsedTime float64 `json:"elapsed"`
	Histogram   LatRecs `json:"histogram"`
	RHistogram  LatRecs `json:"read_histogram"`
	WHistogram  LatRecs `json:"write_histogram"`
	THistogram  LatRecs `json:"trim_histogram"`
	RecSm       LatRecs `json:"-"` // summarized to summary_size records (mean)
	RRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
	WRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
	TRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
}

// Summarizes the LatRecs data into a LatData.
// First argument is the number of samples to put in the summaries.
// Second argument is the number of buckets in the histograms.
// This does all the work in 3 passes, the first getting avg/min/max.
// Then the values are sorted to access the percentiles by index.
// The final pass computes the standard deviation, which requires the average
// from the first pass.
func (lrs LatRecs) Summarize(summary_size int, histogram_size int) (ld LatData) {
	ld.Max = math.SmallestNonzeroFloat64
	ld.Min = math.MaxFloat64
	ld.BeginTs = lrs[0].Time
	ld.EndTs = lrs[len(lrs)-1].Time
	ld.ElapsedTime = math.Abs(ld.BeginTs - ld.EndTs)

	// compute the bucket size, default to 1 if less than 10k samples
	bucket_sz := 1
	if len(lrs) > summary_size {
		bucket_sz = int(math.Ceil(float64(len(lrs)) / float64(summary_size)))
	}
	hgram_bucket_sz := int(math.Ceil(float64(len(lrs)) / float64(histogram_size)))

	// buckets / indexes / counts for summarization
	var arec, rrec, wrec, trec, acnt, rcnt, wcnt, tcnt int
	abkt := make(LatRecs, bucket_sz)
	rbkt := make(LatRecs, bucket_sz)
	wbkt := make(LatRecs, bucket_sz)
	tbkt := make(LatRecs, bucket_sz)

	// summarized samples, a slice is saved to ld after the first pass
	arecsm := make(LatRecs, summary_size)
	// by ddir
	rrecsm := make(LatRecs, summary_size)
	wrecsm := make(LatRecs, summary_size)
	trecsm := make(LatRecs, summary_size)

	// used to build histograms
	var ahgrec, ahgcnt, rhgrec, rhgcnt, whgrec, whgcnt, thgrec, thgcnt int
	ahgbkt := make(LatRecs, hgram_bucket_sz)
	rhgbkt := make(LatRecs, hgram_bucket_sz)
	whgbkt := make(LatRecs, hgram_bucket_sz)
	thgbkt := make(LatRecs, hgram_bucket_sz)
	ahg := make(LatRecs, histogram_size)
	rhg := make(LatRecs, histogram_size)
	whg := make(LatRecs, histogram_size)
	thg := make(LatRecs, histogram_size)

	// list of values, to be sorted for getting percentiles
	lvs := make([]float64, len(lrs))

	for i, lr := range lrs {
		ld.Count++
		ld.Sum += lr.Val

		if lr.Val > ld.Max {
			ld.Max = lr.Val
		}

		if lr.Val < ld.Min {
			ld.Min = lr.Val
		}

		arec, acnt = abkt.updateBucket(arec, acnt, arecsm, lr)
		ahgrec, ahgcnt = ahgbkt.updateBucket(ahgrec, ahgcnt, ahg, lr)
		if lr.Ddir == 0 {
			rrec, rcnt = rbkt.updateBucket(rrec, rcnt, rrecsm, lr)
			rhgrec, rhgcnt = rhgbkt.updateBucket(rhgrec, rhgcnt, rhg, lr)
		} else if lr.Ddir == 1 {
			wrec, wcnt = wbkt.updateBucket(wrec, wcnt, wrecsm, lr)
			whgrec, whgcnt = whgbkt.updateBucket(whgrec, whgcnt, whg, lr)
		} else if lr.Ddir == 2 {
			trec, tcnt = tbkt.updateBucket(trec, tcnt, trecsm, lr)
			thgrec, thgcnt = thgbkt.updateBucket(thgrec, thgcnt, thg, lr)
		}

		lvs[i] = lr.Val // for sorting on value for percentiles
	}

	// there might be less than summary_size samples, save only the populated slice
	ld.RecSm = arecsm[0:acnt]
	ld.RRecSm = rrecsm[0:rcnt]
	ld.WRecSm = wrecsm[0:wcnt]
	ld.TRecSm = trecsm[0:tcnt]
	ld.Histogram = ahg[0:ahgcnt]
	ld.RHistogram = rhg[0:rhgcnt]
	ld.WHistogram = whg[0:whgcnt]
	ld.THistogram = thg[0:thgcnt]

	// sort then assign percentiles
	sort.Float64s(lvs)
	pctl_idx := func(pc float64) int {
		idx := math.Floor(float64(len(lvs))*(pc/100) + 0.5)
		out := int(idx)
		//fmt.Printf("pctl_idx(%f) = %f -> lvs[%d]\n", pc, idx, out)
		return out
	}

	ld.P1 = lvs[pctl_idx(1)]
	ld.P5 = lvs[pctl_idx(5)]
	ld.P10 = lvs[pctl_idx(10)]
	ld.P25 = lvs[pctl_idx(25)]
	ld.P50 = lvs[pctl_idx(50)]
	ld.P75 = lvs[pctl_idx(75)]
	ld.P90 = lvs[pctl_idx(90)]
	ld.P95 = lvs[pctl_idx(95)]
	ld.P99 = lvs[pctl_idx(99)]

	// needed for stddev
	ld.Average = ld.Sum / float64(ld.Count)

	// second pass over values is required to compute the standard deviation
	// use lvs instead of lrs since it was just sorted and might be in cache
	var dsum float64
	for _, v := range lvs {
		dsum += math.Pow((v - ld.Average), 2)
	}
	ld.Variance = dsum / float64(ld.Count)
	ld.Stddev = math.Sqrt(ld.Variance)

	return
}

// write the latdata + CSV summaries to the specified path + filename fragment
// e.g. summary-%s.json, summary-read-%s.csv, summary-write-%s.csv, summary-trim-%s.csv
func (ld *LatData) WriteFiles(fpath string, ffrag string) {
	jsonpath := path.Join(fpath, fmt.Sprintf("summary-%s.json", ffrag))
	fd, err := os.OpenFile(jsonpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Could not open '%s' for write: %s\n", jsonpath, err)
	}
	defer fd.Close()

	enc := json.NewEncoder(fd)

	err = enc.Encode(ld)
	if err != nil {
		log.Fatalf("JSON encoding failed on file '%s': %s\n", jsonpath, err)
	}

	ld.RRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-read-%s.csv", ffrag)))
	ld.WRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-write-%s.csv", ffrag)))
	ld.TRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-trim-%s.csv", ffrag)))
}

// Adds the value to the bucket at index bktidx, with lr. When full
// summarized into smry[smry_idx]. Returns updated indexes.
func (bucket LatRecs) updateBucket(bktidx int, smry_idx int, smry LatRecs, lr *LatRec) (int, int) {
	bucket[bktidx] = lr

	// end of the bucket
	if bktidx+1 == len(bucket) {
		var ptotal, ttotal float64
		for _, v := range bucket {
			ptotal += v.Val
			ttotal += v.Time
		}

		nlr := LatRec{
			Time: math.Floor(ttotal / float64(len(bucket))),
			Val:  ptotal / float64(len(bucket)),
			Ddir: lr.Ddir,
			Bsz:  lr.Bsz,
		}
		//fmt.Printf("smry[%d] = %v (%d)\n", smry_idx, nlr, len(smry))
		smry[smry_idx] = &nlr

		// bucket is now summed & stored, reset the bucket index to 0
		return 0, smry_idx + 1
		// advance the bucket index, stay on the same summary index
	} else {
		return bktidx + 1, smry_idx
	}
}
<commit_msg>update comments<commit_after>package effio

// encoding/csv doesn't strip whitespace and does a fair bit of
// work to handle strings & quoting which are totally unnecessary
// overhead for these files so skip it

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"
)

// the input is ints but plotinum uses float64 so might as well
// start there and avoid the type conversions later
type LatRec struct {
	Time float64 `json:"x"` // time offset from beginning of fio run
	Val  float64 `json:"y"` // latency value
	Ddir uint8   `json:"-"` // 0 = read, 1 = write, 2 = trim
	Bsz  uint16  `json:"-"` // block size
}
type LatRecs []*LatRec

// Loads the CSV output by fio into an LatRecs array of LatRec structs.
func LoadCSV(filename string) LatRecs {
	fmt.Printf("Parsing file: '%s' ... ", filename)

	fd, err := os.Open(filename)
	if err != nil {
		fmt.Printf(" Failed.\nCould not open file '%s' for read: %s\n", filename, err)
		return LatRecs{}
	}
	defer fd.Close()

	started := time.Now()
	records := make(LatRecs, 0)

	var tm, perf float64
	var ddir, bsz int
	bfd := bufio.NewReader(fd)
	var lno int = 0
	for {
		line, _, err := bfd.ReadLine()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("\nRead from file '%s' failed: %s", filename, err)
		}
		lno++

		if lno%10000 == 0 {
			fmt.Printf(".")
		}

		// fio always uses ", " instead of "," as far as I can tell
		r := strings.SplitN(string(line), ", ", 4)
		// probably an impartial record at the end of the file
		if len(r) < 4 || r[0] == "" || r[1] == "" {
			continue
		}

		tm, err = strconv.ParseFloat(r[0], 64)
		if err != nil {
			log.Fatalf("\nParsing field 0 failed in file '%s' at line %d: %s", filename, lno, err)
		}
		perf, err = strconv.ParseFloat(r[1], 64)
		if err != nil {
			log.Fatalf("\nParsing field 1 in file '%s' at line %d: %s", filename, lno, err)
		}
		ddir, err = strconv.Atoi(r[2])
		if err != nil {
			log.Fatalf("\nParsing field 2 failed in file '%s' at line %d: %s", filename, lno, err)
		}
		bsz, err = strconv.Atoi(r[3])
		if err != nil {
			log.Fatalf("\nParsing field 3 failed in file '%s' at line %d: %s", filename, lno, err)
		}

		lr := LatRec{tm, perf, uint8(ddir), uint16(bsz)}
		records = append(records, &lr)
	}

	done := time.Now()
	fmt.Printf(" Done.\nRows: %d Elapsed: %s\n", len(records), done.Sub(started).String())

	return records
}

// implement some plotinum interfaces
func (lrs LatRecs) Len() int {
	return len(lrs)
}

func (lrs LatRecs) XY(i int) (float64, float64) {
	return lrs[i].Time, lrs[i].Val
}

func (lrs LatRecs) Value(i int) float64 {
	return lrs[i].Val
}

func (lrs LatRecs) Values(i int) (vals []float64) {
	for _, l := range lrs {
		vals = append(vals, l.Val)
	}
	return
}

func (lrs LatRecs) DumpCSV(fpath string) {
	fd, err := os.OpenFile(fpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Could not open '%s' for write: %s\n", fpath, err)
	}
	defer fd.Close()

	for _, lr := range lrs {
		// TODO: Something isn't right with the sampling below
		// all the samples should always be full
		if lr == nil {
			break
		}
		fmt.Fprintf(fd, "%f,%f,%d\n", lr.Time, lr.Val, lr.Ddir)
	}
}

type LatData struct {
	Min         float64 `json:"min"`
	Max         float64 `json:"max"`
	Count       int     `json:"count"`
	Sum         float64 `json:"sum"`
	Average     float64 `json:"average"`
	Stddev      float64 `json:"stddev"`
	Variance    float64 `json:"variance"`
	P1          float64 `json:"p1"`
	P5          float64 `json:"p5"`
	P10         float64 `json:"p10"`
	P25         float64 `json:"p25"`
	P50         float64 `json:"p50"`
	P75         float64 `json:"p75"`
	P90         float64 `json:"p90"`
	P95         float64 `json:"p95"`
	P99         float64 `json:"p99"`
	BeginTs     float64 `json:"first_ts"`
	EndTs       float64 `json:"last_ts"`
	ElapsedTime float64 `json:"elapsed"`
	Histogram   LatRecs `json:"histogram"`
	RHistogram  LatRecs `json:"read_histogram"`
	WHistogram  LatRecs `json:"write_histogram"`
	THistogram  LatRecs `json:"trim_histogram"`
	RecSm       LatRecs `json:"-"` // summarized to summary_size records (mean)
	RRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
	WRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
	TRecSm      LatRecs `json:"-"` // summarized to summary_size records (mean)
}

// Summarizes the LatRecs data into a LatData.
// First argument is the number of samples to put in the summaries.
// Second argument is the number of buckets in the histograms.
// This does all the work in 3 passes, the first getting avg/min/max.
// Then the values are sorted to access the percentiles by index.
// The final pass computes the standard deviation, which requires the average
// from the first pass.
func (lrs LatRecs) Summarize(summary_size int, histogram_size int) (ld LatData) {
	ld.Max = math.SmallestNonzeroFloat64
	ld.Min = math.MaxFloat64
	ld.BeginTs = lrs[0].Time
	ld.EndTs = lrs[len(lrs)-1].Time
	ld.ElapsedTime = math.Abs(ld.BeginTs - ld.EndTs)

	// compute the bucket size, default to 1 if less than 10k samples
	bucket_sz := 1
	if len(lrs) > summary_size {
		bucket_sz = int(math.Ceil(float64(len(lrs)) / float64(summary_size)))
	}
	hgram_bucket_sz := int(math.Ceil(float64(len(lrs)) / float64(histogram_size)))

	// buckets / indexes / counts for summarization
	var arec, rrec, wrec, trec, acnt, rcnt, wcnt, tcnt int
	abkt := make(LatRecs, bucket_sz)
	rbkt := make(LatRecs, bucket_sz)
	wbkt := make(LatRecs, bucket_sz)
	tbkt := make(LatRecs, bucket_sz)

	// summarized samples, a slice is saved to ld after the first pass
	arecsm := make(LatRecs, summary_size)
	// by ddir
	rrecsm := make(LatRecs, summary_size)
	wrecsm := make(LatRecs, summary_size)
	trecsm := make(LatRecs, summary_size)

	// used to build histograms, same as summarization, but (usually) much smaller
	var ahgrec, ahgcnt, rhgrec, rhgcnt, whgrec, whgcnt, thgrec, thgcnt int
	ahgbkt := make(LatRecs, hgram_bucket_sz)
	rhgbkt := make(LatRecs, hgram_bucket_sz)
	whgbkt := make(LatRecs, hgram_bucket_sz)
	thgbkt := make(LatRecs, hgram_bucket_sz)
	ahg := make(LatRecs, histogram_size)
	rhg := make(LatRecs, histogram_size)
	whg := make(LatRecs, histogram_size)
	thg := make(LatRecs, histogram_size)

	// list of values, to be sorted for getting percentiles
	lvs := make([]float64, len(lrs))

	for i, lr := range lrs {
		ld.Count++
		ld.Sum += lr.Val

		if lr.Val > ld.Max {
			ld.Max = lr.Val
		}

		if lr.Val < ld.Min {
			ld.Min = lr.Val
		}

		arec, acnt = abkt.updateBucket(arec, acnt, arecsm, lr)
		ahgrec, ahgcnt = ahgbkt.updateBucket(ahgrec, ahgcnt, ahg, lr)
		if lr.Ddir == 0 {
			rrec, rcnt = rbkt.updateBucket(rrec, rcnt, rrecsm, lr)
			rhgrec, rhgcnt = rhgbkt.updateBucket(rhgrec, rhgcnt, rhg, lr)
		} else if lr.Ddir == 1 {
			wrec, wcnt = wbkt.updateBucket(wrec, wcnt, wrecsm, lr)
			whgrec, whgcnt = whgbkt.updateBucket(whgrec, whgcnt, whg, lr)
		} else if lr.Ddir == 2 {
			trec, tcnt = tbkt.updateBucket(trec, tcnt, trecsm, lr)
			thgrec, thgcnt = thgbkt.updateBucket(thgrec, thgcnt, thg, lr)
		}

		lvs[i] = lr.Val // for sorting on value for percentiles
	}

	// there might be less than summary_size samples, save only the populated slice
	ld.RecSm = arecsm[0:acnt]
	ld.RRecSm = rrecsm[0:rcnt]
	ld.WRecSm = wrecsm[0:wcnt]
	ld.TRecSm = trecsm[0:tcnt]
	ld.Histogram = ahg[0:ahgcnt]
	ld.RHistogram = rhg[0:rhgcnt]
	ld.WHistogram = whg[0:whgcnt]
	ld.THistogram = thg[0:thgcnt]

	// sort then assign percentiles
	sort.Float64s(lvs)
	pctl_idx := func(pc float64) int {
		idx := math.Floor(float64(len(lvs))*(pc/100) + 0.5)
		out := int(idx)
		//fmt.Printf("pctl_idx(%f) = %f -> lvs[%d]\n", pc, idx, out)
		return out
	}

	ld.P1 = lvs[pctl_idx(1)]
	ld.P5 = lvs[pctl_idx(5)]
	ld.P10 = lvs[pctl_idx(10)]
	ld.P25 = lvs[pctl_idx(25)]
	ld.P50 = lvs[pctl_idx(50)]
	ld.P75 = lvs[pctl_idx(75)]
	ld.P90 = lvs[pctl_idx(90)]
	ld.P95 = lvs[pctl_idx(95)]
	ld.P99 = lvs[pctl_idx(99)]

	// needed for stddev
	ld.Average = ld.Sum / float64(ld.Count)

	// second pass over values is required to compute the standard deviation
	// use lvs instead of lrs: it's smaller, recently accessed, might be in cache
	var dsum float64
	for _, v := range lvs {
		dsum += math.Pow((v - ld.Average), 2)
	}
	ld.Variance = dsum / float64(ld.Count)
	ld.Stddev = math.Sqrt(ld.Variance)

	return
}

// write the latdata + CSV summaries to the specified path + filename fragment
// e.g. summary-%s.json, summary-read-%s.csv, summary-write-%s.csv, summary-trim-%s.csv
func (ld *LatData) WriteFiles(fpath string, ffrag string) {
	jsonpath := path.Join(fpath, fmt.Sprintf("summary-%s.json", ffrag))
	fd, err := os.OpenFile(jsonpath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("Could not open '%s' for write: %s\n", jsonpath, err)
	}
	defer fd.Close()

	enc := json.NewEncoder(fd)

	err = enc.Encode(ld)
	if err != nil {
		log.Fatalf("JSON encoding failed on file '%s': %s\n", jsonpath, err)
	}

	ld.RRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-read-%s.csv", ffrag)))
	ld.WRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-write-%s.csv", ffrag)))
	ld.TRecSm.DumpCSV(path.Join(fpath, fmt.Sprintf("summary-trim-%s.csv", ffrag)))
}

// Adds the value to the bucket at index bktidx, with lr. When full
// summarized into smry[smry_idx]. Returns updated indexes.
func (bucket LatRecs) updateBucket(bktidx int, smry_idx int, smry LatRecs, lr *LatRec) (int, int) {
	bucket[bktidx] = lr

	// end of the bucket
	if bktidx+1 == len(bucket) {
		var ptotal, ttotal float64
		for _, v := range bucket {
			ptotal += v.Val
			ttotal += v.Time
		}

		nlr := LatRec{
			Time: math.Floor(ttotal / float64(len(bucket))),
			Val:  ptotal / float64(len(bucket)),
			Ddir: lr.Ddir,
			Bsz:  lr.Bsz,
		}
		//fmt.Printf("smry[%d] = %v (%d)\n", smry_idx, nlr, len(smry))
		smry[smry_idx] = &nlr

		// bucket is now summed & stored, reset the bucket index to 0
		return 0, smry_idx + 1
		// advance the bucket index, stay on the same summary index
	} else {
		return bktidx + 1, smry_idx
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 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 errors implements functions to manipulate errors.
package errors

// New returns an error that formats as the given text.
//
// The returned error contains a Frame set to the caller's location and
// implements Formatter to show this information when printed with details.
func New(text string) error {
	return &errorString{text, Caller(1)}
}

// errorString is a trivial implementation of error.
type errorString struct {
	s     string
	frame Frame
}

func (e *errorString) Error() string {
	return e.s
}

func (e *errorString) FormatError(p Printer) (next error) {
	p.Print(e.s)
	e.frame.Format(p)
	return nil
}
<commit_msg>errors: improve performance of New<commit_after>// Copyright 2011 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 errors implements functions to manipulate errors.
package errors

import "runtime"

// New returns an error that formats as the given text.
//
// The returned error contains a Frame set to the caller's location and
// implements Formatter to show this information when printed with details.
func New(text string) error {
	// Inline call to errors.Callers to improve performance.
	var s Frame
	runtime.Callers(2, s.frames[:])
	return &errorString{text, s}
}

// errorString is a trivial implementation of error.
type errorString struct {
	s     string
	frame Frame
}

func (e *errorString) Error() string {
	return e.s
}

func (e *errorString) FormatError(p Printer) (next error) {
	p.Print(e.s)
	e.frame.Format(p)
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2009 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 expvar provides a standardized interface to public variables, such
// as operation counters in servers. It exposes these variables via HTTP at
// /debug/vars in JSON format.
//
// Operations to set or modify these public variables are atomic.
//
// In addition to adding the HTTP handler, this package registers the
// following variables:
//
//	cmdline   os.Args
//	memstats  runtime.Memstats
//
// The package is sometimes only imported for the side effect of
// registering its HTTP handler and the above variables. To use it
// this way, link this package into your program:
//	import _ "expvar"
//
package expvar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
)

// Var is an abstract type for all exported variables.
type Var interface {
	// String returns a valid JSON value for the variable.
	// Types with String methods that do not return valid JSON
	// (such as time.Time) must not be used as a Var.
	String() string
}

// Int is a 64-bit integer variable that satisfies the Var interface.
type Int struct {
	i int64
}

func (v *Int) Value() int64 {
	return atomic.LoadInt64(&v.i)
}

func (v *Int) String() string {
	return strconv.FormatInt(atomic.LoadInt64(&v.i), 10)
}

func (v *Int) Add(delta int64) {
	atomic.AddInt64(&v.i, delta)
}

func (v *Int) Set(value int64) {
	atomic.StoreInt64(&v.i, value)
}

// Float is a 64-bit float variable that satisfies the Var interface.
type Float struct {
	f uint64
}

func (v *Float) Value() float64 {
	return math.Float64frombits(atomic.LoadUint64(&v.f))
}

func (v *Float) String() string {
	return strconv.FormatFloat(
		math.Float64frombits(atomic.LoadUint64(&v.f)), 'g', -1, 64)
}

// Add adds delta to v.
func (v *Float) Add(delta float64) {
	for {
		cur := atomic.LoadUint64(&v.f)
		curVal := math.Float64frombits(cur)
		nxtVal := curVal + delta
		nxt := math.Float64bits(nxtVal)
		if atomic.CompareAndSwapUint64(&v.f, cur, nxt) {
			return
		}
	}
}

// Set sets v to value.
func (v *Float) Set(value float64) {
	atomic.StoreUint64(&v.f, math.Float64bits(value))
}

// Map is a string-to-Var map variable that satisfies the Var interface.
type Map struct {
	m      sync.Map // map[string]Var
	keysMu sync.RWMutex
	keys   []string // sorted
}

// KeyValue represents a single entry in a Map.
type KeyValue struct {
	Key   string
	Value Var
}

func (v *Map) String() string {
	var b bytes.Buffer
	fmt.Fprintf(&b, "{")
	first := true
	v.Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(&b, ", ")
		}
		fmt.Fprintf(&b, "%q: %v", kv.Key, kv.Value)
		first = false
	})
	fmt.Fprintf(&b, "}")
	return b.String()
}

func (v *Map) Init() *Map { return v }

// updateKeys updates the sorted list of keys in v.keys.
func (v *Map) addKey(key string) {
	v.keysMu.Lock()
	defer v.keysMu.Unlock()
	v.keys = append(v.keys, key)
	sort.Strings(v.keys)
}

func (v *Map) Get(key string) Var {
	i, _ := v.m.Load(key)
	av, _ := i.(Var)
	return av
}

func (v *Map) Set(key string, av Var) {
	if _, dup := v.m.LoadOrStore(key, av); dup {
		v.m.Store(key, av)
	} else {
		v.addKey(key)
	}
}

// Add adds delta to the *Int value stored under the given map key.
func (v *Map) Add(key string, delta int64) {
	i, ok := v.m.Load(key)
	if !ok {
		var dup bool
		i, dup = v.m.LoadOrStore(key, new(Int))
		if !dup {
			v.addKey(key)
		}
	}

	// Add to Int; ignore otherwise.
	if iv, ok := i.(*Int); ok {
		iv.Add(delta)
	}
}

// AddFloat adds delta to the *Float value stored under the given map key.
func (v *Map) AddFloat(key string, delta float64) {
	i, ok := v.m.Load(key)
	if !ok {
		var dup bool
		i, dup = v.m.LoadOrStore(key, new(Float))
		if !dup {
			v.addKey(key)
		}
	}

	// Add to Float; ignore otherwise.
	if iv, ok := i.(*Float); ok {
		iv.Add(delta)
	}
}

// Do calls f for each entry in the map.
// The map is locked during the iteration,
// but existing entries may be concurrently updated.
func (v *Map) Do(f func(KeyValue)) {
	v.keysMu.RLock()
	defer v.keysMu.RUnlock()
	for _, k := range v.keys {
		i, _ := v.m.Load(k)
		f(KeyValue{k, i.(Var)})
	}
}

// String is a string variable, and satisfies the Var interface.
type String struct {
	s atomic.Value // string
}

func (v *String) Value() string {
	p, _ := v.s.Load().(string)
	return p
}

// String implements the Val interface. To get the unquoted string
// use Value.
func (v *String) String() string {
	s := v.Value()
	b, _ := json.Marshal(s)
	return string(b)
}

func (v *String) Set(value string) {
	v.s.Store(value)
}

// Func implements Var by calling the function
// and formatting the returned value using JSON.
type Func func() interface{}

func (f Func) Value() interface{} {
	return f()
}

func (f Func) String() string {
	v, _ := json.Marshal(f())
	return string(v)
}

// All published variables.
var (
	vars      sync.Map // map[string]Var
	varKeysMu sync.RWMutex
	varKeys   []string // sorted
)

// Publish declares a named exported variable. This should be called from a
// package's init function when it creates its Vars. If the name is already
// registered then this will log.Panic.
func Publish(name string, v Var) {
	if _, dup := vars.LoadOrStore(name, v); dup {
		log.Panicln("Reuse of exported var name:", name)
	}
	varKeysMu.Lock()
	defer varKeysMu.Unlock()
	varKeys = append(varKeys, name)
	sort.Strings(varKeys)
}

// Get retrieves a named exported variable. It returns nil if the name has
// not been registered.
func Get(name string) Var {
	i, _ := vars.Load(name)
	v, _ := i.(Var)
	return v
}

// Convenience functions for creating new exported variables.

func NewInt(name string) *Int {
	v := new(Int)
	Publish(name, v)
	return v
}

func NewFloat(name string) *Float {
	v := new(Float)
	Publish(name, v)
	return v
}

func NewMap(name string) *Map {
	v := new(Map).Init()
	Publish(name, v)
	return v
}

func NewString(name string) *String {
	v := new(String)
	Publish(name, v)
	return v
}

// Do calls f for each exported variable.
// The global variable map is locked during the iteration,
// but existing entries may be concurrently updated.
func Do(f func(KeyValue)) {
	varKeysMu.RLock()
	defer varKeysMu.RUnlock()
	for _, k := range varKeys {
		val, _ := vars.Load(k)
		f(KeyValue{k, val.(Var)})
	}
}

func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

// Handler returns the expvar HTTP Handler.
//
// This is only needed to install the handler in a non-standard location.
func Handler() http.Handler {
	return http.HandlerFunc(expvarHandler)
}

func cmdline() interface{} {
	return os.Args
}

func memstats() interface{} {
	stats := new(runtime.MemStats)
	runtime.ReadMemStats(stats)
	return *stats
}

func init() {
	http.HandleFunc("/debug/vars", expvarHandler)
	Publish("cmdline", Func(cmdline))
	Publish("memstats", Func(memstats))
}
<commit_msg>expvar: reduce sync.Map overhead for (*Map).Set<commit_after>// Copyright 2009 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 expvar provides a standardized interface to public variables, such
// as operation counters in servers. It exposes these variables via HTTP at
// /debug/vars in JSON format.
//
// Operations to set or modify these public variables are atomic.
//
// In addition to adding the HTTP handler, this package registers the
// following variables:
//
//	cmdline   os.Args
//	memstats  runtime.Memstats
//
// The package is sometimes only imported for the side effect of
// registering its HTTP handler and the above variables. To use it
// this way, link this package into your program:
//	import _ "expvar"
//
package expvar

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
)

// Var is an abstract type for all exported variables.
type Var interface {
	// String returns a valid JSON value for the variable.
	// Types with String methods that do not return valid JSON
	// (such as time.Time) must not be used as a Var.
	String() string
}

// Int is a 64-bit integer variable that satisfies the Var interface.
type Int struct {
	i int64
}

func (v *Int) Value() int64 {
	return atomic.LoadInt64(&v.i)
}

func (v *Int) String() string {
	return strconv.FormatInt(atomic.LoadInt64(&v.i), 10)
}

func (v *Int) Add(delta int64) {
	atomic.AddInt64(&v.i, delta)
}

func (v *Int) Set(value int64) {
	atomic.StoreInt64(&v.i, value)
}

// Float is a 64-bit float variable that satisfies the Var interface.
type Float struct {
	f uint64
}

func (v *Float) Value() float64 {
	return math.Float64frombits(atomic.LoadUint64(&v.f))
}

func (v *Float) String() string {
	return strconv.FormatFloat(
		math.Float64frombits(atomic.LoadUint64(&v.f)), 'g', -1, 64)
}

// Add adds delta to v.
func (v *Float) Add(delta float64) {
	for {
		cur := atomic.LoadUint64(&v.f)
		curVal := math.Float64frombits(cur)
		nxtVal := curVal + delta
		nxt := math.Float64bits(nxtVal)
		if atomic.CompareAndSwapUint64(&v.f, cur, nxt) {
			return
		}
	}
}

// Set sets v to value.
func (v *Float) Set(value float64) {
	atomic.StoreUint64(&v.f, math.Float64bits(value))
}

// Map is a string-to-Var map variable that satisfies the Var interface.
type Map struct {
	m      sync.Map // map[string]Var
	keysMu sync.RWMutex
	keys   []string // sorted
}

// KeyValue represents a single entry in a Map.
type KeyValue struct {
	Key   string
	Value Var
}

func (v *Map) String() string {
	var b bytes.Buffer
	fmt.Fprintf(&b, "{")
	first := true
	v.Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(&b, ", ")
		}
		fmt.Fprintf(&b, "%q: %v", kv.Key, kv.Value)
		first = false
	})
	fmt.Fprintf(&b, "}")
	return b.String()
}

func (v *Map) Init() *Map { return v }

// updateKeys updates the sorted list of keys in v.keys.
func (v *Map) addKey(key string) {
	v.keysMu.Lock()
	defer v.keysMu.Unlock()
	v.keys = append(v.keys, key)
	sort.Strings(v.keys)
}

func (v *Map) Get(key string) Var {
	i, _ := v.m.Load(key)
	av, _ := i.(Var)
	return av
}

func (v *Map) Set(key string, av Var) {
	// Before we store the value, check to see whether the key is new. Try a Load
	// before LoadOrStore: LoadOrStore causes the key interface to escape even on
	// the Load path.
	if _, ok := v.m.Load(key); !ok {
		if _, dup := v.m.LoadOrStore(key, av); !dup {
			v.addKey(key)
			return
		}
	}

	v.m.Store(key, av)
}

// Add adds delta to the *Int value stored under the given map key.
func (v *Map) Add(key string, delta int64) {
	i, ok := v.m.Load(key)
	if !ok {
		var dup bool
		i, dup = v.m.LoadOrStore(key, new(Int))
		if !dup {
			v.addKey(key)
		}
	}

	// Add to Int; ignore otherwise.
	if iv, ok := i.(*Int); ok {
		iv.Add(delta)
	}
}

// AddFloat adds delta to the *Float value stored under the given map key.
func (v *Map) AddFloat(key string, delta float64) {
	i, ok := v.m.Load(key)
	if !ok {
		var dup bool
		i, dup = v.m.LoadOrStore(key, new(Float))
		if !dup {
			v.addKey(key)
		}
	}

	// Add to Float; ignore otherwise.
	if iv, ok := i.(*Float); ok {
		iv.Add(delta)
	}
}

// Do calls f for each entry in the map.
// The map is locked during the iteration,
// but existing entries may be concurrently updated.
func (v *Map) Do(f func(KeyValue)) {
	v.keysMu.RLock()
	defer v.keysMu.RUnlock()
	for _, k := range v.keys {
		i, _ := v.m.Load(k)
		f(KeyValue{k, i.(Var)})
	}
}

// String is a string variable, and satisfies the Var interface.
type String struct {
	s atomic.Value // string
}

func (v *String) Value() string {
	p, _ := v.s.Load().(string)
	return p
}

// String implements the Val interface. To get the unquoted string
// use Value.
func (v *String) String() string {
	s := v.Value()
	b, _ := json.Marshal(s)
	return string(b)
}

func (v *String) Set(value string) {
	v.s.Store(value)
}

// Func implements Var by calling the function
// and formatting the returned value using JSON.
type Func func() interface{}

func (f Func) Value() interface{} {
	return f()
}

func (f Func) String() string {
	v, _ := json.Marshal(f())
	return string(v)
}

// All published variables.
var (
	vars      sync.Map // map[string]Var
	varKeysMu sync.RWMutex
	varKeys   []string // sorted
)

// Publish declares a named exported variable. This should be called from a
// package's init function when it creates its Vars. If the name is already
// registered then this will log.Panic.
func Publish(name string, v Var) {
	if _, dup := vars.LoadOrStore(name, v); dup {
		log.Panicln("Reuse of exported var name:", name)
	}
	varKeysMu.Lock()
	defer varKeysMu.Unlock()
	varKeys = append(varKeys, name)
	sort.Strings(varKeys)
}

// Get retrieves a named exported variable. It returns nil if the name has
// not been registered.
func Get(name string) Var {
	i, _ := vars.Load(name)
	v, _ := i.(Var)
	return v
}

// Convenience functions for creating new exported variables.

func NewInt(name string) *Int {
	v := new(Int)
	Publish(name, v)
	return v
}

func NewFloat(name string) *Float {
	v := new(Float)
	Publish(name, v)
	return v
}

func NewMap(name string) *Map {
	v := new(Map).Init()
	Publish(name, v)
	return v
}

func NewString(name string) *String {
	v := new(String)
	Publish(name, v)
	return v
}

// Do calls f for each exported variable.
// The global variable map is locked during the iteration,
// but existing entries may be concurrently updated.
func Do(f func(KeyValue)) {
	varKeysMu.RLock()
	defer varKeysMu.RUnlock()
	for _, k := range varKeys {
		val, _ := vars.Load(k)
		f(KeyValue{k, val.(Var)})
	}
}

func expvarHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	fmt.Fprintf(w, "{\n")
	first := true
	Do(func(kv KeyValue) {
		if !first {
			fmt.Fprintf(w, ",\n")
		}
		first = false
		fmt.Fprintf(w, "%q: %s", kv.Key, kv.Value)
	})
	fmt.Fprintf(w, "\n}\n")
}

// Handler returns the expvar HTTP Handler.
//
// This is only needed to install the handler in a non-standard location.
func Handler() http.Handler {
	return http.HandlerFunc(expvarHandler)
}

func cmdline() interface{} {
	return os.Args
}

func memstats() interface{} {
	stats := new(runtime.MemStats)
	runtime.ReadMemStats(stats)
	return *stats
}

func init() {
	http.HandleFunc("/debug/vars", expvarHandler)
	Publish("cmdline", Func(cmdline))
	Publish("memstats", Func(memstats))
}
<|endoftext|>


<commit_before>package src

import (
	"fmt"
	"time"

	"github.com/codegangsta/cli"

	log "github.com/Sirupsen/logrus"
)

type RunnerHealth struct {
	failures  int
	lastCheck time.Time
}

type MultiRunner struct {
	config  Config
	jobs    []*Job
	healthy map[string]*RunnerHealth
}

func (mr *MultiRunner) getHealth(runner *RunnerConfig) *RunnerHealth {
	if mr.healthy == nil {
		mr.healthy = map[string]*RunnerHealth{}
	}
	health := mr.healthy[runner.UniqueID()]
	if health == nil {
		health = &RunnerHealth{
			lastCheck: time.Now(),
		}
		mr.healthy[runner.UniqueID()] = health
	}
	return health
}

func (mr *MultiRunner) isHealthy(runner *RunnerConfig) bool {
	health := mr.getHealth(runner)
	if health.failures < HEALTHY_CHECKS {
		return true
	}

	if time.Since(health.lastCheck) > HEALTH_CHECK_INTERVAL*time.Second {
		mr.errorln("Runner", runner.ShortDescription(), "is not healthy, but will be checked!")
		health.failures = 0
		health.lastCheck = time.Now()
		return true
	}

	return false
}

func (mr *MultiRunner) makeHealthy(runner *RunnerConfig) {
	health := mr.getHealth(runner)
	health.failures = 0
	health.lastCheck = time.Now()
}

func (mr *MultiRunner) makeUnhealthy(runner *RunnerConfig) {
	health := mr.getHealth(runner)
	health.failures++

	if health.failures >= HEALTHY_CHECKS {
		mr.errorln("Runner", runner.ShortDescription(), "is not healthy and will be disabled!")
	}
}

func (mr *MultiRunner) requestNewJob() (*GetBuildResponse, *RunnerConfig) {
	for _, runner := range mr.config.Runners {
		if runner == nil {
			continue
		}

		if !mr.isHealthy(runner) {
			continue
		}

		count := 0
		for _, job := range mr.jobs {
			if job.Runner == runner {
				count += 1
			}
		}

		if runner.Limit > 0 && count >= runner.Limit {
			continue
		}

		new_build, healthy := GetBuild(*runner)
		if new_build != nil {
			return new_build, runner
		}

		if healthy {
			mr.makeHealthy(runner)
		} else {
			mr.makeUnhealthy(runner)
		}
	}

	return nil, nil
}

func (mr *MultiRunner) startNewJob(finish chan *Job) *Job {
	if mr.config.Concurrent <= len(mr.jobs) {
		return nil
	}

	log.Debugln(len(mr.jobs), "Requesting a new job...")

	new_build, runner_config := mr.requestNewJob()
	if new_build == nil {
		return nil
	}
	if runner_config == nil {
		// this shouldn't happen
		return nil
	}

	log.Debugln(len(mr.jobs), "Received new job for", runner_config.ShortDescription(), "build", new_build.Id)
	new_job := &Job{
		Build: &Build{
			GetBuildResponse: *new_build,
		},
		Runner: runner_config,
	}

	build_prefix := fmt.Sprintf("runner-%s", runner_config.ShortDescription())

	other_builds := []*Build{}
	for _, other_job := range mr.jobs {
		other_builds = append(other_builds, other_job.Build)
	}
	new_job.Build.GenerateUniqueName(build_prefix, other_builds)

	go func() {
		new_job.Run()
		finish <- new_job
	}()
	return new_job
}

func (mr *MultiRunner) errorln(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Errorln(args...)
}

func (mr *MultiRunner) debugln(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Debugln(args...)
}

func (mr *MultiRunner) println(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Println(args...)
}

func (mr *MultiRunner) addJob(newJob *Job) {
	mr.jobs = append(mr.jobs, newJob)
	mr.debugln("Added a new job", newJob)
}

func (mr *MultiRunner) removeJob(deleteJob *Job) bool {
	for idx, job := range mr.jobs {
		if job == deleteJob {
			mr.jobs[idx] = mr.jobs[len(mr.jobs)-1]
			mr.jobs = mr.jobs[:len(mr.jobs)-1]
			mr.debugln("Removed job", deleteJob)
			return true
		}
	}
	return false
}

func runMulti(c *cli.Context) {
	mr := MultiRunner{}
	err := mr.config.LoadConfig(c.String("config"))
	if err != nil {
		panic(err)
	}

	mr.config.SetChdir()
	mr.println("Starting multi-runner from", c.String("config"), "...")

	job_finish := make(chan *Job)
	reload_config := make(chan Config)
	go ReloadConfig(c.String("config"), mr.config.ModTime, reload_config)

	for {
		new_job := mr.startNewJob(job_finish)
		if new_job != nil {
			mr.addJob(new_job)
		}

		select {
		case finished_job := <-job_finish:
			mr.debugln("Job finished", finished_job)
			mr.removeJob(finished_job)

		case new_config := <-reload_config:
			mr.debugln("Config reloaded.")
			mr.healthy = nil
			mr.config = new_config
			mr.config.SetChdir()

		case <-time.After(CHECK_INTERVAL * time.Second):
			mr.debugln("Check interval fired")
		}
	}
}
<commit_msg>Cleanup code<commit_after>package src

import (
	"fmt"
	"time"

	"github.com/codegangsta/cli"

	log "github.com/Sirupsen/logrus"
)

type RunnerHealth struct {
	failures  int
	lastCheck time.Time
}

type MultiRunner struct {
	config  Config
	jobs    []*Job
	healthy map[string]*RunnerHealth
}

func (mr *MultiRunner) errorln(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Errorln(args...)
}

func (mr *MultiRunner) debugln(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Debugln(args...)
}

func (mr *MultiRunner) println(args ...interface{}) {
	args = append([]interface{}{len(mr.jobs)}, args...)
	log.Println(args...)
}

func (mr *MultiRunner) getHealth(runner *RunnerConfig) *RunnerHealth {
	if mr.healthy == nil {
		mr.healthy = map[string]*RunnerHealth{}
	}
	health := mr.healthy[runner.UniqueID()]
	if health == nil {
		health = &RunnerHealth{
			lastCheck: time.Now(),
		}
		mr.healthy[runner.UniqueID()] = health
	}
	return health
}

func (mr *MultiRunner) isHealthy(runner *RunnerConfig) bool {
	health := mr.getHealth(runner)
	if health.failures < HEALTHY_CHECKS {
		return true
	}

	if time.Since(health.lastCheck) > HEALTH_CHECK_INTERVAL*time.Second {
		mr.errorln("Runner", runner.ShortDescription(), "is not healthy, but will be checked!")
		health.failures = 0
		health.lastCheck = time.Now()
		return true
	}

	return false
}

func (mr *MultiRunner) makeHealthy(runner *RunnerConfig) {
	health := mr.getHealth(runner)
	health.failures = 0
	health.lastCheck = time.Now()
}

func (mr *MultiRunner) makeUnhealthy(runner *RunnerConfig) {
	health := mr.getHealth(runner)
	health.failures++

	if health.failures >= HEALTHY_CHECKS {
		mr.errorln("Runner", runner.ShortDescription(), "is not healthy and will be disabled!")
	}
}

func (mr *MultiRunner) requestNewJob() (*GetBuildResponse, *RunnerConfig) {
	for _, runner := range mr.config.Runners {
		if runner == nil {
			continue
		}

		if !mr.isHealthy(runner) {
			continue
		}

		count := 0
		for _, job := range mr.jobs {
			if job.Runner == runner {
				count += 1
			}
		}

		if runner.Limit > 0 && count >= runner.Limit {
			continue
		}

		new_build, healthy := GetBuild(*runner)
		if new_build != nil {
			return new_build, runner
		}

		if healthy {
			mr.makeHealthy(runner)
		} else {
			mr.makeUnhealthy(runner)
		}
	}

	return nil, nil
}

func (mr *MultiRunner) startNewJob(finish chan *Job) *Job {
	if mr.config.Concurrent <= len(mr.jobs) {
		return nil
	}

	log.Debugln(len(mr.jobs), "Requesting a new job...")

	new_build, runner_config := mr.requestNewJob()
	if new_build == nil {
		return nil
	}
	if runner_config == nil {
		// this shouldn't happen
		return nil
	}

	log.Debugln(len(mr.jobs), "Received new job for", runner_config.ShortDescription(), "build", new_build.Id)
	new_job := &Job{
		Build: &Build{
			GetBuildResponse: *new_build,
		},
		Runner: runner_config,
	}

	build_prefix := fmt.Sprintf("runner-%s", runner_config.ShortDescription())

	other_builds := []*Build{}
	for _, other_job := range mr.jobs {
		other_builds = append(other_builds, other_job.Build)
	}
	new_job.Build.GenerateUniqueName(build_prefix, other_builds)

	go func() {
		new_job.Run()
		finish <- new_job
	}()
	return new_job
}

func (mr *MultiRunner) addJob(newJob *Job) {
	mr.jobs = append(mr.jobs, newJob)
	mr.debugln("Added a new job", newJob)
}

func (mr *MultiRunner) removeJob(deleteJob *Job) bool {
	for idx, job := range mr.jobs {
		if job == deleteJob {
			mr.jobs[idx] = mr.jobs[len(mr.jobs)-1]
			mr.jobs = mr.jobs[:len(mr.jobs)-1]
			mr.debugln("Removed job", deleteJob)
			return true
		}
	}
	return false
}

func runMulti(c *cli.Context) {
	mr := MultiRunner{}
	err := mr.config.LoadConfig(c.String("config"))
	if err != nil {
		panic(err)
	}

	mr.config.SetChdir()
	mr.println("Starting multi-runner from", c.String("config"), "...")

	job_finish := make(chan *Job)
	reload_config := make(chan Config)
	go ReloadConfig(c.String("config"), mr.config.ModTime, reload_config)

	for {
		new_job := mr.startNewJob(job_finish)
		if new_job != nil {
			mr.addJob(new_job)
		}

		select {
		case finished_job := <-job_finish:
			mr.debugln("Job finished", finished_job)
			mr.removeJob(finished_job)

		case new_config := <-reload_config:
			mr.debugln("Config reloaded.")
			mr.healthy = nil
			mr.config = new_config
			mr.config.SetChdir()

		case <-time.After(CHECK_INTERVAL * time.Second):
			mr.debugln("Check interval fired")
		}
	}
}
<|endoftext|>