<commit_before>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
SkipList* is a data structure combining elements of both a skiplist
and the bottom half of a Y-fast trie.  The idea is that you can quickly
search for a sub branch of the skip list and that this branch can
fit entirely within cache, thereby improving the performance characteristics
over a standard skip list.  This also keeps any memcopy operation limited
to O(log M) where M is the size of the desired universe.

Performance vs standard skip list.
BenchmarkInsert-8	 2000000	       949 ns/op
BenchmarkGet-8	 3000000	       516 ns/op
BenchmarkDelete-8	 3000000	       499 ns/op
BenchmarkStarInsert-8	 3000000	       453 ns/op
BenchmarkStarGet-8	 3000000	       524 ns/op
BenchmarkStarDelete-8	 3000000	       469 ns/op
*/
package skip

import "log"

func init() {
	log.Printf(`I HATE THIS.`)
}

// SkipList* implements all methods of a standard skip list but attempts
// to improve performance by ensuring cache locality.
type SkipStarList struct {
	ary uint8
	num uint64
	sl  *SkipList
}

// entryBundle is a helper struct used to define the nodes that
// can be inserted into the SkipList*.
type entryBundle struct {
	key     uint64
	entries Entries
}

// Key will return the key associated with this entity bundle.
// This is required by the Entry interface.
func (eb *entryBundle) Key() uint64 {
	return eb.key
}

func newEntryBundle(key uint64, size uint8) *entryBundle {
	return &entryBundle{
		key:     key,
		entries: make(Entries, 0, size),
	}
}

func (ssl *SkipStarList) init(ifc interface{}) {
	switch ifc.(type) {
	case uint8:
		ssl.ary = 8
	case uint16:
		ssl.ary = 16
	case uint32:
		ssl.ary = 32
	case uint64, uint:
		ssl.ary = 64
	}
	ssl.sl = New(ifc)
}

func (ssl *SkipStarList) getNormalizedKey(key uint64) uint64 {
	key = key/uint64(ssl.ary) + 1
	return key * uint64(ssl.ary)
}

func (ssl *SkipStarList) insert(entry Entry) Entry {
	key := ssl.getNormalizedKey(entry.Key())
	eb, ok := ssl.sl.Get(entry.Key())[0].(*entryBundle)
	if !ok {
		eb = newEntryBundle(key, ssl.ary)
		ssl.sl.Insert(eb)
	}

	e := eb.entries.insert(entry)
	if e == nil {
		ssl.num++
	}
	return e
}

// Insert will insert the provded entries into the SkipList*.  Any
// existing entry with a matching key will be overwritten.  The returned
// list of a list of entries that were overwritten, in order.  A nil
// will be in the in-order position for any non-overwritten entries.
func (ssl *SkipStarList) Insert(entries ...Entry) Entries {
	overwritten := make(Entries, 0, len(entries))
	for _, e := range entries {
		overwritten = append(overwritten, ssl.insert(e))
	}

	return overwritten
}

func (ssl *SkipStarList) get(key uint64) Entry {
	normalizedKey := ssl.getNormalizedKey(key)
	eb, ok := ssl.sl.Get(normalizedKey)[0].(*entryBundle)
	if ok {
		return eb.entries.get(key)
	}
	return nil
}

// Get will return a list of entries associated with the provided keys.
// A nil will be returned for any key not found.
func (ssl *SkipStarList) Get(keys ...uint64) Entries {
	entries := make(Entries, 0, len(keys))
	for _, key := range keys {
		entries = append(entries, ssl.get(key))
	}

	return entries
}

func (ssl *SkipStarList) delete(key uint64) Entry {
	normalizedKey := ssl.getNormalizedKey(key)
	eb, ok := ssl.sl.Get(normalizedKey)[0].(*entryBundle)
	if !ok {
		return nil
	}

	deleted := eb.entries.delete(key)
	if deleted != nil {
		ssl.num--
		if len(eb.entries) == 0 {
			ssl.sl.Delete(eb.key)
		}
	}

	return deleted
}

// Delete will remove the provided keys from the SkipList* and
// return a list of entries that were deleted.
func (ssl *SkipStarList) Delete(keys ...uint64) Entries {
	deleted := make(Entries, 0, len(keys))
	for _, key := range keys {
		deleted = append(deleted, ssl.delete(key))
	}

	return deleted
}

func (ssl *SkipStarList) iter(key uint64) *starIterator {
	normalizedKey := ssl.getNormalizedKey(key)
	iter := ssl.sl.Iter(normalizedKey)
	if !iter.Next() {
		return &starIterator{
			index: iteratorExhausted,
		}
	}

	eb := iter.Value().(*entryBundle)
	return &starIterator{
		index:   eb.entries.search(key) - 1,
		entries: eb.entries,
		iter:    iter,
	}
}

// Iter will return an iterator that will visit every value
// equal to or greater than the provided key.
func (ssl *SkipStarList) Iter(key uint64) Iterator {
	return ssl.iter(key)
}

// Len returns the number of items in the SkipList*.
func (ssl *SkipStarList) Len() uint64 {
	return ssl.num
}

// NewStar will allocate, initialize, and return a new SkipStarList.
// The Skip* list has an node size defined by the provided interface
// parameter.  This parameter must be a uint type (uint8, uint16, etc).
func NewStar(ifc interface{}) *SkipStarList {
	ssl := &SkipStarList{}
	ssl.init(ifc)
	return ssl
}
<commit_msg>Renamed SkipStarList to SkipListStar.<commit_after>/*
Copyright 2014 Workiva, LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
SkipList* is a data structure combining elements of both a skiplist
and the bottom half of a Y-fast trie.  The idea is that you can quickly
search for a sub branch of the skip list and that this branch can
fit entirely within cache, thereby improving the performance characteristics
over a standard skip list.  This also keeps any memcopy operation limited
to O(log M) where M is the size of the desired universe.

Performance vs standard skip list.
BenchmarkInsert-8	 2000000	       949 ns/op
BenchmarkGet-8	 3000000	       516 ns/op
BenchmarkDelete-8	 3000000	       499 ns/op
BenchmarkStarInsert-8	 3000000	       453 ns/op
BenchmarkStarGet-8	 3000000	       524 ns/op
BenchmarkStarDelete-8	 3000000	       469 ns/op
*/
package skip

import "log"

func init() {
	log.Printf(`I HATE THIS.`)
}

// SkipList* implements all methods of a standard skip list but attempts
// to improve performance by ensuring cache locality.
type SkipListStar struct {
	ary uint8
	num uint64
	sl  *SkipList
}

// entryBundle is a helper struct used to define the nodes that
// can be inserted into the SkipList*.
type entryBundle struct {
	key     uint64
	entries Entries
}

// Key will return the key associated with this entity bundle.
// This is required by the Entry interface.
func (eb *entryBundle) Key() uint64 {
	return eb.key
}

func newEntryBundle(key uint64, size uint8) *entryBundle {
	return &entryBundle{
		key:     key,
		entries: make(Entries, 0, size),
	}
}

func (ssl *SkipListStar) init(ifc interface{}) {
	switch ifc.(type) {
	case uint8:
		ssl.ary = 8
	case uint16:
		ssl.ary = 16
	case uint32:
		ssl.ary = 32
	case uint64, uint:
		ssl.ary = 64
	}
	ssl.sl = New(ifc)
}

func (ssl *SkipListStar) getNormalizedKey(key uint64) uint64 {
	key = key/uint64(ssl.ary) + 1
	return key * uint64(ssl.ary)
}

func (ssl *SkipListStar) insert(entry Entry) Entry {
	key := ssl.getNormalizedKey(entry.Key())
	eb, ok := ssl.sl.Get(entry.Key())[0].(*entryBundle)
	if !ok {
		eb = newEntryBundle(key, ssl.ary)
		ssl.sl.Insert(eb)
	}

	e := eb.entries.insert(entry)
	if e == nil {
		ssl.num++
	}
	return e
}

// Insert will insert the provded entries into the SkipList*.  Any
// existing entry with a matching key will be overwritten.  The returned
// list of a list of entries that were overwritten, in order.  A nil
// will be in the in-order position for any non-overwritten entries.
func (ssl *SkipListStar) Insert(entries ...Entry) Entries {
	overwritten := make(Entries, 0, len(entries))
	for _, e := range entries {
		overwritten = append(overwritten, ssl.insert(e))
	}

	return overwritten
}

func (ssl *SkipListStar) get(key uint64) Entry {
	normalizedKey := ssl.getNormalizedKey(key)
	eb, ok := ssl.sl.Get(normalizedKey)[0].(*entryBundle)
	if ok {
		return eb.entries.get(key)
	}
	return nil
}

// Get will return a list of entries associated with the provided keys.
// A nil will be returned for any key not found.
func (ssl *SkipListStar) Get(keys ...uint64) Entries {
	entries := make(Entries, 0, len(keys))
	for _, key := range keys {
		entries = append(entries, ssl.get(key))
	}

	return entries
}

func (ssl *SkipListStar) delete(key uint64) Entry {
	normalizedKey := ssl.getNormalizedKey(key)
	eb, ok := ssl.sl.Get(normalizedKey)[0].(*entryBundle)
	if !ok {
		return nil
	}

	deleted := eb.entries.delete(key)
	if deleted != nil {
		ssl.num--
		if len(eb.entries) == 0 {
			ssl.sl.Delete(eb.key)
		}
	}

	return deleted
}

// Delete will remove the provided keys from the SkipList* and
// return a list of entries that were deleted.
func (ssl *SkipListStar) Delete(keys ...uint64) Entries {
	deleted := make(Entries, 0, len(keys))
	for _, key := range keys {
		deleted = append(deleted, ssl.delete(key))
	}

	return deleted
}

func (ssl *SkipListStar) iter(key uint64) *starIterator {
	normalizedKey := ssl.getNormalizedKey(key)
	iter := ssl.sl.Iter(normalizedKey)
	if !iter.Next() {
		return &starIterator{
			index: iteratorExhausted,
		}
	}

	eb := iter.Value().(*entryBundle)
	return &starIterator{
		index:   eb.entries.search(key) - 1,
		entries: eb.entries,
		iter:    iter,
	}
}

// Iter will return an iterator that will visit every value
// equal to or greater than the provided key.
func (ssl *SkipListStar) Iter(key uint64) Iterator {
	return ssl.iter(key)
}

// Len returns the number of items in the SkipList*.
func (ssl *SkipListStar) Len() uint64 {
	return ssl.num
}

// NewStar will allocate, initialize, and return a new SkipListStar.
// The Skip* list has an node size defined by the provided interface
// parameter.  This parameter must be a uint type (uint8, uint16, etc).
func NewStar(ifc interface{}) *SkipListStar {
	ssl := &SkipListStar{}
	ssl.init(ifc)
	return ssl
}
<|endoftext|>


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

package x509

import (
	"bytes"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"strings"

	"github.com/zigbee-alliance/distributed-compliance-ledger/x/pki/types"
)

type X509Certificate struct {
	Issuer         string
	SerialNumber   string
	Subject        string
	SubjectKeyID   string
	AuthorityKeyID string
	Certificate    *x509.Certificate
}

func DecodeX509Certificate(pemCertificate string) (*X509Certificate, error) {
	block, _ := pem.Decode([]byte(pemCertificate))
	if block == nil {
		return nil, types.NewErrInvalidCertificate("Could not decode pem certificate")
	}

	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, types.NewErrInvalidCertificate(fmt.Sprintf("Could not parse certificate: %v", err.Error()))
	}

	oldVID := []byte("1.3.6.1.4.1.37244.2.1")
	oldPID := []byte("1.3.6.1.4.1.37244.2.2")

	newVIDKey := []byte("vid=0x")
	newPIDKey := []byte("pid=0x")

	issuer := cert.Issuer.String()
	issuer = formatOID([]byte(issuer), oldVID, newVIDKey)
	issuer = formatOID([]byte(issuer), oldPID, newPIDKey)

	subject := cert.Subject.String()
	subject = formatOID([]byte(subject), oldVID, newVIDKey)
	subject = formatOID([]byte(subject), oldPID, newPIDKey)

	certificate := X509Certificate{
		Issuer:         issuer,
		SerialNumber:   cert.SerialNumber.String(),
		Subject:        subject,
		SubjectKeyID:   BytesToHex(cert.SubjectKeyId),
		AuthorityKeyID: BytesToHex(cert.AuthorityKeyId),
		Certificate:    cert,
	}

	return &certificate, nil
}

func formatOID(header []byte, oldKey []byte, newKey []byte) string {
	if i := bytes.Index(header, oldKey); i >= 0 {
		// get value from header
		value := header[i+len(oldKey)+2 : len(string(header))]
		// get last 8 numbers from vidValue
		value = value[len(value)-8:]

		decoded, _ := hex.DecodeString(string(value))

		newKey = []byte(string(newKey) + string(decoded))

		header = header[0:i]
		header = []byte(string(header) + string(newKey))
	}

	return string(header)
}

func BytesToHex(bytes []byte) string {
	if bytes == nil {
		return ""
	}

	bytesHex := make([]string, len(bytes))
	for i, b := range bytes {
		bytesHex[i] = fmt.Sprintf("%X", b)
	}

	return strings.Join(bytesHex, ":")
}

func (c X509Certificate) Verify(parent *X509Certificate) error {
	roots := x509.NewCertPool()
	roots.AddCert(parent.Certificate)

	opts := x509.VerifyOptions{Roots: roots}

	if _, err := c.Certificate.Verify(opts); err != nil {
		return types.NewErrInvalidCertificate(fmt.Sprintf("Certificate verification failed. Error: %v", err))
	}

	return nil
}

func (c X509Certificate) IsSelfSigned() bool {
	if len(c.AuthorityKeyID) > 0 {
		return c.Issuer == c.Subject && c.AuthorityKeyID == c.SubjectKeyID
	} else {
		return c.Issuer == c.Subject
	}
}
<commit_msg>Refactore function FormatOID<commit_after>// Copyright 2020 DSR Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package x509

import (
	"bytes"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"strings"

	"github.com/zigbee-alliance/distributed-compliance-ledger/x/pki/types"
)

type X509Certificate struct {
	Issuer         string
	SerialNumber   string
	Subject        string
	SubjectKeyID   string
	AuthorityKeyID string
	Certificate    *x509.Certificate
}

func DecodeX509Certificate(pemCertificate string) (*X509Certificate, error) {
	block, _ := pem.Decode([]byte(pemCertificate))
	if block == nil {
		return nil, types.NewErrInvalidCertificate("Could not decode pem certificate")
	}

	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, types.NewErrInvalidCertificate(fmt.Sprintf("Could not parse certificate: %v", err.Error()))
	}

	oldVID := []byte("1.3.6.1.4.1.37244.2.1")
	oldPID := []byte("1.3.6.1.4.1.37244.2.2")

	newVIDKey := []byte("vid=0x")
	newPIDKey := []byte("pid=0x")

	issuer := cert.Issuer.String()
	issuer = formatOID([]byte(issuer), oldVID, newVIDKey)
	issuer = formatOID([]byte(issuer), oldPID, newPIDKey)

	subject := cert.Subject.String()
	subject = formatOID([]byte(subject), oldVID, newVIDKey)
	subject = formatOID([]byte(subject), oldPID, newPIDKey)

	certificate := X509Certificate{
		Issuer:         issuer,
		SerialNumber:   cert.SerialNumber.String(),
		Subject:        subject,
		SubjectKeyID:   BytesToHex(cert.SubjectKeyId),
		AuthorityKeyID: BytesToHex(cert.AuthorityKeyId),
		Certificate:    cert,
	}

	return &certificate, nil
}

func FormatOID(header, oldKey, newKey string) string {
	subjectValues := strings.Split(header, ",")

	for index, value := range subjectValues {
		if i := strings.Index(value, oldKey); i >= 0 {
		// get value from header
		value = value[len(value)-8:]

			decoded, _ := hex.DecodeString(value)
			hexStr := "=0x" + string(decoded)

			value = newKey + hexStr
			subjectValues[index] = value
		}
	}

	return strings.Join(subjectValues, ",")
}

func BytesToHex(bytes []byte) string {
	if bytes == nil {
		return ""
	}

	bytesHex := make([]string, len(bytes))
	for i, b := range bytes {
		bytesHex[i] = fmt.Sprintf("%X", b)
	}

	return strings.Join(bytesHex, ":")
}

func (c X509Certificate) Verify(parent *X509Certificate) error {
	roots := x509.NewCertPool()
	roots.AddCert(parent.Certificate)

	opts := x509.VerifyOptions{Roots: roots}

	if _, err := c.Certificate.Verify(opts); err != nil {
		return types.NewErrInvalidCertificate(fmt.Sprintf("Certificate verification failed. Error: %v", err))
	}

	return nil
}

func (c X509Certificate) IsSelfSigned() bool {
	if len(c.AuthorityKeyID) > 0 {
		return c.Issuer == c.Subject && c.AuthorityKeyID == c.SubjectKeyID
	} else {
		return c.Issuer == c.Subject
	}
}
<|endoftext|>


<commit_before>package shell

import (
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"golang.org/x/exp/slices"
	"io"
	"os"
)

func init() {
	Commands = append(Commands, &commandVolumeServerEvacuate{})
}

type commandVolumeServerEvacuate struct {
	targetServer string
}

func (c *commandVolumeServerEvacuate) Name() string {
	return "volumeServer.evacuate"
}

func (c *commandVolumeServerEvacuate) Help() string {
	return `move out all data on a volume server

	volumeServer.evacuate -node <host:port>

	This command moves all data away from the volume server.
	The volumes on the volume servers will be redistributed.

	Usually this is used to prepare to shutdown or upgrade the volume server.

	Sometimes a volume can not be moved because there are no
	good destination to meet the replication requirement.
	E.g. a volume replication 001 in a cluster with 2 volume servers can not be moved.
	You can use "-skipNonMoveable" to move the rest volumes.

`
}

func (c *commandVolumeServerEvacuate) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	vsEvacuateCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeServer := vsEvacuateCommand.String("node", "", "<host>:<port> of the volume server")
	c.targetServer = *vsEvacuateCommand.String("target", "", "<host>:<port> of target volume")
	skipNonMoveable := vsEvacuateCommand.Bool("skipNonMoveable", false, "skip volumes that can not be moved")
	applyChange := vsEvacuateCommand.Bool("force", false, "actually apply the changes")
	retryCount := vsEvacuateCommand.Int("retry", 0, "how many times to retry")
	if err = vsEvacuateCommand.Parse(args); err != nil {
		return nil
	}
	infoAboutSimulationMode(writer, *applyChange, "-force")

	if err = commandEnv.confirmIsLocked(args); err != nil {
		return
	}

	if *volumeServer == "" {
		return fmt.Errorf("need to specify volume server by -node=<host>:<port>")
	}
	for i := 0; i < *retryCount+1; i++ {
		if err = c.volumeServerEvacuate(commandEnv, *volumeServer, *skipNonMoveable, *applyChange, writer); err == nil {
			return nil
		}
	}

	return

}

func (c *commandVolumeServerEvacuate) volumeServerEvacuate(commandEnv *CommandEnv, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) (err error) {
	// 1. confirm the volume server is part of the cluster
	// 2. collect all other volume servers, sort by empty slots
	// 3. move to any other volume server as long as it satisfy the replication requirements

	// list all the volumes
	// collect topology information
	topologyInfo, _, err := collectTopologyInfo(commandEnv, 0)
	if err != nil {
		return err
	}

	if err := c.evacuateNormalVolumes(commandEnv, topologyInfo, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	if err := c.evacuateEcVolumes(commandEnv, topologyInfo, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	return nil
}

func (c *commandVolumeServerEvacuate) evacuateNormalVolumes(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this volume server
	volumeServers := collectVolumeServersByDc(topologyInfo, "")
	thisNode, otherNodes := nodesOtherThan(volumeServers, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster", volumeServer)
	}

	// move away normal volumes
	volumeReplicas, _ := collectVolumeReplicaLocations(topologyInfo)
	for _, diskInfo := range thisNode.info.DiskInfos {
		for _, vol := range diskInfo.VolumeInfos {
			hasMoved, err := moveAwayOneNormalVolume(commandEnv, volumeReplicas, vol, thisNode, otherNodes, applyChange)
			if err != nil {
				return fmt.Errorf("move away volume %d from %s: %v", vol.Id, volumeServer, err)
			}
			if !hasMoved {
				if skipNonMoveable {
					replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(vol.ReplicaPlacement))
					fmt.Fprintf(writer, "skipping non moveable volume %d replication:%s\n", vol.Id, replicaPlacement.String())
				} else {
					return fmt.Errorf("failed to move volume %d from %s", vol.Id, volumeServer)
				}
			}
		}
	}
	return nil
}

func (c *commandVolumeServerEvacuate) evacuateEcVolumes(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this ec volume server
	ecNodes, _ := collectEcVolumeServersByDc(topologyInfo, "")
	thisNode, otherNodes := ecNodesOtherThan(ecNodes, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster\n", volumeServer)
	}

	// move away ec volumes
	for _, diskInfo := range thisNode.info.DiskInfos {
		for _, ecShardInfo := range diskInfo.EcShardInfos {
			hasMoved, err := c.moveAwayOneEcVolume(commandEnv, ecShardInfo, thisNode, otherNodes, applyChange)
			if err != nil {
				return fmt.Errorf("move away volume %d from %s: %v", ecShardInfo.Id, volumeServer, err)
			}
			if !hasMoved {
				if skipNonMoveable {
					fmt.Fprintf(writer, "failed to move away ec volume %d from %s\n", ecShardInfo.Id, volumeServer)
				} else {
					return fmt.Errorf("failed to move away ec volume %d from %s", ecShardInfo.Id, volumeServer)
				}
			}
		}
	}
	return nil
}

func (c *commandVolumeServerEvacuate) moveAwayOneEcVolume(commandEnv *CommandEnv, ecShardInfo *master_pb.VolumeEcShardInformationMessage, thisNode *EcNode, otherNodes []*EcNode, applyChange bool) (hasMoved bool, err error) {

	for _, shardId := range erasure_coding.ShardBits(ecShardInfo.EcIndexBits).ShardIds() {
		slices.SortFunc(otherNodes, func(a, b *EcNode) bool {
			return a.localShardIdCount(ecShardInfo.Id) < b.localShardIdCount(ecShardInfo.Id)
		})
		for i := 0; i < len(otherNodes); i++ {
			emptyNode := otherNodes[i]
			if c.targetServer != "" && c.targetServer != emptyNode.info.Id {
				continue
			}
			collectionPrefix := ""
			if ecShardInfo.Collection != "" {
				collectionPrefix = ecShardInfo.Collection + "_"
			}
			fmt.Fprintf(os.Stdout, "moving ec volume %s%d.%d %s => %s\n", collectionPrefix, ecShardInfo.Id, shardId, thisNode.info.Id, emptyNode.info.Id)
			err = moveMountedShardToEcNode(commandEnv, thisNode, ecShardInfo.Collection, needle.VolumeId(ecShardInfo.Id), shardId, emptyNode, applyChange)
			if err != nil {
				return
			} else {
				hasMoved = true
				break
			}
		}
		if !hasMoved {
			return
		}
	}

	return
}

func moveAwayOneNormalVolume(commandEnv *CommandEnv, volumeReplicas map[uint32][]*VolumeReplica, vol *master_pb.VolumeInformationMessage, thisNode *Node, otherNodes []*Node, applyChange bool) (hasMoved bool, err error) {
	fn := capacityByFreeVolumeCount(types.ToDiskType(vol.DiskType))
	for _, n := range otherNodes {
		n.selectVolumes(func(v *master_pb.VolumeInformationMessage) bool {
			return v.DiskType == vol.DiskType
		})
	}
	slices.SortFunc(otherNodes, func(a, b *Node) bool {
		return a.localVolumeRatio(fn) < b.localVolumeRatio(fn)
	})
	for i := 0; i < len(otherNodes); i++ {
		emptyNode := otherNodes[i]
		hasMoved, err = maybeMoveOneVolume(commandEnv, volumeReplicas, thisNode, vol, emptyNode, applyChange)
		if err != nil {
			return
		}
		if hasMoved {
			break
		}
	}
	return
}

func nodesOtherThan(volumeServers []*Node, thisServer string) (thisNode *Node, otherNodes []*Node) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}

func ecNodesOtherThan(volumeServers []*EcNode, thisServer string) (thisNode *EcNode, otherNodes []*EcNode) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}
<commit_msg>volume server evacuate to target server<commit_after>package shell

import (
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"golang.org/x/exp/slices"
	"io"
	"os"
)

func init() {
	Commands = append(Commands, &commandVolumeServerEvacuate{})
}

type commandVolumeServerEvacuate struct {
	targetServer string
}

func (c *commandVolumeServerEvacuate) Name() string {
	return "volumeServer.evacuate"
}

func (c *commandVolumeServerEvacuate) Help() string {
	return `move out all data on a volume server

	volumeServer.evacuate -node <host:port>

	This command moves all data away from the volume server.
	The volumes on the volume servers will be redistributed.

	Usually this is used to prepare to shutdown or upgrade the volume server.

	Sometimes a volume can not be moved because there are no
	good destination to meet the replication requirement.
	E.g. a volume replication 001 in a cluster with 2 volume servers can not be moved.
	You can use "-skipNonMoveable" to move the rest volumes.

`
}

func (c *commandVolumeServerEvacuate) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	vsEvacuateCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeServer := vsEvacuateCommand.String("node", "", "<host>:<port> of the volume server")
	targetServer := vsEvacuateCommand.String("target", "", "<host>:<port> of target volume")
	skipNonMoveable := vsEvacuateCommand.Bool("skipNonMoveable", false, "skip volumes that can not be moved")
	applyChange := vsEvacuateCommand.Bool("force", false, "actually apply the changes")
	retryCount := vsEvacuateCommand.Int("retry", 0, "how many times to retry")
	if err = vsEvacuateCommand.Parse(args); err != nil {
		return nil
	}
	infoAboutSimulationMode(writer, *applyChange, "-force")

	if err = commandEnv.confirmIsLocked(args); err != nil {
		return
	}

	if *volumeServer == "" {
		return fmt.Errorf("need to specify volume server by -node=<host>:<port>")
	}
	if *targetServer != "" {
		c.targetServer = *targetServer
	}
	for i := 0; i < *retryCount+1; i++ {
		if err = c.volumeServerEvacuate(commandEnv, *volumeServer, *skipNonMoveable, *applyChange, writer); err == nil {
			return nil
		}
	}

	return

}

func (c *commandVolumeServerEvacuate) volumeServerEvacuate(commandEnv *CommandEnv, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) (err error) {
	// 1. confirm the volume server is part of the cluster
	// 2. collect all other volume servers, sort by empty slots
	// 3. move to any other volume server as long as it satisfy the replication requirements

	// list all the volumes
	// collect topology information
	topologyInfo, _, err := collectTopologyInfo(commandEnv, 0)
	if err != nil {
		return err
	}

	if err := c.evacuateNormalVolumes(commandEnv, topologyInfo, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	if err := c.evacuateEcVolumes(commandEnv, topologyInfo, volumeServer, skipNonMoveable, applyChange, writer); err != nil {
		return err
	}

	return nil
}

func (c *commandVolumeServerEvacuate) evacuateNormalVolumes(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this volume server
	volumeServers := collectVolumeServersByDc(topologyInfo, "")
	thisNode, otherNodes := nodesOtherThan(volumeServers, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster", volumeServer)
	}
	if c.targetServer != "" {
		targetServerFound := false
		for _, otherNode := range otherNodes {
			if otherNode.info.Id == c.targetServer {
				otherNodes = []*Node{otherNode}
				targetServerFound = true
				break
			}
		}
		if !targetServerFound {
			return fmt.Errorf("target %s is not found in this cluster", c.targetServer)
		}
	}

	// move away normal volumes
	volumeReplicas, _ := collectVolumeReplicaLocations(topologyInfo)
	for _, diskInfo := range thisNode.info.DiskInfos {
		for _, vol := range diskInfo.VolumeInfos {
			hasMoved, err := c.moveAwayOneNormalVolume(commandEnv, volumeReplicas, vol, thisNode, otherNodes, applyChange)
			if err != nil {
				fmt.Fprintf(writer, "move away volume %d from %s: %v", vol.Id, volumeServer, err)
			}
			if !hasMoved {
				if skipNonMoveable {
					replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(vol.ReplicaPlacement))
					fmt.Fprintf(writer, "skipping non moveable volume %d replication:%s\n", vol.Id, replicaPlacement.String())
				} else {
					return fmt.Errorf("failed to move volume %d from %s", vol.Id, volumeServer)
				}
			}
		}
	}
	return nil
}

func (c *commandVolumeServerEvacuate) evacuateEcVolumes(commandEnv *CommandEnv, topologyInfo *master_pb.TopologyInfo, volumeServer string, skipNonMoveable, applyChange bool, writer io.Writer) error {
	// find this ec volume server
	ecNodes, _ := collectEcVolumeServersByDc(topologyInfo, "")
	thisNode, otherNodes := ecNodesOtherThan(ecNodes, volumeServer)
	if thisNode == nil {
		return fmt.Errorf("%s is not found in this cluster\n", volumeServer)
	}

	// move away ec volumes
	for _, diskInfo := range thisNode.info.DiskInfos {
		for _, ecShardInfo := range diskInfo.EcShardInfos {
			hasMoved, err := c.moveAwayOneEcVolume(commandEnv, ecShardInfo, thisNode, otherNodes, applyChange)
			if err != nil {
				fmt.Fprintf(writer, "move away volume %d from %s: %v", ecShardInfo.Id, volumeServer, err)
			}
			if !hasMoved {
				if skipNonMoveable {
					fmt.Fprintf(writer, "failed to move away ec volume %d from %s\n", ecShardInfo.Id, volumeServer)
				} else {
					return fmt.Errorf("failed to move away ec volume %d from %s", ecShardInfo.Id, volumeServer)
				}
			}
		}
	}
	return nil
}

func (c *commandVolumeServerEvacuate) moveAwayOneEcVolume(commandEnv *CommandEnv, ecShardInfo *master_pb.VolumeEcShardInformationMessage, thisNode *EcNode, otherNodes []*EcNode, applyChange bool) (hasMoved bool, err error) {

	for _, shardId := range erasure_coding.ShardBits(ecShardInfo.EcIndexBits).ShardIds() {
		slices.SortFunc(otherNodes, func(a, b *EcNode) bool {
			return a.localShardIdCount(ecShardInfo.Id) < b.localShardIdCount(ecShardInfo.Id)
		})
		for i := 0; i < len(otherNodes); i++ {
			emptyNode := otherNodes[i]
			collectionPrefix := ""
			if ecShardInfo.Collection != "" {
				collectionPrefix = ecShardInfo.Collection + "_"
			}
			fmt.Fprintf(os.Stdout, "moving ec volume %s%d.%d %s => %s\n", collectionPrefix, ecShardInfo.Id, shardId, thisNode.info.Id, emptyNode.info.Id)
			err = moveMountedShardToEcNode(commandEnv, thisNode, ecShardInfo.Collection, needle.VolumeId(ecShardInfo.Id), shardId, emptyNode, applyChange)
			if err != nil {
				return
			} else {
				hasMoved = true
				break
			}
		}
		if !hasMoved {
			return
		}
	}

	return
}

func (c *commandVolumeServerEvacuate) moveAwayOneNormalVolume(commandEnv *CommandEnv, volumeReplicas map[uint32][]*VolumeReplica, vol *master_pb.VolumeInformationMessage, thisNode *Node, otherNodes []*Node, applyChange bool) (hasMoved bool, err error) {
	fn := capacityByFreeVolumeCount(types.ToDiskType(vol.DiskType))
	for _, n := range otherNodes {
		n.selectVolumes(func(v *master_pb.VolumeInformationMessage) bool {
			return v.DiskType == vol.DiskType
		})
	}
	slices.SortFunc(otherNodes, func(a, b *Node) bool {
		return a.localVolumeRatio(fn) < b.localVolumeRatio(fn)
	})
	for i := 0; i < len(otherNodes); i++ {
		emptyNode := otherNodes[i]
		hasMoved, err = maybeMoveOneVolume(commandEnv, volumeReplicas, thisNode, vol, emptyNode, applyChange)
		if err != nil {
			return
		}
		if hasMoved {
			break
		}
	}
	return
}

func nodesOtherThan(volumeServers []*Node, thisServer string) (thisNode *Node, otherNodes []*Node) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}

func ecNodesOtherThan(volumeServers []*EcNode, thisServer string) (thisNode *EcNode, otherNodes []*EcNode) {
	for _, node := range volumeServers {
		if node.info.Id == thisServer {
			thisNode = node
			continue
		}
		otherNodes = append(otherNodes, node)
	}
	return
}
<|endoftext|>


<commit_before>package voldemort

import (
	"bytes"
	proto "code.google.com/p/goprotobuf/proto"
	"encoding/binary"
	"encoding/xml"
	"errors"
	"fmt"
	vproto "github.com/matzhouse/go-voldemort-protobufs"
	"io"
	"log"
	"math"
	"net"
	"sync"
	"time"
)

// The VoldemortConn struct is used to hold all the data for the Voldemort cluster you need to query
type VoldemortConn struct {
	// TCP connection used to talk to Voldemort instance
	c *net.TCPConn

	// Information about the Voldemort cluster (received from the cluster once you connect to it)
	cl *Cluster

	// A simple mutex to make the conn thread safe
	mu sync.Mutex
}

type Cluster struct {
	Name    string   `xml:"name"`
	Servers []Server `xml:"server"`
}

type Server struct {
	Id         int    `xml:"id"`
	Host       string `xml:"host"`
	Http       int    `xml:"http-port"`
	Socket     int    `xml:"socket-port"`
	Partitions string `xml:"partitions"`
	State      bool
}

func Dial(raddr *net.TCPAddr, proto string) (c *VoldemortConn, err error) {

	conn, err := net.DialTCP("tcp", nil, raddr)

	if err != nil {
		return nil, err
	}

	conn.Write([]byte(proto))

	result := bytes.NewBuffer(nil)
	var buf [2]byte // protocol response only returns a 2 byte response - ok or no

	_, err = conn.Read(buf[0:])

	if err != nil {
		return nil, err
	}

	result.Write(buf[0:2])

	// protocol response doesn't return a new line
	if string(result.Bytes()) != "ok" {
		if string(result.Bytes()) == "no" {
			err = errors.New(fmt.Sprintf("bad protocol set response : %s", result.Bytes()))
			return nil, err
		} else {
			err = errors.New(fmt.Sprintf("unknown protocol response : %s", result.Bytes()))
			return nil, err
		}
	}

	vc := new(VoldemortConn)
	vc.c = conn
	vc.c.SetNoDelay(true)

	cl, err := bootstrap(vc)

	if err != nil {
		log.Println(err)
		log.Println("warning - this client will only be able to talk to the current node")
	}

	go watchcluster()

	vc.cl = cl

	return vc, nil

}

func bootstrap(vc *VoldemortConn) (n *Cluster, err error) {

	n, err = getclusterdata(vc)

	if err != nil {
		return nil, err
	}

	// update the state
	for k, v := range n.Servers {
		saddr := fmt.Sprintf("%s:%d", v.Host, v.Socket)
		if saddr == vc.c.RemoteAddr().String() {
			vc.cl.Servers[k].State = true
		}
	}

	fmt.Println(vc.cl)

	return n, nil

}

func watchcluster() {

}

func get(conn *VoldemortConn, store string, req *vproto.GetRequest, shouldroute bool) (resp *vproto.GetResponse, err error) {

	rt := vproto.RequestType(0)

	vr := &vproto.VoldemortRequest{
		Store:       &store,
		Type:        &rt,
		ShouldRoute: &shouldroute,
	}

	vr.Get = req

	input, err := proto.Marshal(vr)
	if err != nil {
		return nil, err
	}

	output, err := Do(conn, input)

	if err != nil {
		return nil, err
	}

	// no record found
	if output == nil {
		return nil, nil
	}

	resp = new(vproto.GetResponse)

	err = proto.Unmarshal(output, resp)
	if err != nil {
		return nil, err
	}

	if resp.Error != nil {
		err = errors.New(*resp.Error.ErrorMessage)
		return nil, err
	}

	return resp, nil

}

func put(conn *VoldemortConn, store string, req *vproto.PutRequest) (resp *vproto.PutResponse, err error) {

	Routingdecision := true
	rt := vproto.RequestType(2)

	vr := &vproto.VoldemortRequest{
		Store:       &store,
		Type:        &rt,
		ShouldRoute: &Routingdecision,
	}

	vr.Put = req

	input, err := proto.Marshal(vr)
	if err != nil {
		return nil, err
	}

	output, err := Do(conn, input)

	if err != nil {
		return nil, err
	}

	// no record found
	if output == nil {
		return nil, nil
	}

	resp = new(vproto.PutResponse)

	err = proto.Unmarshal(output, resp)
	if err != nil {
		return nil, err
	}

	return resp, nil

}

func getclusterdata(conn *VoldemortConn) (cl *Cluster, err error) {

	req := &vproto.GetRequest{
		Key: []byte("cluster.xml"),
	}

	resp, err := get(conn, "metadata", req, false)

	if err != nil {
		return nil, err
	}

	if resp == nil {
		err = errors.New("metadata not available")
		return nil, err
	}

	cl = &Cluster{}

	err = xml.Unmarshal([]byte(resp.GetVersioned()[0].GetValue()), cl)

	if err != nil {
		return nil, err
	}

	fmt.Println(cl)

	conn.cl = cl

	return cl, nil

}

func getversion(conn *VoldemortConn, store string, key string) (vc *vproto.VectorClock, err error) {

	req := &vproto.GetRequest{
		Key: []byte(key),
	}

	resp, err := get(conn, store, req, true)
	if err != nil {
		return nil, err
	}

	if resp == nil {
		// return new VectorClock
		vc := new(vproto.VectorClock)
		now := int64(time.Now().Unix())
		vc.Timestamp = &now

		id := int32(0)
		v := int64(1) // 1-32767

		ce := &vproto.ClockEntry{
			NodeId:  &id,
			Version: &v,
		}

		vc.Entries = append(vc.Entries, ce)

		return vc, nil
	} else {
		now := int64(time.Now().Unix())

		cv := resp.Versioned[0].Version.Entries[0].Version
		ncv := int64(1) + *cv

		resp.Versioned[0].Version.Entries[0].Version = &ncv

		resp.Versioned[0].Version.Timestamp = &now

		return resp.Versioned[0].Version, nil
	}

}

// Nice getter for a string key returning a string value
func Get(conn *VoldemortConn, store string, key string) (value string, err error) {

	if store == "" || key == "" {
		err = errors.New("store and key cannot be empty")
		return "", err
	}

	req := &vproto.GetRequest{
		Key: []byte(key),
	}

	resp, err := get(conn, store, req, true)

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

	// null response
	if resp == nil {
		return "", nil
	}

	vd := resp.GetVersioned()

	return string(vd[0].GetValue()), nil // simple get newest

}

func Put(conn *VoldemortConn, store string, key string, value string) (b bool, err error) {

	if store == "" || key == "" {
		err = errors.New("store and key cannot be empty")
		return false, err
	}

	vc, err := getversion(conn, store, key)

	if err != nil {
		return false, err
	}

	req := &vproto.PutRequest{
		Key: []byte(key),
		Versioned: &vproto.Versioned{
			Value:   []byte(value),
			Version: vc,
		},
	}

	resp, err := put(conn, store, req)

	if err != nil {
		return false, err
	}

	if resp != nil {
		if resp.Error != nil {
			return false, errors.New(*resp.Error.ErrorMessage)
		}
	}

	return true, nil

}

func Close(conn *VoldemortConn) {
	conn.c.Close()
}

func reconnect(conn *VoldemortConn) {

}

// creates the voldemort request <4 byte length, big endian encoded><message bytes> and receives the same
func Do(conn *VoldemortConn, input []byte) (output []byte, err error) {

	conn.mu.Lock()
	defer conn.mu.Unlock()

	noconn := errors.New("tcp connection not available")

	buf := new(bytes.Buffer)

	var length uint32
	length = uint32(len(input))

	err = binary.Write(buf, binary.BigEndian, length)

	_, err = buf.Write(input)

	if err != nil {
		return nil, err
	}

	// send the command to voldemort
	_, err = conn.c.Write(buf.Bytes())

	if err != nil {
		if err == io.EOF {
			return nil, noconn
		} else {
			return nil, err
		}
	}

	// reset the buffer for the received content
	buf.Reset()

	// Get first 4 bytes from conn to get length of response from voldemort
	var buflen [4]byte
	n, err := conn.c.Read(buflen[0:4])

	if err != nil {
		if err == io.EOF {
			return nil, noconn
		} else {
			return nil, err
		}
	}

	// write correct length of response to buffer
	buf.Write(buflen[0:n])

	var respLen uint32
	err = binary.Read(buf, binary.BigEndian, &respLen)

	if respLen == 0 {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	buf.Reset()
	var tempBuf []byte

	static_len := 512

	rem := math.Mod(float64(respLen), 512)
	loop_count := (int(respLen) - int(rem)) / 512

	for i := 0; i <= int(loop_count); i++ {
		tempBuf = make([]byte, 512)

		n, err = conn.c.Read(tempBuf[0:static_len])

		if err != nil {
			if err == io.EOF {
				return nil, noconn
			} else {
				return nil, err
			}
		}

		buf.Write(tempBuf[0:n])
	}

	return buf.Bytes(), nil

}
<commit_msg>Update client - set Close as vc method<commit_after>package voldemort

import (
	"bytes"
	proto "code.google.com/p/goprotobuf/proto"
	"encoding/binary"
	"encoding/xml"
	"errors"
	"fmt"
	vproto "github.com/matzhouse/go-voldemort-protobufs"
	"io"
	"log"
	"math"
	"net"
	"sync"
	"time"
)

// The VoldemortConn struct is used to hold all the data for the Voldemort cluster you need to query
type VoldemortConn struct {
	// TCP connection used to talk to Voldemort instance
	c *net.TCPConn

	// Information about the Voldemort cluster (received from the cluster once you connect to it)
	cl *Cluster

	// A simple mutex to make the conn thread safe
	mu sync.Mutex
}

// The cluster struct holds all the information taken from the Voldemort cluster when it's first connected
type Cluster struct {
	// Name of Voldemort cluster
	Name string `xml:"name"`

	// Array of Voldemort servers
	Servers []Server `xml:"server"`
}

// The server struct holds all the information about a Voldermort server
type Server struct {
	Id         int    `xml:"id"`
	Host       string `xml:"host"`
	Http       int    `xml:"http-port"`
	Socket     int    `xml:"socket-port"`
	Partitions string `xml:"partitions"`
	State      bool
}

// Returns a VoldemortConn that can be used to talk to a Voldemort cluster
func Dial(raddr *net.TCPAddr, proto string) (c *VoldemortConn, err error) {

	conn, err := net.DialTCP("tcp", nil, raddr)

	//conn, err := net.Dial(network, address)

	if err != nil {
		return nil, err
	}

	err = setProtocol(conn, proto)

	if err != nil {
		return nil, err
	}

	vc := new(VoldemortConn)
	vc.c = conn
	vc.c.SetNoDelay(true)

	cl, err := bootstrap(vc)

	if err != nil {
		log.Println(err)
		log.Println("warning - this client will only be able to talk to the current node")
	}

	vc.cl = cl

	return vc, nil

}

func setProtocol(conn *net.TCPConn, proto string) (err error) {

	_, err = conn.Write([]byte(proto))

	result := bytes.NewBuffer(nil)
	var buf [2]byte // protocol response only returns a 2 byte response - ok or no

	_, err = conn.Read(buf[0:])

	if err != nil {
		return err
	}

	result.Write(buf[0:2])

	// protocol response doesn't return a new line
	if string(result.Bytes()) != "ok" {
		if string(result.Bytes()) == "no" {
			err = errors.New(fmt.Sprintf("bad protocol set response : %s", result.Bytes()))
			return err
		} else {
			err = errors.New(fmt.Sprintf("unknown protocol response : %s", result.Bytes()))
			return err
		}
	}

	return nil

}

func bootstrap(vc *VoldemortConn) (n *Cluster, err error) {

	n, err = getclusterdata(vc)

	if err != nil {
		return nil, err
	}

	// update the state
	for k, v := range n.Servers {
		saddr := fmt.Sprintf("%s:%d", v.Host, v.Socket)
		if saddr == vc.c.RemoteAddr().String() {
			vc.cl.Servers[k].State = true
		}
	}

	return n, nil

}

func get(conn *VoldemortConn, store string, req *vproto.GetRequest, shouldroute bool) (resp *vproto.GetResponse, err error) {

	rt := vproto.RequestType(0)

	vr := &vproto.VoldemortRequest{
		Store:       &store,
		Type:        &rt,
		ShouldRoute: &shouldroute,
	}

	vr.Get = req

	input, err := proto.Marshal(vr)
	if err != nil {
		return nil, err
	}

	output, err := Do(conn, input)

	if err != nil {
		return nil, err
	}

	// no record found
	if output == nil {
		return nil, nil
	}

	resp = new(vproto.GetResponse)

	err = proto.Unmarshal(output, resp)
	if err != nil {
		return nil, err
	}

	if resp.Error != nil {
		err = errors.New(*resp.Error.ErrorMessage)
		return nil, err
	}

	return resp, nil

}

func put(conn *VoldemortConn, store string, req *vproto.PutRequest) (resp *vproto.PutResponse, err error) {

	Routingdecision := true
	rt := vproto.RequestType(2)

	vr := &vproto.VoldemortRequest{
		Store:       &store,
		Type:        &rt,
		ShouldRoute: &Routingdecision,
	}

	vr.Put = req

	input, err := proto.Marshal(vr)
	if err != nil {
		return nil, err
	}

	output, err := Do(conn, input)

	if err != nil {
		return nil, err
	}

	// no record found
	if output == nil {
		return nil, nil
	}

	resp = new(vproto.PutResponse)

	err = proto.Unmarshal(output, resp)
	if err != nil {
		return nil, err
	}

	return resp, nil

}

func getclusterdata(conn *VoldemortConn) (cl *Cluster, err error) {

	req := &vproto.GetRequest{
		Key: []byte("cluster.xml"),
	}

	resp, err := get(conn, "metadata", req, false)

	if err != nil {
		return nil, err
	}

	if resp == nil {
		err = errors.New("metadata not available")
		return nil, err
	}

	cl = &Cluster{}

	err = xml.Unmarshal([]byte(resp.GetVersioned()[0].GetValue()), cl)

	if err != nil {
		return nil, err
	}

	conn.cl = cl

	return cl, nil

}

func getversion(conn *VoldemortConn, store string, key string) (vc *vproto.VectorClock, err error) {

	req := &vproto.GetRequest{
		Key: []byte(key),
	}

	resp, err := get(conn, store, req, true)
	if err != nil {
		return nil, err
	}

	if resp == nil {
		// return new VectorClock
		vc := new(vproto.VectorClock)
		now := int64(time.Now().Unix())
		vc.Timestamp = &now

		id := int32(0)
		v := int64(1) // 1-32767

		ce := &vproto.ClockEntry{
			NodeId:  &id,
			Version: &v,
		}

		vc.Entries = append(vc.Entries, ce)

		return vc, nil
	} else {
		now := int64(time.Now().Unix())

		cv := resp.Versioned[0].Version.Entries[0].Version
		ncv := int64(1) + *cv

		resp.Versioned[0].Version.Entries[0].Version = &ncv

		resp.Versioned[0].Version.Timestamp = &now

		return resp.Versioned[0].Version, nil
	}

}

// Nice getter for a string key returning a string value
func Get(conn *VoldemortConn, store string, key string) (value string, err error) {

	if store == "" || key == "" {
		err = errors.New("store and key cannot be empty")
		return "", err
	}

	req := &vproto.GetRequest{
		Key: []byte(key),
	}

	resp, err := get(conn, store, req, true)

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

	// null response
	if resp == nil {
		return "", nil
	}

	vd := resp.GetVersioned()

	return string(vd[0].GetValue()), nil // simple get newest

}

func Put(conn *VoldemortConn, store string, key string, value string) (b bool, err error) {

	if store == "" || key == "" {
		err = errors.New("store and key cannot be empty")
		return false, err
	}

	vc, err := getversion(conn, store, key)

	if err != nil {
		return false, err
	}

	req := &vproto.PutRequest{
		Key: []byte(key),
		Versioned: &vproto.Versioned{
			Value:   []byte(value),
			Version: vc,
		},
	}

	resp, err := put(conn, store, req)

	if err != nil {
		return false, err
	}

	if resp != nil {
		if resp.Error != nil {
			return false, errors.New(*resp.Error.ErrorMessage)
		}
	}

	return true, nil

}

func (conn *VoldemortConn) Close() {
	conn.mu.Lock()
	defer conn.mu.Unlock()

	conn.c.Close()
}

func reconnect(conn *VoldemortConn) {

}

// creates the voldemort request <4 byte length, big endian encoded><message bytes> and receives the same
func Do(conn *VoldemortConn, input []byte) (output []byte, err error) {

	conn.mu.Lock()
	defer conn.mu.Unlock()

	noconn := errors.New("tcp connection not available")

	buf := new(bytes.Buffer)

	var length uint32
	length = uint32(len(input))

	err = binary.Write(buf, binary.BigEndian, length)

	_, err = buf.Write(input)

	if err != nil {
		return nil, err
	}

	// send the command to voldemort
	_, err = conn.c.Write(buf.Bytes())

	if err != nil {
		if err == io.EOF {
			return nil, noconn
		} else {
			return nil, err
		}
	}

	// reset the buffer for the received content
	buf.Reset()

	// Get first 4 bytes from conn to get length of response from voldemort
	var buflen [4]byte
	n, err := conn.c.Read(buflen[0:4])

	if err != nil {
		if err == io.EOF {
			return nil, noconn
		} else {
			return nil, err
		}
	}

	// write correct length of response to buffer
	buf.Write(buflen[0:n])

	var respLen uint32
	err = binary.Read(buf, binary.BigEndian, &respLen)

	if respLen == 0 {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	buf.Reset()
	var tempBuf []byte

	static_len := 512

	rem := math.Mod(float64(respLen), 512)
	loop_count := (int(respLen) - int(rem)) / 512

	for i := 0; i <= int(loop_count); i++ {
		tempBuf = make([]byte, 512)

		n, err = conn.c.Read(tempBuf[0:static_len])

		if err != nil {
			if err == io.EOF {
				return nil, noconn
			} else {
				return nil, err
			}
		}

		buf.Write(tempBuf[0:n])
	}

	return buf.Bytes(), nil

}
<|endoftext|>


<commit_before>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package vmm provides VMs based on OpenBSD vmm virtualization.
package vmm

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("vmm", ctor, true)
}

type Config struct {
	Count    int    `json:"count"`    // number of VMs to use
	Mem      int    `json:"mem"`      // amount of VM memory in MBs
	Kernel   string `json:"kernel"`   // kernel to boot
	Template string `json:"template"` // vm template
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	os       string
	sshkey   string
	sshuser  string
	sshhost  string
	sshport  int
	merger   *vmimpl.OutputMerger
	vmName   string
	vmm      *exec.Cmd
	consolew io.WriteCloser
}

var ipRegex = regexp.MustCompile(`bound to (([0-9]+\.){3}3)`)

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
		Mem:   512,
	}

	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file '%v' does not exist", env.Image)
	}

	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vmm vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 128 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1-128]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if cfg.Mem < 128 || cfg.Mem > 1048576 {
		return nil, fmt.Errorf("invalid config param mem: %v, want [128-1048576]", cfg.Mem)
	}
	if cfg.Kernel == "" {
		return nil, fmt.Errorf("missing config param kernel")
	}
	if !osutil.IsExist(cfg.Kernel) {
		return nil, fmt.Errorf("kernel '%v' does not exist", cfg.Kernel)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}

	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	image := filepath.Join(workdir, "disk.img")
	if err := osutil.CopyFile(pool.env.Image, image); err != nil {
		return nil, err
	}

	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	inst := &instance{
		cfg:     pool.cfg,
		image:   image,
		debug:   pool.env.Debug,
		os:      pool.env.OS,
		sshkey:  pool.env.SSHKey,
		sshuser: pool.env.SSHUser,
		sshport: 22,
		vmName:  fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:  vmimpl.NewOutputMerger(tee),
	}

	// Stop the instance from the previous run in case it's still running.
	// This is racy even with -w flag, start periodically fails with:
	// vmctl: start vm command failed: Operation already in progress
	// So also sleep for a bit.
	inst.vmctl("stop", inst.vmName, "-f", "-w")
	time.Sleep(3 * time.Second)

	closeInst := inst
	defer func() {
		if closeInst != nil {
			closeInst.Close()
		}
	}()

	if err := inst.Boot(); err != nil {
		return nil, err
	}

	closeInst = nil
	return inst, nil
}

func (inst *instance) Boot() error {
	outr, outw, err := osutil.LongPipe()
	if err != nil {
		return err
	}
	inr, inw, err := osutil.LongPipe()
	if err != nil {
		outr.Close()
		outw.Close()
		return err
	}
	startArgs := []string{
		"start", inst.vmName,
		"-b", inst.cfg.Kernel,
		"-d", inst.image,
		"-m", fmt.Sprintf("%vM", inst.cfg.Mem),
		"-L", // add a local network interface
		"-c", // connect to the console
	}
	if inst.cfg.Template != "" {
		startArgs = append(startArgs, "-t", inst.cfg.Template)
	}
	if inst.debug {
		log.Logf(0, "running command: vmctl %#v", startArgs)
	}
	cmd := osutil.Command("vmctl", startArgs...)
	cmd.Stdin = inr
	cmd.Stdout = outw
	cmd.Stderr = outw
	if err := cmd.Start(); err != nil {
		outr.Close()
		outw.Close()
		inr.Close()
		inw.Close()
		return err
	}
	inst.vmm = cmd
	inst.consolew = inw
	outw.Close()
	inr.Close()
	inst.merger.Add("console", outr)

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	ipch := make(chan string, 1)
	go func() {
		gotip := false
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				bootOutputStop <- true
				return
			}
			if gotip {
				continue
			}
			if ip := parseIP(bootOutput); ip != "" {
				ipch <- ip
				gotip = true
			}
		}
	}()

	select {
	case ip := <-ipch:
		inst.sshhost = ip
	case <-inst.merger.Err:
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "vmm exited", Output: bootOutput}
	case <-time.After(10 * time.Minute):
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "no IP found", Output: bootOutput}
	}

	if err := vmimpl.WaitForSSH(inst.debug, 20*time.Minute, inst.sshhost,
		inst.sshkey, inst.sshuser, inst.os, inst.sshport); err != nil {
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: err.Error(), Output: bootOutput}
	}
	bootOutputStop <- true
	<-bootOutputStop
	return nil
}

func (inst *instance) Close() {
	inst.vmctl("stop", inst.vmName, "-f")
	if inst.consolew != nil {
		inst.consolew.Close()
	}
	if inst.vmm != nil {
		inst.vmm.Process.Kill()
		inst.vmm.Wait()
	}
	inst.merger.Wait()
}

func (inst *instance) Forward(port int) (string, error) {
	octets := strings.Split(inst.sshhost, ".")
	if len(octets) < 3 {
		return "", fmt.Errorf("too few octets in hostname %v", inst.sshhost)
	}
	addr := fmt.Sprintf("%v.%v.%v.2:%v", octets[0], octets[1], octets[2], port)
	return addr, nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	vmDst := filepath.Join("/root", filepath.Base(hostSrc))
	args := append(vmimpl.SCPArgs(inst.debug, inst.sshkey, inst.sshport),
		hostSrc, inst.sshuser+"@"+inst.sshhost+":"+vmDst)
	if inst.debug {
		log.Logf(0, "running command: scp %#v", args)
	}
	_, err := osutil.RunCmd(10*time.Minute, "", "scp", args...)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add("ssh", rpipe)

	args := append(vmimpl.SSHArgs(inst.debug, inst.sshkey, inst.sshport),
		inst.sshuser+"@"+inst.sshhost, command)
	if inst.debug {
		log.Logf(0, "running command: ssh %#v", args)
	}
	cmd := osutil.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) Diagnose() bool {
	// Note: this only works if kernel actually paniced and kernel shows panic console.
	// If kernel just hanged, we've lost connection or detected some non-panic error,
	// console still shows normal login prompt.
	for _, c := range []string{"\n", "trace\n", "show registers\n"} {
		inst.consolew.Write([]byte(c))
		time.Sleep(1 * time.Second)
	}
	return true
}

// Run the given vmctl(8) command and wait for it to finish.
func (inst *instance) vmctl(args ...string) (string, error) {
	if inst.debug {
		log.Logf(0, "running command: vmctl %#v", args)
	}
	out, err := osutil.RunCmd(time.Minute, "", "vmctl", args...)
	if err != nil {
		if inst.debug {
			log.Logf(0, "vmctl failed: %v", err)
		}
		return "", err
	}
	if inst.debug {
		log.Logf(0, "vmctl output: %v", string(out))
	}
	return string(out), nil
}

func parseIP(output []byte) string {
	matches := ipRegex.FindSubmatch(output)
	if len(matches) < 2 {
		return ""
	}
	return string(matches[1])
}
<commit_msg>vm/vmm: use derived disks for VMs<commit_after>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// Package vmm provides VMs based on OpenBSD vmm virtualization.
package vmm

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("vmm", ctor, true)
}

type Config struct {
	Count    int    `json:"count"`    // number of VMs to use
	Mem      int    `json:"mem"`      // amount of VM memory in MBs
	Kernel   string `json:"kernel"`   // kernel to boot
	Template string `json:"template"` // vm template
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	image    string
	debug    bool
	os       string
	sshkey   string
	sshuser  string
	sshhost  string
	sshport  int
	merger   *vmimpl.OutputMerger
	vmName   string
	vmm      *exec.Cmd
	consolew io.WriteCloser
}

var ipRegex = regexp.MustCompile(`bound to (([0-9]+\.){3}3)`)

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
		Mem:   512,
	}

	if !osutil.IsExist(env.Image) {
		return nil, fmt.Errorf("image file '%v' does not exist", env.Image)
	}

	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse vmm vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 128 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1-128]", cfg.Count)
	}
	if env.Debug {
		cfg.Count = 1
	}
	if cfg.Mem < 128 || cfg.Mem > 1048576 {
		return nil, fmt.Errorf("invalid config param mem: %v, want [128-1048576]", cfg.Mem)
	}
	if cfg.Kernel == "" {
		return nil, fmt.Errorf("missing config param kernel")
	}
	if !osutil.IsExist(cfg.Kernel) {
		return nil, fmt.Errorf("kernel '%v' does not exist", cfg.Kernel)
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}

	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	var tee io.Writer
	if pool.env.Debug {
		tee = os.Stdout
	}
	inst := &instance{
		cfg:     pool.cfg,
		image:   filepath.Join(workdir, "disk.qcow2"),
		debug:   pool.env.Debug,
		os:      pool.env.OS,
		sshkey:  pool.env.SSHKey,
		sshuser: pool.env.SSHUser,
		sshport: 22,
		vmName:  fmt.Sprintf("%v-%v", pool.env.Name, index),
		merger:  vmimpl.NewOutputMerger(tee),
	}

	// Stop the instance from the previous run in case it's still running.
	// This is racy even with -w flag, start periodically fails with:
	// vmctl: start vm command failed: Operation already in progress
	// So also sleep for a bit.
	inst.vmctl("stop", inst.vmName, "-f", "-w")
	time.Sleep(3 * time.Second)

	createArgs := []string{
		"create", inst.image,
		"-b", pool.env.Image,
	}
	if _, err := inst.vmctl(createArgs...); err != nil {
		return nil, err
	}

	closeInst := inst
	defer func() {
		if closeInst != nil {
			closeInst.Close()
		}
	}()

	if err := inst.Boot(); err != nil {
		return nil, err
	}

	closeInst = nil
	return inst, nil
}

func (inst *instance) Boot() error {
	outr, outw, err := osutil.LongPipe()
	if err != nil {
		return err
	}
	inr, inw, err := osutil.LongPipe()
	if err != nil {
		outr.Close()
		outw.Close()
		return err
	}
	startArgs := []string{
		"start", inst.vmName,
		"-b", inst.cfg.Kernel,
		"-d", inst.image,
		"-m", fmt.Sprintf("%vM", inst.cfg.Mem),
		"-L", // add a local network interface
		"-c", // connect to the console
	}
	if inst.cfg.Template != "" {
		startArgs = append(startArgs, "-t", inst.cfg.Template)
	}
	if inst.debug {
		log.Logf(0, "running command: vmctl %#v", startArgs)
	}
	cmd := osutil.Command("vmctl", startArgs...)
	cmd.Stdin = inr
	cmd.Stdout = outw
	cmd.Stderr = outw
	if err := cmd.Start(); err != nil {
		outr.Close()
		outw.Close()
		inr.Close()
		inw.Close()
		return err
	}
	inst.vmm = cmd
	inst.consolew = inw
	outw.Close()
	inr.Close()
	inst.merger.Add("console", outr)

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	ipch := make(chan string, 1)
	go func() {
		gotip := false
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				bootOutputStop <- true
				return
			}
			if gotip {
				continue
			}
			if ip := parseIP(bootOutput); ip != "" {
				ipch <- ip
				gotip = true
			}
		}
	}()

	select {
	case ip := <-ipch:
		inst.sshhost = ip
	case <-inst.merger.Err:
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "vmm exited", Output: bootOutput}
	case <-time.After(10 * time.Minute):
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "no IP found", Output: bootOutput}
	}

	if err := vmimpl.WaitForSSH(inst.debug, 20*time.Minute, inst.sshhost,
		inst.sshkey, inst.sshuser, inst.os, inst.sshport); err != nil {
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: err.Error(), Output: bootOutput}
	}
	bootOutputStop <- true
	<-bootOutputStop
	return nil
}

func (inst *instance) Close() {
	inst.vmctl("stop", inst.vmName, "-f")
	if inst.consolew != nil {
		inst.consolew.Close()
	}
	if inst.vmm != nil {
		inst.vmm.Process.Kill()
		inst.vmm.Wait()
	}
	inst.merger.Wait()
}

func (inst *instance) Forward(port int) (string, error) {
	octets := strings.Split(inst.sshhost, ".")
	if len(octets) < 3 {
		return "", fmt.Errorf("too few octets in hostname %v", inst.sshhost)
	}
	addr := fmt.Sprintf("%v.%v.%v.2:%v", octets[0], octets[1], octets[2], port)
	return addr, nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	vmDst := filepath.Join("/root", filepath.Base(hostSrc))
	args := append(vmimpl.SCPArgs(inst.debug, inst.sshkey, inst.sshport),
		hostSrc, inst.sshuser+"@"+inst.sshhost+":"+vmDst)
	if inst.debug {
		log.Logf(0, "running command: scp %#v", args)
	}
	_, err := osutil.RunCmd(10*time.Minute, "", "scp", args...)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add("ssh", rpipe)

	args := append(vmimpl.SSHArgs(inst.debug, inst.sshkey, inst.sshport),
		inst.sshuser+"@"+inst.sshhost, command)
	if inst.debug {
		log.Logf(0, "running command: ssh %#v", args)
	}
	cmd := osutil.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) Diagnose() bool {
	// Note: this only works if kernel actually paniced and kernel shows panic console.
	// If kernel just hanged, we've lost connection or detected some non-panic error,
	// console still shows normal login prompt.
	for _, c := range []string{"\n", "trace\n", "show registers\n"} {
		inst.consolew.Write([]byte(c))
		time.Sleep(1 * time.Second)
	}
	return true
}

// Run the given vmctl(8) command and wait for it to finish.
func (inst *instance) vmctl(args ...string) (string, error) {
	if inst.debug {
		log.Logf(0, "running command: vmctl %#v", args)
	}
	out, err := osutil.RunCmd(time.Minute, "", "vmctl", args...)
	if err != nil {
		if inst.debug {
			log.Logf(0, "vmctl failed: %v", err)
		}
		return "", err
	}
	if inst.debug {
		log.Logf(0, "vmctl output: %v", string(out))
	}
	return string(out), nil
}

func parseIP(output []byte) string {
	matches := ipRegex.FindSubmatch(output)
	if len(matches) < 2 {
		return ""
	}
	return string(matches[1])
}
<|endoftext|>


<commit_before>package vreg

import "testing"

func TestQuery(t *testing.T) {

}
<commit_msg>added lic goreleaser config<commit_after>package vreg

import "testing"

func TestQuery(t *testing.T) {
	Query("JJJ895AX")
}
<|endoftext|>


<commit_before>package pick

import (
	"io"

	"golang.org/x/net/html"
)

type Option struct {
	PageSource io.Reader
	TagName    string
	Attr       *Attr // optional
}

type Attr struct {
	Label string
	Value string
}

func PickAttr(option *Option, AttrLabel string, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		case html.StartTagToken, html.SelfClosingTagToken:
			tagName, attr := z.TagName()

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false
			tmpRes := []string{}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				labelStr := string(label)
				valueStr := string(value)

				// check the attr
				if option.Attr == nil || (option.Attr.Label == labelStr && option.Attr.Value == valueStr) {
					matched = true
				}

				// get the result - even the matched false or true
				if labelStr == AttrLabel {
					tmpRes = append(tmpRes, valueStr)
				}
			}

			// skip the non matched one
			if !matched {
				continue
			}

			// send the result for matched only
			res = append(res, tmpRes...)

			// return when limit
			if limit > 0 && len(res) >= limit {
				return
			}
		}
	}

	return
}

// PickText: get text of elements that matches the Option.TagName (including it's child elements)
func PickText(option *Option, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	depth := 0

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		// get text
		case html.TextToken:
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Text())
			}

		case html.EndTagToken:
			if depth > 0 {
				depth--
			}

		case html.StartTagToken:
			tagName, attr := z.TagName()

			// inside the target
			if depth > 0 && !isSelfClosingTag(tagName) {
				depth++
				continue
			}

			// check limit
			if limit > 0 && len(res) >= limit {
				return
			}

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false

			// empty attr element
			if !attr && option.Attr == nil {
				matched = true
			}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				if option.Attr == nil || (option.Attr.Label == string(label) && option.Attr.Value == string(value)) {
					matched = true
				}
			}

			if !matched {
				continue
			}

			depth++

			// init an empty element
			res = append(res, "")
		}
	}

	return
}

// isSelfClosingTag: return if it's a self closing tag
// html.SelfClosingTagToken is not always correct
// ex:
//		<input /> 	-> html.SelfClosingTagToken
//		<input> 	-> html.StartTagToken
func isSelfClosingTag(tag []byte) bool {
	switch string(tag) {
	case "area", "base", "br", "col", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr":
		return true
	}

	return false
}

// PickHtml: get html of elements that matches the Option.TagName (including it's child elements)
func PickHtml(option *Option, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	depth := 0

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		// get text
		case html.TextToken:
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())
			}

		case html.EndTagToken:
			if depth > 0 {
				depth--
			}

			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())
			}

		case html.SelfClosingTagToken:
			// append to the last element
			res[len(res)-1] = res[len(res)-1] + string(z.Raw())

		case html.StartTagToken:
			tagName, attr := z.TagName()

			// inside the target
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())

				if !isSelfClosingTag(tagName) {
					depth++
				}
				continue
			}

			// check limit
			if limit > 0 && len(res) >= limit {
				return
			}

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false

			// empty attr element
			if !attr && option.Attr == nil {
				matched = true
			}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				if option.Attr == nil || (option.Attr.Label == string(label) && option.Attr.Value == string(value)) {

					matched = true
				}
			}

			if !matched {
				continue
			}

			depth++

			// init an empty element
			res = append(res, "")
		}
	}

	return
}
<commit_msg>add godoc comment [skip ci]<commit_after>package pick

import (
	"io"

	"golang.org/x/net/html"
)

// Option
// PageSource: html reader
// TagName: element's tag. ex: div, p, form, input...
type Option struct {
	PageSource io.Reader
	TagName    string
	Attr       *Attr // optional
}

// define an element's attribute
// ex: <input type="hidden" />
//		Attr{
//			Label: "type"
//			Value: "hidden"
//		}
type Attr struct {
	Label string
	Value string
}

// PickAttr: get attribute Option.Attr of elements that matches the Option.TagName
func PickAttr(option *Option, AttrLabel string, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		case html.StartTagToken, html.SelfClosingTagToken:
			tagName, attr := z.TagName()

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false
			tmpRes := []string{}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				labelStr := string(label)
				valueStr := string(value)

				// check the attr
				if option.Attr == nil || (option.Attr.Label == labelStr && option.Attr.Value == valueStr) {
					matched = true
				}

				// get the result - even the matched false or true
				if labelStr == AttrLabel {
					tmpRes = append(tmpRes, valueStr)
				}
			}

			// skip the non matched one
			if !matched {
				continue
			}

			// send the result for matched only
			res = append(res, tmpRes...)

			// return when limit
			if limit > 0 && len(res) >= limit {
				return
			}
		}
	}

	return
}

// PickText: get text of elements that matches the Option.TagName (including it's child elements)
func PickText(option *Option, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	depth := 0

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		// get text
		case html.TextToken:
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Text())
			}

		case html.EndTagToken:
			if depth > 0 {
				depth--
			}

		case html.StartTagToken:
			tagName, attr := z.TagName()

			// inside the target
			if depth > 0 && !isSelfClosingTag(tagName) {
				depth++
				continue
			}

			// check limit
			if limit > 0 && len(res) >= limit {
				return
			}

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false

			// empty attr element
			if !attr && option.Attr == nil {
				matched = true
			}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				if option.Attr == nil || (option.Attr.Label == string(label) && option.Attr.Value == string(value)) {
					matched = true
				}
			}

			if !matched {
				continue
			}

			depth++

			// init an empty element
			res = append(res, "")
		}
	}

	return
}

// isSelfClosingTag: return if it's a self closing tag
// html.SelfClosingTagToken is not always correct
// ex:
//		<input /> 	-> html.SelfClosingTagToken
//		<input> 	-> html.StartTagToken
func isSelfClosingTag(tag []byte) bool {
	switch string(tag) {
	case "area", "base", "br", "col", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr":
		return true
	}

	return false
}

// PickHtml: get html of elements that matches the Option.TagName (including it's child elements)
func PickHtml(option *Option, limit int) (res []string) {
	if option == nil || option.PageSource == nil {
		return
	}

	z := html.NewTokenizer(option.PageSource)

	depth := 0

	for {
		tokenType := z.Next()

		switch tokenType {

		// ignore the error token
		// quit on eof
		case html.ErrorToken:
			if z.Err() == io.EOF {
				return
			}

		// get text
		case html.TextToken:
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())
			}

		case html.EndTagToken:
			if depth > 0 {
				depth--
			}

			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())
			}

		case html.SelfClosingTagToken:
			// append to the last element
			res[len(res)-1] = res[len(res)-1] + string(z.Raw())

		case html.StartTagToken:
			tagName, attr := z.TagName()

			// inside the target
			if depth > 0 {
				// append to the last element
				res[len(res)-1] = res[len(res)-1] + string(z.Raw())

				if !isSelfClosingTag(tagName) {
					depth++
				}
				continue
			}

			// check limit
			if limit > 0 && len(res) >= limit {
				return
			}

			if string(tagName) != option.TagName {
				continue
			}

			var label, value []byte

			matched := false

			// empty attr element
			if !attr && option.Attr == nil {
				matched = true
			}

			// get attr
			for attr {
				label, value, attr = z.TagAttr()

				if option.Attr == nil || (option.Attr.Label == string(label) && option.Attr.Value == string(value)) {

					matched = true
				}
			}

			if !matched {
				continue
			}

			depth++

			// init an empty element
			res = append(res, "")
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
	"errors"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string,error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if (length == -1) {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

        //Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

        if (err != nil) {

        	target, err := hashRing.Get(metric)
        	if err != nil {
        		log.Panicln(err)
        	}

			// check built packet size and send if metric doesn't fit
        	if packets[target].Len()+size > packetLen {
        		sendPacket(packets[target].Bytes(), target)
        		packets[target].Reset()
        	}
			// add to packet
        	packets[target].Write(buff[offset : offset+size])
        	packets[target].Write(sep)

        	numMetrics++
	    }

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err != nil {
				buff[offset+i] = '\n'
				offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or most host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<commit_msg>revert err check<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
	"errors"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string,error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if (length == -1) {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

        //Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

        if (err != nil) {

        	target, err := hashRing.Get(metric)
        	if err != nil {
        		log.Panicln(err)
        	}

			// check built packet size and send if metric doesn't fit
        	if packets[target].Len()+size > packetLen {
        		sendPacket(packets[target].Bytes(), target)
        		packets[target].Reset()
        	}
			// add to packet
        	packets[target].Write(buff[offset : offset+size])
        	packets[target].Write(sep)

        	numMetrics++
	    }

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err == nil {
				buff[offset+i] = '\n'
				offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or most host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<|endoftext|>