<commit_before>/*
Copyright 2017 Google Inc.

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 messager

import (
	"sort"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// TabletService defines the functions of TabletServer
// that the messager needs for callback.
type TabletService interface {
	CheckMySQL()
	PostponeMessages(ctx context.Context, target *querypb.Target, name string, ids []string) (count int64, err error)
	PurgeMessages(ctx context.Context, target *querypb.Target, name string, timeCutoff int64) (count int64, err error)
}

// Engine is the engine for handling messages.
type Engine struct {
	dbconfigs *dbconfigs.DBConfigs

	mu       sync.Mutex
	isOpen   bool
	managers map[string]*messageManager

	tsv          TabletService
	se           *schema.Engine
	conns        *connpool.Pool
	postponeSema *sync2.Semaphore
}

// NewEngine creates a new Engine.
func NewEngine(tsv TabletService, se *schema.Engine, config tabletenv.TabletConfig) *Engine {
	return &Engine{
		tsv: tsv,
		se:  se,
		conns: connpool.New(
			config.PoolNamePrefix+"MessagerPool",
			config.MessagePoolSize,
			config.MessagePoolPrefillParallelism,
			time.Duration(config.IdleTimeout*1e9),
			tsv,
		),
		postponeSema: sync2.NewSemaphore(config.MessagePostponeCap, 0),
		managers:     make(map[string]*messageManager),
	}
}

// InitDBConfig must be called before Open.
func (me *Engine) InitDBConfig(dbcfgs *dbconfigs.DBConfigs) {
	me.dbconfigs = dbcfgs
}

// Open starts the Engine service.
func (me *Engine) Open() error {
	if me.isOpen {
		return nil
	}
	me.conns.Open(me.dbconfigs.AppWithDB(), me.dbconfigs.DbaWithDB(), me.dbconfigs.AppDebugWithDB())
	me.se.RegisterNotifier("messages", me.schemaChanged)
	me.isOpen = true
	return nil
}

// Close closes the Engine service.
func (me *Engine) Close() {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return
	}
	me.isOpen = false
	me.se.UnregisterNotifier("messages")
	for _, mm := range me.managers {
		mm.Close()
	}
	me.managers = make(map[string]*messageManager)
	me.conns.Close()
}

// Subscribe subscribes to messages from the requested table.
// The function returns a done channel that will be closed when
// the subscription ends, which can be initiated by the send function
// returning io.EOF. The engine can also end a subscription which is
// usually triggered by Close. It's the responsibility of the send
// function to promptly return if the done channel is closed. Otherwise,
// the engine's Close function will hang indefinitely.
func (me *Engine) Subscribe(ctx context.Context, name string, send func(*sqltypes.Result) error) (done <-chan struct{}, err error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "messager engine is closed, probably because this is not a master any more")
	}
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found", name)
	}
	return mm.Subscribe(ctx, send), nil
}

// LockDB obtains db locks for all messages that need to
// be updated and returns the counterpart unlock function.
func (me *Engine) LockDB(newMessages map[string][]*MessageRow, changedMessages map[string][]string) func() {
	// Short-circuit to avoid taking any locks if there's nothing to do.
	if len(newMessages) == 0 && len(changedMessages) == 0 {
		return func() {}
	}

	// Build the set of affected messages tables.
	combined := make(map[string]struct{})
	for name := range newMessages {
		combined[name] = struct{}{}
	}
	for name := range changedMessages {
		combined[name] = struct{}{}
	}

	// Build the list of manager objects (one per table).
	var mms []*messageManager
	// Don't do DBLock while holding lock on mu.
	// It causes deadlocks.
	func() {
		me.mu.Lock()
		defer me.mu.Unlock()
		for name := range combined {
			if mm := me.managers[name]; mm != nil {
				mms = append(mms, mm)
			}
		}
	}()
	if len(mms) > 1 {
		// Always use the same order in which manager objects are locked to avoid deadlocks.
		// The previous order in "mms" is not guaranteed for multiple reasons:
		// - We use a Go map above which does not guarantee an iteration order.
		// - Transactions may not always use the same order when writing to multiple
		//   messages tables.
		sort.Slice(mms, func(i, j int) bool { return mms[i].name.String() < mms[j].name.String() })
	}

	// Lock each manager/messages table.
	for _, mm := range mms {
		mm.DBLock.Lock()
	}
	return func() {
		for _, mm := range mms {
			mm.DBLock.Unlock()
		}
	}
}

// UpdateCaches updates the caches for the committed changes.
func (me *Engine) UpdateCaches(newMessages map[string][]*MessageRow, changedMessages map[string][]string) {
	// Short-circuit to avoid taking any locks if there's nothing to do.
	if len(newMessages) == 0 && len(changedMessages) == 0 {
		return
	}

	me.mu.Lock()
	defer me.mu.Unlock()
	now := time.Now().UnixNano()
	for name, mrs := range newMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		MessageStats.Add([]string{name, "Queued"}, int64(len(mrs)))
		for _, mr := range mrs {
			if mr.TimeNext > now {
				// We don't handle future messages yet.
				continue
			}
			mm.Add(mr)
		}
	}
	for name, ids := range changedMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.cache.Discard(ids)
	}
}

// GenerateLoadMessagesQuery returns the ParsedQuery for loading messages by pk.
// The results of the query can be used in a BuildMessageRow call.
func (me *Engine) GenerateLoadMessagesQuery(name string) (*sqlparser.ParsedQuery, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	return mm.loadMessagesQuery, nil
}

// GenerateAckQuery returns the query and bind vars for acking a message.
func (me *Engine) GenerateAckQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GenerateAckQuery(ids)
	return query, bv, nil
}

// GeneratePostponeQuery returns the query and bind vars for postponing a message.
func (me *Engine) GeneratePostponeQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePostponeQuery(ids)
	return query, bv, nil
}

// GeneratePurgeQuery returns the query and bind vars for purging messages.
func (me *Engine) GeneratePurgeQuery(name string, timeCutoff int64) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePurgeQuery(timeCutoff)
	return query, bv, nil
}

func (me *Engine) schemaChanged(tables map[string]*schema.Table, created, altered, dropped []string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	for _, name := range created {
		t := tables[name]
		if t.Type != schema.Message {
			continue
		}
		if me.managers[name] != nil {
			tabletenv.InternalErrors.Add("Messages", 1)
			log.Errorf("Newly created table alread exists in messages: %s", name)
			continue
		}
		mm := newMessageManager(me.tsv, t, me.conns, me.postponeSema)
		me.managers[name] = mm
		mm.Open()
	}

	// TODO(sougou): Update altered tables.

	for _, name := range dropped {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.Close()
		delete(me.managers, name)
	}
}
<commit_msg>cleanup: correct logging message<commit_after>/*
Copyright 2017 Google Inc.

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 messager

import (
	"sort"
	"sync"
	"time"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// TabletService defines the functions of TabletServer
// that the messager needs for callback.
type TabletService interface {
	CheckMySQL()
	PostponeMessages(ctx context.Context, target *querypb.Target, name string, ids []string) (count int64, err error)
	PurgeMessages(ctx context.Context, target *querypb.Target, name string, timeCutoff int64) (count int64, err error)
}

// Engine is the engine for handling messages.
type Engine struct {
	dbconfigs *dbconfigs.DBConfigs

	mu       sync.Mutex
	isOpen   bool
	managers map[string]*messageManager

	tsv          TabletService
	se           *schema.Engine
	conns        *connpool.Pool
	postponeSema *sync2.Semaphore
}

// NewEngine creates a new Engine.
func NewEngine(tsv TabletService, se *schema.Engine, config tabletenv.TabletConfig) *Engine {
	return &Engine{
		tsv: tsv,
		se:  se,
		conns: connpool.New(
			config.PoolNamePrefix+"MessagerPool",
			config.MessagePoolSize,
			config.MessagePoolPrefillParallelism,
			time.Duration(config.IdleTimeout*1e9),
			tsv,
		),
		postponeSema: sync2.NewSemaphore(config.MessagePostponeCap, 0),
		managers:     make(map[string]*messageManager),
	}
}

// InitDBConfig must be called before Open.
func (me *Engine) InitDBConfig(dbcfgs *dbconfigs.DBConfigs) {
	me.dbconfigs = dbcfgs
}

// Open starts the Engine service.
func (me *Engine) Open() error {
	if me.isOpen {
		return nil
	}
	me.conns.Open(me.dbconfigs.AppWithDB(), me.dbconfigs.DbaWithDB(), me.dbconfigs.AppDebugWithDB())
	me.se.RegisterNotifier("messages", me.schemaChanged)
	me.isOpen = true
	return nil
}

// Close closes the Engine service.
func (me *Engine) Close() {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return
	}
	me.isOpen = false
	me.se.UnregisterNotifier("messages")
	for _, mm := range me.managers {
		mm.Close()
	}
	me.managers = make(map[string]*messageManager)
	me.conns.Close()
}

// Subscribe subscribes to messages from the requested table.
// The function returns a done channel that will be closed when
// the subscription ends, which can be initiated by the send function
// returning io.EOF. The engine can also end a subscription which is
// usually triggered by Close. It's the responsibility of the send
// function to promptly return if the done channel is closed. Otherwise,
// the engine's Close function will hang indefinitely.
func (me *Engine) Subscribe(ctx context.Context, name string, send func(*sqltypes.Result) error) (done <-chan struct{}, err error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "messager engine is closed, probably because this is not a master any more")
	}
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found", name)
	}
	return mm.Subscribe(ctx, send), nil
}

// LockDB obtains db locks for all messages that need to
// be updated and returns the counterpart unlock function.
func (me *Engine) LockDB(newMessages map[string][]*MessageRow, changedMessages map[string][]string) func() {
	// Short-circuit to avoid taking any locks if there's nothing to do.
	if len(newMessages) == 0 && len(changedMessages) == 0 {
		return func() {}
	}

	// Build the set of affected messages tables.
	combined := make(map[string]struct{})
	for name := range newMessages {
		combined[name] = struct{}{}
	}
	for name := range changedMessages {
		combined[name] = struct{}{}
	}

	// Build the list of manager objects (one per table).
	var mms []*messageManager
	// Don't do DBLock while holding lock on mu.
	// It causes deadlocks.
	func() {
		me.mu.Lock()
		defer me.mu.Unlock()
		for name := range combined {
			if mm := me.managers[name]; mm != nil {
				mms = append(mms, mm)
			}
		}
	}()
	if len(mms) > 1 {
		// Always use the same order in which manager objects are locked to avoid deadlocks.
		// The previous order in "mms" is not guaranteed for multiple reasons:
		// - We use a Go map above which does not guarantee an iteration order.
		// - Transactions may not always use the same order when writing to multiple
		//   messages tables.
		sort.Slice(mms, func(i, j int) bool { return mms[i].name.String() < mms[j].name.String() })
	}

	// Lock each manager/messages table.
	for _, mm := range mms {
		mm.DBLock.Lock()
	}
	return func() {
		for _, mm := range mms {
			mm.DBLock.Unlock()
		}
	}
}

// UpdateCaches updates the caches for the committed changes.
func (me *Engine) UpdateCaches(newMessages map[string][]*MessageRow, changedMessages map[string][]string) {
	// Short-circuit to avoid taking any locks if there's nothing to do.
	if len(newMessages) == 0 && len(changedMessages) == 0 {
		return
	}

	me.mu.Lock()
	defer me.mu.Unlock()
	now := time.Now().UnixNano()
	for name, mrs := range newMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		MessageStats.Add([]string{name, "Queued"}, int64(len(mrs)))
		for _, mr := range mrs {
			if mr.TimeNext > now {
				// We don't handle future messages yet.
				continue
			}
			mm.Add(mr)
		}
	}
	for name, ids := range changedMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.cache.Discard(ids)
	}
}

// GenerateLoadMessagesQuery returns the ParsedQuery for loading messages by pk.
// The results of the query can be used in a BuildMessageRow call.
func (me *Engine) GenerateLoadMessagesQuery(name string) (*sqlparser.ParsedQuery, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	return mm.loadMessagesQuery, nil
}

// GenerateAckQuery returns the query and bind vars for acking a message.
func (me *Engine) GenerateAckQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GenerateAckQuery(ids)
	return query, bv, nil
}

// GeneratePostponeQuery returns the query and bind vars for postponing a message.
func (me *Engine) GeneratePostponeQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePostponeQuery(ids)
	return query, bv, nil
}

// GeneratePurgeQuery returns the query and bind vars for purging messages.
func (me *Engine) GeneratePurgeQuery(name string, timeCutoff int64) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePurgeQuery(timeCutoff)
	return query, bv, nil
}

func (me *Engine) schemaChanged(tables map[string]*schema.Table, created, altered, dropped []string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	for _, name := range created {
		t := tables[name]
		if t.Type != schema.Message {
			continue
		}
		if me.managers[name] != nil {
			tabletenv.InternalErrors.Add("Messages", 1)
			log.Errorf("Newly created table already exists in messages: %s", name)
			continue
		}
		mm := newMessageManager(me.tsv, t, me.conns, me.postponeSema)
		me.managers[name] = mm
		mm.Open()
	}

	// TODO(sougou): Update altered tables.

	for _, name := range dropped {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.Close()
		delete(me.managers, name)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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 messager

import (
	"sync"
	"time"

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

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// TabletService defines the functions of TabletServer
// that the messager needs for callback.
type TabletService interface {
	CheckMySQL()
	PostponeMessages(ctx context.Context, target *querypb.Target, name string, ids []string) (count int64, err error)
	PurgeMessages(ctx context.Context, target *querypb.Target, name string, timeCutoff int64) (count int64, err error)
}

// Engine is the engine for handling messages.
type Engine struct {
	dbconfigs dbconfigs.DBConfigs

	mu       sync.Mutex
	isOpen   bool
	managers map[string]*messageManager

	tsv          TabletService
	se           *schema.Engine
	conns        *connpool.Pool
	postponeSema *sync2.Semaphore
}

// NewEngine creates a new Engine.
func NewEngine(tsv TabletService, se *schema.Engine, config tabletenv.TabletConfig) *Engine {
	return &Engine{
		tsv: tsv,
		se:  se,
		conns: connpool.New(
			config.PoolNamePrefix+"MessagerPool",
			config.MessagePoolSize,
			time.Duration(config.IdleTimeout*1e9),
			tsv,
		),
		postponeSema: sync2.NewSemaphore(config.MessagePostponeCap, 0),
		managers:     make(map[string]*messageManager),
	}
}

// InitDBConfig must be called before Open.
func (me *Engine) InitDBConfig(dbcfgs dbconfigs.DBConfigs) {
	me.dbconfigs = dbcfgs
}

// Open starts the Engine service.
func (me *Engine) Open() error {
	if me.isOpen {
		return nil
	}
	me.conns.Open(&me.dbconfigs.App, &me.dbconfigs.Dba, &me.dbconfigs.AppDebug)
	me.se.RegisterNotifier("messages", me.schemaChanged)
	me.isOpen = true
	return nil
}

// Close closes the Engine service.
func (me *Engine) Close() {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return
	}
	me.isOpen = false
	me.se.UnregisterNotifier("messages")
	for _, mm := range me.managers {
		mm.Close()
	}
	me.managers = make(map[string]*messageManager)
	me.conns.Close()
}

// Subscribe subscribes to messages from the requested table.
// The function returns a done channel that will be closed when
// the subscription ends, which can be initiated by the send function
// returning io.EOF. The engine can also end a subscription which is
// usually triggered by Close. It's the responsibility of the send
// function to promptly return if the done channel is closed. Otherwise,
// the engine's Close function will hang indefinitely.
func (me *Engine) Subscribe(ctx context.Context, name string, send func(*sqltypes.Result) error) (done <-chan struct{}, err error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "messager engine is closed, probably because this is not a master any more")
	}
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found", name)
	}
	return mm.Subscribe(ctx, send), nil
}

// LockDB obtains db locks for all messages that need to
// be updated and returns the counterpart unlock function.
func (me *Engine) LockDB(newMessages map[string][]*MessageRow, changedMessages map[string][]string) func() {
	combined := make(map[string]struct{})
	for name := range newMessages {
		combined[name] = struct{}{}
	}
	for name := range changedMessages {
		combined[name] = struct{}{}
	}
	var mms []*messageManager
	// Don't do DBLock while holding lock on mu.
	// It causes deadlocks.
	func() {
		me.mu.Lock()
		defer me.mu.Unlock()
		for name := range combined {
			if mm := me.managers[name]; mm != nil {
				mms = append(mms, mm)
			}
		}
	}()
	for _, mm := range mms {
		mm.DBLock.Lock()
	}
	return func() {
		for _, mm := range mms {
			mm.DBLock.Unlock()
		}
	}
}

// UpdateCaches updates the caches for the committed changes.
func (me *Engine) UpdateCaches(newMessages map[string][]*MessageRow, changedMessages map[string][]string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	now := time.Now().UnixNano()
	for name, mrs := range newMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		MessageStats.Add([]string{name, "Queued"}, int64(len(mrs)))
		for _, mr := range mrs {
			if mr.TimeNext > now {
				// We don't handle future messages yet.
				continue
			}
			mm.Add(mr)
		}
	}
	for name, ids := range changedMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.cache.Discard(ids)
	}
}

// GenerateLoadMessagesQuery returns the ParsedQuery for loading messages by pk.
// The results of the query can be used in a BuildMessageRow call.
func (me *Engine) GenerateLoadMessagesQuery(name string) (*sqlparser.ParsedQuery, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	return mm.loadMessagesQuery, nil
}

// GenerateAckQuery returns the query and bind vars for acking a message.
func (me *Engine) GenerateAckQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GenerateAckQuery(ids)
	return query, bv, nil
}

// GeneratePostponeQuery returns the query and bind vars for postponing a message.
func (me *Engine) GeneratePostponeQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePostponeQuery(ids)
	return query, bv, nil
}

// GeneratePurgeQuery returns the query and bind vars for purging messages.
func (me *Engine) GeneratePurgeQuery(name string, timeCutoff int64) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePurgeQuery(timeCutoff)
	return query, bv, nil
}

func (me *Engine) schemaChanged(tables map[string]*schema.Table, created, altered, dropped []string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	for _, name := range created {
		t := tables[name]
		if t.Type != schema.Message {
			continue
		}
		if me.managers[name] != nil {
			tabletenv.InternalErrors.Add("Messages", 1)
			log.Errorf("Newly created table alread exists in messages: %s", name)
			continue
		}
		mm := newMessageManager(me.tsv, t, me.conns, me.postponeSema)
		me.managers[name] = mm
		mm.Open()
	}

	// TODO(sougou): Update altered tables.

	for _, name := range dropped {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.Close()
		delete(me.managers, name)
	}
}
<commit_msg>messages: LockDB(): Adding comments and newlines to improve readability.<commit_after>/*
Copyright 2017 Google Inc.

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 messager

import (
	"sync"
	"time"

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

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/schema"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
)

// TabletService defines the functions of TabletServer
// that the messager needs for callback.
type TabletService interface {
	CheckMySQL()
	PostponeMessages(ctx context.Context, target *querypb.Target, name string, ids []string) (count int64, err error)
	PurgeMessages(ctx context.Context, target *querypb.Target, name string, timeCutoff int64) (count int64, err error)
}

// Engine is the engine for handling messages.
type Engine struct {
	dbconfigs dbconfigs.DBConfigs

	mu       sync.Mutex
	isOpen   bool
	managers map[string]*messageManager

	tsv          TabletService
	se           *schema.Engine
	conns        *connpool.Pool
	postponeSema *sync2.Semaphore
}

// NewEngine creates a new Engine.
func NewEngine(tsv TabletService, se *schema.Engine, config tabletenv.TabletConfig) *Engine {
	return &Engine{
		tsv: tsv,
		se:  se,
		conns: connpool.New(
			config.PoolNamePrefix+"MessagerPool",
			config.MessagePoolSize,
			time.Duration(config.IdleTimeout*1e9),
			tsv,
		),
		postponeSema: sync2.NewSemaphore(config.MessagePostponeCap, 0),
		managers:     make(map[string]*messageManager),
	}
}

// InitDBConfig must be called before Open.
func (me *Engine) InitDBConfig(dbcfgs dbconfigs.DBConfigs) {
	me.dbconfigs = dbcfgs
}

// Open starts the Engine service.
func (me *Engine) Open() error {
	if me.isOpen {
		return nil
	}
	me.conns.Open(&me.dbconfigs.App, &me.dbconfigs.Dba, &me.dbconfigs.AppDebug)
	me.se.RegisterNotifier("messages", me.schemaChanged)
	me.isOpen = true
	return nil
}

// Close closes the Engine service.
func (me *Engine) Close() {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return
	}
	me.isOpen = false
	me.se.UnregisterNotifier("messages")
	for _, mm := range me.managers {
		mm.Close()
	}
	me.managers = make(map[string]*messageManager)
	me.conns.Close()
}

// Subscribe subscribes to messages from the requested table.
// The function returns a done channel that will be closed when
// the subscription ends, which can be initiated by the send function
// returning io.EOF. The engine can also end a subscription which is
// usually triggered by Close. It's the responsibility of the send
// function to promptly return if the done channel is closed. Otherwise,
// the engine's Close function will hang indefinitely.
func (me *Engine) Subscribe(ctx context.Context, name string, send func(*sqltypes.Result) error) (done <-chan struct{}, err error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	if !me.isOpen {
		return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "messager engine is closed, probably because this is not a master any more")
	}
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found", name)
	}
	return mm.Subscribe(ctx, send), nil
}

// LockDB obtains db locks for all messages that need to
// be updated and returns the counterpart unlock function.
func (me *Engine) LockDB(newMessages map[string][]*MessageRow, changedMessages map[string][]string) func() {
	// Build the set of affected messages tables.
	combined := make(map[string]struct{})
	for name := range newMessages {
		combined[name] = struct{}{}
	}
	for name := range changedMessages {
		combined[name] = struct{}{}
	}

	// Build the list of manager objects (one per table).
	var mms []*messageManager
	// Don't do DBLock while holding lock on mu.
	// It causes deadlocks.
	func() {
		me.mu.Lock()
		defer me.mu.Unlock()
		for name := range combined {
			if mm := me.managers[name]; mm != nil {
				mms = append(mms, mm)
			}
		}
	}()

	// Lock each manager/messages table.
	for _, mm := range mms {
		mm.DBLock.Lock()
	}
	return func() {
		for _, mm := range mms {
			mm.DBLock.Unlock()
		}
	}
}

// UpdateCaches updates the caches for the committed changes.
func (me *Engine) UpdateCaches(newMessages map[string][]*MessageRow, changedMessages map[string][]string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	now := time.Now().UnixNano()
	for name, mrs := range newMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		MessageStats.Add([]string{name, "Queued"}, int64(len(mrs)))
		for _, mr := range mrs {
			if mr.TimeNext > now {
				// We don't handle future messages yet.
				continue
			}
			mm.Add(mr)
		}
	}
	for name, ids := range changedMessages {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.cache.Discard(ids)
	}
}

// GenerateLoadMessagesQuery returns the ParsedQuery for loading messages by pk.
// The results of the query can be used in a BuildMessageRow call.
func (me *Engine) GenerateLoadMessagesQuery(name string) (*sqlparser.ParsedQuery, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	return mm.loadMessagesQuery, nil
}

// GenerateAckQuery returns the query and bind vars for acking a message.
func (me *Engine) GenerateAckQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GenerateAckQuery(ids)
	return query, bv, nil
}

// GeneratePostponeQuery returns the query and bind vars for postponing a message.
func (me *Engine) GeneratePostponeQuery(name string, ids []string) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePostponeQuery(ids)
	return query, bv, nil
}

// GeneratePurgeQuery returns the query and bind vars for purging messages.
func (me *Engine) GeneratePurgeQuery(name string, timeCutoff int64) (string, map[string]*querypb.BindVariable, error) {
	me.mu.Lock()
	defer me.mu.Unlock()
	mm := me.managers[name]
	if mm == nil {
		return "", nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "message table %s not found in schema", name)
	}
	query, bv := mm.GeneratePurgeQuery(timeCutoff)
	return query, bv, nil
}

func (me *Engine) schemaChanged(tables map[string]*schema.Table, created, altered, dropped []string) {
	me.mu.Lock()
	defer me.mu.Unlock()
	for _, name := range created {
		t := tables[name]
		if t.Type != schema.Message {
			continue
		}
		if me.managers[name] != nil {
			tabletenv.InternalErrors.Add("Messages", 1)
			log.Errorf("Newly created table alread exists in messages: %s", name)
			continue
		}
		mm := newMessageManager(me.tsv, t, me.conns, me.postponeSema)
		me.managers[name] = mm
		mm.Open()
	}

	// TODO(sougou): Update altered tables.

	for _, name := range dropped {
		mm := me.managers[name]
		if mm == nil {
			continue
		}
		mm.Close()
		delete(me.managers, name)
	}
}
<|endoftext|>


<commit_before>package goarken

import (
	"encoding/json"
	"errors"
	"github.com/coreos/go-etcd/etcd"
	"github.com/golang/glog"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var (
	serviceRegexp = regexp.MustCompile("/services/(.*)(/.*)*")
)

type Location struct {
	Host string `json:"host"`
	Port int    `json:"port"`
}

func SetServicePrefix(servicePrefix string) {
	serviceRegexp = regexp.MustCompile(servicePrefix + "/(.*)(/.*)*")
}

func (s *Location) Equals(other *Location) bool {
	if s == nil && other == nil {
		return true
	}

	return s != nil && other != nil &&
		s.Host == other.Host &&
		s.Port == other.Port
}

func (s *Location) IsFullyDefined() bool {
	return s.Host != "" && s.Port != 0
}

func getEnvIndexForNode(node *etcd.Node) string {
	return strings.Split(serviceRegexp.FindStringSubmatch(node.Key)[1], "/")[1]
}

func getEnvForNode(node *etcd.Node) string {
	return strings.Split(serviceRegexp.FindStringSubmatch(node.Key)[1], "/")[0]
}

type Service struct {
	Index      string
	NodeKey    string
	Location   *Location
	Domain     string
	Name       string
	Status     *Status
	LastAccess *time.Time
}

func NewService(serviceNode *etcd.Node) (*Service, error) {

	serviceIndex := getEnvIndexForNode(serviceNode)

	if _, err := strconv.Atoi(serviceIndex); err != nil {
		// Don't handle node that are not integer (ie config node)
		return nil, errors.New("Not a service index node")
	}

	service := &Service{}
	service.Location = &Location{}
	service.Index = getEnvIndexForNode(serviceNode)
	service.Name = getEnvForNode(serviceNode)
	service.NodeKey = serviceNode.Key

	for _, node := range serviceNode.Nodes {
		switch node.Key {
		case service.NodeKey + "/location":
			location := &Location{}
			err := json.Unmarshal([]byte(node.Value), location)
			if err == nil {
				service.Location.Host = location.Host
				service.Location.Port = location.Port
			}

		case service.NodeKey + "/domain":
			service.Domain = node.Value
		case service.NodeKey + "/lastAccess":
			lastAccess := node.Value
			lastAccessTime, err := time.Parse(TIME_FORMAT, lastAccess)
			if err != nil {
				glog.Errorf("Error parsing last access date with service %s: %s", service.Name, err)
				break
			}
			service.LastAccess = &lastAccessTime

		case service.NodeKey + "/status":
			service.Status = NewStatus(service, node)
		}
	}
	return service, nil
}

func (s *Service) UnitName() string {
	return "nxio@" + strings.Split(s.Name, "_")[1] + ".service"
}

func (service *Service) Equals(other *Service) bool {
	if service == nil && other == nil {
		return true
	}

	return service != nil && other != nil &&
		service.Location.Equals(other.Location) &&
		service.Status.Equals(other.Status)
}

func (s *Service) StartedSince() *time.Time {
	if s == nil {
		return nil
	}

	if s.Status != nil &&
		s.Status.Current == STARTED_STATUS {
		return s.LastAccess
	} else {
		return nil
	}
}

func (s *Service) Start(client *etcd.Client) error {
	return s.fleetcmd("start", client)
}

func (s *Service) Stop(client *etcd.Client) error {
	return s.fleetcmd("destroy", client)
}

func (s *Service) Passivate(client *etcd.Client) error {
	err := s.fleetcmd("destroy", client)
	if err != nil {
		return err
	}

	statusKey := s.NodeKey + "/status"

	responseCurrent, error := client.Set(statusKey+"/current", PASSIVATED_STATUS, 0)
	if error != nil && responseCurrent == nil {
		glog.Errorf("Setting status current to 'passivated' has failed for Service "+s.Name+": %s", error)
	}

	response, error := client.Set(statusKey+"/expected", PASSIVATED_STATUS, 0)
	if error != nil && response == nil {
		glog.Errorf("Setting status expected to 'passivated' has failed for Service "+s.Name+": %s", error)
	}
	return nil
}

func (s *Service) fleetcmd(command string, client *etcd.Client) error {
	//TODO Use fleet's REST API
	etcdAddress := client.GetCluster()[0]
	cmd := exec.Command("/usr/bin/fleetctl", "--endpoint="+etcdAddress, command, s.UnitName())
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}
<commit_msg>NXIO-477 reporting config/gogeta key in goarken lib<commit_after>package goarken

import (
	"encoding/json"
	"errors"
	"github.com/coreos/go-etcd/etcd"
	"github.com/golang/glog"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

var (
	serviceRegexp = regexp.MustCompile("/services/(.*)(/.*)*")
)

type Location struct {
	Host string `json:"host"`
	Port int    `json:"port"`
}

func SetServicePrefix(servicePrefix string) {
	serviceRegexp = regexp.MustCompile(servicePrefix + "/(.*)(/.*)*")
}

func (s *Location) Equals(other *Location) bool {
	if s == nil && other == nil {
		return true
	}

	return s != nil && other != nil &&
		s.Host == other.Host &&
		s.Port == other.Port
}

func (s *Location) IsFullyDefined() bool {
	return s.Host != "" && s.Port != 0
}

func getEnvIndexForNode(node *etcd.Node) string {
	return strings.Split(serviceRegexp.FindStringSubmatch(node.Key)[1], "/")[1]
}

func getEnvForNode(node *etcd.Node) string {
	return strings.Split(serviceRegexp.FindStringSubmatch(node.Key)[1], "/")[0]
}

type ServiceConfig struct {
	Robots string `json:"robots"`
}

func (config *ServiceConfig) Equals(other *ServiceConfig) bool {
	if config == nil && other == nil {
		return true
	}

	return config != nil && other != nil &&
		config.Robots == other.Robots
}

type Service struct {
	Index      string
	NodeKey    string
	Location   *Location
	Domain     string
	Name       string
	Status     *Status
	LastAccess *time.Time
	Config     *ServiceConfig
}

func NewService(serviceNode *etcd.Node) (*Service, error) {

	serviceIndex := getEnvIndexForNode(serviceNode)

	if _, err := strconv.Atoi(serviceIndex); err != nil {
		// Don't handle node that are not integer (ie config node)
		return nil, errors.New("Not a service index node")
	}

	service := &Service{}
	service.Location = &Location{}
	service.Config = &ServiceConfig{Robots: ""}
	service.Index = getEnvIndexForNode(serviceNode)
	service.Name = getEnvForNode(serviceNode)
	service.NodeKey = serviceNode.Key

	for _, node := range serviceNode.Nodes {
		switch node.Key {
		case service.NodeKey + "/location":
			location := &Location{}
			err := json.Unmarshal([]byte(node.Value), location)
			if err == nil {
				service.Location.Host = location.Host
				service.Location.Port = location.Port
			}

		case service.NodeKey + "/config":
			for _, subNode := range node.Nodes {
				switch subNode.Key {
				case service.NodeKey + "/config/gogeta":
					serviceConfig := &ServiceConfig{}
					err := json.Unmarshal([]byte(subNode.Value), serviceConfig)
					if err == nil {
						service.Config = serviceConfig
					}
				}
			}

		case service.NodeKey + "/domain":
			service.Domain = node.Value
		case service.NodeKey + "/lastAccess":
			lastAccess := node.Value
			lastAccessTime, err := time.Parse(TIME_FORMAT, lastAccess)
			if err != nil {
				glog.Errorf("Error parsing last access date with service %s: %s", service.Name, err)
				break
			}
			service.LastAccess = &lastAccessTime

		case service.NodeKey + "/status":
			service.Status = NewStatus(service, node)
		}
	}
	return service, nil
}

func (s *Service) UnitName() string {
	return "nxio@" + strings.Split(s.Name, "_")[1] + ".service"
}

func (service *Service) Equals(other *Service) bool {
	if service == nil && other == nil {
		return true
	}

	return service != nil && other != nil &&
		service.Location.Equals(other.Location) &&
		service.Status.Equals(other.Status) &&
		service.Config.Equals(other.Config)
}

func (s *Service) StartedSince() *time.Time {
	if s == nil {
		return nil
	}

	if s.Status != nil &&
		s.Status.Current == STARTED_STATUS {
		return s.LastAccess
	} else {
		return nil
	}
}

func (s *Service) Start(client *etcd.Client) error {
	return s.fleetcmd("start", client)
}

func (s *Service) Stop(client *etcd.Client) error {
	return s.fleetcmd("destroy", client)
}

func (s *Service) Passivate(client *etcd.Client) error {
	err := s.fleetcmd("destroy", client)
	if err != nil {
		return err
	}

	statusKey := s.NodeKey + "/status"

	responseCurrent, error := client.Set(statusKey+"/current", PASSIVATED_STATUS, 0)
	if error != nil && responseCurrent == nil {
		glog.Errorf("Setting status current to 'passivated' has failed for Service "+s.Name+": %s", error)
	}

	response, error := client.Set(statusKey+"/expected", PASSIVATED_STATUS, 0)
	if error != nil && response == nil {
		glog.Errorf("Setting status expected to 'passivated' has failed for Service "+s.Name+": %s", error)
	}
	return nil
}

func (s *Service) fleetcmd(command string, client *etcd.Client) error {
	//TODO Use fleet's REST API
	etcdAddress := client.GetCluster()[0]
	cmd := exec.Command("/usr/bin/fleetctl", "--endpoint="+etcdAddress, command, s.UnitName())
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}
<|endoftext|>


<commit_before>package urknall

import (
	"crypto/sha256"
	"fmt"
	"github.com/dynport/gologger"
	"github.com/dynport/gossh"
	"github.com/dynport/urknall/cmd"
	"path"
	"strings"
)

type ProvisionOptions struct {
	LogStdout bool // log stdout of commands
	DryRun    bool
}

type sshClient struct {
	client           *gossh.Client
	host             *Host
	provisionOptions ProvisionOptions
}

func newSSHClient(host *Host, opts *ProvisionOptions) (client *sshClient) {
	if opts == nil {
		opts = &ProvisionOptions{}
	}
	return &sshClient{host: host, client: gossh.New(host.IP, host.user()), provisionOptions: *opts}
}

func (sc *sshClient) provision() (e error) {
	logger.PushPrefix(sc.host.IP)
	defer logger.PopPrefix()

	if e = sc.host.precompileRunlists(); e != nil {
		return e
	}

	return provisionRunlists(sc.host.runlists(), sc.provisionRunlist)
}

func (sc *sshClient) provisionRunlist(rl *Runlist) (e error) {
	tasks := sc.buildTasksForRunlist(rl)

	checksumDir := fmt.Sprintf("/var/cache/urknall/%s", rl.name)

	checksumHash, e := sc.buildChecksumHash(checksumDir)
	if e != nil {
		return fmt.Errorf("failed to build checksum hash: %s", e.Error())
	}

	if sc.host.isSudoRequired() {
		logger.PushPrefix("SUDO")
		defer logger.PopPrefix()
	}

	for i := range tasks {
		task := tasks[i]
		logMsg := task.command.Logging()
		if _, found := checksumHash[task.checksum]; found { // Task is cached.
			logger.Infof("\b[%s][%.8s]%s", gologger.Colorize(33, "CACHED"), task.checksum, logMsg)
			delete(checksumHash, task.checksum) // Delete checksums of cached tasks from hash.
			continue
		}

		if len(checksumHash) > 0 { // All remaining checksums are invalid, as something changed.
			if e = sc.cleanUpRemainingCachedEntries(checksumDir, checksumHash); e != nil {
				return e
			}
			checksumHash = make(map[string]struct{})
		}

		logger.Infof("\b[%s  ][%.8s]%s", gologger.Colorize(34, "EXEC"), task.checksum, logMsg)
		if e = sc.runTask(task, checksumDir); e != nil {
			return e
		}
	}

	return nil
}

func (sc *sshClient) runTask(task *taskData, checksumDir string) (e error) {
	if sc.provisionOptions.DryRun {
		return nil
	}

	checksumFile := fmt.Sprintf("%s/%s", checksumDir, task.checksum)

	stderr := ">(tee " + checksumFile + ".stderr >&2)"
	stdout := checksumFile + ".stdout"

	if sc.provisionOptions.LogStdout {
		sc.client.DebugWriter = logger.Info
		stdout = ">(tee " + stdout + ")"
	}

	sCmd := fmt.Sprintf("bash <<EOF_RUNTASK 2> %s 1> %s\n%s\nEOF_RUNTASK\n", stderr, stdout, task.command.Shell())
	if sc.host.isSudoRequired() {
		sCmd = fmt.Sprintf("sudo bash <<EOF_ZWO_SUDO\n%s\nEOF_ZWO_SUDO\n", sCmd)
	}
	sc.client.ErrorWriter = logger.Error
	rsp, e := sc.client.Execute(sCmd)

	// Write the checksum file (containing information on the command run).
	sc.writeChecksumFile(checksumFile, e != nil, task.command.Logging(), rsp)

	if e != nil {
		return fmt.Errorf("%s (see %s.failed for more information)", e.Error(), checksumFile)
	}
	return nil
}

func (sc *sshClient) executeCommand(cmdRaw string) *gossh.Result {
	if sc.host.isSudoRequired() {
		cmdRaw = fmt.Sprintf("sudo bash <<EOF_ZWO_SUDO\n%s\nEOF_ZWO_SUDO\n", cmdRaw)
	}
	c := &cmd.ShellCommand{Command: cmdRaw}
	result, e := sc.client.Execute(c.Shell())
	if e != nil {
		stderr := ""
		if result != nil {
			stderr = strings.TrimSpace(result.Stderr())
		}
		panic(fmt.Errorf("internal error: %s (%s)", e.Error(), stderr))
	}
	return result
}

func (sc *sshClient) buildChecksumHash(checksumDir string) (checksumMap map[string]struct{}, e error) {
	// Make sure the directory exists.
	sc.executeCommand(fmt.Sprintf("mkdir -p %s", checksumDir))

	checksums := []string{}
	rsp := sc.executeCommand(fmt.Sprintf("ls %s/*.done | xargs", checksumDir))
	for _, checksumFile := range strings.Fields(rsp.Stdout()) {
		checksum := strings.TrimSuffix(path.Base(checksumFile), ".done")
		checksums = append(checksums, checksum)
	}

	checksumMap = make(map[string]struct{})
	for i := range checksums {
		if len(checksums[i]) != 64 {
			return nil, fmt.Errorf("invalid checksum '%s' found in '%s'", checksums[i], checksumDir)
		}
		checksumMap[checksums[i]] = struct{}{}
	}
	return checksumMap, nil
}

func (sc *sshClient) cleanUpRemainingCachedEntries(checksumDir string, checksumHash map[string]struct{}) (e error) {
	invalidCacheEntries := make([]string, 0, len(checksumHash))
	for k, _ := range checksumHash {
		invalidCacheEntries = append(invalidCacheEntries, fmt.Sprintf("%s.done", k))
	}
	if sc.provisionOptions.DryRun {
		logger.Info("invalidated commands:", invalidCacheEntries)
	} else {
		cmd := fmt.Sprintf("cd %s && rm -f *.failed %s", checksumDir, strings.Join(invalidCacheEntries, " "))
		logger.Debug(cmd)
		sc.executeCommand(cmd)
	}
	return nil
}

func (sc *sshClient) writeChecksumFile(checksumFileBase string, failed bool, logMsg string, response *gossh.Result) {
	checksumFile := checksumFileBase
	if failed {
		checksumFile += ".failed"
	} else {
		checksumFile += ".done"
	}

	// Whoa, super hacky stuff to get the command to the checksum file. The command might contain a lot of stuff, like
	// apostrophes and the like, that would totally nuke a quoted string. Though there is a here doc.
	c := []string{
		fmt.Sprintf(`echo "STDOUT #####" >> %s`, checksumFile),
		fmt.Sprintf(`cat %s.stdout >> %s`, checksumFileBase, checksumFile),
		fmt.Sprintf(`echo "STDERR #####" >> %s`, checksumFile),
		fmt.Sprintf(`cat %s.stderr >> %s`, checksumFileBase, checksumFile),
		fmt.Sprintf(`rm -f %s.std*`, checksumFileBase),
	}
	sc.executeCommand(fmt.Sprintf("cat <<EOF_COMMAND > %s && %s\n%s\nEOF_COMMAND\n", checksumFile, strings.Join(c, " && "), logMsg))
}

type taskData struct {
	command  cmd.Command // The command to be executed.
	checksum string      // The checksum of the command.
}

func (sc *sshClient) buildTasksForRunlist(rl *Runlist) (tasks []*taskData) {
	tasks = make([]*taskData, 0, len(rl.commands))

	cmdHash := sha256.New()
	for i := range rl.commands {
		rawCmd := rl.commands[i].Shell()
		cmdHash.Write([]byte(rawCmd))

		task := &taskData{command: rl.commands[i], checksum: fmt.Sprintf("%x", cmdHash.Sum(nil))}
		tasks = append(tasks, task)
	}
	return tasks
}
<commit_msg>improved combined logging of std(err|out)<commit_after>package urknall

import (
	"crypto/sha256"
	"fmt"
	"github.com/dynport/gologger"
	"github.com/dynport/gossh"
	"github.com/dynport/urknall/cmd"
	"path"
	"strings"
)

type ProvisionOptions struct {
	LogStdout bool // log stdout of commands
	DryRun    bool
}

type sshClient struct {
	client           *gossh.Client
	host             *Host
	provisionOptions ProvisionOptions
}

func newSSHClient(host *Host, opts *ProvisionOptions) (client *sshClient) {
	if opts == nil {
		opts = &ProvisionOptions{}
	}
	return &sshClient{host: host, client: gossh.New(host.IP, host.user()), provisionOptions: *opts}
}

func (sc *sshClient) provision() (e error) {
	logger.PushPrefix(sc.host.IP)
	defer logger.PopPrefix()

	if e = sc.host.precompileRunlists(); e != nil {
		return e
	}

	return provisionRunlists(sc.host.runlists(), sc.provisionRunlist)
}

func (sc *sshClient) provisionRunlist(rl *Runlist) (e error) {
	tasks := sc.buildTasksForRunlist(rl)

	checksumDir := fmt.Sprintf("/var/cache/urknall/%s", rl.name)

	checksumHash, e := sc.buildChecksumHash(checksumDir)
	if e != nil {
		return fmt.Errorf("failed to build checksum hash: %s", e.Error())
	}

	if sc.host.isSudoRequired() {
		logger.PushPrefix("SUDO")
		defer logger.PopPrefix()
	}

	for i := range tasks {
		task := tasks[i]
		logMsg := task.command.Logging()
		if _, found := checksumHash[task.checksum]; found { // Task is cached.
			logger.Infof("\b[%s][%.8s]%s", gologger.Colorize(33, "CACHED"), task.checksum, logMsg)
			delete(checksumHash, task.checksum) // Delete checksums of cached tasks from hash.
			continue
		}

		if len(checksumHash) > 0 { // All remaining checksums are invalid, as something changed.
			if e = sc.cleanUpRemainingCachedEntries(checksumDir, checksumHash); e != nil {
				return e
			}
			checksumHash = make(map[string]struct{})
		}

		logger.Infof("\b[%s  ][%.8s]%s", gologger.Colorize(34, "EXEC"), task.checksum, logMsg)
		if e = sc.runTask(task, checksumDir); e != nil {
			return e
		}
	}

	return nil
}

func (sc *sshClient) runTask(task *taskData, checksumDir string) (e error) {
	if sc.provisionOptions.DryRun {
		return nil
	}

	stderr := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stderr:$line"; done | tee /tmp/%s.%s.stderr)`, sc.host.user(), task.checksum)
	stdout := fmt.Sprintf(`>(while read line; do echo "$(date --iso-8601=ns):stdout:$line"; done | tee /tmp/%s.%s.stdout)`, sc.host.user(), task.checksum)

	sc.client.ErrorWriter = logger.Error
	if sc.provisionOptions.LogStdout {
		sc.client.DebugWriter = logger.Info
	}

	sCmd := fmt.Sprintf("bash <<EOF_RUNTASK 1> %s 2> %s\n%s\nEOF_RUNTASK\n", stdout, stderr, task.command.Shell())
	if sc.host.isSudoRequired() {
		sCmd = fmt.Sprintf("sudo %s", sCmd)
	}
	rsp, e := sc.client.Execute(sCmd)

	// Write the checksum file (containing information on the command run).
	sc.writeChecksumFile(checksumDir, task.checksum, e != nil, task.command.Logging(), rsp)

	if e != nil {
		return fmt.Errorf("%s (see %s/%s.failed for more information)", e.Error(), checksumDir, task.checksum)
	}
	return nil
}

func (sc *sshClient) executeCommand(cmdRaw string) *gossh.Result {
	if sc.host.isSudoRequired() {
		cmdRaw = fmt.Sprintf("sudo bash <<EOF_ZWO_SUDO\n%s\nEOF_ZWO_SUDO\n", cmdRaw)
	}
	c := &cmd.ShellCommand{Command: cmdRaw}
	result, e := sc.client.Execute(c.Shell())
	if e != nil {
		stderr := ""
		if result != nil {
			stderr = strings.TrimSpace(result.Stderr())
		}
		panic(fmt.Errorf("internal error: %s (%s)", e.Error(), stderr))
	}
	return result
}

func (sc *sshClient) buildChecksumHash(checksumDir string) (checksumMap map[string]struct{}, e error) {
	// Make sure the directory exists.
	sc.executeCommand(fmt.Sprintf("mkdir -p %s", checksumDir))

	checksums := []string{}
	rsp := sc.executeCommand(fmt.Sprintf("ls %s/*.done | xargs", checksumDir))
	for _, checksumFile := range strings.Fields(rsp.Stdout()) {
		checksum := strings.TrimSuffix(path.Base(checksumFile), ".done")
		checksums = append(checksums, checksum)
	}

	checksumMap = make(map[string]struct{})
	for i := range checksums {
		if len(checksums[i]) != 64 {
			return nil, fmt.Errorf("invalid checksum '%s' found in '%s'", checksums[i], checksumDir)
		}
		checksumMap[checksums[i]] = struct{}{}
	}
	return checksumMap, nil
}

func (sc *sshClient) cleanUpRemainingCachedEntries(checksumDir string, checksumHash map[string]struct{}) (e error) {
	invalidCacheEntries := make([]string, 0, len(checksumHash))
	for k, _ := range checksumHash {
		invalidCacheEntries = append(invalidCacheEntries, fmt.Sprintf("%s.done", k))
	}
	if sc.provisionOptions.DryRun {
		logger.Info("invalidated commands:", invalidCacheEntries)
	} else {
		cmd := fmt.Sprintf("cd %s && rm -f *.failed %s", checksumDir, strings.Join(invalidCacheEntries, " "))
		logger.Debug(cmd)
		sc.executeCommand(cmd)
	}
	return nil
}

func (sc *sshClient) writeChecksumFile(checksumDir, checksum string, failed bool, logMsg string, response *gossh.Result) {
	tmpChecksumFiles := "/tmp/" + sc.host.user() + "." + checksum + ".std*"
	checksumFile := checksumDir + "/" + checksum
	if failed {
		checksumFile += ".failed"
	} else {
		checksumFile += ".done"
	}

	// Whoa, super hacky stuff to get the command to the checksum file. The command might contain a lot of stuff, like
	// apostrophes and the like, that would totally nuke a quoted string. Though there is a here doc.
	c := []string{
		fmt.Sprintf(`cat %s | sort >> %s`, tmpChecksumFiles, checksumFile),
		fmt.Sprintf(`rm -f %s`, tmpChecksumFiles),
	}
	sc.executeCommand(fmt.Sprintf("cat <<EOF_COMMAND > %s && %s\n%s\nEOF_COMMAND\n", checksumFile, strings.Join(c, " && "), logMsg))
}

type taskData struct {
	command  cmd.Command // The command to be executed.
	checksum string      // The checksum of the command.
}

func (sc *sshClient) buildTasksForRunlist(rl *Runlist) (tasks []*taskData) {
	tasks = make([]*taskData, 0, len(rl.commands))

	cmdHash := sha256.New()
	for i := range rl.commands {
		rawCmd := rl.commands[i].Shell()
		cmdHash.Write([]byte(rawCmd))

		task := &taskData{command: rl.commands[i], checksum: fmt.Sprintf("%x", cmdHash.Sum(nil))}
		tasks = append(tasks, task)
	}
	return tasks
}
<|endoftext|>


<commit_before>package xql

import (
    "errors"
    "database/sql"
    "fmt"
    log "github.com/Sirupsen/logrus"
)

type Session struct {
    driverName string
    db *sql.DB
    tx *sql.Tx
    verbose bool
}

func (self *Session) getStatementBuilder() StatementBuilder {
    if s, ok := _statement_builders[self.driverName]; ok {
        return s
    }else{
        panic(fmt.Sprintf("Statement Builder '%s' not registered! ", self.driverName))
    }
    return nil
}

func (self *Session) Close() {
    if self.tx != nil {
        self.tx.Commit()
    }
}

func (self *Session) Query(table *Table, columns ...interface{}) *QuerySet {
    qs := &QuerySet{session:self, offset:-1, limit:-1}
    qs.table = table
    if len(columns) > 0 {
        for _, c := range columns {
            if qc, ok := c.(QueryColumn); ok {
                qs.queries = append(qs.queries, qc)
            }else if qcn, ok := c.(string); ok {
                if col, ok := qs.table.Columns[qcn]; !ok {
                    panic("Invalid column name.")
                }else{
                    qs.queries = append(qs.queries, QueryColumn{FieldName:col.FieldName, Alias:col.FieldName})
                }
            }else{
                panic("Unsupported parameter type!")
            }

        }
    }else{
        //for _, col := range qs.table.Columns {
        //    qs.queries = append(qs.queries, QueryColumn{FieldName:col.FieldName, Alias:col.FieldName})
        //}
    }
    return qs
}

func (self *Session) Begin() error {
    if nil != self.tx {
        return errors.New("Already in Tx!")
    }
    tx, err := self.db.Begin()
    if nil == err {
        self.tx = tx
    }
    return err
}

func (self *Session) Commit() error {
    if nil == self.tx {
        return errors.New("Not open Tx!")
    }
    err := self.tx.Commit()
    if nil == err {
        self.tx = nil
    }
    return err
}

func (self *Session) Rollback() error {
    if nil == self.tx {
        return errors.New("Not open Tx!")
    }
    err := self.tx.Rollback()
    self.tx = nil
    return err
}

func (self *Session) doExec(query string, args ...interface{}) (sql.Result, error) {
    if self.tx != nil {
        if self.verbose { log.Debugln("doExec in Tx: ", query, args) }
        return self.tx.Exec(query, args...)
    }else{
        if self.verbose { log.Debugln("doExec in DB: ", query, args) }
        return self.db.Exec(query, args...)
    }
    return nil, nil
}

func (self *Session) doQuery(query string, args ...interface{}) (*sql.Rows, error) {
    if self.tx != nil {
        if self.verbose { log.Debugln("doQuery in Tx: ", query, args) }
        return self.tx.Query(query, args...)
    }else{
        if self.verbose { log.Debugln("doQuery in DB: ", query, args) }
        return self.db.Query(query, args...)
    }
    return nil, nil
}

func (self *Session) doQueryRow(query string, args ...interface{}) *sql.Row {
    if self.tx != nil {
        if self.verbose { log.Debugln("doQueryRow in Tx: ", query, args) }
        return self.tx.QueryRow(query, args...)
    }else{
        if self.verbose { log.Debugln("doQueryRow in Db: ", query, args) }
        return self.db.QueryRow(query, args...)
    }
    return nil
}<commit_msg>Updated: improved error message.<commit_after>package xql

import (
    "errors"
    "database/sql"
    "fmt"
    log "github.com/Sirupsen/logrus"
)

type Session struct {
    driverName string
    db *sql.DB
    tx *sql.Tx
    verbose bool
}

func (self *Session) getStatementBuilder() StatementBuilder {
    if s, ok := _statement_builders[self.driverName]; ok {
        return s
    }else{
        panic(fmt.Sprintf("Statement Builder '%s' not registered! ", self.driverName))
    }
    return nil
}

func (self *Session) Close() {
    if self.tx != nil {
        self.tx.Commit()
    }
}

func (self *Session) Query(table *Table, columns ...interface{}) *QuerySet {
    qs := &QuerySet{session:self, offset:-1, limit:-1}
    qs.table = table
    if len(columns) > 0 {
        for _, c := range columns {
            if qc, ok := c.(QueryColumn); ok {
                qs.queries = append(qs.queries, qc)
            }else if qcn, ok := c.(string); ok {
                if col, ok := qs.table.Columns[qcn]; !ok {
                    panic("Invalid column name:"+qcn)
                }else{
                    qs.queries = append(qs.queries, QueryColumn{FieldName:col.FieldName, Alias:col.FieldName})
                }
            }else{
                panic("Unsupported parameter type!")
            }

        }
    }else{
        //for _, col := range qs.table.Columns {
        //    qs.queries = append(qs.queries, QueryColumn{FieldName:col.FieldName, Alias:col.FieldName})
        //}
    }
    return qs
}

func (self *Session) Begin() error {
    if nil != self.tx {
        return errors.New("Already in Tx!")
    }
    tx, err := self.db.Begin()
    if nil == err {
        self.tx = tx
    }
    return err
}

func (self *Session) Commit() error {
    if nil == self.tx {
        return errors.New("Not open Tx!")
    }
    err := self.tx.Commit()
    if nil == err {
        self.tx = nil
    }
    return err
}

func (self *Session) Rollback() error {
    if nil == self.tx {
        return errors.New("Not open Tx!")
    }
    err := self.tx.Rollback()
    self.tx = nil
    return err
}

func (self *Session) doExec(query string, args ...interface{}) (sql.Result, error) {
    if self.tx != nil {
        if self.verbose { log.Debugln("doExec in Tx: ", query, args) }
        return self.tx.Exec(query, args...)
    }else{
        if self.verbose { log.Debugln("doExec in DB: ", query, args) }
        return self.db.Exec(query, args...)
    }
    return nil, nil
}

func (self *Session) doQuery(query string, args ...interface{}) (*sql.Rows, error) {
    if self.tx != nil {
        if self.verbose { log.Debugln("doQuery in Tx: ", query, args) }
        return self.tx.Query(query, args...)
    }else{
        if self.verbose { log.Debugln("doQuery in DB: ", query, args) }
        return self.db.Query(query, args...)
    }
    return nil, nil
}

func (self *Session) doQueryRow(query string, args ...interface{}) *sql.Row {
    if self.tx != nil {
        if self.verbose { log.Debugln("doQueryRow in Tx: ", query, args) }
        return self.tx.QueryRow(query, args...)
    }else{
        if self.verbose { log.Debugln("doQueryRow in Db: ", query, args) }
        return self.db.QueryRow(query, args...)
    }
    return nil
}<|endoftext|>


<commit_before>// Copyright (c) 2013 Mathieu Turcotte
// Licensed under the MIT license.

// Package browserchannel provides a server-side browser channel
// implementation. See http://goo.gl/F287G for the client-side API.
package bc

import (
	"encoding/json"
	"errors"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// The browser channel protocol version implemented by this library.
const SupportedProcolVersion = 8

var (
	ErrClosed = errors.New("handler closed")
)

const (
	DefaultBindPath = "bind"
	DefaultTestPath = "test"
)

type queryType int

// Possible values for the query TYPE parameter.
const (
	queryUnset = iota
	queryTerminate
	queryXmlHttp
	queryHtml
	queryTest
)

func parseQueryType(s string) (qtype queryType) {
	switch s {
	case "html":
		qtype = queryHtml
	case "xmlhttp":
		qtype = queryXmlHttp
	case "terminate":
		qtype = queryTerminate
	case "test":
		qtype = queryTest
	}
	return
}

func (qtype queryType) setContentType(rw http.ResponseWriter) {
	if qtype == queryHtml {
		rw.Header().Set("Content-Type", "text/html")
	} else {
		rw.Header().Set("Content-Type", "text/plain")
	}
}

func parseAid(s string) (aid int, err error) {
	aid = -1
	if len(s) == 0 {
		return
	}
	aid, err = strconv.Atoi(s)
	return
}

func parseProtoVersion(s string) (version int) {
	version, err := strconv.Atoi(s)
	if err != nil {
		version = -1
	}
	return
}

type bindParams struct {
	cver    string
	sid     SessionId
	qtype   queryType
	domain  string
	rid     string
	aid     int
	chunked bool
	values  url.Values
	method  string
}

func parseBindParams(req *http.Request, values url.Values) (params *bindParams, err error) {
	cver := req.Form.Get("VER")
	qtype := parseQueryType(req.Form.Get("TYPE"))
	domain := req.Form.Get("DOMAIN")
	rid := req.Form.Get("zx")
	ci := req.Form.Get("CI") == "1"
	sid, err := parseSessionId(req.Form.Get("SID"))
	if err != nil {
		return
	}
	aid, err := parseAid(req.Form.Get("AID"))
	if err != nil {
		return
	}
	params = &bindParams{cver, sid, qtype, domain, rid, aid, ci, values, req.Method}
	return
}

type testParams struct {
	ver    int
	init   bool
	qtype  queryType
	domain string
}

func parseTestParams(req *http.Request) *testParams {
	version := parseProtoVersion(req.Form.Get("VER"))
	qtype := parseQueryType(req.Form.Get("TYPE"))
	domain := req.Form.Get("DOMAIN")
	init := req.Form.Get("MODE") == "init"
	return &testParams{version, init, qtype, domain}
}

var headers = map[string]string{
	"Cache-Control":          "no-cache, no-store, max-age=0, must-revalidate",
	"Expires":                "Fri, 01 Jan 1990 00:00:00 GMT",
	"X-Content-Type-Options": "nosniff",
	"Transfer-Encoding":      "chunked",
	"Pragma":                 "no-cache",
}

type channelMap struct {
	sync.RWMutex
	m map[SessionId]*Channel
}

func (m *channelMap) get(sid SessionId) *Channel {
	m.RLock()
	defer m.RUnlock()
	return m.m[sid]
}

func (m *channelMap) set(sid SessionId, channel *Channel) {
	m.Lock()
	defer m.Unlock()
	m.m[sid] = channel
}

func (m *channelMap) del(sid SessionId) (c *Channel, deleted bool) {
	m.Lock()
	defer m.Unlock()
	c, deleted = m.m[sid]
	delete(m.m, sid)
	return
}

// Contains the browser channel cross domain info for a single domain.
type crossDomainInfo struct {
	hostMatcher *regexp.Regexp
	domain      string
	prefixes    []string
}

func getHostPrefix(info *crossDomainInfo) string {
	if info != nil {
		return info.prefixes[rand.Intn(len(info.prefixes))]
	}
	return ""
}

type ChannelHandler func(*Channel)

type Handler struct {
	corsInfo    *crossDomainInfo
	prefix      string
	channels    *channelMap
	bindPath    string
	testPath    string
	gcChan      chan SessionId
	chanHandler ChannelHandler
}

// Creates a new browser channel handler.
func NewHandler(chanHandler ChannelHandler) (h *Handler) {
	h = new(Handler)
	h.channels = &channelMap{m: make(map[SessionId]*Channel)}
	h.bindPath = DefaultBindPath
	h.testPath = DefaultTestPath
	h.gcChan = make(chan SessionId, 10)
	h.chanHandler = chanHandler
	go h.removeClosedSession()
	return
}

// Sets the cross domain information for this browser channel. The origin is
// used as the Access-Control-Allow-Origin header value and should respect the
// format specified in http://www.w3.org/TR/cors/. The prefix is used to set
// the hostPrefix parameter on the client side.
func (h *Handler) SetCrossDomainPrefix(domain string, prefixes []string) {
	h.corsInfo = &crossDomainInfo{makeOriginMatcher(domain), domain, prefixes}
}

// Removes closed channels from the handler's channel map.
func (h *Handler) removeClosedSession() {
	for {
		sid, ok := <-h.gcChan
		if !ok {
			break
		}

		log.Printf("removing %s from session map\n", sid)

		if channel, ok := h.channels.del(sid); ok {
			channel.dispose()
		} else {
			log.Printf("missing channel for %s in session map\n", sid)
		}
	}
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// The CORS  spec only supports *, null or the exact domain.
	// http://www.w3.org/TR/cors/#access-control-allow-origin-response-header
	// http://tools.ietf.org/html/rfc6454#section-7.1
	origin := req.Header.Get("origin")
	if len(origin) > 0 && h.corsInfo != nil &&
		h.corsInfo.hostMatcher.MatchString(origin) {
		rw.Header().Set("Access-Control-Allow-Origin", origin)
		rw.Header().Set("Access-Control-Allow-Credentials", "true")
	}

	// The body is parsed before calling ParseForm so the values don't get
	// collapsed into a single collection.
	values, err := parseBody(req.Body)
	if err != nil {
		rw.WriteHeader(400)
		return
	}

	req.ParseForm()

	path := req.URL.Path
	if strings.HasSuffix(path, h.testPath) {
		log.Printf("test:%s\n", req.URL)
		h.handleTestRequest(rw, parseTestParams(req))
	} else if strings.HasSuffix(path, h.bindPath) {
		log.Printf("bind:%s:%s\n", req.Method, req.URL)
		params, err := parseBindParams(req, values)
		if err != nil {
			rw.WriteHeader(400)
			return
		}
		h.handleBindRequest(rw, params)
	} else {
		rw.WriteHeader(404)
	}
}

func (h *Handler) handleTestRequest(rw http.ResponseWriter, params *testParams) {
	if params.ver != SupportedProcolVersion {
		rw.WriteHeader(400)
		io.WriteString(rw, "Unsupported protocol version.")
	} else if params.init {
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		io.WriteString(rw, "[\""+getHostPrefix(h.corsInfo)+"\",\"\"]")
	} else {
		params.qtype.setContentType(rw)
		setHeaders(rw, &headers)
		rw.WriteHeader(200)

		if params.qtype == queryHtml {
			writeHtmlHead(rw)
			writeHtmlDomain(rw, params.domain)
			writeHtmlRpc(rw, "11111")
			writeHtmlPadding(rw)
		} else {
			io.WriteString(rw, "11111")
		}

		time.Sleep(2 * time.Second)

		if params.qtype == queryHtml {
			writeHtmlRpc(rw, "2")
			writeHtmlDone(rw)
		} else {
			io.WriteString(rw, "2")
		}
	}
}

func (h *Handler) handleBindRequest(rw http.ResponseWriter, params *bindParams) {
	var channel *Channel
	sid := params.sid

	// If the client has specified a session id, lookup the session object in
	// the sessions map. Lookup failure should be signaled to the client using
	// a 400 status code and a message containing 'Unknown SID'. See
	// goog/net/channelrequest.js for more context on how this error is
	// handled.
	if sid != nullSessionId {
		channel = h.channels.get(sid)
		if channel == nil {
			log.Printf("failed to lookup session %s\n", sid)
			setHeaders(rw, &headers)
			rw.WriteHeader(400)
			io.WriteString(rw, "Unknown SID")
			return
		}
	}

	if channel == nil {
		sid, _ = newSesionId()
		log.Printf("creating session %s\n", sid)
		channel = newChannel(params.cver, sid, h.gcChan, h.corsInfo)
		h.channels.set(sid, channel)
		go h.chanHandler(channel)
		go channel.start()
	}

	if params.aid != -1 {
		channel.acknowledge(params.aid)
	}

	switch params.method {
	case "POST":
		h.handleBindPost(rw, params, channel)
	case "GET":
		h.handleBindGet(rw, params, channel)
	default:
		rw.WriteHeader(400)
	}
}

func (h *Handler) handleBindPost(rw http.ResponseWriter, params *bindParams, channel *Channel) {
	log.Printf("%s: bind parameters: %v\n", channel.Sid, params.values)

	offset, maps, err := parseIncomingMaps(params.values)
	if err != nil {
		rw.WriteHeader(400)
		return
	}

	channel.receiveMaps(offset, maps)

	if channel.state == channelInit {
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		rw.(http.Flusher).Flush()

		// The initial forward request is used as a back channel to send the
		// server configuration: ['c', id, host, version]. This payload has to
		// be sent immediately, so the streaming is disabled on the back
		// channel. Note that the first bind request made by IE<10 does not
		// contain a TYPE=html query parameter and therefore receives the same
		// length prefixed array reply as is sent to the XHR streaming clients.
		backChannel := newBackChannel(channel.Sid, rw, false, "", params.rid)
		channel.setBackChannel(backChannel)
		backChannel.wait()
	} else {
		// On normal forward channel request, the session status is returned
		// to the client. The session status contains 3 pieces of information:
		// does this session has a back channel, the last array id sent to the
		// client and the number of outstanding bytes in the back channel.
		b, _ := json.Marshal(channel.getState())
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		io.WriteString(rw, strconv.FormatInt(int64(len(b)), 10)+"\n")
		rw.Write(b)
	}
}

func (h *Handler) handleBindGet(rw http.ResponseWriter, params *bindParams, channel *Channel) {
	if params.qtype == queryTerminate {
		channel.Close()
	} else {
		params.qtype.setContentType(rw)
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		rw.(http.Flusher).Flush()

		isHtml := params.qtype == queryHtml
		bc := newBackChannel(channel.Sid, rw, isHtml, params.domain, params.rid)
		bc.setChunked(params.chunked)
		channel.setBackChannel(bc)
		bc.wait()
	}
}
<commit_msg>Remove unused ErrClosed.<commit_after>// Copyright (c) 2013 Mathieu Turcotte
// Licensed under the MIT license.

// Package browserchannel provides a server-side browser channel
// implementation. See http://goo.gl/F287G for the client-side API.
package bc

import (
	"encoding/json"
	"errors"
	"io"
	"log"
	"math/rand"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// The browser channel protocol version implemented by this library.
const SupportedProcolVersion = 8

const (
	DefaultBindPath = "bind"
	DefaultTestPath = "test"
)

type queryType int

// Possible values for the query TYPE parameter.
const (
	queryUnset = iota
	queryTerminate
	queryXmlHttp
	queryHtml
	queryTest
)

func parseQueryType(s string) (qtype queryType) {
	switch s {
	case "html":
		qtype = queryHtml
	case "xmlhttp":
		qtype = queryXmlHttp
	case "terminate":
		qtype = queryTerminate
	case "test":
		qtype = queryTest
	}
	return
}

func (qtype queryType) setContentType(rw http.ResponseWriter) {
	if qtype == queryHtml {
		rw.Header().Set("Content-Type", "text/html")
	} else {
		rw.Header().Set("Content-Type", "text/plain")
	}
}

func parseAid(s string) (aid int, err error) {
	aid = -1
	if len(s) == 0 {
		return
	}
	aid, err = strconv.Atoi(s)
	return
}

func parseProtoVersion(s string) (version int) {
	version, err := strconv.Atoi(s)
	if err != nil {
		version = -1
	}
	return
}

type bindParams struct {
	cver    string
	sid     SessionId
	qtype   queryType
	domain  string
	rid     string
	aid     int
	chunked bool
	values  url.Values
	method  string
}

func parseBindParams(req *http.Request, values url.Values) (params *bindParams, err error) {
	cver := req.Form.Get("VER")
	qtype := parseQueryType(req.Form.Get("TYPE"))
	domain := req.Form.Get("DOMAIN")
	rid := req.Form.Get("zx")
	ci := req.Form.Get("CI") == "1"
	sid, err := parseSessionId(req.Form.Get("SID"))
	if err != nil {
		return
	}
	aid, err := parseAid(req.Form.Get("AID"))
	if err != nil {
		return
	}
	params = &bindParams{cver, sid, qtype, domain, rid, aid, ci, values, req.Method}
	return
}

type testParams struct {
	ver    int
	init   bool
	qtype  queryType
	domain string
}

func parseTestParams(req *http.Request) *testParams {
	version := parseProtoVersion(req.Form.Get("VER"))
	qtype := parseQueryType(req.Form.Get("TYPE"))
	domain := req.Form.Get("DOMAIN")
	init := req.Form.Get("MODE") == "init"
	return &testParams{version, init, qtype, domain}
}

var headers = map[string]string{
	"Cache-Control":          "no-cache, no-store, max-age=0, must-revalidate",
	"Expires":                "Fri, 01 Jan 1990 00:00:00 GMT",
	"X-Content-Type-Options": "nosniff",
	"Transfer-Encoding":      "chunked",
	"Pragma":                 "no-cache",
}

type channelMap struct {
	sync.RWMutex
	m map[SessionId]*Channel
}

func (m *channelMap) get(sid SessionId) *Channel {
	m.RLock()
	defer m.RUnlock()
	return m.m[sid]
}

func (m *channelMap) set(sid SessionId, channel *Channel) {
	m.Lock()
	defer m.Unlock()
	m.m[sid] = channel
}

func (m *channelMap) del(sid SessionId) (c *Channel, deleted bool) {
	m.Lock()
	defer m.Unlock()
	c, deleted = m.m[sid]
	delete(m.m, sid)
	return
}

// Contains the browser channel cross domain info for a single domain.
type crossDomainInfo struct {
	hostMatcher *regexp.Regexp
	domain      string
	prefixes    []string
}

func getHostPrefix(info *crossDomainInfo) string {
	if info != nil {
		return info.prefixes[rand.Intn(len(info.prefixes))]
	}
	return ""
}

type ChannelHandler func(*Channel)

type Handler struct {
	corsInfo    *crossDomainInfo
	prefix      string
	channels    *channelMap
	bindPath    string
	testPath    string
	gcChan      chan SessionId
	chanHandler ChannelHandler
}

// Creates a new browser channel handler.
func NewHandler(chanHandler ChannelHandler) (h *Handler) {
	h = new(Handler)
	h.channels = &channelMap{m: make(map[SessionId]*Channel)}
	h.bindPath = DefaultBindPath
	h.testPath = DefaultTestPath
	h.gcChan = make(chan SessionId, 10)
	h.chanHandler = chanHandler
	go h.removeClosedSession()
	return
}

// Sets the cross domain information for this browser channel. The origin is
// used as the Access-Control-Allow-Origin header value and should respect the
// format specified in http://www.w3.org/TR/cors/. The prefix is used to set
// the hostPrefix parameter on the client side.
func (h *Handler) SetCrossDomainPrefix(domain string, prefixes []string) {
	h.corsInfo = &crossDomainInfo{makeOriginMatcher(domain), domain, prefixes}
}

// Removes closed channels from the handler's channel map.
func (h *Handler) removeClosedSession() {
	for {
		sid, ok := <-h.gcChan
		if !ok {
			break
		}

		log.Printf("removing %s from session map\n", sid)

		if channel, ok := h.channels.del(sid); ok {
			channel.dispose()
		} else {
			log.Printf("missing channel for %s in session map\n", sid)
		}
	}
}

func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// The CORS  spec only supports *, null or the exact domain.
	// http://www.w3.org/TR/cors/#access-control-allow-origin-response-header
	// http://tools.ietf.org/html/rfc6454#section-7.1
	origin := req.Header.Get("origin")
	if len(origin) > 0 && h.corsInfo != nil &&
		h.corsInfo.hostMatcher.MatchString(origin) {
		rw.Header().Set("Access-Control-Allow-Origin", origin)
		rw.Header().Set("Access-Control-Allow-Credentials", "true")
	}

	// The body is parsed before calling ParseForm so the values don't get
	// collapsed into a single collection.
	values, err := parseBody(req.Body)
	if err != nil {
		rw.WriteHeader(400)
		return
	}

	req.ParseForm()

	path := req.URL.Path
	if strings.HasSuffix(path, h.testPath) {
		log.Printf("test:%s\n", req.URL)
		h.handleTestRequest(rw, parseTestParams(req))
	} else if strings.HasSuffix(path, h.bindPath) {
		log.Printf("bind:%s:%s\n", req.Method, req.URL)
		params, err := parseBindParams(req, values)
		if err != nil {
			rw.WriteHeader(400)
			return
		}
		h.handleBindRequest(rw, params)
	} else {
		rw.WriteHeader(404)
	}
}

func (h *Handler) handleTestRequest(rw http.ResponseWriter, params *testParams) {
	if params.ver != SupportedProcolVersion {
		rw.WriteHeader(400)
		io.WriteString(rw, "Unsupported protocol version.")
	} else if params.init {
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		io.WriteString(rw, "[\""+getHostPrefix(h.corsInfo)+"\",\"\"]")
	} else {
		params.qtype.setContentType(rw)
		setHeaders(rw, &headers)
		rw.WriteHeader(200)

		if params.qtype == queryHtml {
			writeHtmlHead(rw)
			writeHtmlDomain(rw, params.domain)
			writeHtmlRpc(rw, "11111")
			writeHtmlPadding(rw)
		} else {
			io.WriteString(rw, "11111")
		}

		time.Sleep(2 * time.Second)

		if params.qtype == queryHtml {
			writeHtmlRpc(rw, "2")
			writeHtmlDone(rw)
		} else {
			io.WriteString(rw, "2")
		}
	}
}

func (h *Handler) handleBindRequest(rw http.ResponseWriter, params *bindParams) {
	var channel *Channel
	sid := params.sid

	// If the client has specified a session id, lookup the session object in
	// the sessions map. Lookup failure should be signaled to the client using
	// a 400 status code and a message containing 'Unknown SID'. See
	// goog/net/channelrequest.js for more context on how this error is
	// handled.
	if sid != nullSessionId {
		channel = h.channels.get(sid)
		if channel == nil {
			log.Printf("failed to lookup session %s\n", sid)
			setHeaders(rw, &headers)
			rw.WriteHeader(400)
			io.WriteString(rw, "Unknown SID")
			return
		}
	}

	if channel == nil {
		sid, _ = newSesionId()
		log.Printf("creating session %s\n", sid)
		channel = newChannel(params.cver, sid, h.gcChan, h.corsInfo)
		h.channels.set(sid, channel)
		go h.chanHandler(channel)
		go channel.start()
	}

	if params.aid != -1 {
		channel.acknowledge(params.aid)
	}

	switch params.method {
	case "POST":
		h.handleBindPost(rw, params, channel)
	case "GET":
		h.handleBindGet(rw, params, channel)
	default:
		rw.WriteHeader(400)
	}
}

func (h *Handler) handleBindPost(rw http.ResponseWriter, params *bindParams, channel *Channel) {
	log.Printf("%s: bind parameters: %v\n", channel.Sid, params.values)

	offset, maps, err := parseIncomingMaps(params.values)
	if err != nil {
		rw.WriteHeader(400)
		return
	}

	channel.receiveMaps(offset, maps)

	if channel.state == channelInit {
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		rw.(http.Flusher).Flush()

		// The initial forward request is used as a back channel to send the
		// server configuration: ['c', id, host, version]. This payload has to
		// be sent immediately, so the streaming is disabled on the back
		// channel. Note that the first bind request made by IE<10 does not
		// contain a TYPE=html query parameter and therefore receives the same
		// length prefixed array reply as is sent to the XHR streaming clients.
		backChannel := newBackChannel(channel.Sid, rw, false, "", params.rid)
		channel.setBackChannel(backChannel)
		backChannel.wait()
	} else {
		// On normal forward channel request, the session status is returned
		// to the client. The session status contains 3 pieces of information:
		// does this session has a back channel, the last array id sent to the
		// client and the number of outstanding bytes in the back channel.
		b, _ := json.Marshal(channel.getState())
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		io.WriteString(rw, strconv.FormatInt(int64(len(b)), 10)+"\n")
		rw.Write(b)
	}
}

func (h *Handler) handleBindGet(rw http.ResponseWriter, params *bindParams, channel *Channel) {
	if params.qtype == queryTerminate {
		channel.Close()
	} else {
		params.qtype.setContentType(rw)
		setHeaders(rw, &headers)
		rw.WriteHeader(200)
		rw.(http.Flusher).Flush()

		isHtml := params.qtype == queryHtml
		bc := newBackChannel(channel.Sid, rw, isHtml, params.domain, params.rid)
		bc.setChunked(params.chunked)
		channel.setBackChannel(bc)
		bc.wait()
	}
}
<|endoftext|>