<commit_before>// Copyright 2014, 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 vtgate

// This is a V3 file. Do not intermix with V2.

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/acl"
	"github.com/youtube/vitess/go/cache"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/vtgate/engine"
	"github.com/youtube/vitess/go/vt/vtgate/planbuilder"
	"github.com/youtube/vitess/go/vt/vtgate/vindexes"
)

// Planner is used to compute the plan. It contains
// the vschema, and has a cache of previous computed plans.
type Planner struct {
	serv    topo.SrvTopoServer
	cell    string
	mu      sync.Mutex
	vschema *vindexes.VSchema
	plans   *cache.LRUCache
}

var once sync.Once

// NewPlanner creates a new planner for VTGate.
// It will watch the vschema in the topology until the ctx is closed.
func NewPlanner(ctx context.Context, serv topo.SrvTopoServer, cell string, cacheSize int) *Planner {
	plr := &Planner{
		serv:  serv,
		cell:  cell,
		plans: cache.NewLRUCache(int64(cacheSize)),
	}
	plr.WatchVSchema(ctx)
	once.Do(func() {
		http.Handle("/debug/query_plans", plr)
		http.Handle("/debug/vschema", plr)
	})
	return plr
}

// WatchVSchema watches the VSchema from the topo. The function does
// not return an error. It instead logs warnings on failure.
// We get the list of keyspaces to watch at the beginning (because it's easier)
// and we don't watch that. So when adding a new keyspace, for now,
// vtgate needs to be restarted (after having rebuilt the serving graph).
// This could be fixed by adding a WatchSrvKeyspaceNames API to topo.Server,
// and when it triggers we'd diff the values, and terminate the missing
// SrvKeyspaces, and start watching the new SrvKeyspaces.
func (plr *Planner) WatchVSchema(ctx context.Context) {
	keyspaces, err := plr.serv.GetSrvKeyspaceNames(ctx, plr.cell)
	if err != nil {
		log.Warningf("Error loading vschema: could not read keyspaces: %v", err)
		return
	}

	// we will wait until we get the first value for each keyspace
	// before returning
	var wg sync.WaitGroup

	// mu protects formal
	var mu sync.Mutex
	formal := &vindexes.VSchemaFormal{
		Keyspaces: make(map[string]vindexes.KeyspaceFormal),
	}
	processKeyspace := func(keyspace, kschema string) {
		// Note we use a closure here to allow for defer()
		// to work properly, and also to let the caller
		// do more things even if this fails.

		// unpack the new VSchema, or skip it
		var kformal vindexes.KeyspaceFormal
		if err := json.Unmarshal([]byte(kschema), &kformal); err != nil {
			log.Warningf("Error unmarshalling vschema for keyspace %s: %v", keyspace, err)
			return
		}

		// rebuild the new component
		mu.Lock()
		defer mu.Unlock()
		formal.Keyspaces[keyspace] = kformal
		vschema, err := vindexes.BuildVSchema(formal)
		if err != nil {
			log.Warningf("Error creating VSchema: %v", err)
			return
		}

		plr.mu.Lock()
		plr.vschema = vschema
		plr.mu.Unlock()
	}

	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			gotFirstValue := false

			notifications, err := plr.serv.WatchVSchema(ctx, keyspace)
			if err != nil {
				log.Warningf("Error watching vschema for keyspace %s, will not watch it: %v", keyspace, err)
				return
			}

			for kschema := range notifications {
				processKeyspace(keyspace, kschema)

				if !gotFirstValue {
					gotFirstValue = true
					wg.Done()
				}
			}
		}(keyspace)
	}

	wg.Wait()
}

// VSchema returns the VSchema.
func (plr *Planner) VSchema() *vindexes.VSchema {
	plr.mu.Lock()
	defer plr.mu.Unlock()
	return plr.vschema
}

// GetPlan computes the plan for the given query. If one is in
// the cache, it reuses it.
func (plr *Planner) GetPlan(sql, keyspace string) (*engine.Plan, error) {
	if plr.VSchema() == nil {
		return nil, errors.New("vschema not initialized")
	}
	key := sql
	if keyspace != "" {
		key = keyspace + ":" + sql
	}
	if result, ok := plr.plans.Get(key); ok {
		return result.(*engine.Plan), nil
	}
	plan, err := planbuilder.Build(sql, &wrappedVSchema{
		vschema:  plr.VSchema(),
		keyspace: keyspace,
	})
	if err != nil {
		return nil, err
	}
	plr.plans.Set(sql, plan)
	return plan, nil
}

// ServeHTTP shows the current plans in the query cache.
func (plr *Planner) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if err := acl.CheckAccessHTTP(request, acl.DEBUGGING); err != nil {
		acl.SendError(response, err)
		return
	}
	if request.URL.Path == "/debug/query_plans" {
		keys := plr.plans.Keys()
		response.Header().Set("Content-Type", "text/plain")
		response.Write([]byte(fmt.Sprintf("Length: %d\n", len(keys))))
		for _, v := range keys {
			response.Write([]byte(fmt.Sprintf("%#v\n", v)))
			if plan, ok := plr.plans.Peek(v); ok {
				if b, err := json.MarshalIndent(plan, "", "  "); err != nil {
					response.Write([]byte(err.Error()))
				} else {
					response.Write(b)
				}
				response.Write(([]byte)("\n\n"))
			}
		}
	} else if request.URL.Path == "/debug/vschema" {
		response.Header().Set("Content-Type", "application/json; charset=utf-8")
		b, err := json.MarshalIndent(plr.VSchema(), "", " ")
		if err != nil {
			response.Write([]byte(err.Error()))
			return
		}
		buf := bytes.NewBuffer(nil)
		json.HTMLEscape(buf, b)
		response.Write(buf.Bytes())
	} else {
		response.WriteHeader(http.StatusNotFound)
	}
}

type wrappedVSchema struct {
	vschema  *vindexes.VSchema
	keyspace string
}

func (vs *wrappedVSchema) Find(keyspace, tablename string) (table *vindexes.Table, err error) {
	if keyspace == "" {
		keyspace = vs.keyspace
	}
	return vs.vschema.Find(keyspace, tablename)
}
<commit_msg>Also clearing the plan cache after vschema reload.<commit_after>// Copyright 2014, 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 vtgate

// This is a V3 file. Do not intermix with V2.

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sync"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/acl"
	"github.com/youtube/vitess/go/cache"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/vtgate/engine"
	"github.com/youtube/vitess/go/vt/vtgate/planbuilder"
	"github.com/youtube/vitess/go/vt/vtgate/vindexes"
)

// Planner is used to compute the plan. It contains
// the vschema, and has a cache of previous computed plans.
type Planner struct {
	serv    topo.SrvTopoServer
	cell    string
	mu      sync.Mutex
	vschema *vindexes.VSchema
	plans   *cache.LRUCache
}

var once sync.Once

// NewPlanner creates a new planner for VTGate.
// It will watch the vschema in the topology until the ctx is closed.
func NewPlanner(ctx context.Context, serv topo.SrvTopoServer, cell string, cacheSize int) *Planner {
	plr := &Planner{
		serv:  serv,
		cell:  cell,
		plans: cache.NewLRUCache(int64(cacheSize)),
	}
	plr.WatchVSchema(ctx)
	once.Do(func() {
		http.Handle("/debug/query_plans", plr)
		http.Handle("/debug/vschema", plr)
	})
	return plr
}

// WatchVSchema watches the VSchema from the topo. The function does
// not return an error. It instead logs warnings on failure.
// We get the list of keyspaces to watch at the beginning (because it's easier)
// and we don't watch that. So when adding a new keyspace, for now,
// vtgate needs to be restarted (after having rebuilt the serving graph).
// This could be fixed by adding a WatchSrvKeyspaceNames API to topo.Server,
// and when it triggers we'd diff the values, and terminate the missing
// SrvKeyspaces, and start watching the new SrvKeyspaces.
func (plr *Planner) WatchVSchema(ctx context.Context) {
	keyspaces, err := plr.serv.GetSrvKeyspaceNames(ctx, plr.cell)
	if err != nil {
		log.Warningf("Error loading vschema: could not read keyspaces: %v", err)
		return
	}

	// we will wait until we get the first value for each keyspace
	// before returning
	var wg sync.WaitGroup

	// mu protects formal
	var mu sync.Mutex
	formal := &vindexes.VSchemaFormal{
		Keyspaces: make(map[string]vindexes.KeyspaceFormal),
	}
	processKeyspace := func(keyspace, kschema string) {
		// Note we use a closure here to allow for defer()
		// to work properly, and also to let the caller
		// do more things even if this fails.

		// unpack the new VSchema, or skip it
		var kformal vindexes.KeyspaceFormal
		if err := json.Unmarshal([]byte(kschema), &kformal); err != nil {
			log.Warningf("Error unmarshalling vschema for keyspace %s: %v", keyspace, err)
			return
		}

		// rebuild the new component
		mu.Lock()
		defer mu.Unlock()
		formal.Keyspaces[keyspace] = kformal
		vschema, err := vindexes.BuildVSchema(formal)
		if err != nil {
			log.Warningf("Error creating VSchema: %v", err)
			return
		}

		plr.mu.Lock()
		plr.vschema = vschema
		plr.mu.Unlock()
		plr.plans.Clear()
	}

	for _, keyspace := range keyspaces {
		wg.Add(1)
		go func(keyspace string) {
			gotFirstValue := false

			notifications, err := plr.serv.WatchVSchema(ctx, keyspace)
			if err != nil {
				log.Warningf("Error watching vschema for keyspace %s, will not watch it: %v", keyspace, err)
				return
			}

			for kschema := range notifications {
				processKeyspace(keyspace, kschema)

				if !gotFirstValue {
					gotFirstValue = true
					wg.Done()
				}
			}
		}(keyspace)
	}

	wg.Wait()
}

// VSchema returns the VSchema.
func (plr *Planner) VSchema() *vindexes.VSchema {
	plr.mu.Lock()
	defer plr.mu.Unlock()
	return plr.vschema
}

// GetPlan computes the plan for the given query. If one is in
// the cache, it reuses it.
func (plr *Planner) GetPlan(sql, keyspace string) (*engine.Plan, error) {
	if plr.VSchema() == nil {
		return nil, errors.New("vschema not initialized")
	}
	key := sql
	if keyspace != "" {
		key = keyspace + ":" + sql
	}
	if result, ok := plr.plans.Get(key); ok {
		return result.(*engine.Plan), nil
	}
	plan, err := planbuilder.Build(sql, &wrappedVSchema{
		vschema:  plr.VSchema(),
		keyspace: keyspace,
	})
	if err != nil {
		return nil, err
	}
	plr.plans.Set(sql, plan)
	return plan, nil
}

// ServeHTTP shows the current plans in the query cache.
func (plr *Planner) ServeHTTP(response http.ResponseWriter, request *http.Request) {
	if err := acl.CheckAccessHTTP(request, acl.DEBUGGING); err != nil {
		acl.SendError(response, err)
		return
	}
	if request.URL.Path == "/debug/query_plans" {
		keys := plr.plans.Keys()
		response.Header().Set("Content-Type", "text/plain")
		response.Write([]byte(fmt.Sprintf("Length: %d\n", len(keys))))
		for _, v := range keys {
			response.Write([]byte(fmt.Sprintf("%#v\n", v)))
			if plan, ok := plr.plans.Peek(v); ok {
				if b, err := json.MarshalIndent(plan, "", "  "); err != nil {
					response.Write([]byte(err.Error()))
				} else {
					response.Write(b)
				}
				response.Write(([]byte)("\n\n"))
			}
		}
	} else if request.URL.Path == "/debug/vschema" {
		response.Header().Set("Content-Type", "application/json; charset=utf-8")
		b, err := json.MarshalIndent(plr.VSchema(), "", " ")
		if err != nil {
			response.Write([]byte(err.Error()))
			return
		}
		buf := bytes.NewBuffer(nil)
		json.HTMLEscape(buf, b)
		response.Write(buf.Bytes())
	} else {
		response.WriteHeader(http.StatusNotFound)
	}
}

type wrappedVSchema struct {
	vschema  *vindexes.VSchema
	keyspace string
}

func (vs *wrappedVSchema) Find(keyspace, tablename string) (table *vindexes.Table, err error) {
	if keyspace == "" {
		keyspace = vs.keyspace
	}
	return vs.vschema.Find(keyspace, tablename)
}
<|endoftext|>


<commit_before>// Copyright 2012, 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 wrangler

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	cc "github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/topo"
)

// replaceError replaces original with recent if recent is not nil,
// logging original if it wasn't nil. This should be used in deferred
// cleanup functions if they change the returned error.
func replaceError(original, recent error) error {
	if recent == nil {
		return original
	}
	if original != nil {
		log.Errorf("One of multiple error: %v", original)
	}
	return recent
}

// prepareToSnapshot changes the type of the tablet to backup (when
// the original type is master, it will proceed only if
// forceMasterSnapshot is true). It returns a function that will
// restore the original state.
func (wr *Wrangler) prepareToSnapshot(tabletAlias topo.TabletAlias, forceMasterSnapshot bool) (restoreAfterSnapshot func() error, err error) {
	ti, err := wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}

	originalType := ti.Tablet.Type

	if ti.Tablet.Type == topo.TYPE_MASTER && forceMasterSnapshot {
		// In this case, we don't bother recomputing the serving graph.
		// All queries will have to fail anyway.
		log.Infof("force change type master -> backup: %v", tabletAlias)
		// There is a legitimate reason to force in the case of a single
		// master.
		ti.Tablet.Type = topo.TYPE_BACKUP
		err = topo.UpdateTablet(wr.ts, ti)
	} else {
		err = wr.ChangeType(ti.Alias, topo.TYPE_BACKUP, false)
	}

	if err != nil {
		return
	}

	restoreAfterSnapshot = func() (err error) {
		log.Infof("change type after snapshot: %v %v", tabletAlias, originalType)

		if ti.Tablet.Parent.Uid == topo.NO_TABLET && forceMasterSnapshot {
			log.Infof("force change type backup -> master: %v", tabletAlias)
			ti.Tablet.Type = topo.TYPE_MASTER
			return topo.UpdateTablet(wr.ts, ti)
		}

		return wr.ChangeType(ti.Alias, originalType, false)
	}

	return

}

func (wr *Wrangler) MultiRestore(dstTabletAlias topo.TabletAlias, sources []topo.TabletAlias, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount int, strategy string) error {
	actionPath, err := wr.ai.MultiRestore(dstTabletAlias, &actionnode.MultiRestoreArgs{
		SrcTabletAliases:       sources,
		Concurrency:            concurrency,
		FetchConcurrency:       fetchConcurrency,
		InsertTableConcurrency: insertTableConcurrency,
		FetchRetryCount:        fetchRetryCount,
		Strategy:               strategy})
	if err != nil {
		return err
	}

	return wr.WaitForCompletion(actionPath)
}

func (wr *Wrangler) MultiSnapshot(keyRanges []key.KeyRange, tabletAlias topo.TabletAlias, concurrency int, tables, excludeTables []string, forceMasterSnapshot, skipSlaveRestart bool, maximumFilesize uint64) (manifests []string, parent topo.TabletAlias, err error) {
	restoreAfterSnapshot, err := wr.prepareToSnapshot(tabletAlias, forceMasterSnapshot)
	if err != nil {
		return
	}
	defer func() {
		err = replaceError(err, restoreAfterSnapshot())
	}()

	actionPath, err := wr.ai.MultiSnapshot(tabletAlias, &actionnode.MultiSnapshotArgs{KeyRanges: keyRanges, Concurrency: concurrency, Tables: tables, ExcludeTables: excludeTables, SkipSlaveRestart: skipSlaveRestart, MaximumFilesize: maximumFilesize})
	if err != nil {
		return
	}

	results, err := wr.WaitForCompletionReply(actionPath)
	if err != nil {
		return
	}

	reply := results.(*actionnode.MultiSnapshotReply)

	return reply.ManifestPaths, reply.ParentAlias, nil
}

func (wr *Wrangler) ShardMultiRestore(keyspace, shard string, sources []topo.TabletAlias, tables []string, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount int, strategy string) error {

	// check parameters
	if len(tables) > 0 && len(sources) > 1 {
		return fmt.Errorf("ShardMultiRestore can only handle one source when tables are specified")
	}

	// lock the shard to perform the changes we need done
	actionNode := actionnode.ShardMultiRestore(&actionnode.MultiRestoreArgs{
		SrcTabletAliases:       sources,
		Concurrency:            concurrency,
		FetchConcurrency:       fetchConcurrency,
		InsertTableConcurrency: insertTableConcurrency,
		FetchRetryCount:        fetchRetryCount,
		Strategy:               strategy})
	lockPath, err := wr.lockShard(keyspace, shard, actionNode)
	if err != nil {
		return err
	}

	mrErr := wr.SetSourceShards(keyspace, shard, sources, tables)
	err = wr.unlockShard(keyspace, shard, actionNode, lockPath, mrErr)
	if err != nil {
		if mrErr != nil {
			log.Errorf("unlockShard got error back: %v", err)
			return mrErr
		}
		return err
	}
	if mrErr != nil {
		return mrErr
	}

	// find all tablets in the shard
	destTablets, err := topo.FindAllTabletAliasesInShard(wr.ts, keyspace, shard)
	if err != nil {
		return err
	}

	// now launch MultiRestore on all tablets we need to do
	rec := cc.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, tabletAlias := range destTablets {
		wg.Add(1)
		go func(tabletAlias topo.TabletAlias) {
			log.Infof("Starting multirestore on tablet %v", tabletAlias)
			err := wr.MultiRestore(tabletAlias, sources, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount, strategy)
			log.Infof("Multirestore on tablet %v is done (err=%v)", tabletAlias, err)
			rec.RecordError(err)
			wg.Done()
		}(tabletAlias)
	}
	wg.Wait()

	return rec.Error()
}

func (wr *Wrangler) SetSourceShards(keyspace, shard string, sources []topo.TabletAlias, tables []string) error {
	// read the shard
	shardInfo, err := wr.ts.GetShard(keyspace, shard)
	if err != nil {
		return err
	}

	// If the shard already has sources, maybe it's already been restored,
	// so let's be safe and abort right here.
	if len(shardInfo.SourceShards) > 0 {
		return fmt.Errorf("Shard %v/%v already has SourceShards, not performing the MultiRestore")
	}

	// read the source tablets
	sourceTablets, err := topo.GetTabletMap(wr.TopoServer(), sources)
	if err != nil {
		return err
	}

	// Insert their KeyRange in the SourceShards array.
	// We use a linear 0-based id, that matches what mysqlctld/split.go
	// inserts into _vt.blp_checkpoint.
	shardInfo.SourceShards = make([]topo.SourceShard, len(sourceTablets))
	i := 0
	for _, ti := range sourceTablets {
		shardInfo.SourceShards[i] = topo.SourceShard{
			Uid:      uint32(i),
			Keyspace: ti.Keyspace,
			Shard:    ti.Shard,
			KeyRange: ti.KeyRange,
			Tables:   tables,
		}
		i++
	}

	// and write the shard
	if err = wr.ts.UpdateShard(shardInfo); err != nil {
		return err
	}

	return nil
}
<commit_msg>Fixing error message.<commit_after>// Copyright 2012, 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 wrangler

import (
	"fmt"
	"sync"

	log "github.com/golang/glog"
	cc "github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
	"github.com/youtube/vitess/go/vt/topo"
)

// replaceError replaces original with recent if recent is not nil,
// logging original if it wasn't nil. This should be used in deferred
// cleanup functions if they change the returned error.
func replaceError(original, recent error) error {
	if recent == nil {
		return original
	}
	if original != nil {
		log.Errorf("One of multiple error: %v", original)
	}
	return recent
}

// prepareToSnapshot changes the type of the tablet to backup (when
// the original type is master, it will proceed only if
// forceMasterSnapshot is true). It returns a function that will
// restore the original state.
func (wr *Wrangler) prepareToSnapshot(tabletAlias topo.TabletAlias, forceMasterSnapshot bool) (restoreAfterSnapshot func() error, err error) {
	ti, err := wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return
	}

	originalType := ti.Tablet.Type

	if ti.Tablet.Type == topo.TYPE_MASTER && forceMasterSnapshot {
		// In this case, we don't bother recomputing the serving graph.
		// All queries will have to fail anyway.
		log.Infof("force change type master -> backup: %v", tabletAlias)
		// There is a legitimate reason to force in the case of a single
		// master.
		ti.Tablet.Type = topo.TYPE_BACKUP
		err = topo.UpdateTablet(wr.ts, ti)
	} else {
		err = wr.ChangeType(ti.Alias, topo.TYPE_BACKUP, false)
	}

	if err != nil {
		return
	}

	restoreAfterSnapshot = func() (err error) {
		log.Infof("change type after snapshot: %v %v", tabletAlias, originalType)

		if ti.Tablet.Parent.Uid == topo.NO_TABLET && forceMasterSnapshot {
			log.Infof("force change type backup -> master: %v", tabletAlias)
			ti.Tablet.Type = topo.TYPE_MASTER
			return topo.UpdateTablet(wr.ts, ti)
		}

		return wr.ChangeType(ti.Alias, originalType, false)
	}

	return

}

func (wr *Wrangler) MultiRestore(dstTabletAlias topo.TabletAlias, sources []topo.TabletAlias, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount int, strategy string) error {
	actionPath, err := wr.ai.MultiRestore(dstTabletAlias, &actionnode.MultiRestoreArgs{
		SrcTabletAliases:       sources,
		Concurrency:            concurrency,
		FetchConcurrency:       fetchConcurrency,
		InsertTableConcurrency: insertTableConcurrency,
		FetchRetryCount:        fetchRetryCount,
		Strategy:               strategy})
	if err != nil {
		return err
	}

	return wr.WaitForCompletion(actionPath)
}

func (wr *Wrangler) MultiSnapshot(keyRanges []key.KeyRange, tabletAlias topo.TabletAlias, concurrency int, tables, excludeTables []string, forceMasterSnapshot, skipSlaveRestart bool, maximumFilesize uint64) (manifests []string, parent topo.TabletAlias, err error) {
	restoreAfterSnapshot, err := wr.prepareToSnapshot(tabletAlias, forceMasterSnapshot)
	if err != nil {
		return
	}
	defer func() {
		err = replaceError(err, restoreAfterSnapshot())
	}()

	actionPath, err := wr.ai.MultiSnapshot(tabletAlias, &actionnode.MultiSnapshotArgs{KeyRanges: keyRanges, Concurrency: concurrency, Tables: tables, ExcludeTables: excludeTables, SkipSlaveRestart: skipSlaveRestart, MaximumFilesize: maximumFilesize})
	if err != nil {
		return
	}

	results, err := wr.WaitForCompletionReply(actionPath)
	if err != nil {
		return
	}

	reply := results.(*actionnode.MultiSnapshotReply)

	return reply.ManifestPaths, reply.ParentAlias, nil
}

func (wr *Wrangler) ShardMultiRestore(keyspace, shard string, sources []topo.TabletAlias, tables []string, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount int, strategy string) error {

	// check parameters
	if len(tables) > 0 && len(sources) > 1 {
		return fmt.Errorf("ShardMultiRestore can only handle one source when tables are specified")
	}

	// lock the shard to perform the changes we need done
	actionNode := actionnode.ShardMultiRestore(&actionnode.MultiRestoreArgs{
		SrcTabletAliases:       sources,
		Concurrency:            concurrency,
		FetchConcurrency:       fetchConcurrency,
		InsertTableConcurrency: insertTableConcurrency,
		FetchRetryCount:        fetchRetryCount,
		Strategy:               strategy})
	lockPath, err := wr.lockShard(keyspace, shard, actionNode)
	if err != nil {
		return err
	}

	mrErr := wr.SetSourceShards(keyspace, shard, sources, tables)
	err = wr.unlockShard(keyspace, shard, actionNode, lockPath, mrErr)
	if err != nil {
		if mrErr != nil {
			log.Errorf("unlockShard got error back: %v", err)
			return mrErr
		}
		return err
	}
	if mrErr != nil {
		return mrErr
	}

	// find all tablets in the shard
	destTablets, err := topo.FindAllTabletAliasesInShard(wr.ts, keyspace, shard)
	if err != nil {
		return err
	}

	// now launch MultiRestore on all tablets we need to do
	rec := cc.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, tabletAlias := range destTablets {
		wg.Add(1)
		go func(tabletAlias topo.TabletAlias) {
			log.Infof("Starting multirestore on tablet %v", tabletAlias)
			err := wr.MultiRestore(tabletAlias, sources, concurrency, fetchConcurrency, insertTableConcurrency, fetchRetryCount, strategy)
			log.Infof("Multirestore on tablet %v is done (err=%v)", tabletAlias, err)
			rec.RecordError(err)
			wg.Done()
		}(tabletAlias)
	}
	wg.Wait()

	return rec.Error()
}

func (wr *Wrangler) SetSourceShards(keyspace, shard string, sources []topo.TabletAlias, tables []string) error {
	// read the shard
	shardInfo, err := wr.ts.GetShard(keyspace, shard)
	if err != nil {
		return err
	}

	// If the shard already has sources, maybe it's already been restored,
	// so let's be safe and abort right here.
	if len(shardInfo.SourceShards) > 0 {
		return fmt.Errorf("Shard %v/%v already has SourceShards, not overwriting them", keyspace, shard)
	}

	// read the source tablets
	sourceTablets, err := topo.GetTabletMap(wr.TopoServer(), sources)
	if err != nil {
		return err
	}

	// Insert their KeyRange in the SourceShards array.
	// We use a linear 0-based id, that matches what mysqlctld/split.go
	// inserts into _vt.blp_checkpoint.
	shardInfo.SourceShards = make([]topo.SourceShard, len(sourceTablets))
	i := 0
	for _, ti := range sourceTablets {
		shardInfo.SourceShards[i] = topo.SourceShard{
			Uid:      uint32(i),
			Keyspace: ti.Keyspace,
			Shard:    ti.Shard,
			KeyRange: ti.KeyRange,
			Tables:   tables,
		}
		i++
	}

	// and write the shard
	if err = wr.ts.UpdateShard(shardInfo); err != nil {
		return err
	}

	return nil
}
<|endoftext|>


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

package workflow

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"sqlflow.org/sqlflow/go/codegen/experimental"
	"sqlflow.org/sqlflow/go/database"
	"sqlflow.org/sqlflow/go/ir"
	"sqlflow.org/sqlflow/go/log"
	"sqlflow.org/sqlflow/go/parser"
	pb "sqlflow.org/sqlflow/go/proto"
	"sqlflow.org/sqlflow/go/sql"
	"sqlflow.org/sqlflow/go/workflow/argo"
	"sqlflow.org/sqlflow/go/workflow/couler"
)

// Codegen generates workflow YAML
type Codegen interface {
	GenCode([]ir.SQLFlowStmt, *pb.Session) (string, error)
	GenYAML(string) (string, error)
}

// Workflow submits workflow task and trace step status
type Workflow interface {
	Submit(string) (string, error)
	Fetch(*pb.FetchRequest) (*pb.FetchResponse, error)
}

// New returns Codegen and Workflow instance
func New(backend string) (Codegen, Workflow, error) {
	if backend == "couler" {
		return &couler.Codegen{}, &argo.Workflow{}, nil
	} else if backend == "experimental" {
		return nil, &argo.Workflow{}, nil
	}
	return nil, nil, fmt.Errorf("the specified backend: %s has not support", backend)
}

// Run compiles a SQL program to IRs and submits workflow YAML to Kubernetes
func Run(backend string, sqlProgram string, session *pb.Session, logger *log.Logger) (string, error) {
	if backend != "couler" {
		return "", fmt.Errorf("only couler backend is supported")
	}

	driverName, _, e := database.ParseURL(session.DbConnStr)
	if e != nil {
		return "", e
	}
	sqlProgram, e = parser.RemoveCommentInSQLStatement(sqlProgram)
	if e != nil {
		return "", e
	}

	var yaml string
	var wf Workflow

	useExperimentalCodegen := os.Getenv("SQLFLOW_USE_EXPERIMENTAL_CODEGEN") == "true"
	if !useExperimentalCodegen {
		stmts, e := parser.Parse(driverName, sqlProgram)
		if e != nil {
			return "", e
		}
		sqls := sql.RewriteStatementsWithHints(stmts, driverName)
		spIRs, e := sql.ResolveSQLProgram(sqls, logger)
		if e != nil {
			return "", e
		}
		// New Codegen and workflow operator instance according to the backend identifier
		var cg Codegen
		cg, wf, e = New(backend)
		if e != nil {
			return "", e
		}
		// Generate Fluid/Tekton program
		py, e := cg.GenCode(spIRs, session)
		if e != nil {
			return "", e
		}
		// translate Couler program to workflow YAML
		yaml, e = cg.GenYAML(py)
		if e != nil {
			return "", e
		}
	} else {
		// FIXME(typhoonzero): refactor this later
		wf = &argo.Workflow{}
		py, e := experimental.GenerateCodeCouler(sqlProgram, session)
		if e != nil {
			return "", e
		}
		tmpfile, e := ioutil.TempFile("/tmp", "couler")
		if e != nil {
			return "", e
		}
		defer os.Remove(tmpfile.Name())
		pyFileName := tmpfile.Name()
		if _, e = tmpfile.Write([]byte(py)); e != nil {
			tmpfile.Close()
		}
		cmdline := bytes.Buffer{}
		fmt.Fprintf(&cmdline, "couler run --mode argo --workflow_name sqlflow ")
		if c := os.Getenv("SQLFLOW_WORKFLOW_CLUSTER_CONFIG"); len(c) > 0 {
			fmt.Fprintf(&cmdline, "--cluster_config %s ", c)
		}
		fmt.Fprintf(&cmdline, "--file %s", pyFileName)
		coulerExec := strings.Split(cmdline.String(), " ")
		cmd := exec.Command(coulerExec[0], coulerExec[1:]...)
		cmd.Env = append(os.Environ())
		yamlBytes, e := cmd.CombinedOutput()
		if e != nil {
			return "", e
		}
		yaml = string(yamlBytes)
	}

	return wf.Submit(yaml)
}
<commit_msg>remove experimental backend (#3019)<commit_after>// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package workflow

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"sqlflow.org/sqlflow/go/codegen/experimental"
	"sqlflow.org/sqlflow/go/database"
	"sqlflow.org/sqlflow/go/ir"
	"sqlflow.org/sqlflow/go/log"
	"sqlflow.org/sqlflow/go/parser"
	pb "sqlflow.org/sqlflow/go/proto"
	"sqlflow.org/sqlflow/go/sql"
	"sqlflow.org/sqlflow/go/workflow/argo"
	"sqlflow.org/sqlflow/go/workflow/couler"
)

// Codegen generates workflow YAML
type Codegen interface {
	GenCode([]ir.SQLFlowStmt, *pb.Session) (string, error)
	GenYAML(string) (string, error)
}

// Workflow submits workflow task and trace step status
type Workflow interface {
	Submit(string) (string, error)
	Fetch(*pb.FetchRequest) (*pb.FetchResponse, error)
}

// New returns Codegen and Workflow instance
func New(backend string) (Codegen, Workflow, error) {
	if backend == "couler" {
		return &couler.Codegen{}, &argo.Workflow{}, nil
	}
	return nil, nil, fmt.Errorf("the specified backend: %s has not support", backend)
}

// Run compiles a SQL program to IRs and submits workflow YAML to Kubernetes
func Run(backend string, sqlProgram string, session *pb.Session, logger *log.Logger) (string, error) {
	if backend != "couler" {
		return "", fmt.Errorf("only couler backend is supported")
	}

	driverName, _, e := database.ParseURL(session.DbConnStr)
	if e != nil {
		return "", e
	}
	sqlProgram, e = parser.RemoveCommentInSQLStatement(sqlProgram)
	if e != nil {
		return "", e
	}

	var yaml string
	var wf Workflow

	useExperimentalCodegen := os.Getenv("SQLFLOW_USE_EXPERIMENTAL_CODEGEN") == "true"
	if !useExperimentalCodegen {
		stmts, e := parser.Parse(driverName, sqlProgram)
		if e != nil {
			return "", e
		}
		sqls := sql.RewriteStatementsWithHints(stmts, driverName)
		spIRs, e := sql.ResolveSQLProgram(sqls, logger)
		if e != nil {
			return "", e
		}
		// New Codegen and workflow operator instance according to the backend identifier
		var cg Codegen
		cg, wf, e = New(backend)
		if e != nil {
			return "", e
		}
		// Generate Fluid/Tekton program
		py, e := cg.GenCode(spIRs, session)
		if e != nil {
			return "", e
		}
		// translate Couler program to workflow YAML
		yaml, e = cg.GenYAML(py)
		if e != nil {
			return "", e
		}
	} else {
		// FIXME(typhoonzero): refactor this later
		wf = &argo.Workflow{}
		py, e := experimental.GenerateCodeCouler(sqlProgram, session)
		if e != nil {
			return "", e
		}
		tmpfile, e := ioutil.TempFile("/tmp", "couler")
		if e != nil {
			return "", e
		}
		defer os.Remove(tmpfile.Name())
		pyFileName := tmpfile.Name()
		if _, e = tmpfile.Write([]byte(py)); e != nil {
			tmpfile.Close()
		}
		cmdline := bytes.Buffer{}
		fmt.Fprintf(&cmdline, "couler run --mode argo --workflow_name sqlflow ")
		if c := os.Getenv("SQLFLOW_WORKFLOW_CLUSTER_CONFIG"); len(c) > 0 {
			fmt.Fprintf(&cmdline, "--cluster_config %s ", c)
		}
		fmt.Fprintf(&cmdline, "--file %s", pyFileName)
		coulerExec := strings.Split(cmdline.String(), " ")
		cmd := exec.Command(coulerExec[0], coulerExec[1:]...)
		cmd.Env = append(os.Environ())
		yamlBytes, e := cmd.CombinedOutput()
		if e != nil {
			return "", e
		}
		yaml = string(yamlBytes)
	}

	return wf.Submit(yaml)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015 Joern Rischmueller (joern.rm@gmail.com)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package goldencheetah

import (

	"os"
	"fmt"
	"net/http"

    "github.com/emicklei/go-restful"  // @Version Tag  v1.2
)

// init the Webserver within the GAE framework
func init() {

	ws := new(restful.WebService)

	// ----------------------------------------------------------------------------------
	// setup the charts endpoints - processing see "charts.go"
	// ----------------------------------------------------------------------------------
	ws.
	Path("/v1").
	Doc("Manage Charts").
	Consumes(restful.MIME_JSON).
	Produces(restful.MIME_JSON) // you can specify this per route as well

	ws.Route(ws.POST("/chart").Filter(basicAuthenticate).To(insertChart).
	// docs
	Doc("creates a chart").
	Operation("createChart").
	Reads(ChartAPIv1{})) // from the request

	ws.Route(ws.PUT("/chart").Filter(basicAuthenticate).To(updateChart).
	// docs
	Doc("updates a chart").
	Operation("updatedChart").
	Reads(ChartAPIv1{})) // from the request

	ws.Route(ws.GET("/chart/{id}").Filter(basicAuthenticate).To(getChartById).
	// docs
	Doc("get a chart").
	Operation("getChartbyId").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")).
	Writes(ChartAPIv1{})) // on the response

	ws.Route(ws.DELETE("/chart/{id}").Filter(basicAuthenticate).To(deleteChartById).
	// docs
	Doc("delete a chart by setting the deleted status").
	Operation("deleteChartbyId").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")))

	ws.Route(ws.PUT("/chartcuration/{id}/{newStatus").Filter(basicAuthenticate).To(curateChartById).
	// docs
	Doc("set the curation status of the chart to {newStatus} which must be 'true' or 'false' ").
	Operation("updateChartCurationStatus").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")).
	Param(ws.QueryParameter("newStatus", "true/false curation status").DataType("bool")))

	// Endpoint for ChartHeader only (no JPG or LTMSettings)
	ws.Route(ws.GET("/chartheader").Filter(basicAuthenticate).To(getChartHeader).
	// docs
	Doc("gets a collection of charts header - in buckets of 500 charts - table sort is new to old").
	Operation("getChartHeader").
	Param(ws.QueryParameter("dateFrom", "Date of last change").DataType("string")).
	Writes(ChartAPIv1HeaderOnlyList{})) // on the response

	// Count Chart Headers to be retrieved
	ws.Route(ws.GET("/chartheader/count").Filter(basicAuthenticate).To(getChartHeaderCount).
	// docs
	Doc("gets the number of chart headers for testing,... selection").
	Operation("getChartHeader").
	Param(ws.QueryParameter("dateFrom", "Date of last change").DataType("string")))

	// ----------------------------------------------------------------------------------
	// setup the curator endpoints - processing see "charts.go"
	// ----------------------------------------------------------------------------------
	ws.Route(ws.GET("/curator").Filter(basicAuthenticate).To(getCurator).
	// docs
	Doc("gets a collection of curators").
	Operation("getCurator").
	Param(ws.QueryParameter("curatorId", "UUid of the Curator").DataType("string")).
	Writes(CuratorAPIv1List{})) // on the response

	ws.Route(ws.POST("/curator").Filter(basicAuthenticate).To(insertCurator).
	// docs
	Doc("creates a curator").
	Operation("createCurator").
	Reads(CuratorAPIv1{})) // from the request

	// all routes defined - let's go

	restful.Add(ws)

} // init()

// global declarations
const basicauth = "Basic_Auth"
const authorization = "Authorization"
const dateTimeLayout = "2006-01-02T15:04:05Z"

func basicAuthenticate(req *restful.Request, resp *restful.Response, chain *restful.FilterChain) {
	headerClientId := req.Request.Header.Get(authorization)
	if secretClientId := os.Getenv(basicauth); secretClientId != "" {
		if fmt.Sprint("Basic ",secretClientId) != headerClientId {
			resp.AddHeader("WWW-Authenticate", "Basic realm=Protected Area")
			resp.WriteErrorString(http.StatusUnauthorized, "Not Authorized")
			return
		}
	} else {
		resp.AddHeader("WWW-Authenticate", "Basic realm=Protected Area")
		resp.WriteErrorString(http.StatusInternalServerError, "Authorization configuration missing on Server")
		return
	}

	chain.ProcessFilter(req, resp)
} // basicAuthenticate

// Convenience functions for error handling
func addPlainTextError( r *restful.Response, httpStatus int, errorReason string ) {
	r.AddHeader("Content-Type", "text/plain")
	r.WriteErrorString(httpStatus, errorReason)
}
<commit_msg>Fix API Endpoint nits<commit_after>/*
 * Copyright (c) 2015 Joern Rischmueller (joern.rm@gmail.com)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package goldencheetah

import (

	"os"
	"fmt"
	"net/http"

    "github.com/emicklei/go-restful"  // @Version Tag  v1.2
)

// init the Webserver within the GAE framework
func init() {

	ws := new(restful.WebService)

	// ----------------------------------------------------------------------------------
	// setup the charts endpoints - processing see "charts.go"
	// ----------------------------------------------------------------------------------
	ws.
	Path("/v1").
	Doc("Manage Charts").
	Consumes(restful.MIME_JSON).
	Produces(restful.MIME_JSON) // you can specify this per route as well

	ws.Route(ws.POST("/chart/").Filter(basicAuthenticate).To(insertChart).
	// docs
	Doc("creates a chart").
	Operation("createChart").
	Reads(ChartAPIv1{})) // from the request

	ws.Route(ws.PUT("/chart/").Filter(basicAuthenticate).To(updateChart).
	// docs
	Doc("updates a chart").
	Operation("updatedChart").
	Reads(ChartAPIv1{})) // from the request

	ws.Route(ws.GET("/chart/{id}").Filter(basicAuthenticate).To(getChartById).
	// docs
	Doc("get a chart").
	Operation("getChartbyId").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")).
	Writes(ChartAPIv1{})) // on the response

	ws.Route(ws.DELETE("/chart/{id}").Filter(basicAuthenticate).To(deleteChartById).
	// docs
	Doc("delete a chart by setting the deleted status").
	Operation("deleteChartbyId").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")))

	ws.Route(ws.PUT("/chartcuration/{id}").Filter(basicAuthenticate).To(curateChartById).
	// docs
	Doc("set the curation status of the chart to {newStatus} which must be 'true' or 'false' ").
	Operation("updateChartCurationStatus").
	Param(ws.PathParameter("id", "identifier of the chart").DataType("string")).
	Param(ws.QueryParameter("newStatus", "true/false curation status").DataType("bool")))

	// Endpoint for ChartHeader only (no JPG or LTMSettings)
	ws.Route(ws.GET("/chartheader").Filter(basicAuthenticate).To(getChartHeader).
	// docs
	Doc("gets a collection of charts header - in buckets of 500 charts - table sort is new to old").
	Operation("getChartHeader").
	Param(ws.QueryParameter("dateFrom", "Date of last change").DataType("string")).
	Writes(ChartAPIv1HeaderOnlyList{})) // on the response

	// Count Chart Headers to be retrieved
	ws.Route(ws.GET("/chartheader/count").Filter(basicAuthenticate).To(getChartHeaderCount).
	// docs
	Doc("gets the number of chart headers for testing,... selection").
	Operation("getChartHeader").
	Param(ws.QueryParameter("dateFrom", "Date of last change").DataType("string")))

	// ----------------------------------------------------------------------------------
	// setup the curator endpoints - processing see "charts.go"
	// ----------------------------------------------------------------------------------
	ws.Route(ws.GET("/curator").Filter(basicAuthenticate).To(getCurator).
	// docs
	Doc("gets a collection of curators").
	Operation("getCurator").
	Param(ws.QueryParameter("curatorId", "UUid of the Curator").DataType("string")).
	Writes(CuratorAPIv1List{})) // on the response

	ws.Route(ws.POST("/curator").Filter(basicAuthenticate).To(insertCurator).
	// docs
	Doc("creates a curator").
	Operation("createCurator").
	Reads(CuratorAPIv1{})) // from the request

	// all routes defined - let's go

	restful.Add(ws)

} // init()

// global declarations
const basicauth = "Basic_Auth"
const authorization = "Authorization"
const dateTimeLayout = "2006-01-02T15:04:05Z"

func basicAuthenticate(req *restful.Request, resp *restful.Response, chain *restful.FilterChain) {
	headerClientId := req.Request.Header.Get(authorization)
	if secretClientId := os.Getenv(basicauth); secretClientId != "" {
		if fmt.Sprint("Basic ",secretClientId) != headerClientId {
			resp.AddHeader("WWW-Authenticate", "Basic realm=Protected Area")
			resp.WriteErrorString(http.StatusUnauthorized, "Not Authorized")
			return
		}
	} else {
		resp.AddHeader("WWW-Authenticate", "Basic realm=Protected Area")
		resp.WriteErrorString(http.StatusInternalServerError, "Authorization configuration missing on Server")
		return
	}

	chain.ProcessFilter(req, resp)
} // basicAuthenticate

// Convenience functions for error handling
func addPlainTextError( r *restful.Response, httpStatus int, errorReason string ) {
	r.AddHeader("Content-Type", "text/plain")
	r.WriteErrorString(httpStatus, errorReason)
}
<|endoftext|>


<commit_before>package main

import (
	"../nsq"
	"../util"
	"log"
)

type Topic struct {
	name                 string
	newChannelChan       chan util.ChanReq
	channelMap           map[string]*Channel
	backend              nsq.BackendQueue
	incomingMessageChan  chan *nsq.Message
	msgChan              chan *nsq.Message
	routerSyncChan       chan int
	readSyncChan         chan int
	exitChan             chan int
	channelWriterStarted bool
}

var TopicMap = make(map[string]*Topic)
var newTopicChan = make(chan util.ChanReq)

// Topic constructor
func NewTopic(topicName string, inMemSize int, dataPath string) *Topic {
	topic := &Topic{name: topicName,
		newChannelChan:       make(chan util.ChanReq),
		channelMap:           make(map[string]*Channel),
		backend:              nsq.NewDiskQueue(topicName, dataPath),
		incomingMessageChan:  make(chan *nsq.Message, 5),
		msgChan:              make(chan *nsq.Message, inMemSize),
		routerSyncChan:       make(chan int, 1),
		readSyncChan:         make(chan int),
		exitChan:             make(chan int),
		channelWriterStarted: false}
	go topic.Router(inMemSize, dataPath)
	return topic
}

// GetTopic performs a thread safe operation
// to return a pointer to a Topic object (potentially new)
// see: topicFactory()
func GetTopic(topicName string) *Topic {
	topicChan := make(chan interface{})
	newTopicChan <- util.ChanReq{topicName, topicChan}
	return (<-topicChan).(*Topic)
}

// topicFactory is executed in a goroutine and manages
// the creation/retrieval of Topic objects
func TopicFactory(inMemSize int, dataPath string) {
	var topic *Topic
	var ok bool

	for {
		topicReq := <-newTopicChan
		name := topicReq.Variable.(string)
		if topic, ok = TopicMap[name]; !ok {
			topic = NewTopic(name, inMemSize, dataPath)
			TopicMap[name] = topic
			log.Printf("TOPIC(%s): created", topic.name)
		}
		topicReq.RetChan <- topic
	}
}

// GetChannel performs a thread safe operation
// to return a pointer to a Channel object (potentially new)
// for the given Topic
// see: Topic.Router()
func (t *Topic) GetChannel(channelName string) *Channel {
	channelChan := make(chan interface{})
	t.newChannelChan <- util.ChanReq{channelName, channelChan}
	return (<-channelChan).(*Channel)
}

// PutMessage writes to the appropriate incoming
// message channel
func (t *Topic) PutMessage(msg *nsq.Message) {
	// log.Printf("TOPIC(%s): PutMessage(%s)", t.name, string(msg.Body()))
	t.incomingMessageChan <- msg
}

// MessagePump selects over the in-memory and backend queue and
// writes messages to every channel for this topic, synchronizing
// with the channel router
func (t *Topic) MessagePump() {
	var msg *nsq.Message

	for {
		select {
		case msg = <-t.msgChan:
		case <-t.backend.ReadReadyChan():
			buf, err := t.backend.Get()
			if err != nil {
				log.Printf("ERROR: t.backend.Get() - %s", err.Error())
				continue
			}
			msg = nsq.NewMessage(buf)
		}

		t.readSyncChan <- 1
		log.Printf("TOPIC(%s): channelMap %#v", t.name, t.channelMap)
		for _, channel := range t.channelMap {
			go func(c *Channel) {
				// log.Printf("TOPIC(%s): writing message to channel(%s)", t.name, c.name)
				c.PutMessage(msg)
			}(channel)
		}
		t.routerSyncChan <- 1
	}
}

// Router handles muxing of Topic messages including
// creation of new Channel objects, proxying messages
// to memory or backend, and synchronizing reads
func (t *Topic) Router(inMemSize int, dataPath string) {
	var msg *nsq.Message

	for {
		select {
		case channelReq := <-t.newChannelChan:
			name := channelReq.Variable.(string)
			channel, ok := t.channelMap[name]
			if !ok {
				channel = NewChannel(name, inMemSize, dataPath)
				t.channelMap[name] = channel
				log.Printf("TOPIC(%s): new channel(%s)", t.name, channel.name)
			}
			channelReq.RetChan <- channel
			if !t.channelWriterStarted {
				go t.MessagePump()
				t.channelWriterStarted = true
			}
		case msg = <-t.incomingMessageChan:
			select {
			case t.msgChan <- msg:
				// log.Printf("TOPIC(%s): wrote to messageChan", t.name)
			default:
				err := t.backend.Put(msg.Data)
				if err != nil {
					log.Printf("ERROR: t.backend.Put() - %s", err.Error())
					// TODO: requeue?
				}
				// log.Printf("TOPIC(%s): wrote to backend", t.name)
			}
		case <-t.readSyncChan:
			// log.Printf("TOPIC(%s): read sync START", t.name)
			<-t.routerSyncChan
			// log.Printf("TOPIC(%s): read sync END", t.name)
		case <-t.exitChan:
			return
		}
	}
}

func (t *Topic) Close() error {
	var err error

	log.Printf("TOPIC(%s): closing", t.name)

	t.exitChan <- 1

	for _, channel := range t.channelMap {
		err = channel.Close()
		if err != nil {
			// we need to continue regardless of error to close all the channels
			log.Printf("ERROR: channel(%s) close - %s", channel.name, err.Error())
		}
	}

	err = t.backend.Close()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>resolve elusive bug where channel was accessing variable from outer scope and it could change<commit_after>package main

import (
	"../nsq"
	"../util"
	"log"
)

type Topic struct {
	name                 string
	newChannelChan       chan util.ChanReq
	channelMap           map[string]*Channel
	backend              nsq.BackendQueue
	incomingMessageChan  chan *nsq.Message
	msgChan              chan *nsq.Message
	routerSyncChan       chan int
	readSyncChan         chan int
	exitChan             chan int
	channelWriterStarted bool
}

var TopicMap = make(map[string]*Topic)
var newTopicChan = make(chan util.ChanReq)

// Topic constructor
func NewTopic(topicName string, inMemSize int, dataPath string) *Topic {
	topic := &Topic{name: topicName,
		newChannelChan:       make(chan util.ChanReq),
		channelMap:           make(map[string]*Channel),
		backend:              nsq.NewDiskQueue(topicName, dataPath),
		incomingMessageChan:  make(chan *nsq.Message, 5),
		msgChan:              make(chan *nsq.Message, inMemSize),
		routerSyncChan:       make(chan int, 1),
		readSyncChan:         make(chan int),
		exitChan:             make(chan int),
		channelWriterStarted: false}
	go topic.Router(inMemSize, dataPath)
	return topic
}

// GetTopic performs a thread safe operation
// to return a pointer to a Topic object (potentially new)
// see: topicFactory()
func GetTopic(topicName string) *Topic {
	topicChan := make(chan interface{})
	newTopicChan <- util.ChanReq{topicName, topicChan}
	return (<-topicChan).(*Topic)
}

// topicFactory is executed in a goroutine and manages
// the creation/retrieval of Topic objects
func TopicFactory(inMemSize int, dataPath string) {
	var topic *Topic
	var ok bool

	for {
		topicReq := <-newTopicChan
		name := topicReq.Variable.(string)
		if topic, ok = TopicMap[name]; !ok {
			topic = NewTopic(name, inMemSize, dataPath)
			TopicMap[name] = topic
			log.Printf("TOPIC(%s): created", topic.name)
		}
		topicReq.RetChan <- topic
	}
}

// GetChannel performs a thread safe operation
// to return a pointer to a Channel object (potentially new)
// for the given Topic
// see: Topic.Router()
func (t *Topic) GetChannel(channelName string) *Channel {
	channelChan := make(chan interface{})
	t.newChannelChan <- util.ChanReq{channelName, channelChan}
	return (<-channelChan).(*Channel)
}

// PutMessage writes to the appropriate incoming
// message channel
func (t *Topic) PutMessage(msg *nsq.Message) {
	// log.Printf("TOPIC(%s): PutMessage(%s, %s)", t.name, util.UuidToStr(msg.Uuid()), string(msg.Body()))
	t.incomingMessageChan <- msg
}

// MessagePump selects over the in-memory and backend queue and
// writes messages to every channel for this topic, synchronizing
// with the channel router
func (t *Topic) MessagePump() {
	var msg *nsq.Message

	for {
		select {
		case msg = <-t.msgChan:
		case <-t.backend.ReadReadyChan():
			buf, err := t.backend.Get()
			if err != nil {
				log.Printf("ERROR: t.backend.Get() - %s", err.Error())
				continue
			}
			msg = nsq.NewMessage(buf)
		}

		t.readSyncChan <- 1
		log.Printf("TOPIC(%s): channelMap %#v", t.name, t.channelMap)
		for _, channel := range t.channelMap {
			go func(c *Channel, m *nsq.Message) {
				// log.Printf("TOPIC(%s): writing message to channel(%s)", t.name, c.name)
				c.PutMessage(m)
			}(channel, msg)
		}
		t.routerSyncChan <- 1
	}
}

// Router handles muxing of Topic messages including
// creation of new Channel objects, proxying messages
// to memory or backend, and synchronizing reads
func (t *Topic) Router(inMemSize int, dataPath string) {
	var msg *nsq.Message

	for {
		select {
		case channelReq := <-t.newChannelChan:
			name := channelReq.Variable.(string)
			channel, ok := t.channelMap[name]
			if !ok {
				channel = NewChannel(name, inMemSize, dataPath)
				t.channelMap[name] = channel
				log.Printf("TOPIC(%s): new channel(%s)", t.name, channel.name)
			}
			channelReq.RetChan <- channel
			if !t.channelWriterStarted {
				go t.MessagePump()
				t.channelWriterStarted = true
			}
		case msg = <-t.incomingMessageChan:
			select {
			case t.msgChan <- msg:
				// log.Printf("TOPIC(%s): wrote to messageChan", t.name)
			default:
				err := t.backend.Put(msg.Data)
				if err != nil {
					log.Printf("ERROR: t.backend.Put() - %s", err.Error())
					// TODO: requeue?
				}
				// log.Printf("TOPIC(%s): wrote to backend", t.name)
			}
		case <-t.readSyncChan:
			// log.Printf("TOPIC(%s): read sync START", t.name)
			<-t.routerSyncChan
			// log.Printf("TOPIC(%s): read sync END", t.name)
		case <-t.exitChan:
			return
		}
	}
}

func (t *Topic) Close() error {
	var err error

	log.Printf("TOPIC(%s): closing", t.name)

	t.exitChan <- 1

	for _, channel := range t.channelMap {
		err = channel.Close()
		if err != nil {
			// we need to continue regardless of error to close all the channels
			log.Printf("ERROR: channel(%s) close - %s", channel.name, err.Error())
		}
	}

	err = t.backend.Close()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io"
	"os"
	"regexp"
)

const _FILE_FILTERS = "filters.json"

type ErrInvalidRegexp struct {
	sect, ptrn string // section and pattern
	err        error
}

func (e *ErrInvalidRegexp) Error() string {
	return `invalid regexp filter in "` + e.sect + `": ` +
		e.err.Error() + `; ` + e.ptrn
}

func filter(orders []*Order) []*Order {
	if len(orders) == 0 {
		return orders
	}
	file, err := os.Open(_FILE_FILTERS)
	if err != nil {
		if os.IsExist(err) {
			log.Error.Println("filter(): can't open file with",
				"filter patterns:", err)
		} else {
			log.Warning.Println(_FILE_FILTERS, "was not found")
		}
		return orders
	}
	dec := json.NewDecoder(file)
	var patterns *struct {
		All, OrderName, OKDP, OKPD, OrganisationName []string
	}
	err = dec.Decode(&patterns)
	if err != nil {
		if err != io.EOF {
			log.Error.Println("filter(): can't decode json",
				"from file with filter patterns:", err)
		} else {
			log.Warning.Println("empty filter file", _FILE_FILTERS)
		}
		return orders
	}
	// filter all fields
	for _, pattern := range patterns.All {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrderName) ||
					exp.MatchString(orders[i].OKDP) ||
					exp.MatchString(orders[i].OKPD) ||
					exp.MatchString(orders[i].OrganisationName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"All",
				pattern,
				err,
			})
		}
	}
	// filter each fields
	for _, pattern := range patterns.OrderName {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrderName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OrderName",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OKDP {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OKDP) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OKDP",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OKPD {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OKPD) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OKPD",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OrganisationName {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrganisationName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OrganisationName",
				pattern,
				err,
			})
		}
	}
	return orders
}
<commit_msg>added filtering stats<commit_after>package main

import (
	"encoding/json"
	"io"
	"os"
	"regexp"
)

const _FILE_FILTERS = "filters.json"

type ErrInvalidRegexp struct {
	sect, ptrn string // section and pattern
	err        error
}

func (e *ErrInvalidRegexp) Error() string {
	return `invalid regexp filter in "` + e.sect + `": ` +
		e.err.Error() + `; ` + e.ptrn
}

func filter(orders []*Order) []*Order {
	count := len(orders)
	if count == 0 {
		return orders
	}
	file, err := os.Open(_FILE_FILTERS)
	if err != nil {
		if os.IsExist(err) {
			log.Error.Println("filter(): can't open file with",
				"filter patterns:", err)
		} else {
			log.Warning.Println(_FILE_FILTERS, "was not found")
		}
		return orders
	}
	dec := json.NewDecoder(file)
	var patterns *struct {
		All, OrderName, OKDP, OKPD, OrganisationName []string
	}
	err = dec.Decode(&patterns)
	if err != nil {
		if err != io.EOF {
			log.Error.Println("filter(): can't decode json",
				"from file with filter patterns:", err)
		} else {
			log.Warning.Println("empty filter file", _FILE_FILTERS)
		}
		return orders
	}
	// filter all fields
	for _, pattern := range patterns.All {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrderName) ||
					exp.MatchString(orders[i].OKDP) ||
					exp.MatchString(orders[i].OKPD) ||
					exp.MatchString(orders[i].OrganisationName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"All",
				pattern,
				err,
			})
		}
	}
	// filter each fields
	for _, pattern := range patterns.OrderName {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrderName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OrderName",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OKDP {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OKDP) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OKDP",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OKPD {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OKPD) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OKPD",
				pattern,
				err,
			})
		}
	}
	for _, pattern := range patterns.OrganisationName {
		exp, err := regexp.Compile(pattern)
		if err == nil {
			for i := 0; i < len(orders); {
				if exp.MatchString(orders[i].OrganisationName) {
					orders = append(orders[:i], orders[i+1:]...)
				} else {
					i++
				}
			}
		} else {
			log.Warning.Println(&ErrInvalidRegexp{
				"OrganisationName",
				pattern,
				err,
			})
		}
	}
	filtered := (1 - float32(len(orders))/float32(count))
	log.Warning.Printf("filtered %.1f%%\n", filtered*100)
	return orders
}
<|endoftext|>


<commit_before>package xsd

import (
	"go/ast"
	"go/token"
	"strings"
)

func addBSONTags(decls []ast.Decl) []ast.Decl {
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			for _, field := range structType.Fields.List {
				xmlTag := ParseXMLTag(field.Tag.Value)
				name := camelToSnake(field.Names[0].Name)
				if name == "id" {
					name = "_id"
				}
				jsonTag := JSONTag{
					Name:      name,
					Omitempty: xmlTag.Omitempty,
				}
				if xmlTag.Type == XMLOmitted {
					jsonTag.Type = JSONOmitted
				}
				field.Tag = tag(xmlTag.String() + " " + jsonTag.String())
			}
		}
	}
	return decls
}

func elevateSubArrays(decls []ast.Decl) []ast.Decl {
	arrayTypes := make(map[string]*ast.Field)
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			if len(structType.Fields.List) != 1 {
				continue
			}
			field := structType.Fields.List[0]
			fieldType, ok := field.Type.(*ast.Ident)
			if !ok {
				continue
			}
			if strings.HasPrefix(fieldType.Name, "[]") {
				arrayTypes[typeSpec.Name.Name] = field
			}
		}
	}
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			for _, field := range structType.Fields.List {
				if typ, ok := field.Type.(*ast.Ident); ok {
					typeName := typ.Name
					typeName = strings.TrimPrefix(typeName, "*")
					if f, ok := arrayTypes[typeName]; ok {
						parentTag := ParseXMLTag(field.Tag.Value)
						childTag := ParseXMLTag(f.Tag.Value)
						*field = *f
						field.Tag = tag(XMLTag{
							Name:      parentTag.Name + ">" + childTag.Name,
							Type:      childTag.Type,
							Omitempty: parentTag.Omitempty || childTag.Omitempty,
						}.String())
					}
				}
			}
		}
	}

	var newDecls []ast.Decl
nextDecl:
	for _, decl := range decls {
		if genDecl, ok := decl.(*ast.GenDecl); ok {
			if genDecl.Tok == token.TYPE {
				for _, spec := range genDecl.Specs {
					if typeSpec, ok := spec.(*ast.TypeSpec); ok {
						if _, ok := arrayTypes[typeSpec.Name.Name]; ok {
							continue nextDecl
						}
					}
				}
			}
		}
		newDecls = append(newDecls, decl)
	}
	return newDecls
}
<commit_msg>json tag should not use _id<commit_after>package xsd

import (
	"go/ast"
	"go/token"
	"strings"
)

func addBSONTags(decls []ast.Decl) []ast.Decl {
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			for _, field := range structType.Fields.List {
				xmlTag := ParseXMLTag(field.Tag.Value)
				name := camelToSnake(field.Names[0].Name)
				jsonTag := JSONTag{
					Name:      name,
					Omitempty: xmlTag.Omitempty,
				}
				if xmlTag.Type == XMLOmitted {
					jsonTag.Type = JSONOmitted
				}
				field.Tag = tag(xmlTag.String() + " " + jsonTag.String())
			}
		}
	}
	return decls
}

func elevateSubArrays(decls []ast.Decl) []ast.Decl {
	arrayTypes := make(map[string]*ast.Field)
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			if len(structType.Fields.List) != 1 {
				continue
			}
			field := structType.Fields.List[0]
			fieldType, ok := field.Type.(*ast.Ident)
			if !ok {
				continue
			}
			if strings.HasPrefix(fieldType.Name, "[]") {
				arrayTypes[typeSpec.Name.Name] = field
			}
		}
	}
	for _, decl := range decls {
		genDecl, ok := decl.(*ast.GenDecl)
		if !ok {
			continue
		}
		if genDecl.Tok != token.TYPE {
			continue
		}
		for _, spec := range genDecl.Specs {
			typeSpec, ok := spec.(*ast.TypeSpec)
			if !ok {
				continue
			}
			structType, ok := typeSpec.Type.(*ast.StructType)
			if !ok {
				continue
			}
			for _, field := range structType.Fields.List {
				if typ, ok := field.Type.(*ast.Ident); ok {
					typeName := typ.Name
					typeName = strings.TrimPrefix(typeName, "*")
					if f, ok := arrayTypes[typeName]; ok {
						parentTag := ParseXMLTag(field.Tag.Value)
						childTag := ParseXMLTag(f.Tag.Value)
						*field = *f
						field.Tag = tag(XMLTag{
							Name:      parentTag.Name + ">" + childTag.Name,
							Type:      childTag.Type,
							Omitempty: parentTag.Omitempty || childTag.Omitempty,
						}.String())
					}
				}
			}
		}
	}

	var newDecls []ast.Decl
nextDecl:
	for _, decl := range decls {
		if genDecl, ok := decl.(*ast.GenDecl); ok {
			if genDecl.Tok == token.TYPE {
				for _, spec := range genDecl.Specs {
					if typeSpec, ok := spec.(*ast.TypeSpec); ok {
						if _, ok := arrayTypes[typeSpec.Name.Name]; ok {
							continue nextDecl
						}
					}
				}
			}
		}
		newDecls = append(newDecls, decl)
	}
	return newDecls
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	zkPath         = "/aurora/scheduler"
	zkLeaderPrefix = "singleton_candidate_"
)

type finder interface {
	leaderURL() (string, error)
}

func newFinder(url string) (f finder, err error) {
	if strings.HasPrefix(url, "http://") {
		f = &httpFinder{url: url}
	}

	if strings.HasPrefix(url, "zk://") {
		f = newZkFinder(url)
	}

	if f == nil {
		err = errors.New("finder: bad address")
	}

	return f, err
}

type httpFinder struct {
	url string
}

func (f *httpFinder) leaderURL() (string, error) {
	// This will redirect us to the elected Aurora master
	schedulerURL := fmt.Sprintf("%s/scheduler", f.url)
	rr, err := http.NewRequest("GET", schedulerURL, nil)
	if err != nil {
		return "", err
	}

	rresp, err := httpClient.Transport.RoundTrip(rr)
	if err != nil {
		return "", err
	}
	defer rresp.Body.Close()

	masterLoc := rresp.Header.Get("Location")
	if masterLoc == "" {
		glog.V(6).Info("missing Location header in request")
		masterLoc = schedulerURL
	}

	return strings.TrimRight(masterLoc, "/scheduler"), nil
}

func hostsFromURL(urls string) (hosts []string, err error) {
	for _, s := range strings.Split(urls, ",") {
		u, err := url.Parse(s)
		if err != nil {
			return hosts, err
		}

		hosts = append(hosts, u.Host)
	}

	return hosts, err
}

type zkFinder struct {
	sync.RWMutex
	conn     *zk.Conn
	leaderIP string
}

func newZkFinder(url string) *zkFinder {
	zkSrvs, err := hostsFromURL(url)
	if err != nil {
		panic(err)
	}

	conn, events, err := zk.Connect(zkSrvs, 20*time.Second)
	if err != nil {
		panic(err)
	}

	go func() {
		for ev := range events {
			glog.V(6).Infof("conn: %s server: %s", ev.State, ev.Server)
		}
	}()

	f := zkFinder{conn: conn}
	go f.watch()

	return &f
}

func (f *zkFinder) leaderzNode() (string, error) {
	children, stat, err := f.conn.Children(zkPath)
	if stat == nil {
		err = errors.New("zkFinder: children returned nil stat")
	}
	if err != nil {
		return "", err
	}

	var leaderSeq int
	var leader string
	for _, child := range children {
		path := strings.Split(child, zkLeaderPrefix)
		if len(path) > 1 {
			seq, err := strconv.Atoi(path[1])
			if err != nil {
				return "", err
			}

			if leader == "" {
				leader = child
			}

			if seq <= leaderSeq {
				leaderSeq = seq
				leader = child
			}
		}
	}

	if leader == "" {
		return leader, errors.New("zkFinder: zNode not found")
	}

	return fmt.Sprintf("%s/%s", zkPath, leader), nil
}

func (f *zkFinder) leaderURL() (string, error) {
	f.RLock()
	defer f.RUnlock()

	if f.leaderIP == "" {
		return "", errors.New("zkFinder: no leader found via ZooKeeper")
	}

	return fmt.Sprintf("http://%s:8081", f.leaderIP), nil
}

func (f *zkFinder) watch() {
	for _ = range time.NewTicker(1 * time.Second).C {
		zNode, err := f.leaderzNode()
		if err != nil {
			glog.Warning(err)
			continue
		}

		glog.V(6).Info("leader zNode at: ", zNode)

		data, stat, events, err := f.conn.GetW(zNode)
		if stat == nil {
			err = errors.New("get returned nil stat")
		}
		if err != nil {
			glog.Warning(err)
			continue
		}

		f.Lock()
		f.leaderIP = string(data)
		f.Unlock()

		for ev := range events {
			switch {
			case ev.Err != nil:
				err = fmt.Errorf("watcher error %+v", ev.Err)
			case ev.Type == zk.EventNodeDeleted:
				err = errors.New("leader zNode deleted")
			}

			if err != nil {
				glog.Warning(err)
				break
			}
		}
	}
}
<commit_msg>show what the mutex is protecting<commit_after>package main

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/golang/glog"
	"github.com/samuel/go-zookeeper/zk"
)

const (
	zkPath         = "/aurora/scheduler"
	zkLeaderPrefix = "singleton_candidate_"
)

type finder interface {
	leaderURL() (string, error)
}

func newFinder(url string) (f finder, err error) {
	if strings.HasPrefix(url, "http://") {
		f = &httpFinder{url: url}
	}

	if strings.HasPrefix(url, "zk://") {
		f = newZkFinder(url)
	}

	if f == nil {
		err = errors.New("finder: bad address")
	}

	return f, err
}

type httpFinder struct {
	url string
}

func (f *httpFinder) leaderURL() (string, error) {
	// This will redirect us to the elected Aurora master
	schedulerURL := fmt.Sprintf("%s/scheduler", f.url)
	rr, err := http.NewRequest("GET", schedulerURL, nil)
	if err != nil {
		return "", err
	}

	rresp, err := httpClient.Transport.RoundTrip(rr)
	if err != nil {
		return "", err
	}
	defer rresp.Body.Close()

	masterLoc := rresp.Header.Get("Location")
	if masterLoc == "" {
		glog.V(6).Info("missing Location header in request")
		masterLoc = schedulerURL
	}

	return strings.TrimRight(masterLoc, "/scheduler"), nil
}

func hostsFromURL(urls string) (hosts []string, err error) {
	for _, s := range strings.Split(urls, ",") {
		u, err := url.Parse(s)
		if err != nil {
			return hosts, err
		}

		hosts = append(hosts, u.Host)
	}

	return hosts, err
}

type zkFinder struct {
	conn *zk.Conn

	sync.RWMutex
	leaderIP string
}

func newZkFinder(url string) *zkFinder {
	zkSrvs, err := hostsFromURL(url)
	if err != nil {
		panic(err)
	}

	conn, events, err := zk.Connect(zkSrvs, 20*time.Second)
	if err != nil {
		panic(err)
	}

	go func() {
		for ev := range events {
			glog.V(6).Infof("conn: %s server: %s", ev.State, ev.Server)
		}
	}()

	f := zkFinder{conn: conn}
	go f.watch()

	return &f
}

func (f *zkFinder) leaderzNode() (string, error) {
	children, stat, err := f.conn.Children(zkPath)
	if stat == nil {
		err = errors.New("zkFinder: children returned nil stat")
	}
	if err != nil {
		return "", err
	}

	var leaderSeq int
	var leader string
	for _, child := range children {
		path := strings.Split(child, zkLeaderPrefix)
		if len(path) > 1 {
			seq, err := strconv.Atoi(path[1])
			if err != nil {
				return "", err
			}

			if leader == "" {
				leader = child
			}

			if seq <= leaderSeq {
				leaderSeq = seq
				leader = child
			}
		}
	}

	if leader == "" {
		return leader, errors.New("zkFinder: zNode not found")
	}

	return fmt.Sprintf("%s/%s", zkPath, leader), nil
}

func (f *zkFinder) leaderURL() (string, error) {
	f.RLock()
	defer f.RUnlock()

	if f.leaderIP == "" {
		return "", errors.New("zkFinder: no leader found via ZooKeeper")
	}

	return fmt.Sprintf("http://%s:8081", f.leaderIP), nil
}

func (f *zkFinder) watch() {
	for _ = range time.NewTicker(1 * time.Second).C {
		zNode, err := f.leaderzNode()
		if err != nil {
			glog.Warning(err)
			continue
		}

		glog.V(6).Info("leader zNode at: ", zNode)

		data, stat, events, err := f.conn.GetW(zNode)
		if stat == nil {
			err = errors.New("get returned nil stat")
		}
		if err != nil {
			glog.Warning(err)
			continue
		}

		f.Lock()
		f.leaderIP = string(data)
		f.Unlock()

		for ev := range events {
			switch {
			case ev.Err != nil:
				err = fmt.Errorf("watcher error %+v", ev.Err)
			case ev.Type == zk.EventNodeDeleted:
				err = errors.New("leader zNode deleted")
			}

			if err != nil {
				glog.Warning(err)
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/kylelemons/gousb/usb"
	"log"
	"math/rand"
	"time"
)

type FitbitBase struct {
	ctx                *usb.Context
	dev                *usb.Device
	base               *ANT
	trackerPacketCount int
	currentPacketId    int
	currentBankId      int
}

func (f *FitbitBase) Open() error {
	f.ctx = usb.NewContext()
	devs, err := f.ctx.ListDevices(func(desc *usb.Descriptor) bool {
		if desc.Vendor.String() == "10c4" && desc.Product.String() == "84c4" {

			return true
		}
		return false
	})
	if err == nil {
		if len(devs) == 0 {
			err = fmt.Errorf("no devices found")
		} else {
			f.dev = devs[0]
		}
		for i := 1; i < len(devs); i++ {
			devs[i].Close()
		}
	}
	return err
}

func (f *FitbitBase) SettingUp() error {
	f.dev.Control(64, 0, 65535, 0, []byte{})
	f.dev.Control(64, 1, 8192, 0, []byte{})
	f.dev.Control(64, 0, 0, 0, []byte{})
	f.dev.Control(64, 0, 65535, 0, []byte{})
	f.dev.Control(64, 1, 8192, 0, []byte{})
	f.dev.Control(64, 1, 74, 0, []byte{})
	f.dev.Control(192, 255, 14091, 0, []byte{'\x01'})
	f.dev.Control(64, 3, 2048, 0, []byte{})
	f.dev.Control(64, 19, 0, 0, []byte{
		'\x08', '\x00', '\x00', '\x00',
		'\x40', '\x00', '\x00', '\x00',
		'\x00', '\x00', '\x00', '\x00',
		'\x00', '\x00', '\x00', '\x00'})
	f.dev.Control(64, 18, 12, 0, []byte{})
	err := f.dev.SetConfig(1)
	if err != nil {
		log.Println("setconfig error")
		return err
	}
	err = f.dev.Reset()
	if err != nil {
		log.Println("reset error")
		return err
	}
	err = f.dev.SetConfig(1)
	if err != nil {
		log.Println("setconfig2 error")
		return err
	}
	f.base = &ANT{
		channel: '\x00',
	}
	f.base.reader, err = f.dev.OpenEndpoint(uint8(1), uint8(0), uint8(0), '\x81')
	if err != nil {
		return err
	}
	f.base.writer, err = f.dev.OpenEndpoint(uint8(1), uint8(0), uint8(0), '\x01')
	return err
}

func (f *FitbitBase) Close() {
	f.dev.Close()
	f.ctx.Close()
}

// data transport
func (f *FitbitBase) InitDeviceChannel(channel []byte) error {
	err := f.base.Reset()
	if err != nil {
		log.Println("fitbit reset failed")
		return err
	}
	err = f.base.SetNetworkKey(0, []byte{0, 0, 0, 0, 0, 0, 0, 0})
	if err != nil {
		log.Println("fitbit set network key failed")
		return err
	}
	err = f.base.AssignChannel()
	if err != nil {
		log.Println("fitbit assign channel failed")
		return err
	}
	err = f.base.SetChannelPeriod([]byte{'\x00', '\x10'})
	if err != nil {
		log.Println("fitbit set channel period failed")
		return err
	}
	err = f.base.SetChannelFrequency('\x02')
	if err != nil {
		log.Println("fitbit set channel frequency failed")
		return err
	}
	err = f.base.SetTransmitPower('\x03')
	if err != nil {
		log.Println("fitbit set transmit power failed")
		return err
	}
	err = f.base.SetSearchTimeout('\xff')
	if err != nil {
		log.Println("fitbit set search timeout failed")
		return err
	}
	err = f.base.SetChannelId(channel)
	if err != nil {
		log.Println("fitbit set channel id failed")
		return err
	}
	err = f.base.OpenChannel()
	if err != nil {
		log.Println("fitbit open channel failed")
		return err
	}
	return err
}

func (f *FitbitBase) InitTrackerForTransfer() error {
	err := f.InitDeviceChannel([]byte{'\xff', '\xff', '\x01', '\x01'})
	if err != nil {
		log.Println("fitbit init failed")
		return err
	}
	err = f.WaitForBeacon()
	if err != nil {
		return err
	}
	err = f.ResetTracker()
	if err != nil {
		log.Println("fitbit reset tracker failed")
		return err
	}
	cid := []byte{byte(rand.Intn(254)), byte(rand.Intn(254))}
	err = f.base.SendAcknowledgedData(append(append([]byte{'\x78', '\x02'}, cid...), []byte{'\x00', '\x00', '\x00', '\x00'}...))
	if err != nil {
		log.Println("fitbit ack ack to tracker failed")
		return err
	}
	err = f.base.CloseChannel()
	if err != nil {
		log.Println("fitbit close channel failed")
		return err
	}
	err = f.InitDeviceChannel(append(cid, []byte{0x01, 0x01}...))
	if err != nil {
		log.Println("fitbit reinit channel failed")
		return err
	}
	err = f.WaitForBeacon()
	if err != nil {
		return err
	}
	err = f.PingTracker()
	if err != nil {
		log.Println("fitbit ping tracker failed")
	}
	return err
}

func (f *FitbitBase) ResetTracker() error {
	return f.base.SendAcknowledgedData([]byte{'\x78', '\x01', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'})
}

func (f *FitbitBase) CommandSleep() error {
	return f.base.SendAcknowledgedData([]byte{'\x7f', '\x03', '\x00', '\x00', '\x00', '\x00', '\x00', '\x3c'})
}

func (f *FitbitBase) PingTracker() error {
	return f.base.SendAcknowledgedData([]byte{'\x78', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'})
}

func (f *FitbitBase) WaitForBeacon() error {
	for i := 0; i < 75; i++ {
		data, err := f.base.ReceiveMessage(4096)
		if err == nil && len(data) > 2 && data[2] == '\x4e' {
			return nil
		}
		log.Println("retry beacon", err)
	}
	return fmt.Errorf("Failed to see tracker beacon")
}

func (f *FitbitBase) RunOpcode(opcode, payload []byte) ([]byte, error) {
	for i := 0; i < 4; i++ {
		err := f.SendTrackerPacket(opcode)
		if err != nil {
			log.Println(err)
			continue
		}
		data, err := f.base.ReceiveAcknowledgedReply()
		if err != nil {
			log.Println(err)
			continue
		}
		if data[0] != byte(f.currentPacketId) {
			log.Printf("Tracker Packet IDs don't match! %v %v \n", f.currentPacketId, data[0])
			continue
		}
		if data[1] == '\x42' {
			log.Println("Start DataBank")
			return f.GetDataBank()
		}
		if data[1] == '\x61' {
			if len(payload) > 0 {
				if err := f.SendTrackerPayload(payload); err != nil {
					log.Println("payload failed", err)
					break
				}
				data, err := f.base.ReceiveAcknowledgedReply()
				log.Println("payload", err)
				if len(data) > 0 {
					return data[1:], err
				}
				return data, err
			}
		}
		if data[1] == '\x41' {
			return data[1:], nil
		}
	}
	return []byte{}, fmt.Errorf("failed to run opcode")
}
func (f *FitbitBase) GenPacketId() int {
	f.currentPacketId = '\x38' + f.getTrackerPacketCount()
	return f.currentPacketId
}

func (f *FitbitBase) getTrackerPacketCount() int {
	f.trackerPacketCount++
	if f.trackerPacketCount > 7 {
		f.trackerPacketCount = 0
	}
	return f.trackerPacketCount
}
func (f *FitbitBase) SendTrackerPayload(payload []byte) error {
	p := []byte{'\x00', byte(f.GenPacketId()), '\x80', byte(len(payload)), '\x00', '\x00', '\x00', '\x00'}
	p = append(p, XorSum(payload))
	prefix := []byte{'\x20', '\x40', '\x60'}
	index := 0
	for i := 0; i < len(payload); i += 8 {
		currentPrefix := prefix[index%3]
		var plist []byte
		if (i + 8) > len(payload) {
			plist = append(plist, byte((int(currentPrefix)+'\x80'))|f.base.channel)
			plist = append(plist, payload[i:]...)
		} else {
			plist = append(plist, currentPrefix|f.base.channel)
			plist = append(plist, payload[i:i+8]...)
		}
		for {
			if len(plist) >= 9 {
				break
			}
			plist = append(plist, '\x00')
		}
		p = append(p, plist...)
		index++
	}
	return f.base.SendBurstData(p, 10*time.Millisecond)
}
func (f *FitbitBase) GetDataBank() ([]byte, error) {
	var data []byte
	cmd := byte('\x70')
	for i := 0; i < 2000; i++ {
		bank, err := f.CheckTrackerDataBank(f.currentBankId, cmd)
		if err != nil {
			log.Println(err)
			continue
		}
		f.currentBankId++
		cmd = '\x60' // Send 0x60 on subsequent bursts
		if len(bank) == 0 {
			return data, nil
		}
		data = append(data, bank...)
	}
	return data, fmt.Errorf("Cannot complete data bank")
}

func (f *FitbitBase) CheckTrackerDataBank(index int, cmd byte) ([]byte, error) {
	err := f.SendTrackerPacket([]byte{cmd, '\x00', '\x02', byte(index), '\x00', '\x00', '\x00'})
	var data []byte
	if err == nil {
		data, err = f.GetTrackerBurst()
	}
	return data, err
}
func (f *FitbitBase) SendTrackerPacket(packet []byte) error {
	p := append([]byte{byte(f.GenPacketId())}, packet...)
	return f.base.SendAcknowledgedData(p)
}
func (f *FitbitBase) GetTrackerBurst() ([]byte, error) {
	var data []byte
	d, err := f.base.CheckBurstResponse()
	if err != nil || len(d) < 1 {
		return data, err
	}
	if d[1] != '\x81' {
		return d, fmt.Errorf("Response received is not tracker burst! Got")
	}
	size := (int(d[3]) << 8) | int(d[2])
	if size == 0 {
		return data, err
	}
	datalen := 8 + size
	if datalen < len(d) {
		data = d[8:datalen]
	} else {
		data = d[8:]
	}
	return data, err
}

func (f *FitbitBase) RunDataBankOpcode(index byte) ([]byte, error) {
	return f.RunOpcode([]byte{'\x22', index, '\x00', '\x00', '\x00', '\x00', '\x00'}, []byte{})
}

func (f *FitbitBase) GetTrackerInfo() ([]byte, error) {
	return f.RunOpcode([]byte{'\x24', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'}, []byte{})
}
<commit_msg>bugfix: correct burst data send<commit_after>package main

import (
	"fmt"
	"github.com/kylelemons/gousb/usb"
	"log"
	"math/rand"
	"time"
)

type FitbitBase struct {
	ctx                *usb.Context
	dev                *usb.Device
	base               *ANT
	trackerPacketCount int
	currentPacketId    int
	currentBankId      int
}

func (f *FitbitBase) Open() error {
	f.ctx = usb.NewContext()
	devs, err := f.ctx.ListDevices(func(desc *usb.Descriptor) bool {
		if desc.Vendor.String() == "10c4" && desc.Product.String() == "84c4" {

			return true
		}
		return false
	})
	if err == nil {
		if len(devs) == 0 {
			err = fmt.Errorf("no devices found")
		} else {
			f.dev = devs[0]
		}
		for i := 1; i < len(devs); i++ {
			devs[i].Close()
		}
	}
	return err
}

func (f *FitbitBase) SettingUp() error {
	f.dev.Control(64, 0, 65535, 0, []byte{})
	f.dev.Control(64, 1, 8192, 0, []byte{})
	f.dev.Control(64, 0, 0, 0, []byte{})
	f.dev.Control(64, 0, 65535, 0, []byte{})
	f.dev.Control(64, 1, 8192, 0, []byte{})
	f.dev.Control(64, 1, 74, 0, []byte{})
	f.dev.Control(192, 255, 14091, 0, []byte{'\x01'})
	f.dev.Control(64, 3, 2048, 0, []byte{})
	f.dev.Control(64, 19, 0, 0, []byte{
		'\x08', '\x00', '\x00', '\x00',
		'\x40', '\x00', '\x00', '\x00',
		'\x00', '\x00', '\x00', '\x00',
		'\x00', '\x00', '\x00', '\x00'})
	f.dev.Control(64, 18, 12, 0, []byte{})
	err := f.dev.SetConfig(1)
	if err != nil {
		log.Println("setconfig error")
		return err
	}
	err = f.dev.Reset()
	if err != nil {
		log.Println("reset error")
		return err
	}
	err = f.dev.SetConfig(1)
	if err != nil {
		log.Println("setconfig2 error")
		return err
	}
	f.base = &ANT{
		channel: '\x00',
	}
	f.base.reader, err = f.dev.OpenEndpoint(uint8(1), uint8(0), uint8(0), '\x81')
	if err != nil {
		return err
	}
	f.base.writer, err = f.dev.OpenEndpoint(uint8(1), uint8(0), uint8(0), '\x01')
	return err
}

func (f *FitbitBase) Close() {
	f.dev.Close()
	f.ctx.Close()
}

// data transport
func (f *FitbitBase) InitDeviceChannel(channel []byte) error {
	err := f.base.Reset()
	if err != nil {
		log.Println("fitbit reset failed")
		return err
	}
	err = f.base.SetNetworkKey(0, []byte{0, 0, 0, 0, 0, 0, 0, 0})
	if err != nil {
		log.Println("fitbit set network key failed")
		return err
	}
	err = f.base.AssignChannel()
	if err != nil {
		log.Println("fitbit assign channel failed")
		return err
	}
	err = f.base.SetChannelPeriod([]byte{'\x00', '\x10'})
	if err != nil {
		log.Println("fitbit set channel period failed")
		return err
	}
	err = f.base.SetChannelFrequency('\x02')
	if err != nil {
		log.Println("fitbit set channel frequency failed")
		return err
	}
	err = f.base.SetTransmitPower('\x03')
	if err != nil {
		log.Println("fitbit set transmit power failed")
		return err
	}
	err = f.base.SetSearchTimeout('\xff')
	if err != nil {
		log.Println("fitbit set search timeout failed")
		return err
	}
	err = f.base.SetChannelId(channel)
	if err != nil {
		log.Println("fitbit set channel id failed")
		return err
	}
	err = f.base.OpenChannel()
	if err != nil {
		log.Println("fitbit open channel failed")
		return err
	}
	return err
}

func (f *FitbitBase) InitTrackerForTransfer() error {
	err := f.InitDeviceChannel([]byte{'\xff', '\xff', '\x01', '\x01'})
	if err != nil {
		log.Println("fitbit init failed")
		return err
	}
	err = f.WaitForBeacon()
	if err != nil {
		return err
	}
	err = f.ResetTracker()
	if err != nil {
		log.Println("fitbit reset tracker failed")
		return err
	}
	cid := []byte{byte(rand.Intn(254)), byte(rand.Intn(254))}
	err = f.base.SendAcknowledgedData(append(append([]byte{'\x78', '\x02'}, cid...), []byte{'\x00', '\x00', '\x00', '\x00'}...))
	if err != nil {
		log.Println("fitbit ack ack to tracker failed")
		return err
	}
	err = f.base.CloseChannel()
	if err != nil {
		log.Println("fitbit close channel failed")
		return err
	}
	err = f.InitDeviceChannel(append(cid, []byte{0x01, 0x01}...))
	if err != nil {
		log.Println("fitbit reinit channel failed")
		return err
	}
	err = f.WaitForBeacon()
	if err != nil {
		return err
	}
	err = f.PingTracker()
	if err != nil {
		log.Println("fitbit ping tracker failed")
	}
	return err
}

func (f *FitbitBase) ResetTracker() error {
	return f.base.SendAcknowledgedData([]byte{'\x78', '\x01', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'})
}

func (f *FitbitBase) CommandSleep() error {
	return f.base.SendAcknowledgedData([]byte{'\x7f', '\x03', '\x00', '\x00', '\x00', '\x00', '\x00', '\x3c'})
}

func (f *FitbitBase) PingTracker() error {
	return f.base.SendAcknowledgedData([]byte{'\x78', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'})
}

func (f *FitbitBase) WaitForBeacon() error {
	for i := 0; i < 75; i++ {
		data, err := f.base.ReceiveMessage(4096)
		if err == nil && len(data) > 2 && data[2] == '\x4e' {
			return nil
		}
		log.Println("retry beacon", err)
	}
	return fmt.Errorf("Failed to see tracker beacon")
}

func (f *FitbitBase) RunOpcode(opcode, payload []byte) ([]byte, error) {
	for i := 0; i < 4; i++ {
		err := f.SendTrackerPacket(opcode)
		if err != nil {
			log.Println(err)
			continue
		}
		data, err := f.base.ReceiveAcknowledgedReply()
		if err != nil {
			log.Println(err)
			continue
		}
		if data[0] != byte(f.currentPacketId) {
			log.Printf("Tracker Packet IDs don't match! %v %v \n", f.currentPacketId, data[0])
			continue
		}
		if data[1] == '\x42' {
			log.Println("Start DataBank")
			return f.GetDataBank()
		}
		if data[1] == '\x61' {
			if len(payload) > 0 {
				if err := f.SendTrackerPayload(payload); err != nil {
					log.Println("payload failed", err)
					break
				}
				data, err := f.base.ReceiveAcknowledgedReply()
				log.Println("payload", err)
				if len(data) > 0 {
					return data[1:], err
				}
				return data, err
			}
		}
		if data[1] == '\x41' {
			return data[1:], nil
		}
	}
	return []byte{}, fmt.Errorf("failed to run opcode")
}
func (f *FitbitBase) GenPacketId() int {
	f.currentPacketId = '\x38' + f.getTrackerPacketCount()
	return f.currentPacketId
}

func (f *FitbitBase) getTrackerPacketCount() int {
	f.trackerPacketCount++
	if f.trackerPacketCount > 7 {
		f.trackerPacketCount = 0
	}
	return f.trackerPacketCount
}
func (f *FitbitBase) SendTrackerPayload(payload []byte) error {
	p := []byte{'\x00', byte(f.GenPacketId()), '\x80', byte(len(payload)), '\x00', '\x00', '\x00', '\x00'}
	p = append(p, XorSum(payload))
	prefix := []byte{'\x20', '\x40', '\x60'}
	index := 0
	for i := 0; i < len(payload); i += 8 {
		currentPrefix := prefix[index%3]
		var plist []byte
		if (i + 8) < len(payload) {
			plist = append(plist, currentPrefix|f.base.channel)
			plist = append(plist, payload[i:i+8]...)
		} else {
			plist = append(plist, byte((int(currentPrefix)+'\x80'))|f.base.channel)
			plist = append(plist, payload[i:]...)
		}
		for {
			if len(plist) >= 9 {
				break
			}
			plist = append(plist, '\x00')
		}
		p = append(p, plist...)
		index++
	}
	return f.base.SendBurstData(p, 10*time.Millisecond)
}
func (f *FitbitBase) GetDataBank() ([]byte, error) {
	var data []byte
	cmd := byte('\x70')
	for i := 0; i < 2000; i++ {
		bank, err := f.CheckTrackerDataBank(f.currentBankId, cmd)
		if err != nil {
			log.Println(err)
			continue
		}
		f.currentBankId++
		cmd = '\x60' // Send 0x60 on subsequent bursts
		if len(bank) == 0 {
			return data, nil
		}
		data = append(data, bank...)
	}
	return data, fmt.Errorf("Cannot complete data bank")
}

func (f *FitbitBase) CheckTrackerDataBank(index int, cmd byte) ([]byte, error) {
	err := f.SendTrackerPacket([]byte{cmd, '\x00', '\x02', byte(index), '\x00', '\x00', '\x00'})
	var data []byte
	if err == nil {
		data, err = f.GetTrackerBurst()
	}
	return data, err
}
func (f *FitbitBase) SendTrackerPacket(packet []byte) error {
	p := append([]byte{byte(f.GenPacketId())}, packet...)
	return f.base.SendAcknowledgedData(p)
}
func (f *FitbitBase) GetTrackerBurst() ([]byte, error) {
	var data []byte
	d, err := f.base.CheckBurstResponse()
	if err != nil || len(d) < 1 {
		return data, err
	}
	if d[1] != '\x81' {
		return d, fmt.Errorf("Response received is not tracker burst! Got")
	}
	size := (int(d[3]) << 8) | int(d[2])
	if size == 0 {
		return data, err
	}
	datalen := 8 + size
	if datalen < len(d) {
		data = d[8:datalen]
	} else {
		data = d[8:]
	}
	return data, err
}

func (f *FitbitBase) RunDataBankOpcode(index byte) ([]byte, error) {
	return f.RunOpcode([]byte{'\x22', index, '\x00', '\x00', '\x00', '\x00', '\x00'}, []byte{})
}

func (f *FitbitBase) GetTrackerInfo() ([]byte, error) {
	return f.RunOpcode([]byte{'\x24', '\x00', '\x00', '\x00', '\x00', '\x00', '\x00'}, []byte{})
}
<|endoftext|>


<commit_before>package vmware

import (
	"fmt"
	"runtime"
)

// A driver is able to talk to VMware, control virtual machines, etc.
type Driver interface {
	// CompactDisk compacts a virtual disk.
	CompactDisk(string) error

	// CreateDisk creates a virtual disk with the given size.
	CreateDisk(string, string) error

	// Checks if the VMX file at the given path is running.
	IsRunning(string) (bool, error)

	// Start starts a VM specified by the path to the VMX given.
	Start(string, bool) error

	// Stop stops a VM specified by the path to the VMX given.
	Stop(string) error

	// Get the path to the VMware ISO for the given flavor.
	ToolsIsoPath(string) string

	// Get the path to the DHCP leases file for the given device.
	DhcpLeasesPath(string) string

	// Verify checks to make sure that this driver should function
	// properly. This should check that all the files it will use
	// appear to exist and so on. If everything is okay, this doesn't
	// return an error. Otherwise, this returns an error.
	Verify() error
}

// NewDriver returns a new driver implementation for this operating
// system, or an error if the driver couldn't be initialized.
func NewDriver() (Driver, error) {
	var driver Driver

	switch runtime.GOOS {
	case "darwin":
		driver = &Fusion5Driver{
			AppPath: "/Applications/VMware Fusion.app",
		}
	case "linux":
	case "windows":
		driver = &Workstation9Driver{}
	default:
		return nil, fmt.Errorf("can't find driver for OS: %s", runtime.GOOS)
	}

	if err := driver.Verify(); err != nil {
		return nil, err
	}

	return driver, nil
}
<commit_msg>added missing fallthru in case statement for linux option, fixes #221, #222<commit_after>package vmware

import (
	"fmt"
	"runtime"
)

// A driver is able to talk to VMware, control virtual machines, etc.
type Driver interface {
	// CompactDisk compacts a virtual disk.
	CompactDisk(string) error

	// CreateDisk creates a virtual disk with the given size.
	CreateDisk(string, string) error

	// Checks if the VMX file at the given path is running.
	IsRunning(string) (bool, error)

	// Start starts a VM specified by the path to the VMX given.
	Start(string, bool) error

	// Stop stops a VM specified by the path to the VMX given.
	Stop(string) error

	// Get the path to the VMware ISO for the given flavor.
	ToolsIsoPath(string) string

	// Get the path to the DHCP leases file for the given device.
	DhcpLeasesPath(string) string

	// Verify checks to make sure that this driver should function
	// properly. This should check that all the files it will use
	// appear to exist and so on. If everything is okay, this doesn't
	// return an error. Otherwise, this returns an error.
	Verify() error
}

// NewDriver returns a new driver implementation for this operating
// system, or an error if the driver couldn't be initialized.
func NewDriver() (Driver, error) {
	var driver Driver

	switch runtime.GOOS {
	case "darwin":
		driver = &Fusion5Driver{
			AppPath: "/Applications/VMware Fusion.app",
		}
	case "linux":
		fallthrough
	case "windows":
		driver = &Workstation9Driver{}
	default:
		return nil, fmt.Errorf("can't find driver for OS: %s", runtime.GOOS)
	}

	if err := driver.Verify(); err != nil {
		return nil, err
	}

	return driver, nil
}
<|endoftext|>


<commit_before>package navigate

type NavigateType string

const (
	Left   NavigateType = `left`
	Top    NavigateType = `top`
	Right  NavigateType = `right`
	Bottom NavigateType = `bottom`
)

func NewCollection() *Collection {
	return &Collection{
		Backend:  &Navigates{},
		Frontend: &Navigates{},
	}
}

type Collection struct {
	Backend  *Navigates
	Frontend *Navigates
}

type Navigates map[NavigateType]*List

func (n *Navigates) Add(typ NavigateType, nav *List) {
	(*n)[typ] = nav
}

func (n *Navigates) AddItems(typ NavigateType, index int, items ...*Item) {
	nav := n.Get(typ)
	if nav != nil {
		nav.Add(index, items...)
	}
}

func (n *Navigates) AddTopItems(index int, items ...*Item) {
	n.AddItems(Top, index, items...)
}

func (n *Navigates) AddLeftItems(index int, items ...*Item) {
	n.AddItems(Left, index, items...)
}

func (n *Navigates) AddRightItems(index int, items ...*Item) {
	n.AddItems(Right, index, items...)
}

func (n *Navigates) AddBottomItems(index int, items ...*Item) {
	n.AddItems(Bottom, index, items...)
}

func (n *Navigates) Get(typ NavigateType) (nav *List) {
	nav = (*n)[typ]
	return
}

func (n *Navigates) GetTop() *List {
	return n.Get(Top)
}

func (n *Navigates) GetLeft() *List {
	return n.Get(Left)
}

func (n *Navigates) GetRight() *List {
	return n.Get(Right)
}

func (n *Navigates) GetBottom() *List {
	return n.Get(Bottom)
}

func (n *Navigates) Remove(typ NavigateType) bool {
	_, ok := (*n)[typ]
	if ok {
		delete(*n, typ)
	}
	return ok
}
<commit_msg>update<commit_after>package navigate

type NavigateType string

const (
	Left   NavigateType = `left`
	Top    NavigateType = `top`
	Right  NavigateType = `right`
	Bottom NavigateType = `bottom`
)

func NewCollection() *Collection {
	return &Collection{
		Backend:  &Navigates{},
		Frontend: &Navigates{},
	}
}

type Collection struct {
	Backend  *Navigates
	Frontend *Navigates
}

type Navigates map[NavigateType]*List

func (n *Navigates) Add(typ NavigateType, nav *List) {
	(*n)[typ] = nav
}

func (n *Navigates) AddItems(typ NavigateType, index int, items ...*Item) {
	nav := n.Get(typ)
	if nav == nil {
		nav = &List{}
		(*n)[typ] = nav
	}
	nav.Add(index, items...)
}

func (n *Navigates) AddTopItems(index int, items ...*Item) {
	n.AddItems(Top, index, items...)
}

func (n *Navigates) AddLeftItems(index int, items ...*Item) {
	n.AddItems(Left, index, items...)
}

func (n *Navigates) AddRightItems(index int, items ...*Item) {
	n.AddItems(Right, index, items...)
}

func (n *Navigates) AddBottomItems(index int, items ...*Item) {
	n.AddItems(Bottom, index, items...)
}

func (n *Navigates) Get(typ NavigateType) (nav *List) {
	nav = (*n)[typ]
	return
}

func (n *Navigates) GetTop() *List {
	return n.Get(Top)
}

func (n *Navigates) GetLeft() *List {
	return n.Get(Left)
}

func (n *Navigates) GetRight() *List {
	return n.Get(Right)
}

func (n *Navigates) GetBottom() *List {
	return n.Get(Bottom)
}

func (n *Navigates) Remove(typ NavigateType) bool {
	_, ok := (*n)[typ]
	if ok {
		delete(*n, typ)
	}
	return ok
}
<|endoftext|>