<commit_before>/*
Copyright 2019 The Tekton Authors

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 main

import (
	"context"
	"os"

	defaultconfig "github.com/tektoncd/pipeline/pkg/apis/config"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1alpha1"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1beta1"
	"github.com/tektoncd/pipeline/pkg/contexts"
	"github.com/tektoncd/pipeline/pkg/system"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/sharedmain"
	pkgleaderelection "knative.dev/pkg/leaderelection"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/webhook"
	"knative.dev/pkg/webhook/certificates"
	"knative.dev/pkg/webhook/configmaps"
	"knative.dev/pkg/webhook/resourcesemantics"
	"knative.dev/pkg/webhook/resourcesemantics/conversion"
	"knative.dev/pkg/webhook/resourcesemantics/defaulting"
	"knative.dev/pkg/webhook/resourcesemantics/validation"
)

var types = map[schema.GroupVersionKind]resourcesemantics.GenericCRD{
	// v1alpha1
	v1alpha1.SchemeGroupVersion.WithKind("Pipeline"):         &v1alpha1.Pipeline{},
	v1alpha1.SchemeGroupVersion.WithKind("Task"):             &v1alpha1.Task{},
	v1alpha1.SchemeGroupVersion.WithKind("ClusterTask"):      &v1alpha1.ClusterTask{},
	v1alpha1.SchemeGroupVersion.WithKind("TaskRun"):          &v1alpha1.TaskRun{},
	v1alpha1.SchemeGroupVersion.WithKind("PipelineRun"):      &v1alpha1.PipelineRun{},
	v1alpha1.SchemeGroupVersion.WithKind("Condition"):        &v1alpha1.Condition{},
	v1alpha1.SchemeGroupVersion.WithKind("PipelineResource"): &v1alpha1.PipelineResource{},
	// v1beta1
	v1beta1.SchemeGroupVersion.WithKind("Pipeline"):    &v1beta1.Pipeline{},
	v1beta1.SchemeGroupVersion.WithKind("Task"):        &v1beta1.Task{},
	v1beta1.SchemeGroupVersion.WithKind("ClusterTask"): &v1beta1.ClusterTask{},
	v1beta1.SchemeGroupVersion.WithKind("TaskRun"):     &v1beta1.TaskRun{},
	v1beta1.SchemeGroupVersion.WithKind("PipelineRun"): &v1beta1.PipelineRun{},
}

func newDefaultingAdmissionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// Decorate contexts with the current state of the config.
	store := defaultconfig.NewStore(logging.FromContext(ctx).Named("config-store"))
	store.WatchConfigs(cmw)

	return defaulting.NewAdmissionController(ctx,

		// Name of the resource webhook.
		"webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/defaulting",

		// The resources to validate and default.
		types,

		// A function that infuses the context passed to Validate/SetDefaults with custom metadata.
		func(ctx context.Context) context.Context {
			return contexts.WithUpgradeViaDefaulting(store.ToContext(ctx))
		},

		// Whether to disallow unknown fields.
		true,
	)
}

func newValidationAdmissionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// Decorate contexts with the current state of the config.
	store := defaultconfig.NewStore(logging.FromContext(ctx).Named("config-store"))
	store.WatchConfigs(cmw)
	return validation.NewAdmissionController(ctx,

		// Name of the resource webhook.
		"validation.webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/resource-validation",

		// The resources to validate and default.
		types,

		// A function that infuses the context passed to Validate/SetDefaults with custom metadata.
		func(ctx context.Context) context.Context {
			return contexts.WithUpgradeViaDefaulting(store.ToContext(ctx))
		},

		// Whether to disallow unknown fields.
		true,
	)
}

func newConfigValidationController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	return configmaps.NewAdmissionController(ctx,

		// Name of the configmap webhook.
		"config.webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/config-validation",

		// The configmaps to validate.
		configmap.Constructors{
			logging.ConfigMapName():               logging.NewConfigFromConfigMap,
			defaultconfig.GetDefaultsConfigName(): defaultconfig.NewDefaultsFromConfigMap,
			pkgleaderelection.ConfigMapName():     pkgleaderelection.NewConfigFromConfigMap,
		},
	)
}

func newConversionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// nolint: golint
	var (
		v1alpha1GroupVersion = v1alpha1.SchemeGroupVersion.Version
		v1beta1GroupVersion  = v1beta1.SchemeGroupVersion.Version
	)

	return conversion.NewConversionController(ctx,
		// The path on which to serve the webhook
		"/resource-conversion",

		// Specify the types of custom resource definitions that should be converted
		map[schema.GroupKind]conversion.GroupKindConversion{
			v1beta1.Kind("Task"): {
				DefinitionName: pipeline.TaskResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.Task{},
					v1beta1GroupVersion:  &v1beta1.Task{},
				},
			},
			v1beta1.Kind("ClusterTask"): {
				DefinitionName: pipeline.ClusterTaskResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.ClusterTask{},
					v1beta1GroupVersion:  &v1beta1.ClusterTask{},
				},
			},
			v1beta1.Kind("TaskRun"): {
				DefinitionName: pipeline.TaskRunResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.TaskRun{},
					v1beta1GroupVersion:  &v1beta1.TaskRun{},
				},
			},
			v1beta1.Kind("Pipeline"): {
				DefinitionName: pipeline.PipelineResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.Pipeline{},
					v1beta1GroupVersion:  &v1beta1.Pipeline{},
				},
			},
			v1beta1.Kind("PipelineRun"): {
				DefinitionName: pipeline.PipelineRunResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.PipelineRun{},
					v1beta1GroupVersion:  &v1beta1.PipelineRun{},
				},
			},
		},

		// A function that infuses the context passed to ConvertTo/ConvertFrom/SetDefaults with custom metadata
		func(ctx context.Context) context.Context {
			return ctx
		},
	)
}

func main() {
	serviceName := os.Getenv("WEBHOOK_SERVICE_NAME")
	if serviceName == "" {
		serviceName = "tekton-pipelines-webhook"
	}

	secretName := os.Getenv("WEBHOOK_SECRET_NAME")
	if secretName == "" {
		secretName = "webhook-certs" // #nosec
	}

	// Scope informers to the webhook's namespace instead of cluster-wide
	ctx := injection.WithNamespaceScope(signals.NewContext(), system.GetNamespace())

	// Set up a signal context with our webhook options
	ctx = webhook.WithOptions(ctx, webhook.Options{
		ServiceName: serviceName,
		Port:        8443,
		SecretName:  secretName,
	})

	sharedmain.WebhookMainWithConfig(ctx, "webhook",
		sharedmain.ParseAndGetConfigOrDie(),
		certificates.NewController,
		newDefaultingAdmissionController,
		newValidationAdmissionController,
		newConfigValidationController,
		newConversionController,
	)
}
<commit_msg>Wire up webhook validation for Run objects<commit_after>/*
Copyright 2019 The Tekton Authors

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 main

import (
	"context"
	"os"

	defaultconfig "github.com/tektoncd/pipeline/pkg/apis/config"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1alpha1"
	"github.com/tektoncd/pipeline/pkg/apis/pipeline/v1beta1"
	"github.com/tektoncd/pipeline/pkg/contexts"
	"github.com/tektoncd/pipeline/pkg/system"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/sharedmain"
	pkgleaderelection "knative.dev/pkg/leaderelection"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/webhook"
	"knative.dev/pkg/webhook/certificates"
	"knative.dev/pkg/webhook/configmaps"
	"knative.dev/pkg/webhook/resourcesemantics"
	"knative.dev/pkg/webhook/resourcesemantics/conversion"
	"knative.dev/pkg/webhook/resourcesemantics/defaulting"
	"knative.dev/pkg/webhook/resourcesemantics/validation"
)

var types = map[schema.GroupVersionKind]resourcesemantics.GenericCRD{
	// v1alpha1
	v1alpha1.SchemeGroupVersion.WithKind("Pipeline"):         &v1alpha1.Pipeline{},
	v1alpha1.SchemeGroupVersion.WithKind("Task"):             &v1alpha1.Task{},
	v1alpha1.SchemeGroupVersion.WithKind("ClusterTask"):      &v1alpha1.ClusterTask{},
	v1alpha1.SchemeGroupVersion.WithKind("TaskRun"):          &v1alpha1.TaskRun{},
	v1alpha1.SchemeGroupVersion.WithKind("PipelineRun"):      &v1alpha1.PipelineRun{},
	v1alpha1.SchemeGroupVersion.WithKind("Condition"):        &v1alpha1.Condition{},
	v1alpha1.SchemeGroupVersion.WithKind("PipelineResource"): &v1alpha1.PipelineResource{},
	v1alpha1.SchemeGroupVersion.WithKind("Run"):              &v1alpha1.Run{},
	// v1beta1
	v1beta1.SchemeGroupVersion.WithKind("Pipeline"):    &v1beta1.Pipeline{},
	v1beta1.SchemeGroupVersion.WithKind("Task"):        &v1beta1.Task{},
	v1beta1.SchemeGroupVersion.WithKind("ClusterTask"): &v1beta1.ClusterTask{},
	v1beta1.SchemeGroupVersion.WithKind("TaskRun"):     &v1beta1.TaskRun{},
	v1beta1.SchemeGroupVersion.WithKind("PipelineRun"): &v1beta1.PipelineRun{},
}

func newDefaultingAdmissionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// Decorate contexts with the current state of the config.
	store := defaultconfig.NewStore(logging.FromContext(ctx).Named("config-store"))
	store.WatchConfigs(cmw)

	return defaulting.NewAdmissionController(ctx,

		// Name of the resource webhook.
		"webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/defaulting",

		// The resources to validate and default.
		types,

		// A function that infuses the context passed to Validate/SetDefaults with custom metadata.
		func(ctx context.Context) context.Context {
			return contexts.WithUpgradeViaDefaulting(store.ToContext(ctx))
		},

		// Whether to disallow unknown fields.
		true,
	)
}

func newValidationAdmissionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// Decorate contexts with the current state of the config.
	store := defaultconfig.NewStore(logging.FromContext(ctx).Named("config-store"))
	store.WatchConfigs(cmw)
	return validation.NewAdmissionController(ctx,

		// Name of the resource webhook.
		"validation.webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/resource-validation",

		// The resources to validate and default.
		types,

		// A function that infuses the context passed to Validate/SetDefaults with custom metadata.
		func(ctx context.Context) context.Context {
			return contexts.WithUpgradeViaDefaulting(store.ToContext(ctx))
		},

		// Whether to disallow unknown fields.
		true,
	)
}

func newConfigValidationController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	return configmaps.NewAdmissionController(ctx,

		// Name of the configmap webhook.
		"config.webhook.pipeline.tekton.dev",

		// The path on which to serve the webhook.
		"/config-validation",

		// The configmaps to validate.
		configmap.Constructors{
			logging.ConfigMapName():               logging.NewConfigFromConfigMap,
			defaultconfig.GetDefaultsConfigName(): defaultconfig.NewDefaultsFromConfigMap,
			pkgleaderelection.ConfigMapName():     pkgleaderelection.NewConfigFromConfigMap,
		},
	)
}

func newConversionController(ctx context.Context, cmw configmap.Watcher) *controller.Impl {
	// nolint: golint
	var (
		v1alpha1GroupVersion = v1alpha1.SchemeGroupVersion.Version
		v1beta1GroupVersion  = v1beta1.SchemeGroupVersion.Version
	)

	return conversion.NewConversionController(ctx,
		// The path on which to serve the webhook
		"/resource-conversion",

		// Specify the types of custom resource definitions that should be converted
		map[schema.GroupKind]conversion.GroupKindConversion{
			v1beta1.Kind("Task"): {
				DefinitionName: pipeline.TaskResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.Task{},
					v1beta1GroupVersion:  &v1beta1.Task{},
				},
			},
			v1beta1.Kind("ClusterTask"): {
				DefinitionName: pipeline.ClusterTaskResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.ClusterTask{},
					v1beta1GroupVersion:  &v1beta1.ClusterTask{},
				},
			},
			v1beta1.Kind("TaskRun"): {
				DefinitionName: pipeline.TaskRunResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.TaskRun{},
					v1beta1GroupVersion:  &v1beta1.TaskRun{},
				},
			},
			v1beta1.Kind("Pipeline"): {
				DefinitionName: pipeline.PipelineResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.Pipeline{},
					v1beta1GroupVersion:  &v1beta1.Pipeline{},
				},
			},
			v1beta1.Kind("PipelineRun"): {
				DefinitionName: pipeline.PipelineRunResource.String(),
				HubVersion:     v1alpha1GroupVersion,
				Zygotes: map[string]conversion.ConvertibleObject{
					v1alpha1GroupVersion: &v1alpha1.PipelineRun{},
					v1beta1GroupVersion:  &v1beta1.PipelineRun{},
				},
			},
		},

		// A function that infuses the context passed to ConvertTo/ConvertFrom/SetDefaults with custom metadata
		func(ctx context.Context) context.Context {
			return ctx
		},
	)
}

func main() {
	serviceName := os.Getenv("WEBHOOK_SERVICE_NAME")
	if serviceName == "" {
		serviceName = "tekton-pipelines-webhook"
	}

	secretName := os.Getenv("WEBHOOK_SECRET_NAME")
	if secretName == "" {
		secretName = "webhook-certs" // #nosec
	}

	// Scope informers to the webhook's namespace instead of cluster-wide
	ctx := injection.WithNamespaceScope(signals.NewContext(), system.GetNamespace())

	// Set up a signal context with our webhook options
	ctx = webhook.WithOptions(ctx, webhook.Options{
		ServiceName: serviceName,
		Port:        8443,
		SecretName:  secretName,
	})

	sharedmain.WebhookMainWithConfig(ctx, "webhook",
		sharedmain.ParseAndGetConfigOrDie(),
		certificates.NewController,
		newDefaultingAdmissionController,
		newValidationAdmissionController,
		newConfigValidationController,
		newConversionController,
	)
}
<|endoftext|>


<commit_before>package Mongo

import (
	"log"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)
type Person struct {
	Name string
	Phone string
}
func DB() {
	session, err := mgo.Dial("10.0.0.2:27017")
	if err != nil {
		panic(err)
	}
	defer session.Close()

	// Optional. Switch the session to a monotonic behavior.
	session.SetMode(mgo.Monotonic, true)

	c := session.DB("test").C("people")
	err = c.Insert(&Person{"Ale", "+55 53 8116 9639"},
		&Person{"Cla", "+55 53 8402 8510"})
	if err != nil {
		log.Fatal(err)
	}

	result := Person{}
	err = c.Find(bson.M{"name": "Ale"}).One(&result)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("Phone:", result.Phone)

}
<commit_msg>MongoDB<commit_after>package Mongo

import (
	"gopkg.in/mgo.v2"
)

type M struct {
	DBName   string
	DBServer string
}

var singleton *M

func GetDB() *mgo.Database {
	if singleton == nil {
		singleton = &M{DBName: "test", DBServer: "10.0.0.2:27017"}
	}

	return singleton.connect()
}

func (o *M) connect() *mgo.Database {
	conn, err := mgo.Dial(o.DBServer)
	if err != nil {
		panic(err)
	}

	defer func() {
		if r := recover(); r != nil {
			conn.Clone()
		}
	}()

	// Optional. Switch the session to a monotonic behavior.
	conn.SetMode(mgo.Monotonic, true)
	return conn.DB(o.DBName)
}
<|endoftext|>


<commit_before>package dht

import (
	"errors"
	"fmt"
	"time"

	ds "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore"
	key "github.com/ipfs/go-ipfs/blocks/key"
	pb "github.com/ipfs/go-ipfs/routing/dht/pb"
	lgbl "github.com/ipfs/go-ipfs/thirdparty/loggables"
	proto "gx/ipfs/QmZ4Qi3GaRbjcx28Sme5eMH7RQjGkt8wHxt2a65oLaeFEV/gogo-protobuf/proto"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	peer "gx/ipfs/QmbyvM8zRFDkbFdYyt1MnevUMJ62SiSGbfDFZ3Z8nkrzr4/go-libp2p-peer"
)

// The number of closer peers to send on requests.
var CloserPeerCount = KValue

// dhthandler specifies the signature of functions that handle DHT messages.
type dhtHandler func(context.Context, peer.ID, *pb.Message) (*pb.Message, error)

func (dht *IpfsDHT) handlerForMsgType(t pb.Message_MessageType) dhtHandler {
	switch t {
	case pb.Message_GET_VALUE:
		return dht.handleGetValue
	case pb.Message_PUT_VALUE:
		return dht.handlePutValue
	case pb.Message_FIND_NODE:
		return dht.handleFindPeer
	case pb.Message_ADD_PROVIDER:
		return dht.handleAddProvider
	case pb.Message_GET_PROVIDERS:
		return dht.handleGetProviders
	case pb.Message_PING:
		return dht.handlePing
	default:
		return nil
	}
}

func (dht *IpfsDHT) handleGetValue(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handleGetValue", p).Done()
	log.Debugf("%s handleGetValue for key: %s", dht.self, pmes.GetKey())

	// setup response
	resp := pb.NewMessage(pmes.GetType(), pmes.GetKey(), pmes.GetClusterLevel())

	// first, is there even a key?
	k := key.Key(pmes.GetKey())
	if k == "" {
		return nil, errors.New("handleGetValue but no key was provided")
		// TODO: send back an error response? could be bad, but the other node's hanging.
	}

	rec, err := dht.checkLocalDatastore(k)
	if err != nil {
		return nil, err
	}
	resp.Record = rec

	// Find closest peer on given cluster to desired key and reply with that info
	closer := dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	if len(closer) > 0 {
		closerinfos := peer.PeerInfos(dht.peerstore, closer)
		for _, pi := range closerinfos {
			log.Debugf("handleGetValue returning closer peer: '%s'", pi.ID)
			if len(pi.Addrs) < 1 {
				log.Errorf(`no addresses on peer being sent!
					[local:%s]
					[sending:%s]
					[remote:%s]`, dht.self, pi.ID, p)
			}
		}

		resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), closerinfos)
	}

	return resp, nil
}

func (dht *IpfsDHT) checkLocalDatastore(k key.Key) (*pb.Record, error) {
	log.Debugf("%s handleGetValue looking into ds", dht.self)
	dskey := k.DsKey()
	iVal, err := dht.datastore.Get(dskey)
	log.Debugf("%s handleGetValue looking into ds GOT %v", dht.self, iVal)

	if err == ds.ErrNotFound {
		return nil, nil
	}

	// if we got an unexpected error, bail.
	if err != nil {
		return nil, err
	}

	// if we have the value, send it back
	log.Debugf("%s handleGetValue success!", dht.self)

	byts, ok := iVal.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}

	rec := new(pb.Record)
	err = proto.Unmarshal(byts, rec)
	if err != nil {
		log.Debug("Failed to unmarshal DHT record from datastore.")
		return nil, err
	}

	// if its our record, dont bother checking the times on it
	if peer.ID(rec.GetAuthor()) == dht.self {
		return rec, nil
	}

	var recordIsBad bool
	recvtime, err := u.ParseRFC3339(rec.GetTimeReceived())
	if err != nil {
		log.Info("either no receive time set on record, or it was invalid: ", err)
		recordIsBad = true
	}

	if time.Now().Sub(recvtime) > MaxRecordAge {
		log.Debug("old record found, tossing.")
		recordIsBad = true
	}

	// NOTE: We do not verify the record here beyond checking these timestamps.
	// we put the burden of checking the records on the requester as checking a record
	// may be computationally expensive

	if recordIsBad {
		err := dht.datastore.Delete(dskey)
		if err != nil {
			log.Error("Failed to delete bad record from datastore: ", err)
		}

		return nil, nil // can treat this as not having the record at all
	}

	return rec, nil
}

// Store a value in this peer local storage
func (dht *IpfsDHT) handlePutValue(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handlePutValue", p).Done()
	dskey := key.Key(pmes.GetKey()).DsKey()

	if err := dht.verifyRecordLocally(pmes.GetRecord()); err != nil {
		log.Warningf("Bad dht record in PUT from: %s. %s", key.Key(pmes.GetRecord().GetAuthor()), err)
		return nil, err
	}

	rec := pmes.GetRecord()

	// record the time we receive every record
	rec.TimeReceived = proto.String(u.FormatRFC3339(time.Now()))

	data, err := proto.Marshal(rec)
	if err != nil {
		return nil, err
	}

	err = dht.datastore.Put(dskey, data)
	log.Debugf("%s handlePutValue %v", dht.self, dskey)
	return pmes, err
}

func (dht *IpfsDHT) handlePing(_ context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	log.Debugf("%s Responding to ping from %s!\n", dht.self, p)
	return pmes, nil
}

func (dht *IpfsDHT) handleFindPeer(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handleFindPeer", p).Done()
	resp := pb.NewMessage(pmes.GetType(), "", pmes.GetClusterLevel())
	var closest []peer.ID

	// if looking for self... special case where we send it on CloserPeers.
	if peer.ID(pmes.GetKey()) == dht.self {
		closest = []peer.ID{dht.self}
	} else {
		closest = dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	}

	if closest == nil {
		log.Infof("%s handleFindPeer %s: could not find anything.", dht.self, p)
		return resp, nil
	}

	var withAddresses []peer.PeerInfo
	closestinfos := peer.PeerInfos(dht.peerstore, closest)
	for _, pi := range closestinfos {
		if len(pi.Addrs) > 0 {
			withAddresses = append(withAddresses, pi)
			log.Debugf("handleFindPeer: sending back '%s'", pi.ID)
		}
	}

	resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), withAddresses)
	return resp, nil
}

func (dht *IpfsDHT) handleGetProviders(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	lm := make(lgbl.DeferredMap)
	lm["peer"] = func() interface{} { return p.Pretty() }
	defer log.EventBegin(ctx, "handleGetProviders", lm).Done()

	resp := pb.NewMessage(pmes.GetType(), pmes.GetKey(), pmes.GetClusterLevel())
	key := key.Key(pmes.GetKey())
	lm["key"] = func() interface{} { return key.B58String() }

	// debug logging niceness.
	reqDesc := fmt.Sprintf("%s handleGetProviders(%s, %s): ", dht.self, p, key)
	log.Debugf("%s begin", reqDesc)
	defer log.Debugf("%s end", reqDesc)

	// check if we have this value, to add ourselves as provider.
	has, err := dht.datastore.Has(key.DsKey())
	if err != nil && err != ds.ErrNotFound {
		log.Debugf("unexpected datastore error: %v\n", err)
		has = false
	}

	// setup providers
	providers := dht.providers.GetProviders(ctx, key)
	if has {
		providers = append(providers, dht.self)
		log.Debugf("%s have the value. added self as provider", reqDesc)
	}

	if providers != nil && len(providers) > 0 {
		infos := peer.PeerInfos(dht.peerstore, providers)
		resp.ProviderPeers = pb.PeerInfosToPBPeers(dht.host.Network(), infos)
		log.Debugf("%s have %d providers: %s", reqDesc, len(providers), infos)
	}

	// Also send closer peers.
	closer := dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	if closer != nil {
		infos := peer.PeerInfos(dht.peerstore, closer)
		resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), infos)
		log.Debugf("%s have %d closer peers: %s", reqDesc, len(closer), infos)
	}

	return resp, nil
}

func (dht *IpfsDHT) handleAddProvider(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	lm := make(lgbl.DeferredMap)
	lm["peer"] = func() interface{} { return p.Pretty() }

	defer log.EventBegin(ctx, "handleAddProvider", lm).Done()
	key := key.Key(pmes.GetKey())
	lm["key"] = func() interface{} { return key.B58String() }

	log.Debugf("%s adding %s as a provider for '%s'\n", dht.self, p, key)

	// add provider should use the address given in the message
	pinfos := pb.PBPeersToPeerInfos(pmes.GetProviderPeers())
	for _, pi := range pinfos {
		// We need to disable this check for OpenBazaar pointers to work
		/*if pi.ID != p {
			// we should ignore this provider reccord! not from originator.
			// (we chould sign them and check signature later...)
			log.Debugf("handleAddProvider received provider %s from %s. Ignore.", pi.ID, p)
			continue
		}*/

		if len(pi.Addrs) < 1 {
			log.Debugf("%s got no valid addresses for provider %s. Ignore.", dht.self, p)
			continue
		}

		log.Infof("received provider %s for %s (addrs: %s)", p, key, pi.Addrs)
		if pi.ID != dht.self { // dont add own addrs.
			// add the received addresses to our peerstore.
			dht.peerstore.AddAddrs(pi.ID, pi.Addrs, peer.ProviderAddrTTL)
		}
		dht.providers.AddProvider(ctx, key, p)
	}

	return nil, nil
}
<commit_msg>Use provided ID in AddProvider call<commit_after>package dht

import (
	"errors"
	"fmt"
	"time"

	ds "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/ipfs/go-datastore"
	key "github.com/ipfs/go-ipfs/blocks/key"
	pb "github.com/ipfs/go-ipfs/routing/dht/pb"
	lgbl "github.com/ipfs/go-ipfs/thirdparty/loggables"
	proto "gx/ipfs/QmZ4Qi3GaRbjcx28Sme5eMH7RQjGkt8wHxt2a65oLaeFEV/gogo-protobuf/proto"
	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
	peer "gx/ipfs/QmbyvM8zRFDkbFdYyt1MnevUMJ62SiSGbfDFZ3Z8nkrzr4/go-libp2p-peer"
)

// The number of closer peers to send on requests.
var CloserPeerCount = KValue

// dhthandler specifies the signature of functions that handle DHT messages.
type dhtHandler func(context.Context, peer.ID, *pb.Message) (*pb.Message, error)

func (dht *IpfsDHT) handlerForMsgType(t pb.Message_MessageType) dhtHandler {
	switch t {
	case pb.Message_GET_VALUE:
		return dht.handleGetValue
	case pb.Message_PUT_VALUE:
		return dht.handlePutValue
	case pb.Message_FIND_NODE:
		return dht.handleFindPeer
	case pb.Message_ADD_PROVIDER:
		return dht.handleAddProvider
	case pb.Message_GET_PROVIDERS:
		return dht.handleGetProviders
	case pb.Message_PING:
		return dht.handlePing
	default:
		return nil
	}
}

func (dht *IpfsDHT) handleGetValue(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handleGetValue", p).Done()
	log.Debugf("%s handleGetValue for key: %s", dht.self, pmes.GetKey())

	// setup response
	resp := pb.NewMessage(pmes.GetType(), pmes.GetKey(), pmes.GetClusterLevel())

	// first, is there even a key?
	k := key.Key(pmes.GetKey())
	if k == "" {
		return nil, errors.New("handleGetValue but no key was provided")
		// TODO: send back an error response? could be bad, but the other node's hanging.
	}

	rec, err := dht.checkLocalDatastore(k)
	if err != nil {
		return nil, err
	}
	resp.Record = rec

	// Find closest peer on given cluster to desired key and reply with that info
	closer := dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	if len(closer) > 0 {
		closerinfos := peer.PeerInfos(dht.peerstore, closer)
		for _, pi := range closerinfos {
			log.Debugf("handleGetValue returning closer peer: '%s'", pi.ID)
			if len(pi.Addrs) < 1 {
				log.Errorf(`no addresses on peer being sent!
					[local:%s]
					[sending:%s]
					[remote:%s]`, dht.self, pi.ID, p)
			}
		}

		resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), closerinfos)
	}

	return resp, nil
}

func (dht *IpfsDHT) checkLocalDatastore(k key.Key) (*pb.Record, error) {
	log.Debugf("%s handleGetValue looking into ds", dht.self)
	dskey := k.DsKey()
	iVal, err := dht.datastore.Get(dskey)
	log.Debugf("%s handleGetValue looking into ds GOT %v", dht.self, iVal)

	if err == ds.ErrNotFound {
		return nil, nil
	}

	// if we got an unexpected error, bail.
	if err != nil {
		return nil, err
	}

	// if we have the value, send it back
	log.Debugf("%s handleGetValue success!", dht.self)

	byts, ok := iVal.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}

	rec := new(pb.Record)
	err = proto.Unmarshal(byts, rec)
	if err != nil {
		log.Debug("Failed to unmarshal DHT record from datastore.")
		return nil, err
	}

	// if its our record, dont bother checking the times on it
	if peer.ID(rec.GetAuthor()) == dht.self {
		return rec, nil
	}

	var recordIsBad bool
	recvtime, err := u.ParseRFC3339(rec.GetTimeReceived())
	if err != nil {
		log.Info("either no receive time set on record, or it was invalid: ", err)
		recordIsBad = true
	}

	if time.Now().Sub(recvtime) > MaxRecordAge {
		log.Debug("old record found, tossing.")
		recordIsBad = true
	}

	// NOTE: We do not verify the record here beyond checking these timestamps.
	// we put the burden of checking the records on the requester as checking a record
	// may be computationally expensive

	if recordIsBad {
		err := dht.datastore.Delete(dskey)
		if err != nil {
			log.Error("Failed to delete bad record from datastore: ", err)
		}

		return nil, nil // can treat this as not having the record at all
	}

	return rec, nil
}

// Store a value in this peer local storage
func (dht *IpfsDHT) handlePutValue(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handlePutValue", p).Done()
	dskey := key.Key(pmes.GetKey()).DsKey()

	if err := dht.verifyRecordLocally(pmes.GetRecord()); err != nil {
		log.Warningf("Bad dht record in PUT from: %s. %s", key.Key(pmes.GetRecord().GetAuthor()), err)
		return nil, err
	}

	rec := pmes.GetRecord()

	// record the time we receive every record
	rec.TimeReceived = proto.String(u.FormatRFC3339(time.Now()))

	data, err := proto.Marshal(rec)
	if err != nil {
		return nil, err
	}

	err = dht.datastore.Put(dskey, data)
	log.Debugf("%s handlePutValue %v", dht.self, dskey)
	return pmes, err
}

func (dht *IpfsDHT) handlePing(_ context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	log.Debugf("%s Responding to ping from %s!\n", dht.self, p)
	return pmes, nil
}

func (dht *IpfsDHT) handleFindPeer(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	defer log.EventBegin(ctx, "handleFindPeer", p).Done()
	resp := pb.NewMessage(pmes.GetType(), "", pmes.GetClusterLevel())
	var closest []peer.ID

	// if looking for self... special case where we send it on CloserPeers.
	if peer.ID(pmes.GetKey()) == dht.self {
		closest = []peer.ID{dht.self}
	} else {
		closest = dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	}

	if closest == nil {
		log.Infof("%s handleFindPeer %s: could not find anything.", dht.self, p)
		return resp, nil
	}

	var withAddresses []peer.PeerInfo
	closestinfos := peer.PeerInfos(dht.peerstore, closest)
	for _, pi := range closestinfos {
		if len(pi.Addrs) > 0 {
			withAddresses = append(withAddresses, pi)
			log.Debugf("handleFindPeer: sending back '%s'", pi.ID)
		}
	}

	resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), withAddresses)
	return resp, nil
}

func (dht *IpfsDHT) handleGetProviders(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	lm := make(lgbl.DeferredMap)
	lm["peer"] = func() interface{} { return p.Pretty() }
	defer log.EventBegin(ctx, "handleGetProviders", lm).Done()

	resp := pb.NewMessage(pmes.GetType(), pmes.GetKey(), pmes.GetClusterLevel())
	key := key.Key(pmes.GetKey())
	lm["key"] = func() interface{} { return key.B58String() }

	// debug logging niceness.
	reqDesc := fmt.Sprintf("%s handleGetProviders(%s, %s): ", dht.self, p, key)
	log.Debugf("%s begin", reqDesc)
	defer log.Debugf("%s end", reqDesc)

	// check if we have this value, to add ourselves as provider.
	has, err := dht.datastore.Has(key.DsKey())
	if err != nil && err != ds.ErrNotFound {
		log.Debugf("unexpected datastore error: %v\n", err)
		has = false
	}

	// setup providers
	providers := dht.providers.GetProviders(ctx, key)
	if has {
		providers = append(providers, dht.self)
		log.Debugf("%s have the value. added self as provider", reqDesc)
	}

	if providers != nil && len(providers) > 0 {
		infos := peer.PeerInfos(dht.peerstore, providers)
		resp.ProviderPeers = pb.PeerInfosToPBPeers(dht.host.Network(), infos)
		log.Debugf("%s have %d providers: %s", reqDesc, len(providers), infos)
	}

	// Also send closer peers.
	closer := dht.betterPeersToQuery(pmes, p, CloserPeerCount)
	if closer != nil {
		infos := peer.PeerInfos(dht.peerstore, closer)
		resp.CloserPeers = pb.PeerInfosToPBPeers(dht.host.Network(), infos)
		log.Debugf("%s have %d closer peers: %s", reqDesc, len(closer), infos)
	}

	return resp, nil
}

func (dht *IpfsDHT) handleAddProvider(ctx context.Context, p peer.ID, pmes *pb.Message) (*pb.Message, error) {
	lm := make(lgbl.DeferredMap)
	lm["peer"] = func() interface{} { return p.Pretty() }

	defer log.EventBegin(ctx, "handleAddProvider", lm).Done()
	key := key.Key(pmes.GetKey())
	lm["key"] = func() interface{} { return key.B58String() }

	log.Debugf("%s adding %s as a provider for '%s'\n", dht.self, p, key)

	// add provider should use the address given in the message
	pinfos := pb.PBPeersToPeerInfos(pmes.GetProviderPeers())
	for _, pi := range pinfos {
		// We need to disable this check for OpenBazaar pointers to work
		/*if pi.ID != p {
			// we should ignore this provider reccord! not from originator.
			// (we chould sign them and check signature later...)
			log.Debugf("handleAddProvider received provider %s from %s. Ignore.", pi.ID, p)
			continue
		}*/

		if len(pi.Addrs) < 1 {
			log.Debugf("%s got no valid addresses for provider %s. Ignore.", dht.self, p)
			continue
		}

		log.Infof("received provider %s for %s (addrs: %s)", p, key, pi.Addrs)
		if pi.ID != dht.self { // dont add own addrs.
			// add the received addresses to our peerstore.
			dht.peerstore.AddAddrs(pi.ID, pi.Addrs, peer.ProviderAddrTTL)
		}
		dht.providers.AddProvider(ctx, key, pi.ID)
	}

	return nil, nil
}
<|endoftext|>


<commit_before>package micro

import (
	"context"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/cmd"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/selector"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/transport"
)

type Options struct {
	Broker    broker.Broker
	Cmd       cmd.Cmd
	Client    client.Client
	Server    server.Server
	Registry  registry.Registry
	Transport transport.Transport

	// Register loop interval
	RegisterInterval time.Duration

	// Before and After funcs
	BeforeStart []func() error
	BeforeStop  []func() error
	AfterStart  []func() error
	AfterStop   []func() error

	// Other options for implementations of the interface
	// can be stored in a context
	Context context.Context
}

func newOptions(opts ...Option) Options {
	opt := Options{
		Broker:    broker.DefaultBroker,
		Cmd:       cmd.DefaultCmd,
		Client:    client.DefaultClient,
		Server:    server.DefaultServer,
		Registry:  registry.DefaultRegistry,
		Transport: transport.DefaultTransport,
		Context:   context.Background(),
	}

	for _, o := range opts {
		o(&opt)
	}

	return opt
}

func Broker(b broker.Broker) Option {
	return func(o *Options) {
		o.Broker = b
		// Update Client and Server
		o.Client.Init(client.Broker(b))
		o.Server.Init(server.Broker(b))
	}
}

func Cmd(c cmd.Cmd) Option {
	return func(o *Options) {
		o.Cmd = c
	}
}

func Client(c client.Client) Option {
	return func(o *Options) {
		o.Client = c
	}
}

// Context specifies a context for the service.
// Can be used to signal shutdown of the service.
// Can be used for extra option values.
func Context(ctx context.Context) Option {
	return func(o *Options) {
		o.Context = ctx
	}
}

func Server(s server.Server) Option {
	return func(o *Options) {
		o.Server = s
	}
}

// Registry sets the registry for the service
// and the underlying components
func Registry(r registry.Registry) Option {
	return func(o *Options) {
		o.Registry = r
		// Update Client and Server
		o.Client.Init(client.Registry(r))
		o.Server.Init(server.Registry(r))
		// Update Selector
		o.Client.Options().Selector.Init(selector.Registry(r))
	}
}

// Selector sets the selector for the service client
func Selector(s selector.Selector) Option {
	return func(o *Options) {
		o.Client.Init(client.Selector(s))
	}
}

// Transport sets the transport for the service
// and the underlying components
func Transport(t transport.Transport) Option {
	return func(o *Options) {
		o.Transport = t
		// Update Client and Server
		o.Client.Init(client.Transport(t))
		o.Server.Init(server.Transport(t))
	}
}

// Convenience options

// Name of the service
func Name(n string) Option {
	return func(o *Options) {
		o.Server.Init(server.Name(n))
	}
}

// Version of the service
func Version(v string) Option {
	return func(o *Options) {
		o.Server.Init(server.Version(v))
	}
}

// Metadata associated with the service
func Metadata(md map[string]string) Option {
	return func(o *Options) {
		o.Server.Init(server.Metadata(md))
	}
}

func Flags(flags ...cli.Flag) Option {
	return func(o *Options) {
		o.Cmd.App().Flags = append(o.Cmd.App().Flags, flags...)
	}
}

func Action(a func(*cli.Context)) Option {
	return func(o *Options) {
		o.Cmd.App().Action = a
	}
}

// RegisterTTL specifies the TTL to use when registering the service
func RegisterTTL(t time.Duration) Option {
	return func(o *Options) {
		o.Server.Init(server.RegisterTTL(t))
	}
}

// RegisterInterval specifies the interval on which to re-register
func RegisterInterval(t time.Duration) Option {
	return func(o *Options) {
		o.RegisterInterval = t
	}
}

// WrapClient is a convenience method for wrapping a Client with
// some middleware component. A list of wrappers can be provided.
func WrapClient(w ...client.Wrapper) Option {
	return func(o *Options) {
		// apply in reverse
		for i := len(w); i > 0; i-- {
			o.Client = w[i-1](o.Client)
		}
	}
}

// WrapCall is a convenience method for wrapping a Client CallFunc
func WrapCall(w ...client.CallWrapper) Option {
	return func(o *Options) {
		o.Client.Init(client.WrapCall(w...))
	}
}

// WrapHandler adds a handler Wrapper to a list of options passed into the server
func WrapHandler(w ...server.HandlerWrapper) Option {
	return func(o *Options) {
		var wrappers []server.Option

		for _, wrap := range w {
			wrappers = append(wrappers, server.WrapHandler(wrap))
		}

		// Init once
		o.Server.Init(wrappers...)
	}
}

// WrapSubscriber adds a subscriber Wrapper to a list of options passed into the server
func WrapSubscriber(w ...server.SubscriberWrapper) Option {
	return func(o *Options) {
		var wrappers []server.Option

		for _, wrap := range w {
			wrappers = append(wrappers, server.WrapSubscriber(wrap))
		}

		// Init once
		o.Server.Init(wrappers...)
	}
}

// Before and Afters

func BeforeStart(fn func() error) Option {
	return func(o *Options) {
		o.BeforeStart = append(o.BeforeStart, fn)
	}
}

func BeforeStop(fn func() error) Option {
	return func(o *Options) {
		o.BeforeStop = append(o.BeforeStop, fn)
	}
}

func AfterStart(fn func() error) Option {
	return func(o *Options) {
		o.AfterStart = append(o.AfterStart, fn)
	}
}

func AfterStop(fn func() error) Option {
	return func(o *Options) {
		o.AfterStop = append(o.AfterStop, fn)
	}
}
<commit_msg>Fix broker registry issue<commit_after>package micro

import (
	"context"
	"time"

	"github.com/micro/cli"
	"github.com/micro/go-micro/broker"
	"github.com/micro/go-micro/client"
	"github.com/micro/go-micro/cmd"
	"github.com/micro/go-micro/registry"
	"github.com/micro/go-micro/selector"
	"github.com/micro/go-micro/server"
	"github.com/micro/go-micro/transport"
)

type Options struct {
	Broker    broker.Broker
	Cmd       cmd.Cmd
	Client    client.Client
	Server    server.Server
	Registry  registry.Registry
	Transport transport.Transport

	// Register loop interval
	RegisterInterval time.Duration

	// Before and After funcs
	BeforeStart []func() error
	BeforeStop  []func() error
	AfterStart  []func() error
	AfterStop   []func() error

	// Other options for implementations of the interface
	// can be stored in a context
	Context context.Context
}

func newOptions(opts ...Option) Options {
	opt := Options{
		Broker:    broker.DefaultBroker,
		Cmd:       cmd.DefaultCmd,
		Client:    client.DefaultClient,
		Server:    server.DefaultServer,
		Registry:  registry.DefaultRegistry,
		Transport: transport.DefaultTransport,
		Context:   context.Background(),
	}

	for _, o := range opts {
		o(&opt)
	}

	return opt
}

func Broker(b broker.Broker) Option {
	return func(o *Options) {
		o.Broker = b
		// Update Client and Server
		o.Client.Init(client.Broker(b))
		o.Server.Init(server.Broker(b))
	}
}

func Cmd(c cmd.Cmd) Option {
	return func(o *Options) {
		o.Cmd = c
	}
}

func Client(c client.Client) Option {
	return func(o *Options) {
		o.Client = c
	}
}

// Context specifies a context for the service.
// Can be used to signal shutdown of the service.
// Can be used for extra option values.
func Context(ctx context.Context) Option {
	return func(o *Options) {
		o.Context = ctx
	}
}

func Server(s server.Server) Option {
	return func(o *Options) {
		o.Server = s
	}
}

// Registry sets the registry for the service
// and the underlying components
func Registry(r registry.Registry) Option {
	return func(o *Options) {
		o.Registry = r
		// Update Client and Server
		o.Client.Init(client.Registry(r))
		o.Server.Init(server.Registry(r))
		// Update Selector
		o.Client.Options().Selector.Init(selector.Registry(r))
		// Update Broker
		o.Broker.Init(broker.Registry(r))
	}
}

// Selector sets the selector for the service client
func Selector(s selector.Selector) Option {
	return func(o *Options) {
		o.Client.Init(client.Selector(s))
	}
}

// Transport sets the transport for the service
// and the underlying components
func Transport(t transport.Transport) Option {
	return func(o *Options) {
		o.Transport = t
		// Update Client and Server
		o.Client.Init(client.Transport(t))
		o.Server.Init(server.Transport(t))
	}
}

// Convenience options

// Name of the service
func Name(n string) Option {
	return func(o *Options) {
		o.Server.Init(server.Name(n))
	}
}

// Version of the service
func Version(v string) Option {
	return func(o *Options) {
		o.Server.Init(server.Version(v))
	}
}

// Metadata associated with the service
func Metadata(md map[string]string) Option {
	return func(o *Options) {
		o.Server.Init(server.Metadata(md))
	}
}

func Flags(flags ...cli.Flag) Option {
	return func(o *Options) {
		o.Cmd.App().Flags = append(o.Cmd.App().Flags, flags...)
	}
}

func Action(a func(*cli.Context)) Option {
	return func(o *Options) {
		o.Cmd.App().Action = a
	}
}

// RegisterTTL specifies the TTL to use when registering the service
func RegisterTTL(t time.Duration) Option {
	return func(o *Options) {
		o.Server.Init(server.RegisterTTL(t))
	}
}

// RegisterInterval specifies the interval on which to re-register
func RegisterInterval(t time.Duration) Option {
	return func(o *Options) {
		o.RegisterInterval = t
	}
}

// WrapClient is a convenience method for wrapping a Client with
// some middleware component. A list of wrappers can be provided.
func WrapClient(w ...client.Wrapper) Option {
	return func(o *Options) {
		// apply in reverse
		for i := len(w); i > 0; i-- {
			o.Client = w[i-1](o.Client)
		}
	}
}

// WrapCall is a convenience method for wrapping a Client CallFunc
func WrapCall(w ...client.CallWrapper) Option {
	return func(o *Options) {
		o.Client.Init(client.WrapCall(w...))
	}
}

// WrapHandler adds a handler Wrapper to a list of options passed into the server
func WrapHandler(w ...server.HandlerWrapper) Option {
	return func(o *Options) {
		var wrappers []server.Option

		for _, wrap := range w {
			wrappers = append(wrappers, server.WrapHandler(wrap))
		}

		// Init once
		o.Server.Init(wrappers...)
	}
}

// WrapSubscriber adds a subscriber Wrapper to a list of options passed into the server
func WrapSubscriber(w ...server.SubscriberWrapper) Option {
	return func(o *Options) {
		var wrappers []server.Option

		for _, wrap := range w {
			wrappers = append(wrappers, server.WrapSubscriber(wrap))
		}

		// Init once
		o.Server.Init(wrappers...)
	}
}

// Before and Afters

func BeforeStart(fn func() error) Option {
	return func(o *Options) {
		o.BeforeStart = append(o.BeforeStart, fn)
	}
}

func BeforeStop(fn func() error) Option {
	return func(o *Options) {
		o.BeforeStop = append(o.BeforeStop, fn)
	}
}

func AfterStart(fn func() error) Option {
	return func(o *Options) {
		o.AfterStart = append(o.AfterStart, fn)
	}
}

func AfterStop(fn func() error) Option {
	return func(o *Options) {
		o.AfterStop = append(o.AfterStop, fn)
	}
}
<|endoftext|>


<commit_before>package fst

import (
	"fmt"
	"github.com/balzaczyy/golucene/core/util"
)

// util/fst/Outputs.java

/**
 * Represents the outputs for an FST, providing the basic
 * algebra required for building and traversing the FST.
 *
 *
Note that any operation that returns NO_OUTPUT must
 * return the same singleton object from {@link
 * #getNoOutput}.

 */
type Outputs interface {
	// Eg subtract("foobar", "foo") -> "bar"
	Subtract(output1, output2 interface{}) interface{}
	/** Eg add("foo", "bar") -> "foobar" */
	Add(prefix interface{}, output interface{}) interface{}
	// Encode an output value into a DataOutput.
	Write(interface{}, util.DataOutput) error
	// Encode an final node output value into a DataOutput. By default
	// this just calls write()
	writeFinalOutput(interface{}, util.DataOutput) error
	/** Decode an output value previously written with {@link
	 *  #write(Object, DataOutput)}. */
	Read(in util.DataInput) (e interface{}, err error)
	// Skip the output; defaults to just calling Read() and discarding the result
	SkipOutput(util.DataInput) error
	/** Decode an output value previously written with {@link
	 *  #writeFinalOutput(Object, DataOutput)}.  By default this
	 *  just calls {@link #read(DataInput)}. */
	ReadFinalOutput(in util.DataInput) (e interface{}, err error)
	// Skip the output previously written with WriteFinalOutput;
	// defaults to just calling ReadFinalOutput and discarding the
	// result.
	SkipFinalOutput(util.DataInput) error
	/** NOTE: this output is compared with == so you must
	 *  ensure that all methods return the single object if
	 *  it's really no output */
	NoOutput() interface{}
	outputToString(interface{}) string
	merge(first, second interface{}) interface{}
}

type iOutputsReader interface {
	Read(in util.DataInput) (e interface{}, err error)
	Write(interface{}, util.DataOutput) error
}

type abstractOutputs struct {
	spi iOutputsReader
}

func (out *abstractOutputs) writeFinalOutput(output interface{}, o util.DataOutput) error {
	return out.spi.Write(output, o)
}

func (out *abstractOutputs) SkipOutput(in util.DataInput) error {
	_, err := out.spi.Read(in)
	return err
}

/* Decode an output value previously written with writeFinalOutput(). By default this just calls read(). */
func (out *abstractOutputs) ReadFinalOutput(in util.DataInput) (e interface{}, err error) {
	return out.spi.Read(in)
}

func (out *abstractOutputs) SkipFinalOutput(in util.DataInput) error {
	return out.SkipOutput(in)
}

func (out *abstractOutputs) merge(first, second interface{}) interface{} {
	panic("not supported yet")
}

// fst/NoOutputs.java

var NO_OUTPUT = newNoOutputs()

/* A nil FST Outputs implementation; use this if you just want to build an FSA. */
type NoOutputs struct {
	*abstractOutputs
}

func newNoOutputs() *NoOutputs {
	ans := &NoOutputs{}
	ans.abstractOutputs = &abstractOutputs{ans}
	return ans
}

func (o *NoOutputs) Subtract(output1, output2 interface{}) interface{} {
	assert(output1 == NO_OUTPUT)
	assert(output2 == NO_OUTPUT)
	return NO_OUTPUT
}

func (o *NoOutputs) Add(prefix, output interface{}) interface{} {
	panic("not implemented yet")
}

func (o *NoOutputs) merge(first, second interface{}) interface{} {
	assert(first == NO_OUTPUT)
	assert(second == NO_OUTPUT)
	return NO_OUTPUT
}

func (o *NoOutputs) Write(prefix interface{}, out util.DataOutput) error {
	return nil
}

func (o *NoOutputs) Read(in util.DataInput) (interface{}, error) {
	return NO_OUTPUT, nil
}

func (o *NoOutputs) NoOutput() interface{} {
	return NO_OUTPUT
}

func (o *NoOutputs) outputToString(output interface{}) string {
	return ""
}

// fst/ByteSequenceOutputs.java

/**
 * An FST {@link Outputs} implementation where each output
 * is a sequence of bytes.
 */
type ByteSequenceOutputs struct {
	*abstractOutputs
}

var oneByteSequenceOutputs *ByteSequenceOutputs

func ByteSequenceOutputsSingleton() *ByteSequenceOutputs {
	if oneByteSequenceOutputs == nil {
		oneByteSequenceOutputs = &ByteSequenceOutputs{}
		oneByteSequenceOutputs.abstractOutputs = &abstractOutputs{oneByteSequenceOutputs}
	}
	return oneByteSequenceOutputs
}

func (out *ByteSequenceOutputs) Subtract(output1, output2 interface{}) interface{} {
	panic("not implemented yet")
}

func (out *ByteSequenceOutputs) Add(_prefix interface{}, _output interface{}) interface{} {
	assert(_prefix != nil)
	assert(_output != nil)
	if _prefix == NO_OUTPUT {
		return _output
	} else if _output == NO_OUTPUT {
		return _prefix
	} else {
		prefix, output := _prefix.([]byte), _output.([]byte)
		assert(len(prefix) > 0)
		assert(len(output) > 0)
		result := make([]byte, len(prefix)+len(output))
		copy(result, prefix)
		copy(result[len(prefix):], output)
		return result
	}
}

func (o *ByteSequenceOutputs) Write(obj interface{}, out util.DataOutput) error {
	assert(obj != nil)
	prefix, ok := obj.([]byte)
	assert(ok)
	err := out.WriteVInt(int32(len(prefix)))
	if err == nil {
		err = out.WriteBytes(prefix)
	}
	return err
}

func (out *ByteSequenceOutputs) Read(in util.DataInput) (e interface{}, err error) {
	if length, err := in.ReadVInt(); err == nil {
		// fmt.Printf("Length: %v\n", length)
		if length == 0 {
			e = out.NoOutput()
		} else {
			buf := make([]byte, length)
			e = buf
			err = in.ReadBytes(buf)
		}
	} else {
		fmt.Printf("Failed to read length due to %v", err)
	}
	return e, err
}

func (out *ByteSequenceOutputs) NoOutput() interface{} {
	return NO_OUTPUT
}

func (out *ByteSequenceOutputs) outputToString(output interface{}) string {
	return fmt.Sprintf("%v", output.([]byte))
}

func (out *ByteSequenceOutputs) String() string {
	return "ByteSequenceOutputs"
}

// util/fst/Util.java

/** Looks up the output for this input, or null if the
 *  input is not accepted */
func GetFSTOutput(fst *FST, input []byte) (output interface{}, err error) {
	if fst.inputType != INPUT_TYPE_BYTE1 {
		panic("assert fail")
	}
	fstReader := fst.BytesReader()
	// TODO: would be nice not to alloc this on every lookup
	arc := fst.FirstArc(&Arc{})

	// Accumulate output as we go
	output = fst.outputs.NoOutput()
	for _, v := range input {
		ret, err := fst.FindTargetArc(int(v), arc, arc, fstReader)
		if ret == nil || err != nil {
			return ret, err
		}
		output = fst.outputs.Add(output, arc.Output)
	}

	if arc.IsFinal() {
		return fst.outputs.Add(output, arc.NextFinalOutput), nil
	} else {
		return nil, nil
	}
}
<commit_msg>remove unnecessary cast<commit_after>package fst

import (
	"fmt"
	"github.com/balzaczyy/golucene/core/util"
)

// util/fst/Outputs.java

/**
 * Represents the outputs for an FST, providing the basic
 * algebra required for building and traversing the FST.
 *
 *
Note that any operation that returns NO_OUTPUT must
 * return the same singleton object from {@link
 * #getNoOutput}.

 */
type Outputs interface {
	// Eg subtract("foobar", "foo") -> "bar"
	Subtract(output1, output2 interface{}) interface{}
	/** Eg add("foo", "bar") -> "foobar" */
	Add(prefix interface{}, output interface{}) interface{}
	// Encode an output value into a DataOutput.
	Write(interface{}, util.DataOutput) error
	// Encode an final node output value into a DataOutput. By default
	// this just calls write()
	writeFinalOutput(interface{}, util.DataOutput) error
	/** Decode an output value previously written with {@link
	 *  #write(Object, DataOutput)}. */
	Read(in util.DataInput) (e interface{}, err error)
	// Skip the output; defaults to just calling Read() and discarding the result
	SkipOutput(util.DataInput) error
	/** Decode an output value previously written with {@link
	 *  #writeFinalOutput(Object, DataOutput)}.  By default this
	 *  just calls {@link #read(DataInput)}. */
	ReadFinalOutput(in util.DataInput) (e interface{}, err error)
	// Skip the output previously written with WriteFinalOutput;
	// defaults to just calling ReadFinalOutput and discarding the
	// result.
	SkipFinalOutput(util.DataInput) error
	/** NOTE: this output is compared with == so you must
	 *  ensure that all methods return the single object if
	 *  it's really no output */
	NoOutput() interface{}
	outputToString(interface{}) string
	merge(first, second interface{}) interface{}
}

type iOutputsReader interface {
	Read(in util.DataInput) (e interface{}, err error)
	Write(interface{}, util.DataOutput) error
}

type abstractOutputs struct {
	spi iOutputsReader
}

func (out *abstractOutputs) writeFinalOutput(output interface{}, o util.DataOutput) error {
	return out.spi.Write(output, o)
}

func (out *abstractOutputs) SkipOutput(in util.DataInput) error {
	_, err := out.spi.Read(in)
	return err
}

/* Decode an output value previously written with writeFinalOutput(). By default this just calls read(). */
func (out *abstractOutputs) ReadFinalOutput(in util.DataInput) (e interface{}, err error) {
	return out.spi.Read(in)
}

func (out *abstractOutputs) SkipFinalOutput(in util.DataInput) error {
	return out.SkipOutput(in)
}

func (out *abstractOutputs) merge(first, second interface{}) interface{} {
	panic("not supported yet")
}

// fst/NoOutputs.java

var NO_OUTPUT = newNoOutputs()

/* A nil FST Outputs implementation; use this if you just want to build an FSA. */
type NoOutputs struct {
	*abstractOutputs
}

func newNoOutputs() *NoOutputs {
	ans := &NoOutputs{}
	ans.abstractOutputs = &abstractOutputs{ans}
	return ans
}

func (o *NoOutputs) Subtract(output1, output2 interface{}) interface{} {
	assert(output1 == NO_OUTPUT)
	assert(output2 == NO_OUTPUT)
	return NO_OUTPUT
}

func (o *NoOutputs) Add(prefix, output interface{}) interface{} {
	panic("not implemented yet")
}

func (o *NoOutputs) merge(first, second interface{}) interface{} {
	assert(first == NO_OUTPUT)
	assert(second == NO_OUTPUT)
	return NO_OUTPUT
}

func (o *NoOutputs) Write(prefix interface{}, out util.DataOutput) error {
	return nil
}

func (o *NoOutputs) Read(in util.DataInput) (interface{}, error) {
	return NO_OUTPUT, nil
}

func (o *NoOutputs) NoOutput() interface{} {
	return NO_OUTPUT
}

func (o *NoOutputs) outputToString(output interface{}) string {
	return ""
}

// fst/ByteSequenceOutputs.java

/**
 * An FST {@link Outputs} implementation where each output
 * is a sequence of bytes.
 */
type ByteSequenceOutputs struct {
	*abstractOutputs
}

var oneByteSequenceOutputs *ByteSequenceOutputs

func ByteSequenceOutputsSingleton() *ByteSequenceOutputs {
	if oneByteSequenceOutputs == nil {
		oneByteSequenceOutputs = &ByteSequenceOutputs{}
		oneByteSequenceOutputs.abstractOutputs = &abstractOutputs{oneByteSequenceOutputs}
	}
	return oneByteSequenceOutputs
}

func (out *ByteSequenceOutputs) Subtract(output1, output2 interface{}) interface{} {
	panic("not implemented yet")
}

func (out *ByteSequenceOutputs) Add(_prefix interface{}, _output interface{}) interface{} {
	assert(_prefix != nil)
	assert(_output != nil)
	if _prefix == NO_OUTPUT {
		return _output
	} else if _output == NO_OUTPUT {
		return _prefix
	} else {
		prefix, output := _prefix.([]byte), _output.([]byte)
		assert(len(prefix) > 0)
		assert(len(output) > 0)
		result := make([]byte, len(prefix)+len(output))
		copy(result, prefix)
		copy(result[len(prefix):], output)
		return result
	}
}

func (o *ByteSequenceOutputs) Write(obj interface{}, out util.DataOutput) error {
	assert(obj != nil)
	prefix, ok := obj.([]byte)
	assert(ok)
	err := out.WriteVInt(int32(len(prefix)))
	if err == nil {
		err = out.WriteBytes(prefix)
	}
	return err
}

func (out *ByteSequenceOutputs) Read(in util.DataInput) (e interface{}, err error) {
	if length, err := in.ReadVInt(); err == nil {
		// fmt.Printf("Length: %v\n", length)
		if length == 0 {
			e = out.NoOutput()
		} else {
			buf := make([]byte, length)
			e = buf
			err = in.ReadBytes(buf)
		}
	} else {
		fmt.Printf("Failed to read length due to %v", err)
	}
	return e, err
}

func (out *ByteSequenceOutputs) NoOutput() interface{} {
	return NO_OUTPUT
}

func (out *ByteSequenceOutputs) outputToString(output interface{}) string {
	return fmt.Sprintf("%v", output)
}

func (out *ByteSequenceOutputs) String() string {
	return "ByteSequenceOutputs"
}

// util/fst/Util.java

/** Looks up the output for this input, or null if the
 *  input is not accepted */
func GetFSTOutput(fst *FST, input []byte) (output interface{}, err error) {
	if fst.inputType != INPUT_TYPE_BYTE1 {
		panic("assert fail")
	}
	fstReader := fst.BytesReader()
	// TODO: would be nice not to alloc this on every lookup
	arc := fst.FirstArc(&Arc{})

	// Accumulate output as we go
	output = fst.outputs.NoOutput()
	for _, v := range input {
		ret, err := fst.FindTargetArc(int(v), arc, arc, fstReader)
		if ret == nil || err != nil {
			return ret, err
		}
		output = fst.outputs.Add(output, arc.Output)
	}

	if arc.IsFinal() {
		return fst.outputs.Add(output, arc.NextFinalOutput), nil
	} else {
		return nil, nil
	}
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"testing"
	"time"

	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils/crypto"
)

func TestAddDefaultSubnetDelegatorTxSyntacticVerify(t *testing.T) {
	vm := defaultVM()
	vm.Ctx.Lock.Lock()
	defer func() {
		vm.Shutdown()
		vm.Ctx.Lock.Unlock()
	}()

	// Case 1: tx is nil
	var tx *addDefaultSubnetDelegatorTx
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because tx is nil")
	}

	// Case 2: Tx ID is nil
	tx, err := vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	tx.id = ids.ID{}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because ID is nil")
	}

	// Case 3: Wrong network ID
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID+1,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because the wrong network ID was used")
	}

	// Case 4: Missing Node ID
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	tx.NodeID = ids.ShortID{}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because NodeID is nil")
	}

	// Case 5: Not enough weight
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		MinimumStakeAmount-1,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because of not enough weight")
	}

	// Case 6: Validation length is too short
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateStartTime.Add(MinimumStakingDuration).Unix())-1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	err = tx.SyntacticVerify()
	if err == nil {
		t.Fatal("should have errored because validation length too short")
	}

	// Case 7: Validation length is too long
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateStartTime.Add(MaximumStakingDuration).Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	err = tx.SyntacticVerify()
	if err == nil {
		t.Fatal("should have errored because validation length too long")
	}

	// Case 8: Valid
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err != nil {
		t.Fatal(err)
	}
}

func TestAddDefaultSubnetDelegatorTxSemanticVerify(t *testing.T) {
	vm := defaultVM()
	vm.Ctx.Lock.Lock()
	defer func() {
		vm.Shutdown()
		vm.Ctx.Lock.Unlock()
	}()

	// Case 1: Proposed validator currently validating default subnet
	// but stops validating non-default subnet after stops validating default subnet
	// (note that defaultKey is a genesis validator)
	tx, err := vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator stops validating default subnet earlier than non-default subnet")
	}

	// Case 2: Proposed validator currently validating default subnet
	// and proposed non-default subnet validation period is subset of
	// default subnet validation period
	// (note that defaultKey is a genesis validator)
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatalf("should have failed because the end time is outside the default subnets end time")
	}

	// Add a validator to pending validator set of default subnet
	key, err := vm.factory.NewPrivateKey()
	if err != nil {
		t.Fatal(err)
	}
	pendingDSValidatorID := key.PublicKey().Address()

	// starts validating default subnet 10 seconds after genesis
	DSStartTime := defaultGenesisTime.Add(10 * time.Second)
	DSEndTime := DSStartTime.Add(5 * MinimumStakingDuration)

	addDSTx, err := vm.newAddDefaultSubnetValidatorTx(
		defaultNonce+1,                   // nonce
		defaultStakeAmount,               // stake amount
		uint64(DSStartTime.Unix()),       // start time
		uint64(DSEndTime.Unix()),         // end time
		pendingDSValidatorID,             // node ID
		defaultKey.PublicKey().Address(), // destination
		NumberOfShares,                   // subnet
		testNetworkID,                    // network
		defaultKey,                       // key
	)
	if err != nil {
		t.Fatal(err)
	}

	// Case 3: Proposed validator isn't in pending or current validator sets
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()),
		uint64(DSEndTime.Unix()),
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator not in the current or pending validator sets of the default subnet")
	}

	err = vm.putPendingValidators(
		vm.DB,
		&EventHeap{
			SortByStartTime: true,
			Txs:             []TimedTx{addDSTx},
		},
		DefaultSubnetID,
	)
	if err != nil {
		t.Fatal(err)
	}
	// Node with ID key.PublicKey().Address() now a pending validator for default subnet

	// Case 4: Proposed validator is pending validator of default subnet
	// but starts validating non-default subnet before default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix())-1, // start validating non-default subnet before default subnet
		uint64(DSEndTime.Unix()),
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator starts validating non-default " +
			"subnet before starting to validate default subnet")
	}

	// Case 5: Proposed validator is pending validator of default subnet
	// but stops validating non-default subnet after default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()),
		uint64(DSEndTime.Unix())+1, // stop validating non-default subnet after stopping validating default subnet
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator stops validating non-default " +
			"subnet after stops validating default subnet")
	}

	// Case 6: Proposed validator is pending validator of default subnet
	// and period validating non-default subnet is subset of time validating default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()), // same start time as for default subnet
		uint64(DSEndTime.Unix()),   // same end time as for default subnet
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err != nil {
		t.Fatalf("should have passed verification")
	}

	// Case 7: Proposed validator start validating at/before current timestamp
	// First, advance the timestamp
	newTimestamp := defaultGenesisTime.Add(2 * time.Second)
	if err := vm.putTimestamp(vm.DB, newTimestamp); err != nil {
		t.Fatal(err)
	}

	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,                                          // nonce
		defaultStakeAmount,                                      // weight
		uint64(newTimestamp.Unix()),                             // start time
		uint64(newTimestamp.Add(MinimumStakingDuration).Unix()), // end time
		defaultKey.PublicKey().Address(),                        // node ID
		defaultKey.PublicKey().Address(),                        // destination
		testNetworkID,                                           // network ID
		defaultKey,                                              // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed verification because starts validating at current timestamp")
	}

	// reset the timestamp
	if err := vm.putTimestamp(vm.DB, defaultGenesisTime); err != nil {
		t.Fatal(err)
	}

	// Case 7: Account that pays tx fee doesn't have enough $AVA to pay tx fee
	txFeeSaved := txFee
	txFee = 1 // Do this so test works even when txFee is 0

	// Create new key whose account has no $AVA
	factory := crypto.FactorySECP256K1R{}
	newAcctKey, err := factory.NewPrivateKey()
	if err != nil {
		t.Fatal(err)
	}

	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		1,                                         // nonce (new account has nonce 0 so use nonce 1)
		defaultStakeAmount,                        // weight
		uint64(defaultValidateStartTime.Unix()),   // start time
		uint64(defaultValidateEndTime.Unix()),     // end time
		defaultKey.PublicKey().Address(),          // node ID
		defaultKey.PublicKey().Address(),          // destination
		testNetworkID,                             // network ID
		newAcctKey.(*crypto.PrivateKeySECP256K1R), // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed verification because payer account has no $AVA to pay fee")
	}
	txFee = txFeeSaved // Reset tx fee
}
<commit_msg>Add tests for platform.addDefaultSubnetDelegator for confirming tx fails when attempting to delegate too much as well as confirming balance is correct after delegating.<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package platformvm

import (
	"testing"
	"time"

	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/utils/crypto"
)

func TestAddDefaultSubnetDelegatorTxSyntacticVerify(t *testing.T) {
	vm := defaultVM()
	vm.Ctx.Lock.Lock()
	defer func() {
		vm.Shutdown()
		vm.Ctx.Lock.Unlock()
	}()

	// Case 1: tx is nil
	var tx *addDefaultSubnetDelegatorTx
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because tx is nil")
	}

	// Case 2: Tx ID is nil
	tx, err := vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	tx.id = ids.ID{}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because ID is nil")
	}

	// Case 3: Wrong network ID
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID+1,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because the wrong network ID was used")
	}

	// Case 4: Missing Node ID
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	tx.NodeID = ids.ShortID{}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because NodeID is nil")
	}

	// Case 5: Not enough weight
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		MinimumStakeAmount-1,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err == nil {
		t.Fatal("should have errored because of not enough weight")
	}

	// Case 6: Validation length is too short
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateStartTime.Add(MinimumStakingDuration).Unix())-1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	err = tx.SyntacticVerify()
	if err == nil {
		t.Fatal("should have errored because validation length too short")
	}

	// Case 7: Validation length is too long
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateStartTime.Add(MaximumStakingDuration).Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	err = tx.SyntacticVerify()
	if err == nil {
		t.Fatal("should have errored because validation length too long")
	}

	// Case 8: Valid
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix()),
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := tx.SyntacticVerify(); err != nil {
		t.Fatal(err)
	}
}

func TestAddDefaultSubnetDelegatorTxSemanticVerify(t *testing.T) {
	vm := defaultVM()
	vm.Ctx.Lock.Lock()
	defer func() {
		vm.Shutdown()
		vm.Ctx.Lock.Unlock()
	}()

	// Case 1: Proposed validator currently validating default subnet
	// but stops validating non-default subnet after stops validating default subnet
	// (note that defaultKey is a genesis validator)
	tx, err := vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator stops validating default subnet earlier than non-default subnet")
	}

	// Case 2: Proposed validator currently validating default subnet
	// and proposed non-default subnet validation period is subset of
	// default subnet validation period
	// (note that defaultKey is a genesis validator)
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(defaultValidateStartTime.Unix()),
		uint64(defaultValidateEndTime.Unix())+1,
		defaultKey.PublicKey().Address(),
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatalf("should have failed because the end time is outside the default subnets end time")
	}

	// Add a validator to pending validator set of default subnet
	key, err := vm.factory.NewPrivateKey()
	if err != nil {
		t.Fatal(err)
	}
	pendingDSValidatorID := key.PublicKey().Address()

	// starts validating default subnet 10 seconds after genesis
	DSStartTime := defaultGenesisTime.Add(10 * time.Second)
	DSEndTime := DSStartTime.Add(5 * MinimumStakingDuration)

	addDSTx, err := vm.newAddDefaultSubnetValidatorTx(
		defaultNonce+1,                   // nonce
		defaultStakeAmount,               // stake amount
		uint64(DSStartTime.Unix()),       // start time
		uint64(DSEndTime.Unix()),         // end time
		pendingDSValidatorID,             // node ID
		defaultKey.PublicKey().Address(), // destination
		NumberOfShares,                   // subnet
		testNetworkID,                    // network
		defaultKey,                       // key
	)
	if err != nil {
		t.Fatal(err)
	}

	// Case 3: Proposed validator isn't in pending or current validator sets
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()),
		uint64(DSEndTime.Unix()),
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator not in the current or pending validator sets of the default subnet")
	}

	err = vm.putPendingValidators(
		vm.DB,
		&EventHeap{
			SortByStartTime: true,
			Txs:             []TimedTx{addDSTx},
		},
		DefaultSubnetID,
	)
	if err != nil {
		t.Fatal(err)
	}
	// Node with ID key.PublicKey().Address() now a pending validator for default subnet

	// Case 4: Proposed validator is pending validator of default subnet
	// but starts validating non-default subnet before default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix())-1, // start validating non-default subnet before default subnet
		uint64(DSEndTime.Unix()),
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator starts validating non-default " +
			"subnet before starting to validate default subnet")
	}

	// Case 5: Proposed validator is pending validator of default subnet
	// but stops validating non-default subnet after default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()),
		uint64(DSEndTime.Unix())+1, // stop validating non-default subnet after stopping validating default subnet
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed because validator stops validating non-default " +
			"subnet after stops validating default subnet")
	}

	// Case 6: Proposed validator is pending validator of default subnet
	// and period validating non-default subnet is subset of time validating default subnet
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,
		defaultStakeAmount,
		uint64(DSStartTime.Unix()), // same start time as for default subnet
		uint64(DSEndTime.Unix()),   // same end time as for default subnet
		pendingDSValidatorID,
		defaultKey.PublicKey().Address(),
		testNetworkID,
		defaultKey,
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err != nil {
		t.Fatalf("should have passed verification")
	}

	// Case 7: Proposed validator start validating at/before current timestamp
	// First, advance the timestamp
	newTimestamp := defaultGenesisTime.Add(2 * time.Second)
	if err := vm.putTimestamp(vm.DB, newTimestamp); err != nil {
		t.Fatal(err)
	}

	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		defaultNonce+1,                                          // nonce
		defaultStakeAmount,                                      // weight
		uint64(newTimestamp.Unix()),                             // start time
		uint64(newTimestamp.Add(MinimumStakingDuration).Unix()), // end time
		defaultKey.PublicKey().Address(),                        // node ID
		defaultKey.PublicKey().Address(),                        // destination
		testNetworkID,                                           // network ID
		defaultKey,                                              // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed verification because starts validating at current timestamp")
	}

	// reset the timestamp
	if err := vm.putTimestamp(vm.DB, defaultGenesisTime); err != nil {
		t.Fatal(err)
	}

	// Case 7: Account that pays tx fee doesn't have enough $AVA to pay tx fee
	txFeeSaved := txFee
	txFee = 1 // Do this so test works even when txFee is 0

	// Create new key whose account has no $AVA
	factory := crypto.FactorySECP256K1R{}
	newAcctKey, err := factory.NewPrivateKey()
	if err != nil {
		t.Fatal(err)
	}

	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		1,                                         // nonce (new account has nonce 0 so use nonce 1)
		defaultStakeAmount,                        // weight
		uint64(defaultValidateStartTime.Unix()),   // start time
		uint64(defaultValidateEndTime.Unix()),     // end time
		defaultKey.PublicKey().Address(),          // node ID
		defaultKey.PublicKey().Address(),          // destination
		testNetworkID,                             // network ID
		newAcctKey.(*crypto.PrivateKeySECP256K1R), // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed verification because payer account has no $AVA to pay fee")
	}
	txFee = txFeeSaved // Reset tx fee

	// Case 8: fail verification for spending more funds than it has
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		1,                                         // nonce (new account has nonce 0 so use nonce 1)
		defaultBalance*2,                          // weight
		uint64(defaultValidateStartTime.Unix()),   // start time
		uint64(defaultValidateEndTime.Unix()),     // end time
		defaultKey.PublicKey().Address(),          // node ID
		defaultKey.PublicKey().Address(),          // destination
		testNetworkID,                             // network ID
		newAcctKey.(*crypto.PrivateKeySECP256K1R), // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}
	_, _, _, _, err = tx.SemanticVerify(vm.DB)
	if err == nil {
		t.Fatal("should have failed verification because payer account spent twice the default balance")
	}

	// Case 9: Confirm balance is correct
	tx, err = vm.newAddDefaultSubnetDelegatorTx(
		1,                                         // nonce (new account has nonce 0 so use nonce 1)
		defaultStakeAmount,                        // weight
		uint64(defaultValidateStartTime.Unix()),   // start time
		uint64(defaultValidateEndTime.Unix()),     // end time
		defaultKey.PublicKey().Address(),          // node ID
		defaultKey.PublicKey().Address(),          // destination
		testNetworkID,                             // network ID
		newAcctKey.(*crypto.PrivateKeySECP256K1R), // tx fee payer
	)
	if err != nil {
		t.Fatal(err)
	}

	onCommitDB, _, _, _, err := tx.SemanticVerify(vm.DB)
	if err != nil {
		t.Fatal(err)
	}
	account, err := tx.vm.getAccount(onCommitDB, defaultKey.PublicKey().Address())
	if err != nil {
		t.Fatal(err)
	}
	balance := account.Balance

	if balance == defaultBalance-(defaultStakeAmount+txFee) {
		t.Fatal("")
	}
}
<|endoftext|>


<commit_before>package runner

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/iron-io/titan/common"
	"github.com/iron-io/titan/runner/agent"
	"github.com/iron-io/titan/runner/drivers"
	driverscommon "github.com/iron-io/titan/runner/drivers"
	"github.com/iron-io/titan/runner/drivers/docker"
	"github.com/iron-io/titan/runner/drivers/mock"
)

type Config struct {
	ID      string
	Image   string
	Timeout time.Duration
	AppName string
	Memory  uint64
	Env     map[string]string
	Stdout  io.Writer
	Stderr  io.Writer
}

type Runner struct {
	driver       drivers.Driver
	taskQueue    chan *containerTask
	ml           Logger
	availableMem int64
	usedMem      int64
	usedMemMutex sync.RWMutex
}

var (
	ErrTimeOutNoMemory = errors.New("Task timed out. No available memory.")
	ErrFullQueue       = errors.New("The runner queue is full")

	WaitMemoryTimeout = 10 * time.Second
)

func New(metricLogger Logger) (*Runner, error) {
	// TODO: Is this really required for Titan's driver?
	// Can we remove it?
	env := common.NewEnvironment(func(e *common.Environment) {})

	// TODO: Create a drivers.New(runnerConfig) in Titan
	driver, err := selectDriver("docker", env, &driverscommon.Config{})
	if err != nil {
		return nil, err
	}

	r := &Runner{
		driver:       driver,
		taskQueue:    make(chan *containerTask, 100),
		ml:           metricLogger,
		availableMem: getAvailableMemory(),
		usedMem:      0,
	}

	go r.queueHandler()

	return r, nil
}

// This routine checks for available memory;
// If there's memory then send signal to the task to proceed.
// If there's not available memory to run the task it waits
// If the task waits for more than X seconds it timeouts
func (r *Runner) queueHandler() {
	var task *containerTask
	var waitStart time.Time
	var waitTime time.Duration
	var timedOut bool
	for {
		select {
		case task = <-r.taskQueue:
			waitStart = time.Now()
			timedOut = false
		}

		// Loop waiting for available memory
		canRun := r.checkRequiredMem(task.cfg.Memory)
		for ; !canRun; canRun = r.checkRequiredMem(task.cfg.Memory) {
			waitTime = time.Since(waitStart)
			if waitTime > WaitMemoryTimeout {
				timedOut = true
				break
			}
			time.Sleep(time.Second)
		}

		metricBaseName := fmt.Sprintf("run.%s.", task.cfg.AppName)
		r.ml.LogTime(task.ctx, metricBaseName+"waittime", waitTime)

		if timedOut {
			// Send to a signal to this task saying it cannot run
			r.ml.LogCount(task.ctx, metricBaseName+"timeout", 1)
			task.canRun <- false
			continue
		}

		// Send a signal to this task saying it can run
		task.canRun <- true
	}
}

func (r *Runner) checkRequiredMem(req uint64) bool {
	r.usedMemMutex.RLock()
	defer r.usedMemMutex.RUnlock()
	return r.availableMem-r.usedMem/int64(req)*1024*1024 > 0
}

func (r *Runner) addUsedMem(used int64) {
	r.usedMemMutex.Lock()
	r.usedMem = r.usedMem + used*1024*1024
	if r.usedMem < 0 {
		r.usedMem = 0
	}
	r.usedMemMutex.Unlock()
}

func (r *Runner) checkMemAndUse(req uint64) bool {
	r.usedMemMutex.Lock()
	defer r.usedMemMutex.Unlock()

	used := int64(req) * 1024 * 1024

	if r.availableMem-r.usedMem/used < 0 {
		return false
	}

	r.usedMem += used

	return true
}

func (r *Runner) Run(ctx context.Context, cfg *Config) (drivers.RunResult, error) {
	var err error

	if cfg.Memory == 0 {
		cfg.Memory = 128
	}

	ctask := &containerTask{
		ctx:    ctx,
		cfg:    cfg,
		auth:   &agent.ConfigAuth{},
		canRun: make(chan bool),
	}

	metricBaseName := fmt.Sprintf("run.%s.", cfg.AppName)
	r.ml.LogCount(ctx, metricBaseName+"requests", 1)

	// Check if has enough available memory
	// If available, use it
	if !r.checkMemAndUse(cfg.Memory) {
		// If not, try add task to the queue
		select {
		case r.taskQueue <- ctask:
		default:
			// If queue is full, return error
			r.ml.LogCount(ctx, "queue.full", 1)
			return nil, ErrFullQueue
		}

		// If task was added to the queue, wait for permission
		if ok := <-ctask.canRun; !ok {
			// This task timed out, not available memory
			return nil, ErrTimeOutNoMemory
		}
	} else {
		r.ml.LogTime(ctx, metricBaseName+"waittime", 0)
	}

	closer, err := r.driver.Prepare(ctx, ctask)
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	metricStart := time.Now()

	result, err := r.driver.Run(ctx, ctask)
	if err != nil {
		return nil, err
	}

	r.addUsedMem(-1 * int64(cfg.Memory))

	if result.Status() == "success" {
		r.ml.LogCount(ctx, metricBaseName+"succeeded", 1)
	} else {
		r.ml.LogCount(ctx, metricBaseName+"error", 1)
	}

	metricElapsed := time.Since(metricStart)
	r.ml.LogTime(ctx, metricBaseName+"time", metricElapsed)
	r.ml.LogTime(ctx, "run.exec_time", metricElapsed)

	return result, nil
}

func (r Runner) EnsureImageExists(ctx context.Context, cfg *Config) error {
	ctask := &containerTask{
		cfg:  cfg,
		auth: &agent.ConfigAuth{},
	}

	err := r.driver.EnsureImageExists(ctx, ctask)
	if err != nil {
		return err
	}
	return nil
}

func selectDriver(driver string, env *common.Environment, conf *driverscommon.Config) (drivers.Driver, error) {
	switch driver {
	case "docker":
		docker := docker.NewDocker(env, *conf)
		return docker, nil
	case "mock":
		return mock.New(), nil
	}
	return nil, fmt.Errorf("driver %v not found", driver)
}

func getAvailableMemory() int64 {
	const tooBig = 322122547200 // #300GB or 0, biggest aws instance is 244GB

	var availableMemory uint64
	if os.Getenv("IGNORE_MEMORY") == "1" {
		availableMemory = tooBig
	} else {
		availableMemory, err := checkCgroup()
		if err != nil {
			logrus.WithError(err).Error("Error checking for cgroup memory limits, falling back to host memory available..")
		}
		if availableMemory > tooBig || availableMemory == 0 {
			// Then -m flag probably wasn't set, so use max available on system
			availableMemory, err = checkProc()
			if availableMemory > tooBig || availableMemory == 0 {
				logrus.WithError(err).Fatal("Your Linux version is too old (<3.14) then we can't get the proper information to . You must specify the maximum available memory by passing the -m command with docker run when starting the runner via docker, eg:  `docker run -m 2G ...`")
			}
		}
	}

	return int64(availableMemory)
}

func checkCgroup() (uint64, error) {
	f, err := os.Open("/sys/fs/cgroup/memory/memory.limit_in_bytes")
	if err != nil {
		return 0, err
	}
	defer f.Close()
	b, err := ioutil.ReadAll(f)
	limBytes := string(b)
	limBytes = strings.TrimSpace(limBytes)
	if err != nil {
		return 0, err
	}
	return strconv.ParseUint(limBytes, 10, 64)
}

func checkProc() (uint64, error) {
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return 0, err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		b := scanner.Text()
		if !strings.HasPrefix(b, "MemAvailable") {
			continue
		}

		// expect form:
		// MemAvailable: 1234567890 kB
		tri := strings.Fields(b)
		if len(tri) != 3 {
			return 0, fmt.Errorf("MemAvailable line has unexpected format: %v", b)
		}

		c, err := strconv.ParseUint(tri[1], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse MemAvailable: %v", b)
		}
		switch tri[2] { // convert units to bytes
		case "kB":
			c *= 1024
		case "MB":
			c *= 1024 * 1024
		default:
			return 0, fmt.Errorf("Unexpected units for MemAvailable in /proc/meminfo, need kB or MB, got: %v", tri[2])
		}
		return c, nil
	}

	return 0, fmt.Errorf("Didn't find MemAvailable in /proc/meminfo, kernel is probably < 3.14")
}
<commit_msg>fix sleep time<commit_after>package runner

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/iron-io/titan/common"
	"github.com/iron-io/titan/runner/agent"
	"github.com/iron-io/titan/runner/drivers"
	driverscommon "github.com/iron-io/titan/runner/drivers"
	"github.com/iron-io/titan/runner/drivers/docker"
	"github.com/iron-io/titan/runner/drivers/mock"
)

type Config struct {
	ID      string
	Image   string
	Timeout time.Duration
	AppName string
	Memory  uint64
	Env     map[string]string
	Stdout  io.Writer
	Stderr  io.Writer
}

type Runner struct {
	driver       drivers.Driver
	taskQueue    chan *containerTask
	ml           Logger
	availableMem int64
	usedMem      int64
	usedMemMutex sync.RWMutex
}

var (
	ErrTimeOutNoMemory = errors.New("Task timed out. No available memory.")
	ErrFullQueue       = errors.New("The runner queue is full")

	WaitMemoryTimeout = 10 * time.Second
)

func New(metricLogger Logger) (*Runner, error) {
	// TODO: Is this really required for Titan's driver?
	// Can we remove it?
	env := common.NewEnvironment(func(e *common.Environment) {})

	// TODO: Create a drivers.New(runnerConfig) in Titan
	driver, err := selectDriver("docker", env, &driverscommon.Config{})
	if err != nil {
		return nil, err
	}

	r := &Runner{
		driver:       driver,
		taskQueue:    make(chan *containerTask, 100),
		ml:           metricLogger,
		availableMem: getAvailableMemory(),
		usedMem:      0,
	}

	go r.queueHandler()

	return r, nil
}

// This routine checks for available memory;
// If there's memory then send signal to the task to proceed.
// If there's not available memory to run the task it waits
// If the task waits for more than X seconds it timeouts
func (r *Runner) queueHandler() {
	var task *containerTask
	var waitStart time.Time
	var waitTime time.Duration
	var timedOut bool
	for {
		select {
		case task = <-r.taskQueue:
			waitStart = time.Now()
			timedOut = false
		}

		// Loop waiting for available memory
		canRun := r.checkRequiredMem(task.cfg.Memory)
		for ; !canRun; canRun = r.checkRequiredMem(task.cfg.Memory) {
			waitTime = time.Since(waitStart)
			if waitTime > WaitMemoryTimeout {
				timedOut = true
				break
			}
			time.Sleep(time.Microsecond)
		}

		metricBaseName := fmt.Sprintf("run.%s.", task.cfg.AppName)
		r.ml.LogTime(task.ctx, metricBaseName+"waittime", waitTime)

		if timedOut {
			// Send to a signal to this task saying it cannot run
			r.ml.LogCount(task.ctx, metricBaseName+"timeout", 1)
			task.canRun <- false
			continue
		}

		// Send a signal to this task saying it can run
		task.canRun <- true
	}
}

func (r *Runner) checkRequiredMem(req uint64) bool {
	r.usedMemMutex.RLock()
	defer r.usedMemMutex.RUnlock()
	return r.availableMem-r.usedMem/int64(req)*1024*1024 > 0
}

func (r *Runner) addUsedMem(used int64) {
	r.usedMemMutex.Lock()
	r.usedMem = r.usedMem + used*1024*1024
	if r.usedMem < 0 {
		r.usedMem = 0
	}
	r.usedMemMutex.Unlock()
}

func (r *Runner) checkMemAndUse(req uint64) bool {
	r.usedMemMutex.Lock()
	defer r.usedMemMutex.Unlock()

	used := int64(req) * 1024 * 1024

	if r.availableMem-r.usedMem/used < 0 {
		return false
	}

	r.usedMem += used

	return true
}

func (r *Runner) Run(ctx context.Context, cfg *Config) (drivers.RunResult, error) {
	var err error

	if cfg.Memory == 0 {
		cfg.Memory = 128
	}

	ctask := &containerTask{
		ctx:    ctx,
		cfg:    cfg,
		auth:   &agent.ConfigAuth{},
		canRun: make(chan bool),
	}

	metricBaseName := fmt.Sprintf("run.%s.", cfg.AppName)
	r.ml.LogCount(ctx, metricBaseName+"requests", 1)

	// Check if has enough available memory
	// If available, use it
	if !r.checkMemAndUse(cfg.Memory) {
		// If not, try add task to the queue
		select {
		case r.taskQueue <- ctask:
		default:
			// If queue is full, return error
			r.ml.LogCount(ctx, "queue.full", 1)
			return nil, ErrFullQueue
		}

		// If task was added to the queue, wait for permission
		if ok := <-ctask.canRun; !ok {
			// This task timed out, not available memory
			return nil, ErrTimeOutNoMemory
		}
	} else {
		r.ml.LogTime(ctx, metricBaseName+"waittime", 0)
	}

	closer, err := r.driver.Prepare(ctx, ctask)
	if err != nil {
		return nil, err
	}
	defer closer.Close()

	metricStart := time.Now()

	result, err := r.driver.Run(ctx, ctask)
	if err != nil {
		return nil, err
	}

	r.addUsedMem(-1 * int64(cfg.Memory))

	if result.Status() == "success" {
		r.ml.LogCount(ctx, metricBaseName+"succeeded", 1)
	} else {
		r.ml.LogCount(ctx, metricBaseName+"error", 1)
	}

	metricElapsed := time.Since(metricStart)
	r.ml.LogTime(ctx, metricBaseName+"time", metricElapsed)
	r.ml.LogTime(ctx, "run.exec_time", metricElapsed)

	return result, nil
}

func (r Runner) EnsureImageExists(ctx context.Context, cfg *Config) error {
	ctask := &containerTask{
		cfg:  cfg,
		auth: &agent.ConfigAuth{},
	}

	err := r.driver.EnsureImageExists(ctx, ctask)
	if err != nil {
		return err
	}
	return nil
}

func selectDriver(driver string, env *common.Environment, conf *driverscommon.Config) (drivers.Driver, error) {
	switch driver {
	case "docker":
		docker := docker.NewDocker(env, *conf)
		return docker, nil
	case "mock":
		return mock.New(), nil
	}
	return nil, fmt.Errorf("driver %v not found", driver)
}

func getAvailableMemory() int64 {
	const tooBig = 322122547200 // #300GB or 0, biggest aws instance is 244GB

	var availableMemory uint64
	if os.Getenv("IGNORE_MEMORY") == "1" {
		availableMemory = tooBig
	} else {
		availableMemory, err := checkCgroup()
		if err != nil {
			logrus.WithError(err).Error("Error checking for cgroup memory limits, falling back to host memory available..")
		}
		if availableMemory > tooBig || availableMemory == 0 {
			// Then -m flag probably wasn't set, so use max available on system
			availableMemory, err = checkProc()
			if availableMemory > tooBig || availableMemory == 0 {
				logrus.WithError(err).Fatal("Your Linux version is too old (<3.14) then we can't get the proper information to . You must specify the maximum available memory by passing the -m command with docker run when starting the runner via docker, eg:  `docker run -m 2G ...`")
			}
		}
	}

	return int64(availableMemory)
}

func checkCgroup() (uint64, error) {
	f, err := os.Open("/sys/fs/cgroup/memory/memory.limit_in_bytes")
	if err != nil {
		return 0, err
	}
	defer f.Close()
	b, err := ioutil.ReadAll(f)
	limBytes := string(b)
	limBytes = strings.TrimSpace(limBytes)
	if err != nil {
		return 0, err
	}
	return strconv.ParseUint(limBytes, 10, 64)
}

func checkProc() (uint64, error) {
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return 0, err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		b := scanner.Text()
		if !strings.HasPrefix(b, "MemAvailable") {
			continue
		}

		// expect form:
		// MemAvailable: 1234567890 kB
		tri := strings.Fields(b)
		if len(tri) != 3 {
			return 0, fmt.Errorf("MemAvailable line has unexpected format: %v", b)
		}

		c, err := strconv.ParseUint(tri[1], 10, 64)
		if err != nil {
			return 0, fmt.Errorf("Could not parse MemAvailable: %v", b)
		}
		switch tri[2] { // convert units to bytes
		case "kB":
			c *= 1024
		case "MB":
			c *= 1024 * 1024
		default:
			return 0, fmt.Errorf("Unexpected units for MemAvailable in /proc/meminfo, need kB or MB, got: %v", tri[2])
		}
		return c, nil
	}

	return 0, fmt.Errorf("Didn't find MemAvailable in /proc/meminfo, kernel is probably < 3.14")
}
<|endoftext|>