<commit_before>package node

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/Dataman-Cloud/swan/src/agent"
	"github.com/Dataman-Cloud/swan/src/apiserver"
	"github.com/Dataman-Cloud/swan/src/config"
	"github.com/Dataman-Cloud/swan/src/event"
	"github.com/Dataman-Cloud/swan/src/manager"
	"github.com/Dataman-Cloud/swan/src/swancontext"
	"github.com/Dataman-Cloud/swan/src/types"
	"github.com/Dataman-Cloud/swan/src/utils/httpclient"

	"github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
	"github.com/coreos/etcd/pkg/fileutil"
	"github.com/twinj/uuid"
	"golang.org/x/net/context"
)

const (
	NodeIDFileName    = "/ID"
	JoinRetryInterval = 5
)

type Node struct {
	ID                string
	agent             *agent.Agent     // hold reference to agent, take function when in agent mode
	manager           *manager.Manager // hold a instance of manager, make logic taking place
	ctx               context.Context
	joinRetryInterval time.Duration
	RaftID            uint64
	stopC             chan struct{}
}

func NewNode(config config.SwanConfig) (*Node, error) {
	nodeID, err := loadOrCreateNodeID(config)
	if err != nil {
		return nil, err
	}

	// init swanconfig instance
	config.NodeID = nodeID
	_ = swancontext.NewSwanContext(config, event.New())

	if !swancontext.IsManager() && !swancontext.IsAgent() {
		return nil, errors.New("node must be started with at least one role in [manager,agent]")
	}

	node := &Node{
		ID:                nodeID,
		joinRetryInterval: time.Second * JoinRetryInterval,
		stopC:             make(chan struct{}, 1),
	}

	err = os.MkdirAll(config.DataDir+"/"+nodeID, 0644)
	if err != nil {
		logrus.Errorf("os.MkdirAll got error: %s", err)
		return nil, err
	}

	db, err := bolt.Open(config.DataDir+"/"+nodeID+"/swan.db", 0644, nil)
	if err != nil {
		logrus.Errorf("Init bolt store failed:%s", err)
		return nil, err
	}

	raftID, err := loadOrCreateRaftID(db)
	if err != nil {
		logrus.Errorf("Init raft ID failed:%s", err)
		return nil, err
	}
	node.RaftID = raftID

	if swancontext.IsManager() {
		m, err := manager.New(db)
		if err != nil {
			return nil, err
		}
		node.manager = m
	}

	if swancontext.IsAgent() {
		a, err := agent.New()
		if err != nil {
			return nil, err
		}
		node.agent = a
	}

	nodeApi := &NodeApi{node}
	apiserver.Install(swancontext.Instance().ApiServer, nodeApi)

	return node, nil
}

func loadOrCreateNodeID(swanConfig config.SwanConfig) (string, error) {
	nodeIDFile := swanConfig.DataDir + NodeIDFileName
	if !fileutil.Exist(nodeIDFile) {
		os.MkdirAll(swanConfig.DataDir, 0700)

		nodeID := uuid.NewV4().String()
		idFile, err := os.OpenFile(nodeIDFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return "", err
		}

		if _, err = idFile.WriteString(nodeID); err != nil {
			return "", err
		}

		logrus.Infof("starting swan node, ID file was not found started with  new ID: %s", nodeID)
		return nodeID, nil

	} else {
		idFile, err := os.Open(nodeIDFile)
		if err != nil {
			return "", err
		}

		nodeID, err := ioutil.ReadAll(idFile)
		if err != nil {
			return "", err
		}

		logrus.Infof("starting swan node, ID file was found started with ID: %s", string(nodeID))
		return string(nodeID), nil
	}
}

// node start from here
// - 1, start manager if needed
// - 2, start agent if needed
// - 3, agent join to managers if needed
// - 4, start the API server, both for agent and client
// - 5, enter loop, wait for error or ctx.Done
func (n *Node) Start(ctx context.Context) error {
	errChan := make(chan error, 1)

	swanConfig := swancontext.Instance().Config
	nodeInfo := types.Node{
		ID:                n.ID,
		AdvertiseAddr:     swanConfig.AdvertiseAddr,
		ListenAddr:        swanConfig.ListenAddr,
		RaftListenAddr:    swanConfig.RaftListenAddr,
		RaftAdvertiseAddr: swanConfig.RaftAdvertiseAddr,
		Role:              types.NodeRole(swanConfig.Mode),
		RaftID:            n.RaftID,
	}

	if swancontext.IsManager() {
		go func() {
			if swancontext.IsNewCluster() {
				errChan <- n.runManager(ctx, n.RaftID, []types.Node{nodeInfo}, true)
			} else {
				existedNodes, err := n.JoinAsManager(nodeInfo)
				if err != nil {
					errChan <- err
				}

				errChan <- n.runManager(ctx, n.RaftID, existedNodes, false)
			}
		}()
	}

	if swancontext.IsAgent() {
		go func() {
			errChan <- n.runAgent(ctx)
		}()

		go func() {
			err := n.JoinAsAgent(nodeInfo)
			if err != nil {
				errChan <- err
			}
		}()
	}

	go func() {
		errChan <- swancontext.Instance().ApiServer.Start()
	}()

	for {
		select {
		case err := <-errChan:
			return err
		case <-n.stopC:
			n.Stop()
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (n *Node) runAgent(ctx context.Context) error {
	agentCtx, cancel := context.WithCancel(ctx)
	n.agent.CancelFunc = cancel
	return n.agent.Start(agentCtx)
}

func (n *Node) runManager(ctx context.Context, raftID uint64, peers []types.Node, isNewCluster bool) error {
	managerCtx, cancel := context.WithCancel(ctx)
	n.manager.CancelFunc = cancel
	return n.manager.Start(managerCtx, raftID, peers, isNewCluster)
}

// node stop
func (n *Node) Stop() {
	n.agent.Stop()
	n.manager.Stop()
}

func (n *Node) JoinAsAgent(nodeInfo types.Node) error {
	swanConfig := swancontext.Instance().Config
	if len(swanConfig.JoinAddrs) == 0 {
		return errors.New("start agent failed. Error: joinAddrs must be no empty")
	}

	for _, managerAddr := range swanConfig.JoinAddrs {
		registerAddr := "http://" + managerAddr + config.API_PREFIX + "/nodes"
		_, err := httpclient.NewDefaultClient().POST(context.TODO(), registerAddr, nil, nodeInfo, nil)
		if err != nil {
			logrus.Infof("register to %s got error: %s, try again after %d seconds", registerAddr, err.Error(), JoinRetryInterval)
		}

		if err == nil {
			logrus.Infof("agent register to manager success with managerAddr: %s", managerAddr)
			return nil
		}
	}

	time.Sleep(n.joinRetryInterval)
	n.JoinAsAgent(nodeInfo)
	return nil
}

func (n *Node) JoinAsManager(nodeInfo types.Node) ([]types.Node, error) {
	swanConfig := swancontext.Instance().Config
	if len(swanConfig.JoinAddrs) == 0 {
		return nil, errors.New("start agent failed. Error: joinAddrs must be no empty")
	}

	for _, managerAddr := range swanConfig.JoinAddrs {
		registerAddr := "http://" + managerAddr + config.API_PREFIX + "/nodes"
		resp, err := httpclient.NewDefaultClient().POST(context.TODO(), registerAddr, nil, nodeInfo, nil)
		if err != nil {
			logrus.Errorf("register to %s got error: %s", registerAddr, err.Error())
			continue
		}

		var nodes []types.Node
		if err := json.Unmarshal(resp, &nodes); err != nil {
			logrus.Errorf("register to %s got error: %s", registerAddr, err.Error())
			continue
		}

		var managerNodes []types.Node
		for _, existedNode := range nodes {
			if existedNode.IsManager() {
				managerNodes = append(managerNodes, existedNode)
			}
		}

		return managerNodes, nil
	}

	return nil, errors.New("add manager failed")
}

func loadOrCreateRaftID(db *bolt.DB) (uint64, error) {
	var raftID uint64
	tx, err := db.Begin(true)
	if err != nil {
		return raftID, err
	}
	defer tx.Commit()

	var (
		raftIDBukctName = []byte("raftnode")
		raftIDDataKey   = []byte("raftid")
	)
	raftIDBkt := tx.Bucket(raftIDBukctName)
	if raftIDBkt == nil {
		raftIDBkt, err = tx.CreateBucketIfNotExists(raftIDBukctName)
		if err != nil {
			return raftID, err
		}

		raftID = uint64(rand.Int63()) + 1
		if err := raftIDBkt.Put(raftIDDataKey, []byte(strconv.FormatUint(raftID, 10))); err != nil {
			return raftID, err
		}
		logrus.Infof("raft ID was not found create a new raftID %x", raftID)
		return raftID, nil
	} else {
		raftID_ := raftIDBkt.Get(raftIDDataKey)
		raftID, err = strconv.ParseUint(string(raftID_), 10, 64)
		if err != nil {
			return raftID, err
		}

		return raftID, nil
	}
}
<commit_msg>fix bug: panic when remove manager node<commit_after>package node

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/Dataman-Cloud/swan/src/agent"
	"github.com/Dataman-Cloud/swan/src/apiserver"
	"github.com/Dataman-Cloud/swan/src/config"
	"github.com/Dataman-Cloud/swan/src/event"
	"github.com/Dataman-Cloud/swan/src/manager"
	"github.com/Dataman-Cloud/swan/src/swancontext"
	"github.com/Dataman-Cloud/swan/src/types"
	"github.com/Dataman-Cloud/swan/src/utils/httpclient"

	"github.com/Sirupsen/logrus"
	"github.com/boltdb/bolt"
	"github.com/coreos/etcd/pkg/fileutil"
	"github.com/twinj/uuid"
	"golang.org/x/net/context"
)

const (
	NodeIDFileName    = "/ID"
	JoinRetryInterval = 5
)

type Node struct {
	ID                string
	agent             *agent.Agent     // hold reference to agent, take function when in agent mode
	manager           *manager.Manager // hold a instance of manager, make logic taking place
	ctx               context.Context
	joinRetryInterval time.Duration
	RaftID            uint64
	stopC             chan struct{}
}

func NewNode(config config.SwanConfig) (*Node, error) {
	nodeID, err := loadOrCreateNodeID(config)
	if err != nil {
		return nil, err
	}

	// init swanconfig instance
	config.NodeID = nodeID
	_ = swancontext.NewSwanContext(config, event.New())

	if !swancontext.IsManager() && !swancontext.IsAgent() {
		return nil, errors.New("node must be started with at least one role in [manager,agent]")
	}

	node := &Node{
		ID:                nodeID,
		joinRetryInterval: time.Second * JoinRetryInterval,
		stopC:             make(chan struct{}, 1),
	}

	err = os.MkdirAll(config.DataDir+"/"+nodeID, 0644)
	if err != nil {
		logrus.Errorf("os.MkdirAll got error: %s", err)
		return nil, err
	}

	db, err := bolt.Open(config.DataDir+"/"+nodeID+"/swan.db", 0644, nil)
	if err != nil {
		logrus.Errorf("Init bolt store failed:%s", err)
		return nil, err
	}

	raftID, err := loadOrCreateRaftID(db)
	if err != nil {
		logrus.Errorf("Init raft ID failed:%s", err)
		return nil, err
	}
	node.RaftID = raftID

	if swancontext.IsManager() {
		m, err := manager.New(db)
		if err != nil {
			return nil, err
		}
		node.manager = m
	}

	if swancontext.IsAgent() {
		a, err := agent.New()
		if err != nil {
			return nil, err
		}
		node.agent = a
	}

	nodeApi := &NodeApi{node}
	apiserver.Install(swancontext.Instance().ApiServer, nodeApi)

	return node, nil
}

func loadOrCreateNodeID(swanConfig config.SwanConfig) (string, error) {
	nodeIDFile := swanConfig.DataDir + NodeIDFileName
	if !fileutil.Exist(nodeIDFile) {
		os.MkdirAll(swanConfig.DataDir, 0700)

		nodeID := uuid.NewV4().String()
		idFile, err := os.OpenFile(nodeIDFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			return "", err
		}

		if _, err = idFile.WriteString(nodeID); err != nil {
			return "", err
		}

		logrus.Infof("starting swan node, ID file was not found started with  new ID: %s", nodeID)
		return nodeID, nil

	} else {
		idFile, err := os.Open(nodeIDFile)
		if err != nil {
			return "", err
		}

		nodeID, err := ioutil.ReadAll(idFile)
		if err != nil {
			return "", err
		}

		logrus.Infof("starting swan node, ID file was found started with ID: %s", string(nodeID))
		return string(nodeID), nil
	}
}

// node start from here
// - 1, start manager if needed
// - 2, start agent if needed
// - 3, agent join to managers if needed
// - 4, start the API server, both for agent and client
// - 5, enter loop, wait for error or ctx.Done
func (n *Node) Start(ctx context.Context) error {
	errChan := make(chan error, 1)

	swanConfig := swancontext.Instance().Config
	nodeInfo := types.Node{
		ID:                n.ID,
		AdvertiseAddr:     swanConfig.AdvertiseAddr,
		ListenAddr:        swanConfig.ListenAddr,
		RaftListenAddr:    swanConfig.RaftListenAddr,
		RaftAdvertiseAddr: swanConfig.RaftAdvertiseAddr,
		Role:              types.NodeRole(swanConfig.Mode),
		RaftID:            n.RaftID,
	}

	if swancontext.IsManager() {
		go func() {
			if swancontext.IsNewCluster() {
				errChan <- n.runManager(ctx, n.RaftID, []types.Node{nodeInfo}, true)
			} else {
				existedNodes, err := n.JoinAsManager(nodeInfo)
				if err != nil {
					errChan <- err
				}

				errChan <- n.runManager(ctx, n.RaftID, existedNodes, false)
			}
		}()
	}

	if swancontext.IsAgent() {
		go func() {
			errChan <- n.runAgent(ctx)
		}()

		go func() {
			err := n.JoinAsAgent(nodeInfo)
			if err != nil {
				errChan <- err
			}
		}()
	}

	go func() {
		errChan <- swancontext.Instance().ApiServer.Start()
	}()

	for {
		select {
		case err := <-errChan:
			return err
		case <-n.stopC:
			n.Stop()
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (n *Node) runAgent(ctx context.Context) error {
	agentCtx, cancel := context.WithCancel(ctx)
	n.agent.CancelFunc = cancel
	return n.agent.Start(agentCtx)
}

func (n *Node) runManager(ctx context.Context, raftID uint64, peers []types.Node, isNewCluster bool) error {
	managerCtx, cancel := context.WithCancel(ctx)
	n.manager.CancelFunc = cancel
	return n.manager.Start(managerCtx, raftID, peers, isNewCluster)
}

// node stop
func (n *Node) Stop() {
	if swancontext.IsAgent() {
		n.agent.Stop()
	}

	if swancontext.IsManager() {
		n.manager.Stop()
	}
}

func (n *Node) JoinAsAgent(nodeInfo types.Node) error {
	swanConfig := swancontext.Instance().Config
	if len(swanConfig.JoinAddrs) == 0 {
		return errors.New("start agent failed. Error: joinAddrs must be no empty")
	}

	for _, managerAddr := range swanConfig.JoinAddrs {
		registerAddr := "http://" + managerAddr + config.API_PREFIX + "/nodes"
		_, err := httpclient.NewDefaultClient().POST(context.TODO(), registerAddr, nil, nodeInfo, nil)
		if err != nil {
			logrus.Infof("register to %s got error: %s, try again after %d seconds", registerAddr, err.Error(), JoinRetryInterval)
		}

		if err == nil {
			logrus.Infof("agent register to manager success with managerAddr: %s", managerAddr)
			return nil
		}
	}

	time.Sleep(n.joinRetryInterval)
	n.JoinAsAgent(nodeInfo)
	return nil
}

func (n *Node) JoinAsManager(nodeInfo types.Node) ([]types.Node, error) {
	swanConfig := swancontext.Instance().Config
	if len(swanConfig.JoinAddrs) == 0 {
		return nil, errors.New("start agent failed. Error: joinAddrs must be no empty")
	}

	for _, managerAddr := range swanConfig.JoinAddrs {
		registerAddr := "http://" + managerAddr + config.API_PREFIX + "/nodes"
		resp, err := httpclient.NewDefaultClient().POST(context.TODO(), registerAddr, nil, nodeInfo, nil)
		if err != nil {
			logrus.Errorf("register to %s got error: %s", registerAddr, err.Error())
			continue
		}

		var nodes []types.Node
		if err := json.Unmarshal(resp, &nodes); err != nil {
			logrus.Errorf("register to %s got error: %s", registerAddr, err.Error())
			continue
		}

		var managerNodes []types.Node
		for _, existedNode := range nodes {
			if existedNode.IsManager() {
				managerNodes = append(managerNodes, existedNode)
			}
		}

		return managerNodes, nil
	}

	return nil, errors.New("add manager failed")
}

func loadOrCreateRaftID(db *bolt.DB) (uint64, error) {
	var raftID uint64
	tx, err := db.Begin(true)
	if err != nil {
		return raftID, err
	}
	defer tx.Commit()

	var (
		raftIDBukctName = []byte("raftnode")
		raftIDDataKey   = []byte("raftid")
	)
	raftIDBkt := tx.Bucket(raftIDBukctName)
	if raftIDBkt == nil {
		raftIDBkt, err = tx.CreateBucketIfNotExists(raftIDBukctName)
		if err != nil {
			return raftID, err
		}

		raftID = uint64(rand.Int63()) + 1
		if err := raftIDBkt.Put(raftIDDataKey, []byte(strconv.FormatUint(raftID, 10))); err != nil {
			return raftID, err
		}
		logrus.Infof("raft ID was not found create a new raftID %x", raftID)
		return raftID, nil
	} else {
		raftID_ := raftIDBkt.Get(raftIDDataKey)
		raftID, err = strconv.ParseUint(string(raftID_), 10, 64)
		if err != nil {
			return raftID, err
		}

		return raftID, nil
	}
}
<|endoftext|>


<commit_before>package metadata

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
)

// TODO: rename to SchemaMan Type?
type DatamanType string

// DatamanType is a method for describing the golang type in schema
// This allows us to treat everything as interfaces{} in most of the code yet
// still be in a strongly typed language

const (
	Document DatamanType = "document"
	String               = "string" // max len 4096
	Text                 = "text"
	// We should support converting anything to an int that doesn't lose data
	Int = "int"
	// TODO: int64
	// TODO: uint
	// TODO: uint64
	Bool = "bool"
	// TODO: actually implement
	DateTime = "datetime"
)

// TODO: have this register the type? Right now this assumes this is in-sync with field_type_internal.go (which is bad to do)
func (f DatamanType) ToFieldType() *FieldType {
	return &FieldType{
		Name:        "_" + string(f),
		DatamanType: f,
	}
}

// Normalize the given interface into what we want/expect
func (f DatamanType) Normalize(val interface{}) (interface{}, error) {
	switch f {
	case Document:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case map[string]interface{}:
			return typedVal, nil
		case string:
			mapVal := make(map[string]interface{})
			if err := json.Unmarshal([]byte(typedVal), mapVal); err == nil {
				return mapVal, nil
			} else {
				return nil, err
			}
		default:
			return nil, fmt.Errorf("Not a document")
		}

	case String:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case string:
			// TODO: default, code this out somewhere
			if len(typedVal) > 4096 {
				return nil, fmt.Errorf("String too long!")
			}
			return typedVal, nil
		default:
			return nil, fmt.Errorf("Not a string")
		}

	case Text:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case string:
			return typedVal, nil
		default:
			return nil, fmt.Errorf("Not text")
		}
	case Int:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		// TODO: remove? Or error if we would lose precision
		case int32:
			return int(typedVal), nil
		case int64:
			return int(typedVal), nil
		case int:
			return typedVal, nil
		case float64:
			return int(typedVal), nil
		case string:
			if typedVal == "" {
				return nil, nil
			} else {
				return strconv.ParseInt(typedVal, 10, 64)
			}
		default:
			return nil, fmt.Errorf("Unknown Int type: %s", reflect.TypeOf(val))
		}
	case Bool:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case bool:
			return typedVal, nil
		case string:
			return strconv.ParseBool(typedVal)
		default:
			return nil, fmt.Errorf("Not a bool")
		}
	// TODO: implement
	case DateTime:
		return nil, fmt.Errorf("DateTime currently unimplemented")
	}
	return nil, fmt.Errorf("Unknown type \"%s\" defined", f)
}

// TODO: have method which will reflect type to determine dataman type
// then we can have the datasources just call the method with the largest thing
// they can store in a given field type to determine the closest dataman_type
<commit_msg>Add TODO<commit_after>package metadata

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
)

// TODO: rename to SchemaMan Type?
type DatamanType string

// DatamanType is a method for describing the golang type in schema
// This allows us to treat everything as interfaces{} in most of the code yet
// still be in a strongly typed language

const (
	Document DatamanType = "document"
	String               = "string" // max len 4096
	Text                 = "text"
	// We should support converting anything to an int that doesn't lose data
	Int = "int"
	// TODO: int64
	// TODO: uint
	// TODO: uint64
	Bool = "bool"
	// TODO: actually implement
	DateTime = "datetime"
)

// TODO: have this register the type? Right now this assumes this is in-sync with field_type_internal.go (which is bad to do)
func (f DatamanType) ToFieldType() *FieldType {
	return &FieldType{
		Name:        "_" + string(f),
		DatamanType: f,
	}
}

// Normalize the given interface into what we want/expect
func (f DatamanType) Normalize(val interface{}) (interface{}, error) {
	switch f {
	case Document:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case map[string]interface{}:
			return typedVal, nil
		// TODO: put this behind a switch to enforce strictness
		case string:
			mapVal := make(map[string]interface{})
			if err := json.Unmarshal([]byte(typedVal), mapVal); err == nil {
				return mapVal, nil
			} else {
				return nil, err
			}
		default:
			return nil, fmt.Errorf("Not a document")
		}

	case String:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case string:
			// TODO: default, code this out somewhere
			if len(typedVal) > 4096 {
				return nil, fmt.Errorf("String too long!")
			}
			return typedVal, nil
		default:
			return nil, fmt.Errorf("Not a string")
		}

	case Text:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case string:
			return typedVal, nil
		default:
			return nil, fmt.Errorf("Not text")
		}
	case Int:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		// TODO: remove? Or error if we would lose precision
		case int32:
			return int(typedVal), nil
		case int64:
			return int(typedVal), nil
		case int:
			return typedVal, nil
		case float64:
			return int(typedVal), nil
		case string:
			if typedVal == "" {
				return nil, nil
			} else {
				return strconv.ParseInt(typedVal, 10, 64)
			}
		default:
			return nil, fmt.Errorf("Unknown Int type: %s", reflect.TypeOf(val))
		}
	case Bool:
		switch typedVal := val.(type) {
		case nil:
			return nil, nil
		case bool:
			return typedVal, nil
		case string:
			return strconv.ParseBool(typedVal)
		default:
			return nil, fmt.Errorf("Not a bool")
		}
	// TODO: implement
	case DateTime:
		return nil, fmt.Errorf("DateTime currently unimplemented")
	}
	return nil, fmt.Errorf("Unknown type \"%s\" defined", f)
}

// TODO: have method which will reflect type to determine dataman type
// then we can have the datasources just call the method with the largest thing
// they can store in a given field type to determine the closest dataman_type
<|endoftext|>


<commit_before>// Copyright (C) 2015 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package sync

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/syncthing/syncthing/lib/logger"
)

const (
	logThreshold = 100 * time.Millisecond
	shortWait    = 5 * time.Millisecond
	longWait     = 125 * time.Millisecond
)

func TestTypes(t *testing.T) {
	debug = false
	l.SetDebug("sync", false)

	if _, ok := NewMutex().(*sync.Mutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewRWMutex().(*sync.RWMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewWaitGroup().(*sync.WaitGroup); !ok {
		t.Error("Wrong type")
	}

	debug = true
	l.SetDebug("sync", true)

	if _, ok := NewMutex().(*loggedMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewRWMutex().(*loggedRWMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewWaitGroup().(*loggedWaitGroup); !ok {
		t.Error("Wrong type")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestMutex(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	mut := NewMutex()
	mut.Lock()
	clock.wind(shortWait)
	mut.Unlock()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	mut.Lock()
	clock.wind(longWait)
	mut.Unlock()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestRWMutex(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	mut := NewRWMutex()
	mut.Lock()
	clock.wind(shortWait)
	mut.Unlock()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	mut.Lock()
	clock.wind(longWait)
	mut.Unlock()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	// Testing rlocker logging
	wait := make(chan struct{})
	locking := make(chan struct{})

	mut.RLock()
	go func() {
		close(locking)
		mut.Lock()
		close(wait)
	}()

	<-locking
	clock.wind(longWait)
	mut.RUnlock()
	<-wait

	mut.Unlock()

	if len(messages) != 2 {
		t.Errorf("Unexpected message count")
	}
	if !strings.Contains(messages[1], "RUnlockers while locking:\nat sync") || !strings.Contains(messages[1], "sync_test.go:") {
		t.Error("Unexpected message")
	}

	// Testing multiple rlockers
	mut.RLock()
	mut.RLock()
	mut.RLock()
	_ = 1 // skip empty critical section check
	mut.RUnlock()
	mut.RUnlock()
	mut.RUnlock()

	debug = false
	l.SetDebug("sync", false)
}

func TestWaitGroup(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	wg := NewWaitGroup()
	wg.Add(1)
	waiting := make(chan struct{})

	go func() {
		<-waiting
		clock.wind(shortWait)
		wg.Done()
	}()

	close(waiting)
	wg.Wait()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	wg = NewWaitGroup()
	waiting = make(chan struct{})

	wg.Add(1)
	go func() {
		<-waiting
		clock.wind(longWait)
		wg.Done()
	}()

	close(waiting)
	wg.Wait()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestTimeoutCond(t *testing.T) {
	// WARNING this test relies heavily on threads not being stalled at particular points.
	// As such, it's pretty unstable on the build server. It has been left in as it still
	// exercises the deadlock detector, and one of the two things it tests is still functional.
	// See the comments in runLocks

	const (
		// Low values to avoid being intrusive in continuous testing. Can be
		// increased significantly for stress testing.
		iterations = 100
		routines   = 10

		timeMult = 2
	)

	c := NewTimeoutCond(NewMutex())

	// Start a routine to periodically broadcast on the cond.

	go func() {
		d := time.Duration(routines) * timeMult * time.Millisecond / 2
		t.Log("Broadcasting every", d)
		for i := 0; i < iterations; i++ {
			time.Sleep(d)

			c.L.Lock()
			c.Broadcast()
			c.L.Unlock()
		}
	}()

	// Start several routines that wait on it with different timeouts.

	var results [routines][2]int
	var wg sync.WaitGroup
	for i := 0; i < routines; i++ {
		i := i
		wg.Add(1)
		go func() {
			d := time.Duration(i) * timeMult * time.Millisecond
			t.Logf("Routine %d waits for %v\n", i, d)
			succ, fail := runLocks(t, iterations, c, d)
			results[i][0] = succ
			results[i][1] = fail
			wg.Done()
		}()
	}

	wg.Wait()

	// Print a table of routine number: successes, failures.

	for i, v := range results {
		t.Logf("%4d: %4d %4d\n", i, v[0], v[1])
	}
}

func runLocks(t *testing.T, iterations int, c *TimeoutCond, d time.Duration) (succ, fail int) {
	for i := 0; i < iterations; i++ {
		c.L.Lock()

		// The thread may be stalled, so we can't test the 'succeeded late' case reliably.
		// Therefore make sure that we start t0 before starting the timeout, and only test
		// the 'failed early' case.

		t0 := time.Now()
		w := c.SetupWait(d)

		res := w.Wait()
		waited := time.Since(t0)

		// Allow 20% slide in either direction, and a five milliseconds of
		// scheduling delay... In tweaking these it was clear that things
		// worked like the should, so if this becomes a spurious failure
		// kind of thing feel free to remove or give significantly more
		// slack.

		if !res && waited < d*8/10 {
			t.Errorf("Wait failed early, %v < %v", waited, d)
		}
		if res && waited > d*11/10+5*time.Millisecond {
			// Ideally this would be t.Errorf
			t.Logf("WARNING: Wait succeeded late, %v > %v. This is probably a thread scheduling issue", waited, d)
		}

		w.Stop()

		if res {
			succ++
		} else {
			fail++
		}
		c.L.Unlock()
	}
	return
}

type testClock struct {
	time time.Time
	mut  sync.Mutex
}

func newTestClock() *testClock {
	return &testClock{
		time: time.Now(),
	}
}

func (t *testClock) Now() time.Time {
	t.mut.Lock()
	now := t.time
	t.time = t.time.Add(time.Nanosecond)
	t.mut.Unlock()
	return now
}

func (t *testClock) wind(d time.Duration) {
	t.mut.Lock()
	t.time = t.time.Add(d)
	t.mut.Unlock()
}
<commit_msg>lib/sync: Cleanly fail instead of panic in tests (#6088)<commit_after>// Copyright (C) 2015 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

package sync

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/syncthing/syncthing/lib/logger"
)

const (
	logThreshold = 100 * time.Millisecond
	shortWait    = 5 * time.Millisecond
	longWait     = 125 * time.Millisecond
)

func TestTypes(t *testing.T) {
	debug = false
	l.SetDebug("sync", false)

	if _, ok := NewMutex().(*sync.Mutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewRWMutex().(*sync.RWMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewWaitGroup().(*sync.WaitGroup); !ok {
		t.Error("Wrong type")
	}

	debug = true
	l.SetDebug("sync", true)

	if _, ok := NewMutex().(*loggedMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewRWMutex().(*loggedRWMutex); !ok {
		t.Error("Wrong type")
	}

	if _, ok := NewWaitGroup().(*loggedWaitGroup); !ok {
		t.Error("Wrong type")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestMutex(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	mut := NewMutex()
	mut.Lock()
	clock.wind(shortWait)
	mut.Unlock()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	mut.Lock()
	clock.wind(longWait)
	mut.Unlock()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestRWMutex(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	mut := NewRWMutex()
	mut.Lock()
	clock.wind(shortWait)
	mut.Unlock()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	mut.Lock()
	clock.wind(longWait)
	mut.Unlock()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	// Testing rlocker logging
	wait := make(chan struct{})
	locking := make(chan struct{})

	mut.RLock()
	go func() {
		close(locking)
		mut.Lock()
		close(wait)
	}()

	<-locking
	clock.wind(longWait)
	mut.RUnlock()
	<-wait

	mut.Unlock()

	if len(messages) != 2 {
		t.Errorf("Unexpected message count")
	} else if !strings.Contains(messages[1], "RUnlockers while locking:\nat sync") || !strings.Contains(messages[1], "sync_test.go:") {
		t.Error("Unexpected message")
	}

	// Testing multiple rlockers
	mut.RLock()
	mut.RLock()
	mut.RLock()
	_ = 1 // skip empty critical section check
	mut.RUnlock()
	mut.RUnlock()
	mut.RUnlock()

	debug = false
	l.SetDebug("sync", false)
}

func TestWaitGroup(t *testing.T) {
	oldClock := defaultClock
	clock := newTestClock()
	defaultClock = clock
	defer func() { defaultClock = oldClock }()

	debug = true
	l.SetDebug("sync", true)
	threshold = logThreshold

	msgmut := sync.Mutex{}
	var messages []string

	l.AddHandler(logger.LevelDebug, func(_ logger.LogLevel, message string) {
		msgmut.Lock()
		messages = append(messages, message)
		msgmut.Unlock()
	})

	wg := NewWaitGroup()
	wg.Add(1)
	waiting := make(chan struct{})

	go func() {
		<-waiting
		clock.wind(shortWait)
		wg.Done()
	}()

	close(waiting)
	wg.Wait()

	if len(messages) > 0 {
		t.Errorf("Unexpected message count")
	}

	wg = NewWaitGroup()
	waiting = make(chan struct{})

	wg.Add(1)
	go func() {
		<-waiting
		clock.wind(longWait)
		wg.Done()
	}()

	close(waiting)
	wg.Wait()

	if len(messages) != 1 {
		t.Errorf("Unexpected message count")
	}

	debug = false
	l.SetDebug("sync", false)
}

func TestTimeoutCond(t *testing.T) {
	// WARNING this test relies heavily on threads not being stalled at particular points.
	// As such, it's pretty unstable on the build server. It has been left in as it still
	// exercises the deadlock detector, and one of the two things it tests is still functional.
	// See the comments in runLocks

	const (
		// Low values to avoid being intrusive in continuous testing. Can be
		// increased significantly for stress testing.
		iterations = 100
		routines   = 10

		timeMult = 2
	)

	c := NewTimeoutCond(NewMutex())

	// Start a routine to periodically broadcast on the cond.

	go func() {
		d := time.Duration(routines) * timeMult * time.Millisecond / 2
		t.Log("Broadcasting every", d)
		for i := 0; i < iterations; i++ {
			time.Sleep(d)

			c.L.Lock()
			c.Broadcast()
			c.L.Unlock()
		}
	}()

	// Start several routines that wait on it with different timeouts.

	var results [routines][2]int
	var wg sync.WaitGroup
	for i := 0; i < routines; i++ {
		i := i
		wg.Add(1)
		go func() {
			d := time.Duration(i) * timeMult * time.Millisecond
			t.Logf("Routine %d waits for %v\n", i, d)
			succ, fail := runLocks(t, iterations, c, d)
			results[i][0] = succ
			results[i][1] = fail
			wg.Done()
		}()
	}

	wg.Wait()

	// Print a table of routine number: successes, failures.

	for i, v := range results {
		t.Logf("%4d: %4d %4d\n", i, v[0], v[1])
	}
}

func runLocks(t *testing.T, iterations int, c *TimeoutCond, d time.Duration) (succ, fail int) {
	for i := 0; i < iterations; i++ {
		c.L.Lock()

		// The thread may be stalled, so we can't test the 'succeeded late' case reliably.
		// Therefore make sure that we start t0 before starting the timeout, and only test
		// the 'failed early' case.

		t0 := time.Now()
		w := c.SetupWait(d)

		res := w.Wait()
		waited := time.Since(t0)

		// Allow 20% slide in either direction, and a five milliseconds of
		// scheduling delay... In tweaking these it was clear that things
		// worked like the should, so if this becomes a spurious failure
		// kind of thing feel free to remove or give significantly more
		// slack.

		if !res && waited < d*8/10 {
			t.Errorf("Wait failed early, %v < %v", waited, d)
		}
		if res && waited > d*11/10+5*time.Millisecond {
			// Ideally this would be t.Errorf
			t.Logf("WARNING: Wait succeeded late, %v > %v. This is probably a thread scheduling issue", waited, d)
		}

		w.Stop()

		if res {
			succ++
		} else {
			fail++
		}
		c.L.Unlock()
	}
	return
}

type testClock struct {
	time time.Time
	mut  sync.Mutex
}

func newTestClock() *testClock {
	return &testClock{
		time: time.Now(),
	}
}

func (t *testClock) Now() time.Time {
	t.mut.Lock()
	now := t.time
	t.time = t.time.Add(time.Nanosecond)
	t.mut.Unlock()
	return now
}

func (t *testClock) wind(d time.Duration) {
	t.mut.Lock()
	t.time = t.time.Add(d)
	t.mut.Unlock()
}
<|endoftext|>


<commit_before>package middleware

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/DeedleFake/Go-PhysicsFS/physfs"
	"github.com/carbonsrv/carbon/modules/glue"
	"github.com/carbonsrv/carbon/modules/helpers"
	"github.com/carbonsrv/carbon/modules/scheduler"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pmylund/go-cache"
	"github.com/vifino/golua/lua"
	"github.com/vifino/luar"
	"net/http"
	"time"
)

// Cache
var kvstore *cache.Cache
var cbc *cache.Cache
var cfe *cache.Cache
var LDumper *lua.State

func cacheDump(file string) (string, error, bool) {
	data_tmp, found := cbc.Get(file)
	if found == false {
		data, err := fileRead(file)
		if err != nil {
			return "", err, false
		}
		res, err := bcdump(data)
		if err != nil {
			return "", err, true
		}
		cbc.Set(file, res, cache.DefaultExpiration)
		return res, nil, false
	} else {
		//debug("Using Bytecode-cache for " + file)
		return data_tmp.(string), nil, false
	}
}
func bcdump(data string) (string, error) {
	if LDumper.LoadString(data) != 0 {
		return "", errors.New(LDumper.ToString(-1))
	}
	defer LDumper.Pop(1)
	return LDumper.FDump(), nil
}

// FS
var filesystem http.FileSystem

func fileExists(file string) bool {
	data_tmp, found := cfe.Get(file)
	if found == false {
		exists := physfs.Exists(file)
		cfe.Set(file, exists, cache.DefaultExpiration)
		return exists
	} else {
		return data_tmp.(bool)
	}
}

func fileRead(file string) (string, error) {
	f, err := filesystem.Open(file)
	defer f.Close()
	if err != nil {
		return "", err
	}
	fi, err := f.Stat()
	if err != nil {
		return "", err
	}
	r := bufio.NewReader(f)
	buf := make([]byte, fi.Size())
	_, err = r.Read(buf)
	if err != nil {
		return "", err
	}
	return string(buf), err
}

// Preloader/Starter
var jobs int
var Preloaded chan *lua.State

func Preloader() {
	Preloaded = make(chan *lua.State, jobs)
	for {
		//fmt.Println("preloading")
		L := luar.Init()
		Bind(L)
		err := L.DoString(glue.MainGlue())
		if err != nil {
			panic(err)
		}
		err = L.DoString(glue.RouteGlue())
		if err != nil {
			panic(err)
		}
		Preloaded <- L
	}
}
func GetInstance() *lua.State {
	//fmt.Println("grabbing instance")
	L := <-Preloaded
	//fmt.Println("Done")
	return L
}

// Init
func Init(j int, cfe_new *cache.Cache, kvstore_new *cache.Cache) {
	cfe = cfe_new
	kvstore = kvstore_new
	jobs = j
	filesystem = physfs.FileSystem()
	cbc = cache.New(5*time.Minute, 30*time.Second) // Initialize cache with 5 minute lifetime and purge every 30 seconds
	LDumper = luar.Init()
}

// PHP-like lua scripts
func Lua() func(*gin.Context) {
	//LDumper := luar.Init()
	return func(context *gin.Context) {
		file := context.Request.URL.Path
		if fileExists(file) {
			//fmt.Println("start")
			L := GetInstance()
			//fmt.Println("after start")
			defer scheduler.Add(func() {
				L.Close()
			})
			//fmt.Println("after after start")
			BindContext(L, context)
			//fmt.Println("before cache")
			code, err, lerr := cacheDump(file)
			//fmt.Println("after cache")
			if err != nil {
				if lerr == false {
					context.Next()
					return
				} else {
					helpers.HTMLString(context, http.StatusInternalServerError, `
					<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in file `+context.Request.URL.Path+`:

						`+string(err.Error())+`

					`)
					context.Abort()
					return
				}
			}
			//fmt.Println("before loadbuffer")
			L.LoadBuffer(code, len(code), file) // This shouldn't error, was checked earlier.
			if L.Pcall(0, 0, 0) != 0 {          // != 0 means error in execution
				helpers.HTMLString(context, http.StatusInternalServerError, `
				<title>Runtime Error in `+context.Request.URL.Path+`</title>

Runtime Error in file `+context.Request.URL.Path+`:

					`+L.ToString(-1)+`

				`)
				context.Abort()
				return
			}
			/*L.DoString("return CONTENT_TO_RETURN")
			v := luar.CopyTableToMap(L, nil, -1)
			m := v.(map[string]interface{})
			i := int(m["code"].(float64))
			if err != nil {
				i = http.StatusOK
			}*/
			//helpers.HTMLString(context, i, m["content"].(string))
		} else {
			context.Next()
		}
	}
}

// Route creation by lua
func DLR_NS(bcode string, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) {
	/*code, err := bcdump(code)
	if err != nil {
		return func(*gin.Context) {}, err
	}*/
	return func(context *gin.Context) {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		defer scheduler.Add(func() {
			L.Close()
		})
		BindContext(L, context)
		//fmt.Println("before loadbuffer")
		/*if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}*/
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
	}, nil
}
func DLR_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but reuses states. Much faster. Higher memory use though, because more states.
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func DLRWS_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but for websockets. Not working because?!
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		conn, err := upgrader.Upgrade(context.Writer, context.Request, nil)
		if err != nil {
			fmt.Println("Websocket error: " + err.Error()) // silent error.
		}
		luar.Register(L, "ws", luar.Map{
			"BinaryMessage": websocket.BinaryMessage,
			"TextMessage":   websocket.TextMessage,
			//"read":          conn.ReadMessage,
			//"send":          conn.SendMessage,
			"read": (func() (int, string, error) {
				messageType, p, err := conn.ReadMessage()
				if err != nil {
					return -1, "", err
				}
				return messageType, string(p), nil
			}),
			"send": (func(t int, cnt string) error {
				return conn.WriteMessage(t, []byte(cnt))
			}),
		})
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			fmt.Println("Websocket error: " + L.ToString(-1))
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}
<commit_msg>A little bit inefficient, but it works!<commit_after>package middleware

import (
	"bufio"
	"errors"
	"fmt"
	"github.com/DeedleFake/Go-PhysicsFS/physfs"
	"github.com/carbonsrv/carbon/modules/glue"
	"github.com/carbonsrv/carbon/modules/helpers"
	"github.com/carbonsrv/carbon/modules/scheduler"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pmylund/go-cache"
	"github.com/vifino/golua/lua"
	"github.com/vifino/luar"
	"net/http"
	"time"
)

// Cache
var kvstore *cache.Cache
var cbc *cache.Cache
var cfe *cache.Cache
var LDumper *lua.State

func cacheDump(file string) (string, error, bool) {
	data_tmp, found := cbc.Get(file)
	if found == false {
		data, err := fileRead(file)
		if err != nil {
			return "", err, false
		}
		res, err := bcdump(data)
		if err != nil {
			return "", err, true
		}
		cbc.Set(file, res, cache.DefaultExpiration)
		return res, nil, false
	} else {
		//debug("Using Bytecode-cache for " + file)
		return data_tmp.(string), nil, false
	}
}
func bcdump(data string) (string, error) {
	if LDumper.LoadString(data) != 0 {
		return "", errors.New(LDumper.ToString(-1))
	}
	defer LDumper.Pop(1)
	return LDumper.FDump(), nil
}

// FS
var filesystem http.FileSystem

func fileExists(file string) bool {
	data_tmp, found := cfe.Get(file)
	if found == false {
		exists := physfs.Exists(file)
		cfe.Set(file, exists, cache.DefaultExpiration)
		return exists
	} else {
		return data_tmp.(bool)
	}
}

func fileRead(file string) (string, error) {
	f, err := filesystem.Open(file)
	defer f.Close()
	if err != nil {
		return "", err
	}
	fi, err := f.Stat()
	if err != nil {
		return "", err
	}
	r := bufio.NewReader(f)
	buf := make([]byte, fi.Size())
	_, err = r.Read(buf)
	if err != nil {
		return "", err
	}
	return string(buf), err
}

// Preloader/Starter
var jobs int
var Preloaded chan *lua.State

func Preloader() {
	Preloaded = make(chan *lua.State, jobs)
	for {
		//fmt.Println("preloading")
		L := luar.Init()
		Bind(L)
		err := L.DoString(glue.MainGlue())
		if err != nil {
			panic(err)
		}
		err = L.DoString(glue.RouteGlue())
		if err != nil {
			panic(err)
		}
		Preloaded <- L
	}
}
func GetInstance() *lua.State {
	//fmt.Println("grabbing instance")
	L := <-Preloaded
	//fmt.Println("Done")
	return L
}

// Init
func Init(j int, cfe_new *cache.Cache, kvstore_new *cache.Cache) {
	cfe = cfe_new
	kvstore = kvstore_new
	jobs = j
	filesystem = physfs.FileSystem()
	cbc = cache.New(5*time.Minute, 30*time.Second) // Initialize cache with 5 minute lifetime and purge every 30 seconds
	LDumper = luar.Init()
}

// PHP-like lua scripts
func Lua() func(*gin.Context) {
	//LDumper := luar.Init()
	return func(context *gin.Context) {
		file := context.Request.URL.Path
		if fileExists(file) {
			//fmt.Println("start")
			L := GetInstance()
			//fmt.Println("after start")
			defer scheduler.Add(func() {
				L.Close()
			})
			//fmt.Println("after after start")
			BindContext(L, context)
			//fmt.Println("before cache")
			code, err, lerr := cacheDump(file)
			//fmt.Println("after cache")
			if err != nil {
				if lerr == false {
					context.Next()
					return
				} else {
					helpers.HTMLString(context, http.StatusInternalServerError, `
					<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in file `+context.Request.URL.Path+`:

						`+string(err.Error())+`

					`)
					context.Abort()
					return
				}
			}
			//fmt.Println("before loadbuffer")
			L.LoadBuffer(code, len(code), file) // This shouldn't error, was checked earlier.
			if L.Pcall(0, 0, 0) != 0 {          // != 0 means error in execution
				helpers.HTMLString(context, http.StatusInternalServerError, `
				<title>Runtime Error in `+context.Request.URL.Path+`</title>

Runtime Error in file `+context.Request.URL.Path+`:

					`+L.ToString(-1)+`

				`)
				context.Abort()
				return
			}
			/*L.DoString("return CONTENT_TO_RETURN")
			v := luar.CopyTableToMap(L, nil, -1)
			m := v.(map[string]interface{})
			i := int(m["code"].(float64))
			if err != nil {
				i = http.StatusOK
			}*/
			//helpers.HTMLString(context, i, m["content"].(string))
		} else {
			context.Next()
		}
	}
}

// Route creation by lua
func DLR_NS(bcode string, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) {
	/*code, err := bcdump(code)
	if err != nil {
		return func(*gin.Context) {}, err
	}*/
	return func(context *gin.Context) {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		defer scheduler.Add(func() {
			L.Close()
		})
		BindContext(L, context)
		//fmt.Println("before loadbuffer")
		/*if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Syntax Error in `+context.Request.URL.Path+`</title>

Syntax Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}*/
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
	}, nil
}
func DLR_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but reuses states. Much faster. Higher memory use though, because more states.
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		if L.LoadBuffer(bcode, len(bcode), "route") != 0 {
			return func(context *gin.Context) {}, errors.New(L.ToString(-1))
		}
		L.PushValue(-1)
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			helpers.HTMLString(context, http.StatusInternalServerError, `
			<title>Runtime Error on `+context.Request.URL.Path+`</title>

Runtime Error in Lua Route on `+context.Request.URL.Path+`:

				`+L.ToString(-1)+`

			`)
			context.Abort()
			return
		}
		L.PushValue(-1)
		schan <- L
	}, nil
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func DLRWS_RUS(bcode string, instances int, dobind bool, vals map[string]interface{}) (func(*gin.Context), error) { // Same as above, but for websockets. Not working because?!
	insts := instances
	if instances < 0 {
		insts = 2
		if jobs/2 > 1 {
			insts = jobs
		}
	}
	schan := make(chan *lua.State, insts)
	for i := 0; i < jobs/2; i++ {
		L := GetInstance()
		if dobind {
			luar.Register(L, "", vals)
		}
		schan <- L
	}
	return func(context *gin.Context) {
		L := <-schan
		BindContext(L, context)
		conn, err := upgrader.Upgrade(context.Writer, context.Request, nil)
		if err != nil {
			fmt.Println("Websocket error: " + err.Error()) // silent error.
		}
		luar.Register(L, "ws", luar.Map{
			"BinaryMessage": websocket.BinaryMessage,
			"TextMessage":   websocket.TextMessage,
			//"read":          conn.ReadMessage,
			//"send":          conn.SendMessage,
			"read": (func() (int, string, error) {
				messageType, p, err := conn.ReadMessage()
				if err != nil {
					return -1, "", err
				}
				return messageType, string(p), nil
			}),
			"send": (func(t int, cnt string) error {
				return conn.WriteMessage(t, []byte(cnt))
			}),
		})
		L.LoadBuffer(bcode, len(bcode), "route")
		if L.Pcall(0, 0, 0) != 0 { // != 0 means error in execution
			fmt.Println("Websocket Lua error: " + L.ToString(-1))
			context.Abort()
			return
		}
		schan <- L
	}, nil
}
<|endoftext|>


<commit_before>package renter

import (
	"time"
)

// dropChunk will remove a worker from the responsibility of tracking a chunk.
func (w *worker) dropChunk(uc *unfinishedChunk) {
	uc.mu.Lock()
	uc.workersRemaining--
	uc.mu.Unlock()
	w.renter.managedReleaseIdleChunkPieces(uc)
	w.renter.heapWG.Done()
}

// dropUploadChunks will release all of the upload chunks that the worker has
// received.
func (w *worker) dropUploadChunks() {
	for i := 0; i < len(w.unprocessedChunks); i++ {
		w.dropChunk(w.unprocessedChunks[i])
	}
	w.unprocessedChunks = w.unprocessedChunks[:0]
}

// managedKillUploading will disable all uploading for the worker.
func (w *worker) managedKillUploading() {
	w.mu.Lock()
	w.dropUploadChunks()
	w.uploadTerminated = true
	w.mu.Unlock()
}

// managedNextUploadChunk will pull the next potential chunk out of the worker's
// work queue for uploading.
func (w *worker) managedNextUploadChunk() (nextChunk *unfinishedChunk, pieceIndex uint64) {
	w.mu.Lock()
	defer w.mu.Unlock()

	// Loop through the unprocessed chunks and find some work to do.
	for range w.unprocessedChunks {
		// Pull a chunk off of the unprocessed chunks stack.
		chunk := w.unprocessedChunks[0]
		w.unprocessedChunks = w.unprocessedChunks[1:]
		nextChunk, pieceIndex := w.processUploadChunk(chunk)
		if nextChunk != nil {
			return nextChunk, pieceIndex
		}
	}
	return nil, 0 // no work found
}

// managedQueueUploadChunk will take a chunk and add it to the worker's repair
// stack.
func (w *worker) managedQueueUploadChunk(uc *unfinishedChunk) {
	// Check that the worker is allowed to be uploading before grabbing the
	// worker lock.
	utility, exists := w.renter.hostContractor.ContractUtility(w.contract.ID)
	goodForUpload := exists && utility.GoodForUpload
	w.mu.Lock()
	defer w.mu.Unlock()

	if !goodForUpload || w.uploadTerminated || w.onUploadCooldown() {
		// The worker should not be uploading, remove the chunk.
		w.dropChunk(uc)
		return
	}
	w.unprocessedChunks = append(w.unprocessedChunks, uc)

	// Send a signal informing the work thread that there is work.
	select {
	case w.uploadChan <- struct{}{}:
	default:
	}
}

// managedUpload will perform some upload work.
func (w *worker) managedUpload(uc *unfinishedChunk, pieceIndex uint64) {
	// Open an editing connection to the host.
	e, err := w.renter.hostContractor.Editor(w.contract.ID, w.renter.tg.StopChan())
	if err != nil {
		w.renter.log.Debugln("Worker failed to acquire an editor:", err)
		w.mu.Lock()
		w.uploadFailed(uc, pieceIndex)
		w.mu.Unlock()
		return
	}
	defer e.Close()

	// Perform the upload, and update the failure stats based on the success of
	// the upload attempt.
	root, err := e.Upload(uc.physicalChunkData[pieceIndex])
	if err != nil {
		w.renter.log.Debugln("Worker failed to upload via the editor:", err)
		w.mu.Lock()
		w.uploadFailed(uc, pieceIndex)
		w.mu.Unlock()
		return
	}
	w.mu.Lock()
	w.uploadConsecutiveFailures = 0
	w.mu.Unlock()

	// Update the renter metadata.
	addr := e.Address()
	endHeight := e.EndHeight()
	id := w.renter.mu.Lock()
	uc.renterFile.mu.Lock()
	contract, exists := uc.renterFile.contracts[w.contract.ID]
	if !exists {
		contract = fileContract{
			ID:          w.contract.ID,
			IP:          addr,
			WindowStart: endHeight,
		}
	}
	contract.Pieces = append(contract.Pieces, pieceData{
		Chunk:      uc.index,
		Piece:      pieceIndex,
		MerkleRoot: root,
	})
	uc.renterFile.contracts[w.contract.ID] = contract
	w.renter.saveFile(uc.renterFile)
	uc.renterFile.mu.Unlock()
	w.renter.mu.Unlock(id)

	// Upload is complete. Update the state of the chunk and the renter's memory
	// available to reflect the completed upload.
	uc.mu.Lock()
	releaseSize := len(uc.physicalChunkData[pieceIndex])
	uc.piecesRegistered--
	uc.piecesCompleted++
	uc.physicalChunkData[pieceIndex] = nil
	uc.memoryReleased += uint64(releaseSize)
	uc.mu.Unlock()
	w.renter.memoryManager.Return(uint64(releaseSize))
	w.dropChunk(uc)
}

// onUploadCooldown returns true if the worker is on cooldown from failed
// uploads.
func (w *worker) onUploadCooldown() bool {
	requiredCooldown := uploadFailureCooldown
	for i := 0; i < w.uploadConsecutiveFailures && i < maxConsecutivePenalty; i++ {
		requiredCooldown *= 2
	}
	return time.Now().Before(w.uploadRecentFailure.Add(requiredCooldown))
}

// processUploadChunk will process a chunk from the worker chunk queue.
func (w *worker) processUploadChunk(uc *unfinishedChunk) (nextChunk *unfinishedChunk, pieceIndex uint64) {
	// Determine the usability value of this worker.
	utility, exists := w.renter.hostContractor.ContractUtility(w.contract.ID)
	if !exists {
		return nil, 0
	}
	goodForUpload := utility.GoodForUpload

	// Determine what sort of help this chunk needs.
	uc.mu.Lock()
	// Determine that this worker is not on cooldonw.
	_, candidateHost := uc.unusedHosts[w.hostPubKey.String()]
	chunkComplete := uc.piecesNeeded <= uc.piecesCompleted
	needsHelp := uc.piecesNeeded > uc.piecesCompleted+uc.piecesRegistered

	// If the chunk does not need help from this worker, release the chunk.
	if chunkComplete || !candidateHost || !goodForUpload || w.onUploadCooldown() {
		// This worker no longer needs to track this chunk.
		uc.mu.Unlock()
		w.dropChunk(uc)
		return nil, 0
	}

	// If the chunk needs help from this worker, find a piece to upload and
	// return the stats for that piece.
	index := 0
	if needsHelp {
		// Select a piece and mark that a piece has been selected.
		for i := 0; i < len(uc.pieceUsage); i++ {
			if !uc.pieceUsage[i] {
				index = i
				uc.pieceUsage[i] = true
				break
			}
		}
		delete(uc.unusedHosts, w.hostPubKey.String())
		uc.piecesRegistered++
		uc.mu.Unlock()
		return uc, uint64(index)
	}
	// Add this worker to the set of standby workers for this chunk.
	uc.workersStandby = append(uc.workersStandby, w)
	uc.mu.Unlock()

	// The chunk could need help from this worker, but only if other workers who
	// are performing uploads experience failures. Put this chunk on standby.
	return nil, 0
}

// uploadFailed is called if a worker failed to upload part of an unfinished
// chunk.
func (w *worker) uploadFailed(uc *unfinishedChunk, pieceIndex uint64) {
	// Mark the failure in the worker if the gateway says we are online. It's
	// not the worker's fault if we are offline.
	if w.renter.g.Online() {
		w.uploadRecentFailure = time.Now()
		w.uploadConsecutiveFailures++
	}
	uc.mu.Lock()
	uc.piecesRegistered--
	uc.pieceUsage[pieceIndex] = false
	uc.notifyStandbyWorkers()
	uc.mu.Unlock()
	w.dropChunk(uc)
	w.dropUploadChunks()
}
<commit_msg>fix up upload concurrency<commit_after>package renter

import (
	"time"
)

// managedDropChunk will remove a worker from the responsibility of tracking a chunk.
//
// This function is managed instead of static because it is against convention
// to be calling functions on other objects (in this case, the renter) while
// holding a lock.
func (w *worker) managedDropChunk(uc *unfinishedChunk) {
	uc.mu.Lock()
	uc.workersRemaining--
	uc.mu.Unlock()
	w.renter.managedReleaseIdleChunkPieces(uc)
	w.renter.heapWG.Done()
}

// managedDropUploadChunks will release all of the upload chunks that the worker
// has received.
func (w *worker) managedDropUploadChunks() {
	// Make a copy of the slice under lock, clear the slice, then drop the
	// chunks without a lock (managed function).
	var chunksToDrop []*unfinishedChunk
	w.mu.Lock()
	for i := 0; i < len(w.unprocessedChunks); i++ {
		chunksToDrop = append(chunksToDrop, w.unprocessedChunks[i])
	}
	w.unprocessedChunks = w.unprocessedChunks[:0]
	w.mu.Unlock()

	for i := 0; i < len(chunksToDrop); i++ {
		w.managedDropChunk(chunksToDrop[i])
	}
}

// managedKillUploading will disable all uploading for the worker.
func (w *worker) managedKillUploading() {
	// Mark the worker as disabled so that incoming chunks are rejected.
	w.mu.Lock()
	w.uploadTerminated = true
	w.mu.Unlock()

	// After the worker is marked as disabled, clear out all of the chunks.
	w.managedDropUploadChunks()
}

// managedNextUploadChunk will pull the next potential chunk out of the worker's
// work queue for uploading.
func (w *worker) managedNextUploadChunk() (nextChunk *unfinishedChunk, pieceIndex uint64) {
	// Loop through the unprocessed chunks and find some work to do.
	for {
		// Pull a chunk off of the unprocessed chunks stack.
		w.mu.Lock()
		if len(w.unprocessedChunks) <= 0 {
			w.mu.Unlock()
			break
		}
		chunk := w.unprocessedChunks[0]
		w.unprocessedChunks = w.unprocessedChunks[1:]
		w.mu.Unlock()

		// Process the chunk and return it if valid.
		nextChunk, pieceIndex := w.managedProcessUploadChunk(chunk)
		if nextChunk != nil {
			return nextChunk, pieceIndex
		}
	}
	return nil, 0 // no work found
}

// managedQueueUploadChunk will take a chunk and add it to the worker's repair
// stack.
func (w *worker) managedQueueUploadChunk(uc *unfinishedChunk) {
	// Check that the worker is allowed to be uploading before grabbing the
	// worker lock.
	utility, exists := w.renter.hostContractor.ContractUtility(w.contract.ID)
	goodForUpload := exists && utility.GoodForUpload
	w.mu.Lock()
	if !goodForUpload || w.uploadTerminated || w.onUploadCooldown() {
		// The worker should not be uploading, remove the chunk.
		w.mu.Unlock()
		w.managedDropChunk(uc)
		return
	}
	w.unprocessedChunks = append(w.unprocessedChunks, uc)
	w.mu.Unlock()

	// Send a signal informing the work thread that there is work.
	select {
	case w.uploadChan <- struct{}{}:
	default:
	}
}

// managedUpload will perform some upload work.
func (w *worker) managedUpload(uc *unfinishedChunk, pieceIndex uint64) {
	// Open an editing connection to the host.
	e, err := w.renter.hostContractor.Editor(w.contract.ID, w.renter.tg.StopChan())
	if err != nil {
		w.renter.log.Debugln("Worker failed to acquire an editor:", err)
		w.managedUploadFailed(uc, pieceIndex)
		return
	}
	defer e.Close()

	// Perform the upload, and update the failure stats based on the success of
	// the upload attempt.
	root, err := e.Upload(uc.physicalChunkData[pieceIndex])
	if err != nil {
		w.renter.log.Debugln("Worker failed to upload via the editor:", err)
		w.managedUploadFailed(uc, pieceIndex)
		return
	}
	w.mu.Lock()
	w.uploadConsecutiveFailures = 0
	w.mu.Unlock()

	// Update the renter metadata.
	addr := e.Address()
	endHeight := e.EndHeight()
	id := w.renter.mu.Lock()
	uc.renterFile.mu.Lock()
	contract, exists := uc.renterFile.contracts[w.contract.ID]
	if !exists {
		contract = fileContract{
			ID:          w.contract.ID,
			IP:          addr,
			WindowStart: endHeight,
		}
	}
	contract.Pieces = append(contract.Pieces, pieceData{
		Chunk:      uc.index,
		Piece:      pieceIndex,
		MerkleRoot: root,
	})
	uc.renterFile.contracts[w.contract.ID] = contract
	w.renter.saveFile(uc.renterFile)
	uc.renterFile.mu.Unlock()
	w.renter.mu.Unlock(id)

	// Upload is complete. Update the state of the chunk and the renter's memory
	// available to reflect the completed upload.
	uc.mu.Lock()
	releaseSize := len(uc.physicalChunkData[pieceIndex])
	uc.piecesRegistered--
	uc.piecesCompleted++
	uc.physicalChunkData[pieceIndex] = nil
	uc.memoryReleased += uint64(releaseSize)
	uc.mu.Unlock()
	w.renter.memoryManager.Return(uint64(releaseSize))
	w.managedDropChunk(uc)
}

// onUploadCooldown returns true if the worker is on cooldown from failed
// uploads.
func (w *worker) onUploadCooldown() bool {
	requiredCooldown := uploadFailureCooldown
	for i := 0; i < w.uploadConsecutiveFailures && i < maxConsecutivePenalty; i++ {
		requiredCooldown *= 2
	}
	return time.Now().Before(w.uploadRecentFailure.Add(requiredCooldown))
}

// managedProcessUploadChunk will process a chunk from the worker chunk queue.
func (w *worker) managedProcessUploadChunk(uc *unfinishedChunk) (nextChunk *unfinishedChunk, pieceIndex uint64) {
	// Determine the usability value of this worker.
	utility, exists := w.renter.hostContractor.ContractUtility(w.contract.ID)
	if !exists {
		return nil, 0
	}
	goodForUpload := utility.GoodForUpload

	// Determine what sort of help this chunk needs.
	w.mu.Lock()
	onCooldown := w.onUploadCooldown()
	w.mu.Unlock()
	uc.mu.Lock()
	_, candidateHost := uc.unusedHosts[w.hostPubKey.String()]
	chunkComplete := uc.piecesNeeded <= uc.piecesCompleted
	needsHelp := uc.piecesNeeded > uc.piecesCompleted+uc.piecesRegistered
	// If the chunk does not need help from this worker, release the chunk.
	if chunkComplete || !candidateHost || !goodForUpload || onCooldown {
		// This worker no longer needs to track this chunk.
		uc.mu.Unlock()
		w.managedDropChunk(uc)
		return nil, 0
	}

	// If the chunk needs help from this worker, find a piece to upload and
	// return the stats for that piece.
	index := 0
	if needsHelp {
		// Select a piece and mark that a piece has been selected.
		for i := 0; i < len(uc.pieceUsage); i++ {
			if !uc.pieceUsage[i] {
				index = i
				uc.pieceUsage[i] = true
				break
			}
		}
		delete(uc.unusedHosts, w.hostPubKey.String())
		uc.piecesRegistered++
		uc.mu.Unlock()
		return uc, uint64(index)
	}
	// Add this worker to the set of standby workers for this chunk.
	uc.workersStandby = append(uc.workersStandby, w)
	uc.mu.Unlock()

	// The chunk could need help from this worker, but only if other workers who
	// are performing uploads experience failures. Put this chunk on standby.
	return nil, 0
}

// managedUploadFailed is called if a worker failed to upload part of an unfinished
// chunk.
func (w *worker) managedUploadFailed(uc *unfinishedChunk, pieceIndex uint64) {
	// Mark the failure in the worker if the gateway says we are online. It's
	// not the worker's fault if we are offline.
	if w.renter.g.Online() {
		w.mu.Lock()
		w.uploadRecentFailure = time.Now()
		w.uploadConsecutiveFailures++
		w.mu.Unlock()
	}

	// Unregister the piece from the chunk and hunt for a replacement.
	uc.mu.Lock()
	uc.piecesRegistered--
	uc.pieceUsage[pieceIndex] = false
	uc.notifyStandbyWorkers()
	uc.mu.Unlock()

	// Drop this chunk from the worker. Because the worker is currently on
	// cooldown, drop all remaining chunks from the worker as well.
	w.managedDropChunk(uc)
	w.managedDropUploadChunks()
}
<|endoftext|>


<commit_before>package queue

import (
	"log"

	"github.com/streadway/amqp"
)

type QueueConnection struct {
	Connection *amqp.Connection
	Channel    *amqp.Channel

	ack  chan uint64
	nack chan uint64
}

func NewQueueConnection(amqpURI string) (*QueueConnection, error) {
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return nil, err
	}

	channel, err := connection.Channel()
	if err != nil {
		return nil, err
	}

	err = channel.Confirm(false)
	if err != nil {
		return nil, err
	}

	ack, nack := channel.NotifyConfirm(make(chan uint64, 1), make(chan uint64, 1))

	return &QueueConnection{
		Connection: connection,
		Channel:    channel,
		ack:        ack,
		nack:       nack,
	}, nil
}

func (c *QueueConnection) Close() error {
	err := c.Channel.Close()
	if err != nil {
		return err
	}

	return c.Connection.Close()
}

func (c *QueueConnection) Consume(queueName string) (<-chan amqp.Delivery, error) {
	return c.Channel.Consume(
		queueName,
		"",
		false, // autoAck
		false, // this won't be the sole consumer
		true,  // don't deliver messages from same connection
		false, // the broker owns when consumption can begin
		nil)   // arguments
}

func (c *QueueConnection) ExchangeDeclare(exchangeName string, exchangeType string) error {
	return c.Channel.ExchangeDeclare(
		exchangeName, // name of the exchange
		exchangeType, // type
		true,         // durable
		false,        // delete when complete
		false,        // internal
		false,        // noWait
		nil,          // arguments
	)
}

func (c *QueueConnection) QueueDeclare(queueName string) (amqp.Queue, error) {
	queue, err := c.Channel.QueueDeclare(
		queueName, // name of the queue
		true,      // durable
		false,     // delete when usused
		false,     // exclusive
		false,     // noWait
		nil)       // arguments
	if err != nil {
		return amqp.Queue{
			Name: queueName,
		}, err
	}

	return queue, nil
}

func (c *QueueConnection) BindQueueToExchange(queueName string, exchangeName string) error {
	return c.Channel.QueueBind(
		queueName,
		"#", // key to marshall with
		exchangeName,
		true, // noWait
		nil)  // arguments
}

func (c *QueueConnection) Publish(exchangeName string, routingKey string, contentType string, body string) error {
	defer publisherConfirm(c)

	return c.Channel.Publish(
		exchangeName, // publish to an exchange
		routingKey,   // routing to 0 or more queues
		false,        // mandatory
		false,        // immediate
		amqp.Publishing{
			Headers:         amqp.Table{},
			ContentType:     contentType,
			ContentEncoding: "",
			Body:            []byte(body),
			DeliveryMode:    amqp.Persistent,
			Priority:        0, // 0-9
		})
}

func publisherConfirm(c *QueueConnection) {
	select {
	case tag := <-c.ack:
		log.Println("Acknowledge message publish:", tag)
		err := c.Channel.Ack(tag, false)
		if err != nil {
			log.Fatal("Couldn't ack:", tag, err)
		}
	case tag := <-c.nack:
		log.Println("Couldn't acknowledge message publish:", tag)
		err := c.Channel.Nack(tag, false, true)
		if err != nil {
			log.Fatal("Couldn't nack:", tag, err)
		}
	}
}
<commit_msg>Remove publisher confirms for now<commit_after>package queue

import (
	"github.com/streadway/amqp"
)

type QueueConnection struct {
	Connection *amqp.Connection
	Channel    *amqp.Channel
}

func NewQueueConnection(amqpURI string) (*QueueConnection, error) {
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return nil, err
	}

	channel, err := connection.Channel()
	if err != nil {
		return nil, err
	}

	return &QueueConnection{
		Connection: connection,
		Channel:    channel,
	}, nil
}

func (c *QueueConnection) Close() error {
	err := c.Channel.Close()
	if err != nil {
		return err
	}

	return c.Connection.Close()
}

func (c *QueueConnection) Consume(queueName string) (<-chan amqp.Delivery, error) {
	return c.Channel.Consume(
		queueName,
		"",
		false, // autoAck
		false, // this won't be the sole consumer
		true,  // don't deliver messages from same connection
		false, // the broker owns when consumption can begin
		nil)   // arguments
}

func (c *QueueConnection) ExchangeDeclare(exchangeName string, exchangeType string) error {
	return c.Channel.ExchangeDeclare(
		exchangeName, // name of the exchange
		exchangeType, // type
		true,         // durable
		false,        // delete when complete
		false,        // internal
		false,        // noWait
		nil,          // arguments
	)
}

func (c *QueueConnection) QueueDeclare(queueName string) (amqp.Queue, error) {
	queue, err := c.Channel.QueueDeclare(
		queueName, // name of the queue
		true,      // durable
		false,     // delete when usused
		false,     // exclusive
		false,     // noWait
		nil)       // arguments
	if err != nil {
		return amqp.Queue{
			Name: queueName,
		}, err
	}

	return queue, nil
}

func (c *QueueConnection) BindQueueToExchange(queueName string, exchangeName string) error {
	return c.Channel.QueueBind(
		queueName,
		"#", // key to marshall with
		exchangeName,
		true, // noWait
		nil)  // arguments
}

func (c *QueueConnection) Publish(exchangeName string, routingKey string, contentType string, body string) error {
	return c.Channel.Publish(
		exchangeName, // publish to an exchange
		routingKey,   // routing to 0 or more queues
		false,        // mandatory
		false,        // immediate
		amqp.Publishing{
			Headers:         amqp.Table{},
			ContentType:     contentType,
			ContentEncoding: "",
			Body:            []byte(body),
			DeliveryMode:    amqp.Persistent,
			Priority:        0, // 0-9
		})
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
	"gopkg.in/yaml.v2"
)

var commandDashboards = cli.Command{
	Name: "dashboards",
	Subcommands: []cli.Command{
		{
			Name:  "generate",
			Usage: "Generate custom dashboard",
			Description: `
    A custom dashboard is registered from a yaml file..
    Requests "POST /api/v0/dashboards". See https://mackerel.io/ja/api-docs/entry/dashboards#create.
`,
			Action: doGenerateDashboards,
			Flags: []cli.Flag{
				cli.BoolFlag{Name: "print, p", Usage: "markdown is output in standard output."},
			},
		},
	},
}

type graphsConfig struct {
	Title           string             `yaml:"title"`
	URLPath         string             `yaml:"url_path"`
	GraphType       string             `yaml:"graph_type"`
	Height          int                `yaml:"height"`
	Width           int                `yaml:"width"`
	HostGraphFormat []*hostGraphFormat `yaml:"host_graphs"`
	GraphFormat     []*graphFormat     `yaml:"graphs"`
}

type hostGraphFormat struct {
	Headline   string   `yaml:"headline"`
	HostIDs    []string `yaml:"host_ids"`
	GraphNames []string `yaml:"graph_names"`
	Period     string   `yaml:"period"`
}
type graphFormat struct {
	Headline    string      `yaml:"headline"`
	ColumnCount int         `yaml:"column_count"`
	GraphDefs   []*graphDef `yaml:"graph_def"`
}
type graphDef struct {
	HostID      string `yaml:"host_id"`
	ServiceName string `yaml:"service_name"`
	RoleName    string `yaml:"role_name"`
	Query       string `yaml:"query"`
	GraphName   string `yaml:"graph_name"`
	Period      string `yaml:"period"`
	Stacked     bool   `yaml:"stacked"`
	Simplified  bool   `yaml:"simplified"`
}

func (g graphDef) isHostGraph() bool {
	return g.HostID != ""
}
func (g graphDef) isServiceGraph() bool {
	return g.ServiceName != "" && g.RoleName == ""
}
func (g graphDef) isRoleGraph() bool {
	return g.ServiceName != "" && g.RoleName != ""
}
func (g graphDef) isExpressionGraph() bool {
	return g.Query != ""
}

type baseGraph interface {
	getURL(string, bool) string
	getHeight() int
	getWidth() int
}

type hostGraph struct {
	HostID    string
	GraphType string
	Graph     string
	Period    string
	height    int
	width     int
}

func (h hostGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/hosts/%s"+extension, orgName, h.HostID))
	param := url.Values{}
	param.Add("graph", h.Graph)
	param.Add("period", h.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (h hostGraph) generateGraphString(orgName string) string {
	if h.GraphType == "iframe" {
		return makeIframeTag(orgName, h)
	}
	return makeImageMarkdown(orgName, h)
}
func (h hostGraph) getHeight() int {
	return h.height
}
func (h hostGraph) getWidth() int {
	return h.width
}

type serviceGraph struct {
	ServiceName string
	GraphType   string
	Graph       string
	Period      string
	height      int
	width       int
}

func (s serviceGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/services/%s"+extension, orgName, s.ServiceName))
	param := url.Values{}
	param.Add("graph", s.Graph)
	param.Add("period", s.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (s serviceGraph) generateGraphString(orgName string) string {
	if s.GraphType == "iframe" {
		return makeIframeTag(orgName, s)
	}
	return makeImageMarkdown(orgName, s)
}
func (s serviceGraph) getHeight() int {
	return s.height
}
func (s serviceGraph) getWidth() int {
	return s.width
}

type roleGraph struct {
	ServiceName string
	RoleName    string
	GraphType   string
	Graph       string
	Period      string
	Stacked     bool
	Simplified  bool
	height      int
	width       int
}

func (r roleGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/services/%s/%s"+extension, orgName, r.ServiceName, r.RoleName))
	param := url.Values{}
	param.Add("graph", r.Graph)
	param.Add("stacked", strconv.FormatBool(r.Stacked))
	param.Add("simplified", strconv.FormatBool(r.Simplified))
	param.Add("period", r.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (r roleGraph) generateGraphString(orgName string) string {
	if r.GraphType == "iframe" {
		return makeIframeTag(orgName, r)
	}
	return makeImageMarkdown(orgName, r)
}
func (r roleGraph) getHeight() int {
	return r.height
}
func (r roleGraph) getWidth() int {
	return r.width
}

type expressionGraph struct {
	Query     string
	GraphType string
	Period    string
	height    int
	width     int
}

func (e expressionGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/advanced-graph"+extension, orgName))
	param := url.Values{}
	param.Add("query", e.Query)
	param.Add("period", e.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (e expressionGraph) generateGraphString(orgName string) string {
	if e.GraphType == "iframe" {
		return makeIframeTag(orgName, e)
	}
	return makeImageMarkdown(orgName, e)
}
func (e expressionGraph) getHeight() int {
	return e.height
}
func (e expressionGraph) getWidth() int {
	return e.width
}

func makeIframeTag(orgName string, g baseGraph) string {
	return fmt.Sprintf(`<iframe src="%s" height="%d" width="%d" frameborder="0"></iframe>`, g.getURL(orgName, false), g.getHeight(), g.getWidth())
}

func makeImageMarkdown(orgName string, g baseGraph) string {
	return fmt.Sprintf("[![graph](%s)](%s)", g.getURL(orgName, true), g.getURL(orgName, true))
}

func doGenerateDashboards(c *cli.Context) error {
	conffile := c.GlobalString("conf")

	isStdout := c.Bool("print")

	argFilePath := c.Args()
	if len(argFilePath) < 1 {
		logger.Log("error", "at least one argumet is required.")
		cli.ShowCommandHelp(c, "generate")
		os.Exit(1)
	}

	buf, err := ioutil.ReadFile(argFilePath[0])
	logger.DieIf(err)

	yml := graphsConfig{}
	err = yaml.Unmarshal(buf, &yml)
	logger.DieIf(err)

	client := newMackerel(conffile)

	org, err := client.GetOrg()
	logger.DieIf(err)

	if yml.Title == "" {
		logger.Log("error", "title is required in yaml.")
		os.Exit(1)
	}
	if yml.URLPath == "" {
		logger.Log("error", "url_path is required in yaml.")
		os.Exit(1)
	}
	if yml.GraphType == "" {
		yml.GraphType = "iframe"
	}
	if yml.GraphType != "iframe" && yml.GraphType != "image" {
		logger.Log("error", "graph_type should be 'iframe' or 'image'.")
		os.Exit(1)
	}
	if yml.Height == 0 {
		yml.Height = 200
	}
	if yml.Width == 0 {
		yml.Width = 400
	}

	if yml.HostGraphFormat != nil && yml.GraphFormat != nil {
		logger.Log("error", "you cannot specify both 'graphs' and host_graphs'.")
	}

	var markdown string
	for _, h := range yml.HostGraphFormat {
		markdown += generateHostGraphsMarkdown(org.Name, h, yml.GraphType, yml.Height, yml.Width, client)
	}
	for _, g := range yml.GraphFormat {
		markdown += generateGraphsMarkdown(org.Name, g, yml.GraphType, yml.Height, yml.Width)
	}

	if isStdout {
		fmt.Println(markdown)
	} else {
		updateDashboard := &mackerel.Dashboard{
			Title:        yml.Title,
			BodyMarkDown: markdown,
			URLPath:      yml.URLPath,
		}

		dashboards, fetchError := client.FindDashboards()
		logger.DieIf(fetchError)

		dashboardID := ""
		for _, ds := range dashboards {
			if ds.URLPath == yml.URLPath {
				dashboardID = ds.ID
			}
		}

		if dashboardID == "" {
			_, createError := client.CreateDashboard(updateDashboard)
			logger.DieIf(createError)
		} else {
			_, updateError := client.UpdateDashboard(dashboardID, updateDashboard)
			logger.DieIf(updateError)
		}
	}

	return nil
}

func generateHostGraphsMarkdown(orgName string, hostGraphs *hostGraphFormat, graphType string, height int, width int, client *mackerel.Client) string {

	var markdown string

	if hostGraphs.Headline != "" {
		markdown += fmt.Sprintf("## %s\n", hostGraphs.Headline)
	}
	markdown += generateHostGraphsTableHeader(hostGraphs.HostIDs, client)

	if hostGraphs.Period == "" {
		hostGraphs.Period = "1h"
	}

	for _, graphName := range hostGraphs.GraphNames {
		currentColumnCount := 0
		for _, hostID := range hostGraphs.HostIDs {
			h := &hostGraph{
				hostID,
				graphType,
				graphName,
				hostGraphs.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName), len(hostGraphs.HostIDs))
			currentColumnCount++
			if currentColumnCount >= len(hostGraphs.HostIDs) {
				markdown += "|\n"
				currentColumnCount = 0
			}
		}
	}

	return markdown
}

func generateHostGraphsTableHeader(hostIDs []string, client *mackerel.Client) string {
	var header string
	for _, hostID := range hostIDs {
		host, err := client.FindHost(hostID)
		logger.DieIf(err)

		var hostName string
		if host.DisplayName != "" {
			hostName = host.DisplayName
		} else {
			hostName = host.Name
		}

		header += "|" + hostName
	}

	header += "|\n" + generateGraphTableHeader(len(hostIDs))

	return header
}

func generateGraphsMarkdown(orgName string, graphs *graphFormat, graphType string, height int, width int) string {

	var markdown string
	if graphs.ColumnCount == 0 {
		graphs.ColumnCount = 1
	}
	currentColumnCount := 0

	if graphs.Headline != "" {
		markdown += fmt.Sprintf("## %s\n", graphs.Headline)
	}
	markdown += generateGraphTableHeader(graphs.ColumnCount)

	for i, gd := range graphs.GraphDefs {
		var graphDefCount = 0
		if gd.isHostGraph() {
			graphDefCount++
		}
		if gd.isServiceGraph() {
			graphDefCount++
		}
		if gd.isRoleGraph() {
			graphDefCount++
		}
		if gd.isExpressionGraph() {
			graphDefCount++
		}
		if graphDefCount != 1 {
			logger.Log("error", "at least one between hostId, service_name and query is required.")
			os.Exit(1)
		}

		if gd.Period == "" {
			gd.Period = "1h"
		}

		if gd.isHostGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for host graph.")
				os.Exit(1)
			}

			h := &hostGraph{
				gd.HostID,
				graphType,
				gd.GraphName,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName), graphs.ColumnCount)
		}

		if gd.isServiceGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for service graph.")
				os.Exit(1)
			}

			h := &serviceGraph{
				gd.ServiceName,
				graphType,
				gd.GraphName,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName), graphs.ColumnCount)
		}

		if gd.isRoleGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for role graph.")
				os.Exit(1)
			}

			r := &roleGraph{
				gd.ServiceName,
				gd.RoleName,
				graphType,
				gd.GraphName,
				gd.Period,
				gd.Stacked,
				gd.Simplified,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, r.generateGraphString(orgName), graphs.ColumnCount)
		}

		if gd.isExpressionGraph() {
			e := &expressionGraph{
				gd.Query,
				graphType,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, e.generateGraphString(orgName), graphs.ColumnCount)
		}

		currentColumnCount++
		if currentColumnCount >= graphs.ColumnCount || i >= len(graphs.GraphDefs)-1 {
			markdown += "|\n"
			currentColumnCount = 0
		}
	}

	return markdown
}

func generateGraphTableHeader(confColumnCount int) string {
	return strings.Repeat("|:-:", confColumnCount) + "|\n"
}

func appendMarkdown(markdown string, addItem string, confColumnCount int) string {
	return markdown + "|" + addItem
}
<commit_msg>remove unused parameter<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/mackerelio/mackerel-client-go"
	"github.com/mackerelio/mkr/logger"
	"gopkg.in/yaml.v2"
)

var commandDashboards = cli.Command{
	Name: "dashboards",
	Subcommands: []cli.Command{
		{
			Name:  "generate",
			Usage: "Generate custom dashboard",
			Description: `
    A custom dashboard is registered from a yaml file..
    Requests "POST /api/v0/dashboards". See https://mackerel.io/ja/api-docs/entry/dashboards#create.
`,
			Action: doGenerateDashboards,
			Flags: []cli.Flag{
				cli.BoolFlag{Name: "print, p", Usage: "markdown is output in standard output."},
			},
		},
	},
}

type graphsConfig struct {
	Title           string             `yaml:"title"`
	URLPath         string             `yaml:"url_path"`
	GraphType       string             `yaml:"graph_type"`
	Height          int                `yaml:"height"`
	Width           int                `yaml:"width"`
	HostGraphFormat []*hostGraphFormat `yaml:"host_graphs"`
	GraphFormat     []*graphFormat     `yaml:"graphs"`
}

type hostGraphFormat struct {
	Headline   string   `yaml:"headline"`
	HostIDs    []string `yaml:"host_ids"`
	GraphNames []string `yaml:"graph_names"`
	Period     string   `yaml:"period"`
}
type graphFormat struct {
	Headline    string      `yaml:"headline"`
	ColumnCount int         `yaml:"column_count"`
	GraphDefs   []*graphDef `yaml:"graph_def"`
}
type graphDef struct {
	HostID      string `yaml:"host_id"`
	ServiceName string `yaml:"service_name"`
	RoleName    string `yaml:"role_name"`
	Query       string `yaml:"query"`
	GraphName   string `yaml:"graph_name"`
	Period      string `yaml:"period"`
	Stacked     bool   `yaml:"stacked"`
	Simplified  bool   `yaml:"simplified"`
}

func (g graphDef) isHostGraph() bool {
	return g.HostID != ""
}
func (g graphDef) isServiceGraph() bool {
	return g.ServiceName != "" && g.RoleName == ""
}
func (g graphDef) isRoleGraph() bool {
	return g.ServiceName != "" && g.RoleName != ""
}
func (g graphDef) isExpressionGraph() bool {
	return g.Query != ""
}

type baseGraph interface {
	getURL(string, bool) string
	getHeight() int
	getWidth() int
}

type hostGraph struct {
	HostID    string
	GraphType string
	Graph     string
	Period    string
	height    int
	width     int
}

func (h hostGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/hosts/%s"+extension, orgName, h.HostID))
	param := url.Values{}
	param.Add("graph", h.Graph)
	param.Add("period", h.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (h hostGraph) generateGraphString(orgName string) string {
	if h.GraphType == "iframe" {
		return makeIframeTag(orgName, h)
	}
	return makeImageMarkdown(orgName, h)
}
func (h hostGraph) getHeight() int {
	return h.height
}
func (h hostGraph) getWidth() int {
	return h.width
}

type serviceGraph struct {
	ServiceName string
	GraphType   string
	Graph       string
	Period      string
	height      int
	width       int
}

func (s serviceGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/services/%s"+extension, orgName, s.ServiceName))
	param := url.Values{}
	param.Add("graph", s.Graph)
	param.Add("period", s.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (s serviceGraph) generateGraphString(orgName string) string {
	if s.GraphType == "iframe" {
		return makeIframeTag(orgName, s)
	}
	return makeImageMarkdown(orgName, s)
}
func (s serviceGraph) getHeight() int {
	return s.height
}
func (s serviceGraph) getWidth() int {
	return s.width
}

type roleGraph struct {
	ServiceName string
	RoleName    string
	GraphType   string
	Graph       string
	Period      string
	Stacked     bool
	Simplified  bool
	height      int
	width       int
}

func (r roleGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/services/%s/%s"+extension, orgName, r.ServiceName, r.RoleName))
	param := url.Values{}
	param.Add("graph", r.Graph)
	param.Add("stacked", strconv.FormatBool(r.Stacked))
	param.Add("simplified", strconv.FormatBool(r.Simplified))
	param.Add("period", r.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (r roleGraph) generateGraphString(orgName string) string {
	if r.GraphType == "iframe" {
		return makeIframeTag(orgName, r)
	}
	return makeImageMarkdown(orgName, r)
}
func (r roleGraph) getHeight() int {
	return r.height
}
func (r roleGraph) getWidth() int {
	return r.width
}

type expressionGraph struct {
	Query     string
	GraphType string
	Period    string
	height    int
	width     int
}

func (e expressionGraph) getURL(orgName string, isImage bool) string {
	extension := ""
	if isImage {
		extension = ".png"
	}
	u, _ := url.Parse(fmt.Sprintf("https://mackerel.io/embed/orgs/%s/advanced-graph"+extension, orgName))
	param := url.Values{}
	param.Add("query", e.Query)
	param.Add("period", e.Period)
	u.RawQuery = param.Encode()
	return u.String()
}
func (e expressionGraph) generateGraphString(orgName string) string {
	if e.GraphType == "iframe" {
		return makeIframeTag(orgName, e)
	}
	return makeImageMarkdown(orgName, e)
}
func (e expressionGraph) getHeight() int {
	return e.height
}
func (e expressionGraph) getWidth() int {
	return e.width
}

func makeIframeTag(orgName string, g baseGraph) string {
	return fmt.Sprintf(`<iframe src="%s" height="%d" width="%d" frameborder="0"></iframe>`, g.getURL(orgName, false), g.getHeight(), g.getWidth())
}

func makeImageMarkdown(orgName string, g baseGraph) string {
	return fmt.Sprintf("[![graph](%s)](%s)", g.getURL(orgName, true), g.getURL(orgName, true))
}

func doGenerateDashboards(c *cli.Context) error {
	conffile := c.GlobalString("conf")

	isStdout := c.Bool("print")

	argFilePath := c.Args()
	if len(argFilePath) < 1 {
		logger.Log("error", "at least one argumet is required.")
		cli.ShowCommandHelp(c, "generate")
		os.Exit(1)
	}

	buf, err := ioutil.ReadFile(argFilePath[0])
	logger.DieIf(err)

	yml := graphsConfig{}
	err = yaml.Unmarshal(buf, &yml)
	logger.DieIf(err)

	client := newMackerel(conffile)

	org, err := client.GetOrg()
	logger.DieIf(err)

	if yml.Title == "" {
		logger.Log("error", "title is required in yaml.")
		os.Exit(1)
	}
	if yml.URLPath == "" {
		logger.Log("error", "url_path is required in yaml.")
		os.Exit(1)
	}
	if yml.GraphType == "" {
		yml.GraphType = "iframe"
	}
	if yml.GraphType != "iframe" && yml.GraphType != "image" {
		logger.Log("error", "graph_type should be 'iframe' or 'image'.")
		os.Exit(1)
	}
	if yml.Height == 0 {
		yml.Height = 200
	}
	if yml.Width == 0 {
		yml.Width = 400
	}

	if yml.HostGraphFormat != nil && yml.GraphFormat != nil {
		logger.Log("error", "you cannot specify both 'graphs' and host_graphs'.")
	}

	var markdown string
	for _, h := range yml.HostGraphFormat {
		markdown += generateHostGraphsMarkdown(org.Name, h, yml.GraphType, yml.Height, yml.Width, client)
	}
	for _, g := range yml.GraphFormat {
		markdown += generateGraphsMarkdown(org.Name, g, yml.GraphType, yml.Height, yml.Width)
	}

	if isStdout {
		fmt.Println(markdown)
	} else {
		updateDashboard := &mackerel.Dashboard{
			Title:        yml.Title,
			BodyMarkDown: markdown,
			URLPath:      yml.URLPath,
		}

		dashboards, fetchError := client.FindDashboards()
		logger.DieIf(fetchError)

		dashboardID := ""
		for _, ds := range dashboards {
			if ds.URLPath == yml.URLPath {
				dashboardID = ds.ID
			}
		}

		if dashboardID == "" {
			_, createError := client.CreateDashboard(updateDashboard)
			logger.DieIf(createError)
		} else {
			_, updateError := client.UpdateDashboard(dashboardID, updateDashboard)
			logger.DieIf(updateError)
		}
	}

	return nil
}

func generateHostGraphsMarkdown(orgName string, hostGraphs *hostGraphFormat, graphType string, height int, width int, client *mackerel.Client) string {

	var markdown string

	if hostGraphs.Headline != "" {
		markdown += fmt.Sprintf("## %s\n", hostGraphs.Headline)
	}
	markdown += generateHostGraphsTableHeader(hostGraphs.HostIDs, client)

	if hostGraphs.Period == "" {
		hostGraphs.Period = "1h"
	}

	for _, graphName := range hostGraphs.GraphNames {
		currentColumnCount := 0
		for _, hostID := range hostGraphs.HostIDs {
			h := &hostGraph{
				hostID,
				graphType,
				graphName,
				hostGraphs.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName))
			currentColumnCount++
			if currentColumnCount >= len(hostGraphs.HostIDs) {
				markdown += "|\n"
				currentColumnCount = 0
			}
		}
	}

	return markdown
}

func generateHostGraphsTableHeader(hostIDs []string, client *mackerel.Client) string {
	var header string
	for _, hostID := range hostIDs {
		host, err := client.FindHost(hostID)
		logger.DieIf(err)

		var hostName string
		if host.DisplayName != "" {
			hostName = host.DisplayName
		} else {
			hostName = host.Name
		}

		header += "|" + hostName
	}

	header += "|\n" + generateGraphTableHeader(len(hostIDs))

	return header
}

func generateGraphsMarkdown(orgName string, graphs *graphFormat, graphType string, height int, width int) string {

	var markdown string
	if graphs.ColumnCount == 0 {
		graphs.ColumnCount = 1
	}
	currentColumnCount := 0

	if graphs.Headline != "" {
		markdown += fmt.Sprintf("## %s\n", graphs.Headline)
	}
	markdown += generateGraphTableHeader(graphs.ColumnCount)

	for i, gd := range graphs.GraphDefs {
		var graphDefCount = 0
		if gd.isHostGraph() {
			graphDefCount++
		}
		if gd.isServiceGraph() {
			graphDefCount++
		}
		if gd.isRoleGraph() {
			graphDefCount++
		}
		if gd.isExpressionGraph() {
			graphDefCount++
		}
		if graphDefCount != 1 {
			logger.Log("error", "at least one between hostId, service_name and query is required.")
			os.Exit(1)
		}

		if gd.Period == "" {
			gd.Period = "1h"
		}

		if gd.isHostGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for host graph.")
				os.Exit(1)
			}

			h := &hostGraph{
				gd.HostID,
				graphType,
				gd.GraphName,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName))
		}

		if gd.isServiceGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for service graph.")
				os.Exit(1)
			}

			h := &serviceGraph{
				gd.ServiceName,
				graphType,
				gd.GraphName,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, h.generateGraphString(orgName))
		}

		if gd.isRoleGraph() {
			if gd.GraphName == "" {
				logger.Log("error", "graph_name is required for role graph.")
				os.Exit(1)
			}

			r := &roleGraph{
				gd.ServiceName,
				gd.RoleName,
				graphType,
				gd.GraphName,
				gd.Period,
				gd.Stacked,
				gd.Simplified,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, r.generateGraphString(orgName))
		}

		if gd.isExpressionGraph() {
			e := &expressionGraph{
				gd.Query,
				graphType,
				gd.Period,
				height,
				width,
			}
			markdown = appendMarkdown(markdown, e.generateGraphString(orgName))
		}

		currentColumnCount++
		if currentColumnCount >= graphs.ColumnCount || i >= len(graphs.GraphDefs)-1 {
			markdown += "|\n"
			currentColumnCount = 0
		}
	}

	return markdown
}

func generateGraphTableHeader(confColumnCount int) string {
	return strings.Repeat("|:-:", confColumnCount) + "|\n"
}

func appendMarkdown(markdown string, addItem string) string {
	return markdown + "|" + addItem
}
<|endoftext|>


<commit_before>package data

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"math"
)

type Block struct {
	/*
	   magic - 6
	   block_len - 2
	   num_links - 2
	   offset - 4
	   content_len - 4
	   links - num_links * 768
	   content - content_len
	   padding
	*/
	magic       [6]byte
	block_len   int8
	num_links   int8
	offset      int32
	content_len int32
	links       []Link
	content     []byte
	padding     []byte
}

func NewBlock() *Block {
	return &Block{}
}

func (blk *Block) Create(block_size int, links []Link, content []byte) {

	blk.magic = [6]byte{0xF0, 0x07, 0xDA, 0x7A, '\r', '\n'}
	blk.block_len = int8(math.Log2(float64(block_size)))
	blk.num_links = int8(len(links))
	blk.offset = int32(blk.num_links)*96 + 16
	blk.content_len = int32(len(content))

	blk.links = links
	blk.content = content

	padding_len := block_size - 16 - int(blk.content_len) - (int(blk.num_links) * 96)

	padding := make([]byte, padding_len)
	_, err := rand.Read(padding)

	if err != nil {
		fmt.Println(err)
	}
}

func (blk *Block) GetBytes() []byte {
	var bin_buf bytes.Buffer
	binary.Write(&bin_buf, binary.BigEndian, blk)

	return bin_buf.Bytes()
}

func (blk *Block) GetKey() []byte {

	bin_blk := blk.GetBytes()

	// Create hash and feed data into ongoing hash
	hash := sha256.New()
	hash.Write(bin_blk)

	// Return the hash and append to nil string
	return hash.Sum(nil)
}

func (blk *Block) Encrypt(key []byte) (ciphertext []byte) {

	plaintext := blk.GetBytes()

	block, err := aes.NewCipher(key[16:])
	if err != nil {
		panic(err)
	}

	ciphertext = make([]byte, len(plaintext))
	iv := key[:16]

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, plaintext)

	return
}

func (blk *Block) ContentSize() int32 {
	return blk.content_len
}

func (blk *Block) Offset() int32 {
	return blk.offset
}

func (blk *Block) NumLinks() int8 {
	return blk.num_links
}
<commit_msg>Cleans up naming.<commit_after>package data

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"math"
)

type Block struct {
	/*
	   magic - 6
	   block_len - 2
	   num_links - 2
	   offset - 4
	   content_len - 4
	   links - num_links * 768
	   content - content_len
	   padding
	*/
	magic       [6]byte
	block_len   int8
	num_links   int8
	offset      int32
	content_len int32
	links       []Link
	content     []byte
	padding     []byte
}

func NewBlock() *Block {
	return &Block{}
}

func (blk *Block) Create(block_size int, links []Link, content []byte) {

	blk.magic = [6]byte{0xF0, 0x07, 0xDA, 0x7A, '\r', '\n'}
	blk.block_len = int8(math.Log2(float64(block_size)))
	blk.num_links = int8(len(links))
	blk.offset = int32(blk.num_links)*96 + 16
	blk.content_len = int32(len(content))

	blk.links = links
	blk.content = content

	padding_len := block_size - 16 - int(blk.content_len) - (int(blk.num_links) * 96)

	padding := make([]byte, padding_len)
	_, err := rand.Read(padding)

	if err != nil {
		fmt.Println(err)
	}
}

func (blk *Block) GetBytes() []byte {
	var bin_buf bytes.Buffer
	binary.Write(&bin_buf, binary.BigEndian, blk)

	return bin_buf.Bytes()
}

func (blk *Block) GetHash() []byte {

	bin_blk := blk.GetBytes()

	// Create hash and feed data into ongoing hash
	hash := sha256.New()
	hash.Write(bin_blk)

	// Return the hash and append to nil string
	return hash.Sum(nil)
}

func (blk *Block) Encrypt(hash []byte) (ciphertext []byte) {

	plaintext := blk.GetBytes()
	key := hash[16:]
	iv := hash[:16]

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	ciphertext = make([]byte, len(plaintext))

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, plaintext)

	return
}

func (blk *Block) ContentSize() int32 {
	return blk.content_len
}

func (blk *Block) Offset() int32 {
	return blk.offset
}

func (blk *Block) NumLinks() int8 {
	return blk.num_links
}
<|endoftext|>