<commit_before>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/logutil"
	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/vterrors"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// Instance encapsulate the execution state of vtworker.
type Instance struct {
	// Default wrangler for all operations.
	// Users can specify their own in RunCommand() e.g. the gRPC server does this.
	wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error
	lastRunStopTime     time.Time

	// backgroundContext is context.Background() from main() which has to be plumbed through.
	backgroundContext      context.Context
	topoServer             topo.Server
	cell                   string
	commandDisplayInterval time.Duration
}

// NewInstance creates a new Instance.
func NewInstance(ctx context.Context, ts topo.Server, cell string, commandDisplayInterval time.Duration) *Instance {
	wi := &Instance{backgroundContext: ctx, topoServer: ts, cell: cell, commandDisplayInterval: commandDisplayInterval}
	// Note: setAndStartWorker() also adds a MemoryLogger for the webserver.
	wi.wr = wi.CreateWrangler(logutil.NewConsoleLogger())
	return wi
}

// CreateWrangler creates a new wrangler using the instance specific configuration.
func (wi *Instance) CreateWrangler(logger logutil.Logger) *wrangler.Wrangler {
	return wrangler.New(logger, wi.topoServer, tmclient.NewTabletManagerClient())
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *Instance) setAndStartWorker(wrk Worker, wr *wrangler.Wrangler) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentContext != nil {
		return nil, vterrors.FromError(vtrpcpb.ErrorCode_TRANSIENT_ERROR,
			fmt.Errorf("A worker job is already in progress: %v", wi.currentWorker))
	}

	if wi.currentWorker != nil {
		// During the grace period, we answer with a retryable error.
		const gracePeriod = 1 * time.Minute
		gracePeriodEnd := time.Now().Add(gracePeriod)
		if wi.lastRunStopTime.Before(gracePeriodEnd) {
			return nil, vterrors.FromError(vtrpcpb.ErrorCode_TRANSIENT_ERROR,
				fmt.Errorf("A worker job was recently stopped (%f seconds ago): %v",
					time.Now().Sub(wi.lastRunStopTime).Seconds(),
					wi.currentWorker))
		}

		// QUERY_NOT_SERVED = FailedPrecondition => manual resolution required.
		return nil, vterrors.FromError(vtrpcpb.ErrorCode_QUERY_NOT_SERVED,
			fmt.Errorf("The worker job was stopped %.1f minutes ago, but not reset. You have to reset it manually. Job: %v",
				time.Now().Sub(wi.lastRunStopTime).Minutes(),
				wi.currentWorker))
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(wi.backgroundContext)
	wi.lastRunError = nil
	wi.lastRunStopTime = time.Unix(0, 0)
	done := make(chan struct{})
	wranglerLogger := wr.Logger()
	if wr == wi.wr {
		// If it's the default wrangler, do not reuse its logger because it may have been set before.
		// Resuing it would result into an endless recursion.
		wranglerLogger = logutil.NewConsoleLogger()
	}
	wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, wranglerLogger))

	// one go function runs the worker, changes state when done
	go func() {
		log.Infof("Starting worker...")
		var err error

		// Catch all panics and always save the execution state at the end.
		defer func() {
			// The recovery code is a copy of servenv.HandlePanic().
			if x := recover(); x != nil {
				err = fmt.Errorf("uncaught %v panic: %v", "vtworker", x)
			}

			wi.currentWorkerMutex.Lock()
			wi.currentContext = nil
			wi.currentCancelFunc = nil
			wi.lastRunError = err
			wi.lastRunStopTime = time.Now()
			wi.currentWorkerMutex.Unlock()
			close(done)
		}()

		// run will take a long time
		err = wrk.Run(wi.currentContext)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *Instance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		for s := range sigChan {
			// We got a signal, notify our modules.
			// Use an extra function to properly unlock using defer.
			func() {
				wi.currentWorkerMutex.Lock()
				defer wi.currentWorkerMutex.Unlock()
				if wi.currentCancelFunc != nil {
					log.Infof("Trying to cancel current worker after receiving signal: %v", s)
					wi.currentCancelFunc()
				} else {
					log.Infof("Shutting down idle worker after receiving signal: %v", s)
					os.Exit(0)
				}
			}()
		}
	}()
}

// Reset resets the state of a finished worker. It returns an error if the worker is still running.
func (wi *Instance) Reset() error {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentWorker == nil {
		return nil
	}

	// check the worker is really done
	if wi.currentContext == nil {
		wi.currentWorker = nil
		wi.currentMemoryLogger = nil
		return nil
	}

	return errors.New("worker still executing")
}

// Cancel calls the cancel function of the current vtworker job.
// It returns true, if a job was running. False otherwise.
// NOTE: Cancel won't reset the state as well. Use Reset() to do so.
func (wi *Instance) Cancel() bool {
	wi.currentWorkerMutex.Lock()

	if wi.currentWorker == nil || wi.currentCancelFunc == nil {
		wi.currentWorkerMutex.Unlock()
		return false
	}

	cancel := wi.currentCancelFunc
	wi.currentWorkerMutex.Unlock()

	cancel()

	return true
}
<commit_msg>vtworker: Log head of stacktrace in case of a panic.<commit_after>// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/tb"
	"github.com/youtube/vitess/go/vt/logutil"
	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/vterrors"
	"github.com/youtube/vitess/go/vt/wrangler"
	"golang.org/x/net/context"
)

// Instance encapsulate the execution state of vtworker.
type Instance struct {
	// Default wrangler for all operations.
	// Users can specify their own in RunCommand() e.g. the gRPC server does this.
	wr *wrangler.Wrangler

	// mutex is protecting all the following variables
	// 3 states here:
	// - no job ever ran (or reset was run): currentWorker is nil,
	// currentContext/currentCancelFunc is nil, lastRunError is nil
	// - one worker running: currentWorker is set,
	//   currentContext/currentCancelFunc is set, lastRunError is nil
	// - (at least) one worker already ran, none is running atm:
	//   currentWorker is set, currentContext is nil, lastRunError
	//   has the error returned by the worker.
	currentWorkerMutex  sync.Mutex
	currentWorker       Worker
	currentMemoryLogger *logutil.MemoryLogger
	currentContext      context.Context
	currentCancelFunc   context.CancelFunc
	lastRunError        error
	lastRunStopTime     time.Time

	// backgroundContext is context.Background() from main() which has to be plumbed through.
	backgroundContext      context.Context
	topoServer             topo.Server
	cell                   string
	commandDisplayInterval time.Duration
}

// NewInstance creates a new Instance.
func NewInstance(ctx context.Context, ts topo.Server, cell string, commandDisplayInterval time.Duration) *Instance {
	wi := &Instance{backgroundContext: ctx, topoServer: ts, cell: cell, commandDisplayInterval: commandDisplayInterval}
	// Note: setAndStartWorker() also adds a MemoryLogger for the webserver.
	wi.wr = wi.CreateWrangler(logutil.NewConsoleLogger())
	return wi
}

// CreateWrangler creates a new wrangler using the instance specific configuration.
func (wi *Instance) CreateWrangler(logger logutil.Logger) *wrangler.Wrangler {
	return wrangler.New(logger, wi.topoServer, tmclient.NewTabletManagerClient())
}

// setAndStartWorker will set the current worker.
// We always log to both memory logger (for display on the web) and
// console logger (for records / display of command line worker).
func (wi *Instance) setAndStartWorker(wrk Worker, wr *wrangler.Wrangler) (chan struct{}, error) {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentContext != nil {
		return nil, vterrors.FromError(vtrpcpb.ErrorCode_TRANSIENT_ERROR,
			fmt.Errorf("A worker job is already in progress: %v", wi.currentWorker))
	}

	if wi.currentWorker != nil {
		// During the grace period, we answer with a retryable error.
		const gracePeriod = 1 * time.Minute
		gracePeriodEnd := time.Now().Add(gracePeriod)
		if wi.lastRunStopTime.Before(gracePeriodEnd) {
			return nil, vterrors.FromError(vtrpcpb.ErrorCode_TRANSIENT_ERROR,
				fmt.Errorf("A worker job was recently stopped (%f seconds ago): %v",
					time.Now().Sub(wi.lastRunStopTime).Seconds(),
					wi.currentWorker))
		}

		// QUERY_NOT_SERVED = FailedPrecondition => manual resolution required.
		return nil, vterrors.FromError(vtrpcpb.ErrorCode_QUERY_NOT_SERVED,
			fmt.Errorf("The worker job was stopped %.1f minutes ago, but not reset. You have to reset it manually. Job: %v",
				time.Now().Sub(wi.lastRunStopTime).Minutes(),
				wi.currentWorker))
	}

	wi.currentWorker = wrk
	wi.currentMemoryLogger = logutil.NewMemoryLogger()
	wi.currentContext, wi.currentCancelFunc = context.WithCancel(wi.backgroundContext)
	wi.lastRunError = nil
	wi.lastRunStopTime = time.Unix(0, 0)
	done := make(chan struct{})
	wranglerLogger := wr.Logger()
	if wr == wi.wr {
		// If it's the default wrangler, do not reuse its logger because it may have been set before.
		// Resuing it would result into an endless recursion.
		wranglerLogger = logutil.NewConsoleLogger()
	}
	wr.SetLogger(logutil.NewTeeLogger(wi.currentMemoryLogger, wranglerLogger))

	// one go function runs the worker, changes state when done
	go func() {
		log.Infof("Starting worker...")
		var err error

		// Catch all panics and always save the execution state at the end.
		defer func() {
			// The recovery code is a copy of servenv.HandlePanic().
			if x := recover(); x != nil {
				log.Errorf("uncaught vtworker panic: %v\n%s", x, tb.Stack(4))
				err = fmt.Errorf("uncaught vtworker panic: %v", x)
			}

			wi.currentWorkerMutex.Lock()
			wi.currentContext = nil
			wi.currentCancelFunc = nil
			wi.lastRunError = err
			wi.lastRunStopTime = time.Now()
			wi.currentWorkerMutex.Unlock()
			close(done)
		}()

		// run will take a long time
		err = wrk.Run(wi.currentContext)
	}()

	return done, nil
}

// InstallSignalHandlers installs signal handler which exit vtworker gracefully.
func (wi *Instance) InstallSignalHandlers() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		for s := range sigChan {
			// We got a signal, notify our modules.
			// Use an extra function to properly unlock using defer.
			func() {
				wi.currentWorkerMutex.Lock()
				defer wi.currentWorkerMutex.Unlock()
				if wi.currentCancelFunc != nil {
					log.Infof("Trying to cancel current worker after receiving signal: %v", s)
					wi.currentCancelFunc()
				} else {
					log.Infof("Shutting down idle worker after receiving signal: %v", s)
					os.Exit(0)
				}
			}()
		}
	}()
}

// Reset resets the state of a finished worker. It returns an error if the worker is still running.
func (wi *Instance) Reset() error {
	wi.currentWorkerMutex.Lock()
	defer wi.currentWorkerMutex.Unlock()

	if wi.currentWorker == nil {
		return nil
	}

	// check the worker is really done
	if wi.currentContext == nil {
		wi.currentWorker = nil
		wi.currentMemoryLogger = nil
		return nil
	}

	return errors.New("worker still executing")
}

// Cancel calls the cancel function of the current vtworker job.
// It returns true, if a job was running. False otherwise.
// NOTE: Cancel won't reset the state as well. Use Reset() to do so.
func (wi *Instance) Cancel() bool {
	wi.currentWorkerMutex.Lock()

	if wi.currentWorker == nil || wi.currentCancelFunc == nil {
		wi.currentWorkerMutex.Unlock()
		return false
	}

	cancel := wi.currentCancelFunc
	wi.currentWorkerMutex.Unlock()

	cancel()

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"reflect"
	"testing"
)

func TestScanConcurrently(t *testing.T) {
	const noRows = 10
	f := fakeReader(noRows)
	ps := scanConcurrently(f)
	fmt.Println(ps)
	if got := len(ps); got != noRows {
		t.Errorf("len(Persons) = %v; want %v", got, noRows)
	}
	want := Person{
		ID:   4,
		Name: "name4",
		Age:  noRows - 4}
	if got := ps[3]; !reflect.DeepEqual(got, want) {
		t.Errorf("p[3] = %v; want %v", got, want)
	}
}
<commit_msg>file_scanner: clean concurrently tests<commit_after>package main

import "testing"

func TestScanConcurrently(t *testing.T) {
	testScan(t, scanConcurrently)
}
<|endoftext|>


<commit_before>// Copyright 2018 The Terraformer 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 aws

import (
	"context"
	"fmt"
	"log"
	"strings"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/route53"
)

var route53AllowEmptyValues = []string{}

var route53AdditionalFields = map[string]interface{}{}

type Route53Generator struct {
	AWSService
}

func (g *Route53Generator) createZonesResources(svc *route53.Client) []terraformutils.Resource {
	var resources []terraformutils.Resource
	p := route53.NewListHostedZonesPaginator(svc, &route53.ListHostedZonesInput{})
	for p.HasMorePages() {
		page, err := p.NextPage(context.TODO())
		if err != nil {
			log.Println(err)
			return resources
		}
		for _, zone := range page.HostedZones {
			zoneID := cleanZoneID(StringValue(zone.Id))
			resources = append(resources, terraformutils.NewResource(
				zoneID,
				zoneID+"_"+strings.TrimSuffix(StringValue(zone.Name), "."),
				"aws_route53_zone",
				"aws",
				map[string]string{
					"name":          StringValue(zone.Name),
					"force_destroy": "false",
				},
				route53AllowEmptyValues,
				route53AdditionalFields,
			))
			records := g.createRecordsResources(svc, zoneID)
			resources = append(resources, records...)
		}
	}
	return resources
}

func (Route53Generator) createRecordsResources(svc *route53.Client, zoneID string) []terraformutils.Resource {
	var resources []terraformutils.Resource
	var startRecordIdentifier *string
	var sets *route53.ListResourceRecordSetsOutput
	var err error
	listParams := &route53.ListResourceRecordSetsInput{
		HostedZoneId:          aws.String(zoneID),
		StartRecordIdentifier: startRecordIdentifier,
	}

	sets, err = svc.ListResourceRecordSets(context.TODO(), listParams)
	if err != nil {
		log.Println(err)
		return resources
	}

	for sets == nil || sets.NextRecordIdentifier != nil {
		for _, record := range sets.ResourceRecordSets {
			recordName := wildcardUnescape(StringValue(record.Name))
			typeString := string(record.Type)
			resources = append(resources, terraformutils.NewResource(
				fmt.Sprintf("%s_%s_%s_%s", zoneID, recordName, typeString, StringValue(record.SetIdentifier)),
				fmt.Sprintf("%s_%s_%s_%s", zoneID, recordName, typeString, StringValue(record.SetIdentifier)),
				"aws_route53_record",
				"aws",
				map[string]string{
					"name":           strings.TrimSuffix(recordName, "."),
					"zone_id":        zoneID,
					"type":           typeString,
					"set_identifier": StringValue(record.SetIdentifier),
				},
				route53AllowEmptyValues,
				route53AdditionalFields,
			))
		}
	}
	return resources
}

// Generate TerraformResources from AWS API,
// create terraform resource for each zone + each record
func (g *Route53Generator) InitResources() error {
	config, e := g.generateConfig()
	if e != nil {
		return e
	}
	svc := route53.NewFromConfig(config)

	g.Resources = g.createZonesResources(svc)
	return nil
}

func (g *Route53Generator) PostConvertHook() error {
	for i, resourceRecord := range g.Resources {
		if resourceRecord.InstanceInfo.Type == "aws_route53_zone" {
			continue
		}
		item := resourceRecord.Item
		zoneID := item["zone_id"].(string)
		for _, resourceZone := range g.Resources {
			if resourceZone.InstanceInfo.Type != "aws_route53_zone" {
				continue
			}
			if zoneID == resourceZone.InstanceState.ID {
				g.Resources[i].Item["zone_id"] = "${aws_route53_zone." + resourceZone.ResourceName + ".zone_id}"
			}
		}
		if _, aliasExist := resourceRecord.Item["alias"]; aliasExist {
			if _, ttlExist := resourceRecord.Item["ttl"]; ttlExist {
				delete(g.Resources[i].Item, "ttl")
			}
		}
	}
	return nil
}

func wildcardUnescape(s string) string {
	if strings.Contains(s, "\\052") {
		s = strings.Replace(s, "\\052", "*", 1)
	}
	return s
}

// cleanZoneID is used to remove the leading /hostedzone/
func cleanZoneID(id string) string {
	return cleanPrefix(id, "/hostedzone/")
}

// cleanPrefix removes a string prefix from an ID
func cleanPrefix(id, prefix string) string {
	return strings.TrimPrefix(id, prefix)
}
<commit_msg>Fixed route53 (#875)<commit_after>// Copyright 2018 The Terraformer 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 aws

import (
	"context"
	"fmt"
	"log"
	"strings"

	"github.com/GoogleCloudPlatform/terraformer/terraformutils"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/route53"
)

var route53AllowEmptyValues = []string{}

var route53AdditionalFields = map[string]interface{}{}

type Route53Generator struct {
	AWSService
}

func (g *Route53Generator) createZonesResources(svc *route53.Client) []terraformutils.Resource {
	var resources []terraformutils.Resource
	p := route53.NewListHostedZonesPaginator(svc, &route53.ListHostedZonesInput{})
	for p.HasMorePages() {
		page, err := p.NextPage(context.TODO())
		if err != nil {
			log.Println(err)
			return resources
		}
		for _, zone := range page.HostedZones {
			zoneID := cleanZoneID(StringValue(zone.Id))
			resources = append(resources, terraformutils.NewResource(
				zoneID,
				zoneID+"_"+strings.TrimSuffix(StringValue(zone.Name), "."),
				"aws_route53_zone",
				"aws",
				map[string]string{
					"name":          StringValue(zone.Name),
					"force_destroy": "false",
				},
				route53AllowEmptyValues,
				route53AdditionalFields,
			))
			records := g.createRecordsResources(svc, zoneID)
			resources = append(resources, records...)
		}
	}
	return resources
}

func (Route53Generator) createRecordsResources(svc *route53.Client, zoneID string) []terraformutils.Resource {
	var resources []terraformutils.Resource
	var sets *route53.ListResourceRecordSetsOutput
	var err error
	listParams := &route53.ListResourceRecordSetsInput{
		HostedZoneId: aws.String(zoneID),
	}

	for {
		sets, err = svc.ListResourceRecordSets(context.TODO(), listParams)
		if err != nil {
			log.Println(err)
			return resources
		}
		for _, record := range sets.ResourceRecordSets {
			recordName := wildcardUnescape(StringValue(record.Name))
			typeString := string(record.Type)
			resources = append(resources, terraformutils.NewResource(
				fmt.Sprintf("%s_%s_%s_%s", zoneID, recordName, typeString, StringValue(record.SetIdentifier)),
				fmt.Sprintf("%s_%s_%s_%s", zoneID, recordName, typeString, StringValue(record.SetIdentifier)),
				"aws_route53_record",
				"aws",
				map[string]string{
					"name":           strings.TrimSuffix(recordName, "."),
					"zone_id":        zoneID,
					"type":           typeString,
					"set_identifier": StringValue(record.SetIdentifier),
				},
				route53AllowEmptyValues,
				route53AdditionalFields,
			))
		}

		if sets.IsTruncated {
			listParams.StartRecordName = sets.NextRecordName
			listParams.StartRecordType = sets.NextRecordType
			listParams.StartRecordIdentifier = sets.NextRecordIdentifier
		} else {
			break
		}
	}
	return resources
}

// Generate TerraformResources from AWS API,
// create terraform resource for each zone + each record
func (g *Route53Generator) InitResources() error {
	config, e := g.generateConfig()
	if e != nil {
		return e
	}
	svc := route53.NewFromConfig(config)

	g.Resources = g.createZonesResources(svc)
	return nil
}

func (g *Route53Generator) PostConvertHook() error {
	for i, resourceRecord := range g.Resources {
		if resourceRecord.InstanceInfo.Type == "aws_route53_zone" {
			continue
		}
		item := resourceRecord.Item
		zoneID := item["zone_id"].(string)
		for _, resourceZone := range g.Resources {
			if resourceZone.InstanceInfo.Type != "aws_route53_zone" {
				continue
			}
			if zoneID == resourceZone.InstanceState.ID {
				g.Resources[i].Item["zone_id"] = "${aws_route53_zone." + resourceZone.ResourceName + ".zone_id}"
			}
		}
		if _, aliasExist := resourceRecord.Item["alias"]; aliasExist {
			if _, ttlExist := resourceRecord.Item["ttl"]; ttlExist {
				delete(g.Resources[i].Item, "ttl")
			}
		}
	}
	return nil
}

func wildcardUnescape(s string) string {
	if strings.Contains(s, "\\052") {
		s = strings.Replace(s, "\\052", "*", 1)
	}
	return s
}

// cleanZoneID is used to remove the leading /hostedzone/
func cleanZoneID(id string) string {
	return cleanPrefix(id, "/hostedzone/")
}

// cleanPrefix removes a string prefix from an ID
func cleanPrefix(id, prefix string) string {
	return strings.TrimPrefix(id, prefix)
}
<|endoftext|>


<commit_before>// Package kubernetes implements kubernetes micro runtime
package kubernetes

import (
	"fmt"
	"sync"
	"time"

	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/util/kubernetes/client"
	"github.com/micro/go-micro/util/log"
)

// action to take on runtime service
type action int

const (
	start action = iota
	update
	stop
)

type kubernetes struct {
	sync.RWMutex
	// options configure runtime
	options runtime.Options
	// indicates if we're running
	running bool
	// used to stop the runtime
	closed chan bool
	// client is kubernetes client
	client client.Client
}

// getService queries kubernetes for micro service
// NOTE: this function is not thread-safe
func (k *kubernetes) getService(labels map[string]string) ([]*runtime.Service, error) {
	// get the service status
	serviceList := new(client.ServiceList)
	r := &client.Resource{
		Kind:  "service",
		Value: serviceList,
	}

	// get the service from k8s
	if err := k.client.Get(r, labels); err != nil {
		return nil, err
	}

	// get the deployment status
	depList := new(client.DeploymentList)
	d := &client.Resource{
		Kind:  "deployment",
		Value: depList,
	}

	// get the deployment from k8s
	if err := k.client.Get(d, labels); err != nil {
		return nil, err
	}

	// service map
	svcMap := make(map[string]*runtime.Service)

	// collect info from kubernetes service
	for _, kservice := range serviceList.Items {
		// name of the service
		name := kservice.Metadata.Labels["name"]
		// version of the service
		version := kservice.Metadata.Labels["version"]

		// save as service
		svcMap[name+version] = &runtime.Service{
			Name:     name,
			Version:  version,
			Metadata: make(map[string]string),
		}

		// copy annotations metadata into service metadata
		for k, v := range kservice.Metadata.Annotations {
			svcMap[name+version].Metadata[k] = v
		}
	}

	// collect additional info from kubernetes deployment
	for _, kdep := range depList.Items {
		// name of the service
		name := kdep.Metadata.Labels["name"]
		// versio of the service
		version := kdep.Metadata.Labels["version"]

		// access existing service map based on name + version
		if svc, ok := svcMap[name+version]; ok {
			// we're expecting our own service name in metadata
			if _, ok := kdep.Metadata.Annotations["name"]; !ok {
				continue
			}

			// set the service name, version and source
			// based on existing annotations we stored
			svc.Name = kdep.Metadata.Annotations["name"]
			svc.Version = kdep.Metadata.Annotations["version"]
			svc.Source = kdep.Metadata.Annotations["source"]

			// delete from metadata
			delete(kdep.Metadata.Annotations, "name")
			delete(kdep.Metadata.Annotations, "version")
			delete(kdep.Metadata.Annotations, "source")

			// copy all annotations metadata into service metadata
			for k, v := range kdep.Metadata.Annotations {
				svc.Metadata[k] = v
			}

			// parse out deployment status and inject into service metadata
			if len(kdep.Status.Conditions) > 0 {
				status := kdep.Status.Conditions[0].Type
				// pick the last known condition type and mark the service status with it
				log.Debugf("Runtime setting %s service deployment status: %v", name, status)
				svc.Metadata["status"] = status
			}

			// parse out deployment build
			if build, ok := kdep.Spec.Template.Metadata.Annotations["build"]; ok {
				buildTime, err := time.Parse(time.RFC3339, build)
				if err != nil {
					log.Debugf("Runtime failed parsing build time for %s: %v", name, err)
					continue
				}
				svc.Metadata["build"] = fmt.Sprintf("%d", buildTime.Unix())
				continue
			}
			// if no build annotation is found, set it to current time
			svc.Metadata["build"] = fmt.Sprintf("%d", time.Now().Unix())
		}
	}

	// collect all the services and return
	services := make([]*runtime.Service, 0, len(serviceList.Items))

	for _, service := range svcMap {
		services = append(services, service)
	}

	return services, nil
}

// run runs the runtime management loop
func (k *kubernetes) run(events <-chan runtime.Event) {
	t := time.NewTicker(time.Second * 10)
	defer t.Stop()

	for {
		select {
		case <-t.C:
			// TODO: figure out what to do here
			// - do we even need the ticker for k8s services?
		case event := <-events:
			// NOTE: we only handle Update events for now
			log.Debugf("Runtime received notification event: %v", event)
			switch event.Type {
			case runtime.Update:
				// only process if there's an actual service
				// we do not update all the things individually
				if len(event.Service) == 0 {
					continue
				}

				// format the name
				name := client.Format(event.Service)

				// set the default labels
				labels := map[string]string{
					"micro": k.options.Type,
					"name":  name,
				}

				if len(event.Version) > 0 {
					labels["version"] = event.Version
				}

				// get the deployment status
				deployed := new(client.DeploymentList)

				// get the existing service rather than creating a new one
				err := k.client.Get(&client.Resource{
					Kind:  "deployment",
					Value: deployed,
				}, labels)

				if err != nil {
					log.Debugf("Runtime update failed to get service %s: %v", event.Service, err)
					continue
				}

				// technically we should not receive multiple versions but hey ho
				for _, service := range deployed.Items {
					// check the name matches
					if service.Metadata.Name != name {
						continue
					}

					// update build time annotation
					if service.Spec.Template.Metadata.Annotations == nil {
						service.Spec.Template.Metadata.Annotations = make(map[string]string)

					}

					// check the existing build timestamp
					if build, ok := service.Spec.Template.Metadata.Annotations["build"]; ok {
						buildTime, err := time.Parse(time.RFC3339, build)
						if err == nil && !event.Timestamp.After(buildTime) {
							continue
						}
					}

					// update the build time
					service.Spec.Template.Metadata.Annotations["build"] = event.Timestamp.Format(time.RFC3339)

					log.Debugf("Runtime updating service: %s deployment: %s", event.Service, service.Metadata.Name)
					if err := k.client.Update(deploymentResource(&service)); err != nil {
						log.Debugf("Runtime failed to update service %s: %v", event.Service, err)
						continue
					}
				}
			}
		case <-k.closed:
			log.Debugf("Runtime stopped")
			return
		}
	}
}

// Init initializes runtime options
func (k *kubernetes) Init(opts ...runtime.Option) error {
	k.Lock()
	defer k.Unlock()

	for _, o := range opts {
		o(&k.options)
	}

	return nil
}

// Creates a service
func (k *kubernetes) Create(s *runtime.Service, opts ...runtime.CreateOption) error {
	k.Lock()
	defer k.Unlock()

	options := runtime.CreateOptions{
		Type: k.options.Type,
	}
	for _, o := range opts {
		o(&options)
	}

	// quickly prevalidate the name and version
	name := s.Name
	if len(s.Version) > 0 {
		name = name + "-" + s.Version
	}

	// format as we'll format in the deployment
	name = client.Format(name)

	// create new kubernetes micro service
	service := newService(s, options)

	// start the service
	return service.Start(k.client)
}

// Read returns all instances of given service
func (k *kubernetes) Read(opts ...runtime.ReadOption) ([]*runtime.Service, error) {
	k.Lock()
	defer k.Unlock()

	// set the default labels
	labels := map[string]string{
		"micro": k.options.Type,
	}

	var options runtime.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	if len(options.Service) > 0 {
		labels["name"] = client.Format(options.Service)
	}

	// add version to labels if a version has been supplied
	if len(options.Version) > 0 {
		labels["version"] = options.Version
	}

	if len(options.Type) > 0 {
		labels["micro"] = options.Type
	}

	return k.getService(labels)
}

// List the managed services
func (k *kubernetes) List() ([]*runtime.Service, error) {
	k.Lock()
	defer k.Unlock()

	labels := map[string]string{
		"micro": k.options.Type,
	}

	log.Debugf("Runtime listing all micro services")

	return k.getService(labels)
}

// Update the service in place
func (k *kubernetes) Update(s *runtime.Service) error {
	// create new kubernetes micro service
	service := newService(s, runtime.CreateOptions{
		Type: k.options.Type,
	})

	// update build time annotation
	service.kdeploy.Spec.Template.Metadata.Annotations["build"] = time.Now().Format(time.RFC3339)

	return service.Update(k.client)
}

// Delete removes a service
func (k *kubernetes) Delete(s *runtime.Service) error {
	k.Lock()
	defer k.Unlock()

	// create new kubernetes micro service
	service := newService(s, runtime.CreateOptions{
		Type: k.options.Type,
	})

	return service.Stop(k.client)
}

// Start starts the runtime
func (k *kubernetes) Start() error {
	k.Lock()
	defer k.Unlock()

	// already running
	if k.running {
		return nil
	}

	// set running
	k.running = true
	k.closed = make(chan bool)

	var events <-chan runtime.Event
	if k.options.Notifier != nil {
		var err error
		events, err = k.options.Notifier.Notify()
		if err != nil {
			// TODO: should we bail here?
			log.Debugf("Runtime failed to start update notifier")
		}
	}

	go k.run(events)

	return nil
}

// Stop shuts down the runtime
func (k *kubernetes) Stop() error {
	k.Lock()
	defer k.Unlock()

	if !k.running {
		return nil
	}

	select {
	case <-k.closed:
		return nil
	default:
		close(k.closed)
		// set not running
		k.running = false
		// stop the notifier too
		if k.options.Notifier != nil {
			return k.options.Notifier.Close()
		}
	}

	return nil
}

// String implements stringer interface
func (k *kubernetes) String() string {
	return "kubernetes"
}

// NewRuntime creates new kubernetes runtime
func NewRuntime(opts ...runtime.Option) runtime.Runtime {
	// get default options
	options := runtime.Options{
		// Create labels with type "micro": "service"
		Type: "service",
	}

	// apply requested options
	for _, o := range opts {
		o(&options)
	}

	// kubernetes client
	client := client.NewClusterClient()

	return &kubernetes{
		options: options,
		closed:  make(chan bool),
		client:  client,
	}
}
<commit_msg>runtime/kubernetes: remove unused name variable (#1078)<commit_after>// Package kubernetes implements kubernetes micro runtime
package kubernetes

import (
	"fmt"
	"sync"
	"time"

	"github.com/micro/go-micro/runtime"
	"github.com/micro/go-micro/util/kubernetes/client"
	"github.com/micro/go-micro/util/log"
)

// action to take on runtime service
type action int

const (
	start action = iota
	update
	stop
)

type kubernetes struct {
	sync.RWMutex
	// options configure runtime
	options runtime.Options
	// indicates if we're running
	running bool
	// used to stop the runtime
	closed chan bool
	// client is kubernetes client
	client client.Client
}

// getService queries kubernetes for micro service
// NOTE: this function is not thread-safe
func (k *kubernetes) getService(labels map[string]string) ([]*runtime.Service, error) {
	// get the service status
	serviceList := new(client.ServiceList)
	r := &client.Resource{
		Kind:  "service",
		Value: serviceList,
	}

	// get the service from k8s
	if err := k.client.Get(r, labels); err != nil {
		return nil, err
	}

	// get the deployment status
	depList := new(client.DeploymentList)
	d := &client.Resource{
		Kind:  "deployment",
		Value: depList,
	}

	// get the deployment from k8s
	if err := k.client.Get(d, labels); err != nil {
		return nil, err
	}

	// service map
	svcMap := make(map[string]*runtime.Service)

	// collect info from kubernetes service
	for _, kservice := range serviceList.Items {
		// name of the service
		name := kservice.Metadata.Labels["name"]
		// version of the service
		version := kservice.Metadata.Labels["version"]

		// save as service
		svcMap[name+version] = &runtime.Service{
			Name:     name,
			Version:  version,
			Metadata: make(map[string]string),
		}

		// copy annotations metadata into service metadata
		for k, v := range kservice.Metadata.Annotations {
			svcMap[name+version].Metadata[k] = v
		}
	}

	// collect additional info from kubernetes deployment
	for _, kdep := range depList.Items {
		// name of the service
		name := kdep.Metadata.Labels["name"]
		// versio of the service
		version := kdep.Metadata.Labels["version"]

		// access existing service map based on name + version
		if svc, ok := svcMap[name+version]; ok {
			// we're expecting our own service name in metadata
			if _, ok := kdep.Metadata.Annotations["name"]; !ok {
				continue
			}

			// set the service name, version and source
			// based on existing annotations we stored
			svc.Name = kdep.Metadata.Annotations["name"]
			svc.Version = kdep.Metadata.Annotations["version"]
			svc.Source = kdep.Metadata.Annotations["source"]

			// delete from metadata
			delete(kdep.Metadata.Annotations, "name")
			delete(kdep.Metadata.Annotations, "version")
			delete(kdep.Metadata.Annotations, "source")

			// copy all annotations metadata into service metadata
			for k, v := range kdep.Metadata.Annotations {
				svc.Metadata[k] = v
			}

			// parse out deployment status and inject into service metadata
			if len(kdep.Status.Conditions) > 0 {
				status := kdep.Status.Conditions[0].Type
				// pick the last known condition type and mark the service status with it
				log.Debugf("Runtime setting %s service deployment status: %v", name, status)
				svc.Metadata["status"] = status
			}

			// parse out deployment build
			if build, ok := kdep.Spec.Template.Metadata.Annotations["build"]; ok {
				buildTime, err := time.Parse(time.RFC3339, build)
				if err != nil {
					log.Debugf("Runtime failed parsing build time for %s: %v", name, err)
					continue
				}
				svc.Metadata["build"] = fmt.Sprintf("%d", buildTime.Unix())
				continue
			}
			// if no build annotation is found, set it to current time
			svc.Metadata["build"] = fmt.Sprintf("%d", time.Now().Unix())
		}
	}

	// collect all the services and return
	services := make([]*runtime.Service, 0, len(serviceList.Items))

	for _, service := range svcMap {
		services = append(services, service)
	}

	return services, nil
}

// run runs the runtime management loop
func (k *kubernetes) run(events <-chan runtime.Event) {
	t := time.NewTicker(time.Second * 10)
	defer t.Stop()

	for {
		select {
		case <-t.C:
			// TODO: figure out what to do here
			// - do we even need the ticker for k8s services?
		case event := <-events:
			// NOTE: we only handle Update events for now
			log.Debugf("Runtime received notification event: %v", event)
			switch event.Type {
			case runtime.Update:
				// only process if there's an actual service
				// we do not update all the things individually
				if len(event.Service) == 0 {
					continue
				}

				// format the name
				name := client.Format(event.Service)

				// set the default labels
				labels := map[string]string{
					"micro": k.options.Type,
					"name":  name,
				}

				if len(event.Version) > 0 {
					labels["version"] = event.Version
				}

				// get the deployment status
				deployed := new(client.DeploymentList)

				// get the existing service rather than creating a new one
				err := k.client.Get(&client.Resource{
					Kind:  "deployment",
					Value: deployed,
				}, labels)

				if err != nil {
					log.Debugf("Runtime update failed to get service %s: %v", event.Service, err)
					continue
				}

				// technically we should not receive multiple versions but hey ho
				for _, service := range deployed.Items {
					// check the name matches
					if service.Metadata.Name != name {
						continue
					}

					// update build time annotation
					if service.Spec.Template.Metadata.Annotations == nil {
						service.Spec.Template.Metadata.Annotations = make(map[string]string)

					}

					// check the existing build timestamp
					if build, ok := service.Spec.Template.Metadata.Annotations["build"]; ok {
						buildTime, err := time.Parse(time.RFC3339, build)
						if err == nil && !event.Timestamp.After(buildTime) {
							continue
						}
					}

					// update the build time
					service.Spec.Template.Metadata.Annotations["build"] = event.Timestamp.Format(time.RFC3339)

					log.Debugf("Runtime updating service: %s deployment: %s", event.Service, service.Metadata.Name)
					if err := k.client.Update(deploymentResource(&service)); err != nil {
						log.Debugf("Runtime failed to update service %s: %v", event.Service, err)
						continue
					}
				}
			}
		case <-k.closed:
			log.Debugf("Runtime stopped")
			return
		}
	}
}

// Init initializes runtime options
func (k *kubernetes) Init(opts ...runtime.Option) error {
	k.Lock()
	defer k.Unlock()

	for _, o := range opts {
		o(&k.options)
	}

	return nil
}

// Creates a service
func (k *kubernetes) Create(s *runtime.Service, opts ...runtime.CreateOption) error {
	k.Lock()
	defer k.Unlock()

	options := runtime.CreateOptions{
		Type: k.options.Type,
	}
	for _, o := range opts {
		o(&options)
	}

	// quickly prevalidate the name and version
	name := s.Name
	if len(s.Version) > 0 {
		name = name + "-" + s.Version
	}

	// create new kubernetes micro service
	service := newService(s, options)

	// start the service
	return service.Start(k.client)
}

// Read returns all instances of given service
func (k *kubernetes) Read(opts ...runtime.ReadOption) ([]*runtime.Service, error) {
	k.Lock()
	defer k.Unlock()

	// set the default labels
	labels := map[string]string{
		"micro": k.options.Type,
	}

	var options runtime.ReadOptions
	for _, o := range opts {
		o(&options)
	}

	if len(options.Service) > 0 {
		labels["name"] = client.Format(options.Service)
	}

	// add version to labels if a version has been supplied
	if len(options.Version) > 0 {
		labels["version"] = options.Version
	}

	if len(options.Type) > 0 {
		labels["micro"] = options.Type
	}

	return k.getService(labels)
}

// List the managed services
func (k *kubernetes) List() ([]*runtime.Service, error) {
	k.Lock()
	defer k.Unlock()

	labels := map[string]string{
		"micro": k.options.Type,
	}

	log.Debugf("Runtime listing all micro services")

	return k.getService(labels)
}

// Update the service in place
func (k *kubernetes) Update(s *runtime.Service) error {
	// create new kubernetes micro service
	service := newService(s, runtime.CreateOptions{
		Type: k.options.Type,
	})

	// update build time annotation
	service.kdeploy.Spec.Template.Metadata.Annotations["build"] = time.Now().Format(time.RFC3339)

	return service.Update(k.client)
}

// Delete removes a service
func (k *kubernetes) Delete(s *runtime.Service) error {
	k.Lock()
	defer k.Unlock()

	// create new kubernetes micro service
	service := newService(s, runtime.CreateOptions{
		Type: k.options.Type,
	})

	return service.Stop(k.client)
}

// Start starts the runtime
func (k *kubernetes) Start() error {
	k.Lock()
	defer k.Unlock()

	// already running
	if k.running {
		return nil
	}

	// set running
	k.running = true
	k.closed = make(chan bool)

	var events <-chan runtime.Event
	if k.options.Notifier != nil {
		var err error
		events, err = k.options.Notifier.Notify()
		if err != nil {
			// TODO: should we bail here?
			log.Debugf("Runtime failed to start update notifier")
		}
	}

	go k.run(events)

	return nil
}

// Stop shuts down the runtime
func (k *kubernetes) Stop() error {
	k.Lock()
	defer k.Unlock()

	if !k.running {
		return nil
	}

	select {
	case <-k.closed:
		return nil
	default:
		close(k.closed)
		// set not running
		k.running = false
		// stop the notifier too
		if k.options.Notifier != nil {
			return k.options.Notifier.Close()
		}
	}

	return nil
}

// String implements stringer interface
func (k *kubernetes) String() string {
	return "kubernetes"
}

// NewRuntime creates new kubernetes runtime
func NewRuntime(opts ...runtime.Option) runtime.Runtime {
	// get default options
	options := runtime.Options{
		// Create labels with type "micro": "service"
		Type: "service",
	}

	// apply requested options
	for _, o := range opts {
		o(&options)
	}

	// kubernetes client
	client := client.NewClusterClient()

	return &kubernetes{
		options: options,
		closed:  make(chan bool),
		client:  client,
	}
}
<|endoftext|>


<commit_before>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"

	"github.com/couchbaselabs/cbfs/client"
	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

// TODO: this is currently reading the entire tar stream into a buffer before
// handing it to cbfs.  Should be using an io.Pipe.
func (d DatasetSplitter) Run() {

	logg.LogTo("DATASET_SPLITTER", "Datasetsplitter.run()!.  Config: %+v Dataset: %+v", d.Configuration, d.Dataset)

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile with id: %v.  Error: %v", d.Dataset.DatafileID, err)
		logg.LogError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip
	tr1, tr2, err := d.openTwoTarGzStreams(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		logg.LogError(errMsg)
		return
	}

	prTrain, pwTrain := io.Pipe()
	prTest, pwTest := io.Pipe()

	// bufferTesting := &bytes.Buffer{}
	// bufferTraining := &bytes.Buffer{}

	tarWriterTesting := tar.NewWriter(pwTest)
	tarWriterTraining := tar.NewWriter(pwTrain)

	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	logg.LogTo("DATASET_SPLITTER", "Created cbfs client: %v", cbfs)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	logg.LogTo("DATASET_SPLITTER", "options: %v", options)

	destTraining := fmt.Sprintf("%v/training.tar.gz", d.Dataset.Id)
	destTesting := fmt.Sprintf("%v/testing.tar.gz", d.Dataset.Id)
	logg.LogTo("DATASET_SPLITTER", " %v %v ", destTesting, destTraining)

	go func() {
		logg.LogTo("DATASET_SPLITTER", "Calling transform")
		err = d.transform(tr1, tr2, tarWriterTraining, tarWriterTesting)
		if err != nil {
			errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
			logg.LogError(errMsg)
			return
		}
		pwTest.Close()
		pwTrain.Close()
	}()

	// At this point bufferTesting has all data
	// logg.LogTo("DATASET_SPLITTER", "bufferTesting size: %d", bufferTesting.Len())
	// logg.LogTo("DATASET_SPLITTER", "bufferTraining size: %d", bufferTraining.Len())

	/*
		go func() {
			if err := cbfs.Put("", destTesting, prTest, options); err != nil {
				errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTesting, err)
				logg.LogError(errMsg)
				return

			}
			logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTesting)

		}()

		if err := cbfs.Put("", destTraining, prTrain, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTraining, err)
			logg.LogError(errMsg)
			return
		}
	*/

	wg := &sync.WaitGroup{}
	wg.Add(1)
	wg.Add(1)
	go func() {
		defer wg.Done()
		logg.LogTo("DATASET_SPLITTER", "prTest ReadAll()")
		bytes, err := ioutil.ReadAll(prTest)
		logg.LogTo("DATASET_SPLITTER", "prTest ReadAll: byteslen: %v, err: %v", len(bytes), err)
	}()

	go func() {
		defer wg.Done()
		logg.LogTo("DATASET_SPLITTER", "prTrain ReadAll()")
		bytes, err := ioutil.ReadAll(prTrain)
		logg.LogTo("DATASET_SPLITTER", "prTrain ReadAll: byteslen: %v, err: %v", len(bytes), err)

	}()
	wg.Wait()

	// logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTraining)

}

// Opens to tar.gz streams to the same url.  The reason this is done twice is due to
// ugly hack, which is documented in the transform() method.
func (d DatasetSplitter) openTwoTarGzStreams(url string) (*tar.Reader, *tar.Reader, error) {

	resp1, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader1, err := gzip.NewReader(resp1.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader1 := tar.NewReader(gzipReader1)

	resp2, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader2, err := gzip.NewReader(resp2.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader2 := tar.NewReader(gzipReader2)

	return tarReader1, tarReader2, nil

}

// Read from source tar stream and write training and test to given tar writers
//
// TODO: fix ugly hack.  Since I'm trying to read from the source stream *twice*, which
// doesn't work, the workaround is to expect *two* source streams: source1 and source2.
// That way after source1 is read, source2 is ready for reading from the beginning
func (d DatasetSplitter) transform(source1, source2 *tar.Reader, train, test *tar.Writer) error {

	// build a map from the source
	sourceMap, err := d.createMap(source1)
	if err != nil {
		return err
	}
	logg.Log("sourceMap: %+v", sourceMap)

	// split the map into training and test
	trainMap, testMap, err := d.splitMap(sourceMap)
	if err != nil {
		return err
	}

	// iterate over the source
	logg.LogTo("DATASET_SPLITTER", "iterate over source")
	for {
		hdr, err := source2.Next()

		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return err
		}

		// distribute to writers based on training and test maps
		var twToAdd *tar.Writer

		// if strings.HasPrefix(hdr.Name, "foo") {
		if trainMap.hasPath(hdr.Name) {
			// add to training tar writer
			twToAdd = train
		} else if testMap.hasPath(hdr.Name) {
			// add to testing tar writer
			twToAdd = test
		} else {
			logg.LogPanic("File not present in either test/train: %v", hdr.Name)
		}

		// TODO: is there a more efficient way to do this?
		bytes, err := ioutil.ReadAll(source2)
		if err != nil {
			return err
		}

		if err := twToAdd.WriteHeader(hdr); err != nil {
			return err
		}
		if _, err := twToAdd.Write(bytes); err != nil {
			return err
		}

	}

	logg.LogTo("DATASET_SPLITTER", "done iterating over source")

	// close writers
	logg.LogTo("DATASET_SPLITTER", "Closing writers")
	if err := train.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar writer: %v", err)
		logg.LogError(errMsg)
		return err
	}
	if err := test.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar reader: %v", err)
		logg.LogError(errMsg)
		return err
	}
	logg.LogTo("DATASET_SPLITTER", "Closed writers")

	return nil
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}

// Create a map of folder -> []filename for all entries in the archive
func (d DatasetSplitter) createMap(source *tar.Reader) (filemap, error) {

	resultMap := filemap{}
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return nil, err
		}

		pathComponents := strings.Split(hdr.Name, "/")

		if len(pathComponents) != 2 {
			return nil, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

		directory := pathComponents[0]
		filename := pathComponents[1]

		resultMap.addFileToDirectory(directory, filename)

	}

	return resultMap, nil

}

// Split map into training and testing disjoint subsets based on values
// of DatasetSplitter's Dataset
func (d DatasetSplitter) splitMap(source filemap) (training filemap, testing filemap, err error) {

	training = filemap{}
	testing = filemap{}

	// iterate over source keys
	for directory, files := range source {

		numTraining := int(float64(len(files)) * d.Dataset.TrainingDataset.SplitPercentage)

		numTest := len(files) - int(numTraining)

		// split files into subsets based on ratios in dataset
		trainingFiles, testFiles, err := splitFiles(files, numTraining, numTest)

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

		// add to respective maps
		training[directory] = trainingFiles
		testing[directory] = testFiles

	}

	return training, testing, nil

}

func splitFiles(files []string, numTraining, numTest int) (training []string, test []string, err error) {
	for i, file := range files {
		if i < numTraining {
			training = append(training, file)
		} else {
			test = append(test, file)
		}
	}
	return
}
<commit_msg>Code cleanup and docs<commit_after>package elasticthought

import (
	"archive/tar"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"sync"

	"github.com/couchbaselabs/cbfs/client"
	"github.com/couchbaselabs/logg"
)

// Worker job that splits a dataset into training/test set
type DatasetSplitter struct {
	Configuration Configuration
	Dataset       Dataset
}

// Run this job
func (d DatasetSplitter) Run() {

	// Find the datafile object associated with dataset
	db := d.Configuration.DbConnection()
	datafile, err := d.Dataset.GetDatafile(db)
	if err != nil {
		errMsg := fmt.Errorf("Error looking up datafile with id: %v.  Error: %v", d.Dataset.DatafileID, err)
		logg.LogError(errMsg)
		return
	}

	// Open the url -- content type should be application/x-gzip
	tr1, tr2, err := d.openTwoTarGzStreams(datafile.Url)
	if err != nil {
		errMsg := fmt.Errorf("Error opening tar.gz streams: %v", err)
		logg.LogError(errMsg)
		return
	}

	// Create pipes
	prTrain, pwTrain := io.Pipe()
	prTest, pwTest := io.Pipe()

	// Create tar writers on the write end of the pipes
	tarWriterTesting := tar.NewWriter(pwTest)
	tarWriterTraining := tar.NewWriter(pwTrain)

	// Create a cbfs client
	cbfs, err := cbfsclient.New(d.Configuration.CbfsUrl)
	options := cbfsclient.PutOptions{
		ContentType: "application/x-gzip",
	}
	if err != nil {
		errMsg := fmt.Errorf("Error creating cbfs client: %v", err)
		logg.LogError(errMsg)
		return
	}

	// Figure out where to store these on cbfs
	destTraining := fmt.Sprintf("%v/training.tar.gz", d.Dataset.Id)
	destTesting := fmt.Sprintf("%v/testing.tar.gz", d.Dataset.Id)

	// Spawn a goroutine that will read from tar.gz reader coming from url data
	// and write to the training and test tar writers (which are on write ends of pipe)
	go func() {
		logg.LogTo("DATASET_SPLITTER", "Calling transform")
		err = d.transform(tr1, tr2, tarWriterTraining, tarWriterTesting)
		if err != nil {
			errMsg := fmt.Errorf("Error transforming tar stream: %v", err)
			logg.LogError(errMsg)
			return
		}

		// Must close _underlying_ piped writers, or the piped readers will
		// never get an EOF.  Closing just the tar writers that wrap the underlying
		// piped writers is not enough.
		pwTest.Close()
		pwTrain.Close()
	}()

	// Spawn goroutines to read off the read ends of the pipe and store in cbfs
	wg := &sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()

		if err := cbfs.Put("", destTesting, prTest, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTesting, err)
			logg.LogError(errMsg)
			return

		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTesting)
	}()
	go func() {
		defer wg.Done()
		if err := cbfs.Put("", destTraining, prTrain, options); err != nil {
			errMsg := fmt.Errorf("Error writing %v to cbfs: %v", destTraining, err)
			logg.LogError(errMsg)
			return
		}
		logg.LogTo("DATASET_SPLITTER", "Wrote %v to cbfs", destTraining)
	}()

	// Wait for the piped readers to finish
	wg.Wait()

}

// Opens to tar.gz streams to the same url.  The reason this is done twice is due to
// ugly hack, which is documented in the transform() method.
func (d DatasetSplitter) openTwoTarGzStreams(url string) (*tar.Reader, *tar.Reader, error) {

	resp1, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader1, err := gzip.NewReader(resp1.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader1 := tar.NewReader(gzipReader1)

	resp2, err := http.Get(url)
	if err != nil {
		return nil, nil, err
	}
	gzipReader2, err := gzip.NewReader(resp2.Body)
	if err != nil {
		return nil, nil, err
	}
	tarReader2 := tar.NewReader(gzipReader2)

	return tarReader1, tarReader2, nil

}

// Read from source tar stream and write training and test to given tar writers
//
// TODO: fix ugly hack.  Since I'm trying to read from the source stream *twice*, which
// doesn't work, the workaround is to expect *two* source streams: source1 and source2.
// That way after source1 is read, source2 is ready for reading from the beginning
func (d DatasetSplitter) transform(source1, source2 *tar.Reader, train, test *tar.Writer) error {

	// build a map from the source
	sourceMap, err := d.createMap(source1)
	if err != nil {
		return err
	}
	logg.Log("sourceMap: %+v", sourceMap)

	// split the map into training and test
	trainMap, testMap, err := d.splitMap(sourceMap)
	if err != nil {
		return err
	}

	// iterate over the source
	logg.LogTo("DATASET_SPLITTER", "iterate over source")
	for {
		hdr, err := source2.Next()

		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return err
		}

		// distribute to writers based on training and test maps
		var twToAdd *tar.Writer

		// if strings.HasPrefix(hdr.Name, "foo") {
		if trainMap.hasPath(hdr.Name) {
			// add to training tar writer
			twToAdd = train
		} else if testMap.hasPath(hdr.Name) {
			// add to testing tar writer
			twToAdd = test
		} else {
			logg.LogPanic("File not present in either test/train: %v", hdr.Name)
		}

		// TODO: is there a more efficient way to do this?
		bytes, err := ioutil.ReadAll(source2)
		if err != nil {
			return err
		}

		if err := twToAdd.WriteHeader(hdr); err != nil {
			return err
		}
		if _, err := twToAdd.Write(bytes); err != nil {
			return err
		}

	}

	logg.LogTo("DATASET_SPLITTER", "done iterating over source")

	// close writers
	logg.LogTo("DATASET_SPLITTER", "Closing writers")
	if err := train.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar writer: %v", err)
		logg.LogError(errMsg)
		return err
	}
	if err := test.Close(); err != nil {
		errMsg := fmt.Errorf("Error closing tar reader: %v", err)
		logg.LogError(errMsg)
		return err
	}
	logg.LogTo("DATASET_SPLITTER", "Closed writers")

	return nil
}

// Validate that the source tar stream conforms to expected specs
func (d DatasetSplitter) validate(source *tar.Reader) (bool, error) {

	// validation rules:
	// 1. has at least 2 files
	// 2. the depth of each file is 2 (folder/filename.xxx)

	numFiles := 0
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return false, err
		}
		numFiles += 1

		pathComponents := strings.Split(hdr.Name, "/")
		if len(pathComponents) != 2 {
			return false, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

	}

	if numFiles < 2 {
		return false, fmt.Errorf("Archive must contain at least 2 files")
	}

	return true, nil
}

// Create a map of folder -> []filename for all entries in the archive
func (d DatasetSplitter) createMap(source *tar.Reader) (filemap, error) {

	resultMap := filemap{}
	for {
		hdr, err := source.Next()
		if err == io.EOF {
			// end of tar archive
			break
		}
		if err != nil {
			return nil, err
		}

		pathComponents := strings.Split(hdr.Name, "/")

		if len(pathComponents) != 2 {
			return nil, fmt.Errorf("Path does not have 2 components: %v", hdr.Name)
		}

		directory := pathComponents[0]
		filename := pathComponents[1]

		resultMap.addFileToDirectory(directory, filename)

	}

	return resultMap, nil

}

// Split map into training and testing disjoint subsets based on values
// of DatasetSplitter's Dataset
func (d DatasetSplitter) splitMap(source filemap) (training filemap, testing filemap, err error) {

	training = filemap{}
	testing = filemap{}

	// iterate over source keys
	for directory, files := range source {

		numTraining := int(float64(len(files)) * d.Dataset.TrainingDataset.SplitPercentage)

		numTest := len(files) - int(numTraining)

		// split files into subsets based on ratios in dataset
		trainingFiles, testFiles, err := splitFiles(files, numTraining, numTest)

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

		// add to respective maps
		training[directory] = trainingFiles
		testing[directory] = testFiles

	}

	return training, testing, nil

}

func splitFiles(files []string, numTraining, numTest int) (training []string, test []string, err error) {
	for i, file := range files {
		if i < numTraining {
			training = append(training, file)
		} else {
			test = append(test, file)
		}
	}
	return
}
<|endoftext|>


<commit_before>package tests

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/common"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/order"
	"github.com/bitfinexcom/bitfinex-api-go/v2"
	"github.com/bitfinexcom/bitfinex-api-go/v2/websocket"
)

func TestWebsocketOrder(t *testing.T) {
	if !auth {
		t.Skip("no credentials, skipping order creation")
	}

	wg := sync.WaitGroup{}
	wg.Add(1) // 1. Authentication event

	c := websocket.New().Credentials(key, secret)

	errch := make(chan error)

	err := c.Connect()
	if err != nil {
		t.Fatalf("connecting to websocket service: %s", err)
	}
	defer c.Close()

	go func() {
		for ev := range c.Listen() {
			switch e := ev.(type) {
			case *bitfinex.Notification:
				if e.Status == "ERROR" && e.Type == "on-req" {
					t.Errorf("failed to create order: %s", e.Text)
				}
			case *order.New:
				wg.Done()
			case *order.Cancel:
				wg.Done()
			case error:
				t.Logf("Listen() error: %s", ev)
				errch <- ev.(error)
				wg.Done()
			}
		}
	}()
	/*
		err = c.Authenticate(context.Background())
		if err != nil {
			t.Fatalf("authenticating with websocket service: %s", err)
		}
		if err := wait(&wg, errch, 2*time.Second); err != nil {
			t.Fatalf("failed to authenticate with websocket service: %s", err)
		}
	*/
	wg.Add(1)
	n := time.Now()
	cid := n.Unix()
	cidDate := n.Format("2006-01-02")
	o := &order.NewRequest{
		CID:    cid,
		Type:   common.OrderTypeExchangeLimit,
		Symbol: common.TradingPrefix + bitfinex.BTCUSD,
		Amount: 1.0,
		Price:  28.5,
	}

	ctx, cxl1 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cxl1()
	err = c.SubmitOrder(ctx, o)
	if err != nil {
		t.Fatalf("failed to send OrderNewRequest: %s", err)
	}
	if err := wait(&wg, errch, 2*time.Second); err != nil {
		t.Fatalf("failed to create order: %s", err)
	}

	oc := &order.CancelRequest{
		CID:     cid,
		CIDDate: cidDate,
	}

	wg.Add(1)
	ctx, cxl2 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cxl2()
	err = c.SubmitCancel(ctx, oc)
	if err != nil {
		t.Fatalf("failed to send OrderCancelRequest: %s", err)
	}
	if err := wait(&wg, errch, 2*time.Second); err != nil {
		t.Fatalf("failed to cancel order: %s", err)
	}
}
<commit_msg>tests/integration/v2/live_websocket_private_test.go putting notification package to work<commit_after>package tests

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/common"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/notification"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/order"
	"github.com/bitfinexcom/bitfinex-api-go/v2"
	"github.com/bitfinexcom/bitfinex-api-go/v2/websocket"
)

func TestWebsocketOrder(t *testing.T) {
	if !auth {
		t.Skip("no credentials, skipping order creation")
	}

	wg := sync.WaitGroup{}
	wg.Add(1) // 1. Authentication event

	c := websocket.New().Credentials(key, secret)

	errch := make(chan error)

	err := c.Connect()
	if err != nil {
		t.Fatalf("connecting to websocket service: %s", err)
	}
	defer c.Close()

	go func() {
		for ev := range c.Listen() {
			switch e := ev.(type) {
			case *notification.Notification:
				if e.Status == "ERROR" && e.Type == "on-req" {
					t.Errorf("failed to create order: %s", e.Text)
				}
			case *order.New:
				wg.Done()
			case *order.Cancel:
				wg.Done()
			case error:
				t.Logf("Listen() error: %s", ev)
				errch <- ev.(error)
				wg.Done()
			}
		}
	}()
	/*
		err = c.Authenticate(context.Background())
		if err != nil {
			t.Fatalf("authenticating with websocket service: %s", err)
		}
		if err := wait(&wg, errch, 2*time.Second); err != nil {
			t.Fatalf("failed to authenticate with websocket service: %s", err)
		}
	*/
	wg.Add(1)
	n := time.Now()
	cid := n.Unix()
	cidDate := n.Format("2006-01-02")
	o := &order.NewRequest{
		CID:    cid,
		Type:   common.OrderTypeExchangeLimit,
		Symbol: common.TradingPrefix + bitfinex.BTCUSD,
		Amount: 1.0,
		Price:  28.5,
	}

	ctx, cxl1 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cxl1()
	err = c.SubmitOrder(ctx, o)
	if err != nil {
		t.Fatalf("failed to send OrderNewRequest: %s", err)
	}
	if err := wait(&wg, errch, 2*time.Second); err != nil {
		t.Fatalf("failed to create order: %s", err)
	}

	oc := &order.CancelRequest{
		CID:     cid,
		CIDDate: cidDate,
	}

	wg.Add(1)
	ctx, cxl2 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cxl2()
	err = c.SubmitCancel(ctx, oc)
	if err != nil {
		t.Fatalf("failed to send OrderCancelRequest: %s", err)
	}
	if err := wait(&wg, errch, 2*time.Second); err != nil {
		t.Fatalf("failed to cancel order: %s", err)
	}
}
<|endoftext|>


<commit_before>package framework

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

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

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

const (
	NumOfIterations uint64 = uint64(10)
)

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

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

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

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

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

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

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) SetEpoch(epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.testablyFail("SetEpoch", strconv.FormatUint(epoch, 10)) {
		return
	}

	t.param = &dummyData{}
	t.gradient = &dummyData{}

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

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

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

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

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

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

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

		// In real ML, we modify the gradient first. But here it is noop.
		// Notice that we only
		if t.epoch == NumOfIterations {
			t.framework.ShutdownJob()
			close(t.finishChan)
		} else {
			t.logger.Printf("master finished current epoch, task: %d, epoch: %d", t.taskID, t.epoch)
			t.framework.IncEpoch()
		}
	}
}

func (t *dummyMaster) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if len(args) >= 1 && t.config["failepoch"] != "" {
		// we need to care about fail at specific epoch
		if t.config["failepoch"] != args[0] {
			return false
		}
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("master task %d testably fail, method: %s\n", t.taskID, method)
	t.framework.(*framework).stop()
	t.NodeProducer <- true
	return true
}

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

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

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

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

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

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

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

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

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

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

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

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

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

func (t *dummySlave) ChildDataReady(childID uint64, req string, resp []byte) {
	t.logger.Printf("slave ChildDataReady, task: %d, epoch: %d, child: %d\n", t.taskID, t.epoch, childID)
	d := new(dummyData)
	json.Unmarshal(resp, d)
	if _, ok := t.fromChildren[childID]; ok {
		return
	}
	t.fromChildren[childID] = d
	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		// In real ML, we add the gradient first.
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		// If this failure happens, a new node will redo computing again.
		if t.testablyFail("ChildDataReady") {
			return
		}

		t.framework.FlagMetaToParent("GradientReady")

		// if this failure happens, the parent could
		// 1. not have the data yet. In such case, the parent could
		//   1.1 not request the data before a new node restarts. This will cause
		//       double requests since we provide at-least-once semantics.
		//   1.2 request the data with a failed host (request should fail or be
		//       responded with error message).
		// 2. already get the data.
		if t.testablyFail("ChildDataReady") {
			return
		}
	}
}

func (t *dummySlave) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("slave task %d testably fail, method: %s\n", t.taskID, method)
	t.framework.(*framework).stop()
	t.NodeProducer <- true
	return true
}

func probablyFail(levelStr string) bool {
	level, err := strconv.Atoi(levelStr)
	if err != nil {
		return false
	}
	if level < rand.Intn(100)+1 {
		return false
	}
	return true
}

// used for testing
type SimpleTaskBuilder struct {
	GDataChan    chan int32
	FinishChan   chan struct{}
	NodeProducer chan bool
	MasterConfig map[string]string
	SlaveConfig  map[string]string
}

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

import (
	"encoding/json"
	"log"
	"math/rand"
	"os"
	"strconv"
	"sync"

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

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

const (
	NumOfIterations uint64 = uint64(10)
)

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

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

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

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

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

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

// This give the task an opportunity to cleanup and regroup.
func (t *dummyMaster) SetEpoch(epoch uint64) {
	t.logger.Printf("master SetEpoch, task: %d, epoch: %d\n", t.taskID, epoch)
	if t.testablyFail("SetEpoch", strconv.FormatUint(epoch, 10)) {
		return
	}

	t.param = &dummyData{}
	t.gradient = &dummyData{}

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

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

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

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

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

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

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

		// In real ML, we modify the gradient first. But here it is noop.
		// Notice that we only
		if t.epoch == NumOfIterations {
			t.framework.ShutdownJob()
			close(t.finishChan)
		} else {
			t.logger.Printf("master finished current epoch, task: %d, epoch: %d", t.taskID, t.epoch)
			t.framework.IncEpoch()
		}
	}
}

func (t *dummyMaster) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if len(args) >= 1 && t.config["failepoch"] != "" {
		// we need to care about fail at specific epoch
		if t.config["failepoch"] != args[0] {
			return false
		}
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("master task %d testably fail, method: %s\n", t.taskID, method)
	t.framework.(*framework).stop()
	t.NodeProducer <- true
	return true
}

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

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

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

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

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

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

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

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

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

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

func (t *dummySlave) ParentDataReady(parentID uint64, req string, resp []byte) {
	t.logger.Printf("slave ParentDataReady, task: %d, epoch: %d, parent: %d\n", t.taskID, t.epoch, parentID)
	if t.testablyFail("ParentDataReady") {
		return
	}
	if t.gradientReady.IsDone() {
		return
	}
	t.param = new(dummyData)
	json.Unmarshal(resp, t.param)
	// We need to carry out local compuation.
	t.gradient.Value = t.param.Value * int32(t.framework.GetTaskID())
	t.gradientReady.Done()

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

func (t *dummySlave) ChildDataReady(childID uint64, req string, resp []byte) {
	t.logger.Printf("slave ChildDataReady, task: %d, epoch: %d, child: %d\n", t.taskID, t.epoch, childID)
	d := new(dummyData)
	json.Unmarshal(resp, d)
	if _, ok := t.fromChildren[childID]; ok {
		return
	}
	t.fromChildren[childID] = d
	// This is a weak form of checking. We can also check the task ids.
	// But this really means that we get all the events from children, we
	// should go into the next epoch now.
	if len(t.fromChildren) == len(t.framework.GetTopology().GetChildren(t.epoch)) {
		// If a new node restart and find out both parent and child meta ready, it will
		// simultaneously request both data. We need to wait until gradient data is there.
		t.gradientReady.Wait()
		// In real ML, we add the gradient first.
		for _, g := range t.fromChildren {
			t.gradient.Value += g.Value
		}

		// If this failure happens, a new node will redo computing again.
		if t.testablyFail("ChildDataReady") {
			return
		}

		t.framework.FlagMetaToParent("GradientReady")

		// if this failure happens, the parent could
		// 1. not have the data yet. In such case, the parent could
		//   1.1 not request the data before a new node restarts. This will cause
		//       double requests since we provide at-least-once semantics.
		//   1.2 request the data with a failed host (request should fail or be
		//       responded with error message).
		// 2. already get the data.
		if t.testablyFail("ChildDataReady") {
			return
		}
	}
}

func (t *dummySlave) testablyFail(method string, args ...string) bool {
	if t.config == nil {
		return false
	}
	if t.config[method] != "fail" {
		return false
	}
	if !probablyFail(t.config["faillevel"]) {
		return false
	}
	t.logger.Printf("slave task %d testably fail, method: %s\n", t.taskID, method)
	t.framework.(*framework).stop()
	t.NodeProducer <- true
	return true
}

func probablyFail(levelStr string) bool {
	level, err := strconv.Atoi(levelStr)
	if err != nil {
		return false
	}
	if level < rand.Intn(100)+1 {
		return false
	}
	return true
}

// used for testing
type SimpleTaskBuilder struct {
	GDataChan    chan int32
	FinishChan   chan struct{}
	NodeProducer chan bool
	MasterConfig map[string]string
	SlaveConfig  map[string]string
}

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

// I am writing this count down latch because sync.WaitGroup doesn't support
// decrementing counter when it's 0.
type countDownLatch struct {
	sync.Mutex
	cond    *sync.Cond
	counter int
}

func newCountDownLatch() *countDownLatch {
	c := new(countDownLatch)
	c.cond = sync.NewCond(c)
	return c
}

func (c *countDownLatch) ResetCounter(count int) {
	c.counter = count
}

func (c *countDownLatch) IsDone() bool {
	c.Lock()
	defer c.Unlock()
	return c.counter == 0
}

func (c *countDownLatch) Done() {
	c.Lock()
	defer c.Unlock()
	if c.counter == 0 {
		return
	}
	c.counter--
	if c.counter == 0 {
		c.cond.Broadcast()
	}
}

func (c *countDownLatch) Wait() {
	c.Lock()
	defer c.Unlock()
	if c.counter == 0 {
		return
	}
	c.cond.Wait()
}
<|endoftext|>


<commit_before>package query

import (
	"context"
	"fmt"
	"sync"

	"github.com/rs/rest-layer/schema"
)

// Resource represents type that can be queried by Projection.Eval.
type Resource interface {
	// Find executes the query and returns the matching items.
	Find(ctx context.Context, query *Query) ([]map[string]interface{}, error)

	// MultiGet get some items by their id and return them in the same order. If one
	// or more item(s) is not found, their slot in the response is set to nil.
	MultiGet(ctx context.Context, ids []interface{}) ([]map[string]interface{}, error)

	// SubResource returns the sub-resource at path. If path starts with a
	// dot, the lookup is performed relative to the current resource.
	SubResource(ctx context.Context, path string) (Resource, error)

	// Validator returns the schema.Validator associated with the resource.
	Validator() schema.Validator
}

// Eval evaluate the projection on the given payload with the help of the
// validator. The resolver is used to fetch payload of references outside of the
// provided payload.
func (p Projection) Eval(ctx context.Context, payload map[string]interface{}, rsc Resource) (map[string]interface{}, error) {
	rbr := &referenceBatchResolver{}
	validator := rsc.Validator()
	payload, err := evalProjection(ctx, p, payload, validator, rbr)
	if err == nil {
		// Execute the batched reference resolutions.
		err = rbr.execute(ctx, rsc)
	}
	return payload, err
}

func evalProjection(ctx context.Context, p Projection, payload map[string]interface{}, validator schema.Validator, rbr *referenceBatchResolver) (map[string]interface{}, error) {
	res := map[string]interface{}{}
	resMu := sync.Mutex{} // XXX use sync.Map once Go 1.9 is out
	if len(p) == 0 {
		// When the Projection is empty, it's like saying "all fields".
		// This allows notations like id,user{} to embed all fields of the user
		// sub-resource.
		for fn := range payload {
			p = append(p, ProjectionField{Name: fn})
		}
	}
	for _, pf := range p {
		name := pf.Name
		// Handle aliasing
		if pf.Alias != "" {
			name = pf.Alias
		}
		def := validator.GetField(pf.Name)
		// Skip hidden fields
		if def != nil && def.Hidden {
			continue
		}
		if val, found := payload[pf.Name]; found {
			// Handle sub field selection (if field has a value)
			if len(pf.Children) > 0 && val != nil {
				if def != nil && def.Schema != nil {
					subval, ok := val.(map[string]interface{})
					if !ok {
						return nil, fmt.Errorf("%s: invalid value: not a dict", pf.Name)
					}
					var err error
					if subval, err = evalProjection(ctx, pf.Children, subval, def.Schema, rbr); err != nil {
						return nil, fmt.Errorf("%s.%v", pf.Name, err)
					}
					if res[name], err = resolveFieldHandler(ctx, pf, def, subval); err != nil {
						return nil, err
					}
				} else if ref, ok := def.Validator.(*schema.Reference); ok {
					// Execute sub-request in batch
					q := &Query{Predicate: Predicate{Equal{Field: "id", Value: val}}}
					rbr.request(ref.Path, q, func(payloads []map[string]interface{}, validator schema.Validator) error {
						var v interface{}
						if len(payloads) == 1 {
							payload, err := evalProjection(ctx, pf.Children, payloads[0], validator, rbr)
							if err != nil {
								return fmt.Errorf("%s: error applying Projection on sub-field: %v", name, err)
							}
							if v, err = resolveFieldHandler(ctx, pf, def, payload); err != nil {
								return fmt.Errorf("%s: error resolving field handler on sub-field: %v", name, err)
							}
						}
						resMu.Lock()
						res[name] = v
						resMu.Unlock()
						return nil
					})
				} else {
					return nil, fmt.Errorf("%s: field as no children", pf.Name)
				}
			} else {
				var err error
				if res[name], err = resolveFieldHandler(ctx, pf, def, val); err != nil {
					return nil, err
				}
			}
		} else if def != nil {
			// If field is not found, it may be a connection
			if ref, ok := def.Validator.(*schema.Connection); ok {
				id, ok := payload["id"]
				if !ok {
					return nil, fmt.Errorf("%s: error applying projection on sub-resource: item lacks ID field", pf.Name)
				}
				q, err := connectionQuery(pf, ref.Field, id)
				if err != nil {
					return nil, err
				}
				rbr.request(ref.Path, q, func(payloads []map[string]interface{}, validator schema.Validator) (err error) {
					for i := range payloads {
						if payloads[i], err = evalProjection(ctx, pf.Children, payloads[i], validator, rbr); err != nil {
							return fmt.Errorf("%s: error applying projection on sub-resource item #%d: %v", pf.Name, i, err)
						}
					}
					var v interface{}
					if v, err = resolveFieldHandler(ctx, pf, def, payloads); err != nil {
						return fmt.Errorf("%s: error resolving field handler on sub-resource: %v", name, err)
					}
					resMu.Lock()
					res[name] = v
					resMu.Unlock()
					return nil
				})
			}
		}
	}
	return res, nil
}

// connectionQuery builds a query from a projection field on a schema.Connection type field.
func connectionQuery(pf ProjectionField, field string, id interface{}) (*Query, error) {
	q := &Query{
		Predicate: Predicate{Equal{Field: field, Value: id}},
	}
	if filter, ok := pf.Params["filter"].(string); ok {
		p, err := ParsePredicate(filter)
		if err != nil {
			return nil, fmt.Errorf("%s: invalid filter: %v", pf.Name, err)
		}
		q.Predicate = append(q.Predicate, p)
	}
	if sort, ok := pf.Params["sort"].(string); ok {
		s, err := ParseSort(sort)
		if err != nil {
			return nil, fmt.Errorf("%s: invalid sort: %v", pf.Name, err)
		}
		q.Sort = s
	}
	skip := 0
	if v, ok := pf.Params["skip"].(int); ok {
		skip = v
	}
	page := 1
	if v, ok := pf.Params["page"].(int); ok {
		page = v
	}
	limit := 20
	if v, ok := pf.Params["limit"].(int); ok {
		limit = v
	}
	q.Window = Page(page, limit, skip)
	return q, nil
}

// resolveFieldHandler calls the field handler with the provided params (if any).
func resolveFieldHandler(ctx context.Context, pf ProjectionField, def *schema.Field, val interface{}) (interface{}, error) {
	if def == nil {
		return val, nil
	}
	var err error
	if def.Handler != nil && len(pf.Params) > 0 {
		val, err = def.Handler(ctx, val, pf.Params)
		if err != nil {
			return nil, fmt.Errorf("%s: %v", pf.Name, err)
		}
	}
	if s, ok := def.Validator.(schema.FieldSerializer); ok {
		val, err = s.Serialize(val)
		if err != nil {
			return nil, fmt.Errorf("%s: %v", pf.Name, err)
		}
	}
	return val, nil
}
<commit_msg>Fix projection resolver data-race (#154)<commit_after>package query

import (
	"context"
	"fmt"
	"sync"

	"github.com/rs/rest-layer/schema"
)

// Resource represents type that can be queried by Projection.Eval.
type Resource interface {
	// Find executes the query and returns the matching items.
	Find(ctx context.Context, query *Query) ([]map[string]interface{}, error)

	// MultiGet get some items by their id and return them in the same order. If one
	// or more item(s) is not found, their slot in the response is set to nil.
	MultiGet(ctx context.Context, ids []interface{}) ([]map[string]interface{}, error)

	// SubResource returns the sub-resource at path. If path starts with a
	// dot, the lookup is performed relative to the current resource.
	SubResource(ctx context.Context, path string) (Resource, error)

	// Validator returns the schema.Validator associated with the resource.
	Validator() schema.Validator
}

// Eval evaluate the projection on the given payload with the help of the
// validator. The resolver is used to fetch payload of references outside of the
// provided payload.
func (p Projection) Eval(ctx context.Context, payload map[string]interface{}, rsc Resource) (map[string]interface{}, error) {
	rbr := &referenceBatchResolver{}
	validator := rsc.Validator()
	payload, err := evalProjection(ctx, p, payload, validator, rbr)
	if err == nil {
		// Execute the batched reference resolutions.
		err = rbr.execute(ctx, rsc)
	}
	return payload, err
}

func evalProjection(ctx context.Context, p Projection, payload map[string]interface{}, validator schema.Validator, rbr *referenceBatchResolver) (map[string]interface{}, error) {
	res := map[string]interface{}{}
	resMu := sync.Mutex{} // XXX use sync.Map once Go 1.9 is out
	if len(p) == 0 {
		// When the Projection is empty, it's like saying "all fields".
		// This allows notations like id,user{} to embed all fields of the user
		// sub-resource.
		for fn := range payload {
			p = append(p, ProjectionField{Name: fn})
		}
	}
	for i := range p {
		pf := p[i]
		name := pf.Name
		// Handle aliasing
		if pf.Alias != "" {
			name = pf.Alias
		}
		def := validator.GetField(pf.Name)
		// Skip hidden fields
		if def != nil && def.Hidden {
			continue
		}
		if val, found := payload[pf.Name]; found {
			// Handle sub field selection (if field has a value)
			if len(pf.Children) > 0 && val != nil {
				if def != nil && def.Schema != nil {
					subval, ok := val.(map[string]interface{})
					if !ok {
						return nil, fmt.Errorf("%s: invalid value: not a dict", pf.Name)
					}
					var err error
					if subval, err = evalProjection(ctx, pf.Children, subval, def.Schema, rbr); err != nil {
						return nil, fmt.Errorf("%s.%v", pf.Name, err)
					}
					if res[name], err = resolveFieldHandler(ctx, pf, def, subval); err != nil {
						return nil, err
					}
				} else if ref, ok := def.Validator.(*schema.Reference); ok {
					// Execute sub-request in batch
					q := &Query{Predicate: Predicate{Equal{Field: "id", Value: val}}}
					rbr.request(ref.Path, q, func(payloads []map[string]interface{}, validator schema.Validator) error {
						var v interface{}
						if len(payloads) == 1 {
							payload, err := evalProjection(ctx, pf.Children, payloads[0], validator, rbr)
							if err != nil {
								return fmt.Errorf("%s: error applying Projection on sub-field: %v", name, err)
							}
							if v, err = resolveFieldHandler(ctx, pf, def, payload); err != nil {
								return fmt.Errorf("%s: error resolving field handler on sub-field: %v", name, err)
							}
						}
						resMu.Lock()
						res[name] = v
						resMu.Unlock()
						return nil
					})
				} else {
					return nil, fmt.Errorf("%s: field as no children", pf.Name)
				}
			} else {
				var err error
				if res[name], err = resolveFieldHandler(ctx, pf, def, val); err != nil {
					return nil, err
				}
			}
		} else if def != nil {
			// If field is not found, it may be a connection
			if ref, ok := def.Validator.(*schema.Connection); ok {
				id, ok := payload["id"]
				if !ok {
					return nil, fmt.Errorf("%s: error applying projection on sub-resource: item lacks ID field", pf.Name)
				}
				q, err := connectionQuery(pf, ref.Field, id)
				if err != nil {
					return nil, err
				}
				rbr.request(ref.Path, q, func(payloads []map[string]interface{}, validator schema.Validator) (err error) {
					for i := range payloads {
						if payloads[i], err = evalProjection(ctx, pf.Children, payloads[i], validator, rbr); err != nil {
							return fmt.Errorf("%s: error applying projection on sub-resource item #%d: %v", pf.Name, i, err)
						}
					}
					var v interface{}
					if v, err = resolveFieldHandler(ctx, pf, def, payloads); err != nil {
						return fmt.Errorf("%s: error resolving field handler on sub-resource: %v", name, err)
					}
					resMu.Lock()
					res[name] = v
					resMu.Unlock()
					return nil
				})
			}
		}
	}
	return res, nil
}

// connectionQuery builds a query from a projection field on a schema.Connection type field.
func connectionQuery(pf ProjectionField, field string, id interface{}) (*Query, error) {
	q := &Query{
		Predicate: Predicate{Equal{Field: field, Value: id}},
	}
	if filter, ok := pf.Params["filter"].(string); ok {
		p, err := ParsePredicate(filter)
		if err != nil {
			return nil, fmt.Errorf("%s: invalid filter: %v", pf.Name, err)
		}
		q.Predicate = append(q.Predicate, p)
	}
	if sort, ok := pf.Params["sort"].(string); ok {
		s, err := ParseSort(sort)
		if err != nil {
			return nil, fmt.Errorf("%s: invalid sort: %v", pf.Name, err)
		}
		q.Sort = s
	}
	skip := 0
	if v, ok := pf.Params["skip"].(int); ok {
		skip = v
	}
	page := 1
	if v, ok := pf.Params["page"].(int); ok {
		page = v
	}
	limit := 20
	if v, ok := pf.Params["limit"].(int); ok {
		limit = v
	}
	q.Window = Page(page, limit, skip)
	return q, nil
}

// resolveFieldHandler calls the field handler with the provided params (if any).
func resolveFieldHandler(ctx context.Context, pf ProjectionField, def *schema.Field, val interface{}) (interface{}, error) {
	if def == nil {
		return val, nil
	}
	var err error
	if def.Handler != nil && len(pf.Params) > 0 {
		val, err = def.Handler(ctx, val, pf.Params)
		if err != nil {
			return nil, fmt.Errorf("%s: %v", pf.Name, err)
		}
	}
	if s, ok := def.Validator.(schema.FieldSerializer); ok {
		val, err = s.Serialize(val)
		if err != nil {
			return nil, fmt.Errorf("%s: %v", pf.Name, err)
		}
	}
	return val, nil
}
<|endoftext|>