<commit_before>/*
Copyright 2020 The Vitess Authors.

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

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

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

package tabletserver

import (
	"flag"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/sqlparser"

	"vitess.io/vitess/go/vt/dbconfigs"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/timer"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"

	"context"

	"github.com/golang/protobuf/proto"

	"vitess.io/vitess/go/history"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager/vreplication"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
)

var (
	// blpFunc is a legaacy feature.
	// TODO(sougou): remove after legacy resharding worflows are removed.
	blpFunc = vreplication.StatusSummary

	errUnintialized = "tabletserver uninitialized"

	streamHealthBufferSize = flag.Uint("stream_health_buffer_size", 20, "max streaming health entries to buffer per streaming health client")
)

// healthStreamer streams health information to callers.
type healthStreamer struct {
	stats              *tabletenv.Stats
	degradedThreshold  time.Duration
	unhealthyThreshold sync2.AtomicDuration

	mu      sync.Mutex
	ctx     context.Context
	cancel  context.CancelFunc
	clients map[chan *querypb.StreamHealthResponse]struct{}
	state   *querypb.StreamHealthResponse

	history *history.History

	ticks       *timer.Timer
	dbConfig    dbconfigs.Connector
	conns       *connpool.Pool
	initSuccess bool
}

func newHealthStreamer(env tabletenv.Env, alias topodatapb.TabletAlias) *healthStreamer {
	var newTimer *timer.Timer
	var pool *connpool.Pool
	if env.Config().SignalWhenSchemaChange {
		reloadTime := env.Config().SchemaReloadIntervalSeconds.Get()
		newTimer = timer.NewTimer(reloadTime)
		// We need one connection for the reloader.
		pool = connpool.NewPool(env, "", tabletenv.ConnPoolConfig{
			Size:               1,
			IdleTimeoutSeconds: env.Config().OltpReadPool.IdleTimeoutSeconds,
		})
	}
	return &healthStreamer{
		stats:              env.Stats(),
		degradedThreshold:  env.Config().Healthcheck.DegradedThresholdSeconds.Get(),
		unhealthyThreshold: sync2.NewAtomicDuration(env.Config().Healthcheck.UnhealthyThresholdSeconds.Get()),
		clients:            make(map[chan *querypb.StreamHealthResponse]struct{}),

		state: &querypb.StreamHealthResponse{
			Target:      &querypb.Target{},
			TabletAlias: &alias,
			RealtimeStats: &querypb.RealtimeStats{
				HealthError: errUnintialized,
			},
		},

		history: history.New(5),
		ticks:   newTimer,
		conns:   pool,
	}
}

func (hs *healthStreamer) InitDBConfig(target querypb.Target, cp dbconfigs.Connector) {
	// Weird test failures happen if we don't instantiate
	// a separate variable.
	inner := target
	hs.state.Target = &inner
	hs.dbConfig = cp
}

func (hs *healthStreamer) Open() {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel != nil {
		return
	}
	hs.ctx, hs.cancel = context.WithCancel(context.Background())
	if hs.conns != nil {
		// if we don't have a live conns object, it means we are not configured to signal when the schema changes
		hs.conns.Open(hs.dbConfig, hs.dbConfig, hs.dbConfig)
		hs.ticks.Start(func() {
			if err := hs.reload(); err != nil {
				log.Errorf("periodic schema reload failed in health stream: %v", err)
			}
		})
	}

}

func (hs *healthStreamer) Close() {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel != nil {
		if hs.ticks != nil {
			hs.ticks.Stop()
			hs.conns.Close()
		}
		hs.cancel()
		hs.cancel = nil
	}
}

func (hs *healthStreamer) Stream(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
	ch, hsCtx := hs.register()
	if hsCtx == nil {
		return vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "tabletserver is shutdown")
	}
	defer hs.unregister(ch)

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-hsCtx.Done():
			return vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "tabletserver is shutdown")
		case shr, ok := <-ch:
			if !ok {
				return vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "stream health buffer overflowed. client should reconnect for up-to-date status")
			}
			if err := callback(shr); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
	}
}

func (hs *healthStreamer) register() (chan *querypb.StreamHealthResponse, context.Context) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel == nil {
		return nil, nil
	}

	ch := make(chan *querypb.StreamHealthResponse, *streamHealthBufferSize)
	hs.clients[ch] = struct{}{}

	// Send the current state immediately.
	ch <- proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	return ch, hs.ctx
}

func (hs *healthStreamer) unregister(ch chan *querypb.StreamHealthResponse) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	delete(hs.clients, ch)
}

func (hs *healthStreamer) ChangeState(tabletType topodatapb.TabletType, terTimestamp time.Time, lag time.Duration, err error, serving bool) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	hs.state.Target.TabletType = tabletType
	if tabletType == topodatapb.TabletType_MASTER {
		hs.state.TabletExternallyReparentedTimestamp = terTimestamp.Unix()
	} else {
		hs.state.TabletExternallyReparentedTimestamp = 0
	}
	if err != nil {
		hs.state.RealtimeStats.HealthError = err.Error()
	} else {
		hs.state.RealtimeStats.HealthError = ""
	}
	hs.state.RealtimeStats.SecondsBehindMaster = uint32(lag.Seconds())
	hs.state.Serving = serving

	hs.state.RealtimeStats.SecondsBehindMasterFilteredReplication, hs.state.RealtimeStats.BinlogPlayersCount = blpFunc()
	hs.state.RealtimeStats.Qps = hs.stats.QPSRates.TotalRate()

	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)

	hs.broadCastToClients(shr)
	hs.history.Add(&historyRecord{
		Time:       time.Now(),
		serving:    shr.Serving,
		tabletType: shr.Target.TabletType,
		lag:        lag,
		err:        err,
	})
}

func (hs *healthStreamer) broadCastToClients(shr *querypb.StreamHealthResponse) {
	for ch := range hs.clients {
		select {
		case ch <- shr:
		default:
			// We can't block this state change on broadcasting to a streaming health client, but we
			// also don't want to silently fail to inform a streaming health client of a state change
			// because it can allow a vtgate to get wedged in a state where it's wrong about whether
			// a tablet is healthy and can't automatically recover (see
			//  https://github.com/vitessio/vitess/issues/5445). If we can't send a health update
			// to this client we'll close() the channel which will ultimate fail the streaming health
			// RPC and cause vtgates to reconnect.
			//
			// An alternative approach for streaming health would be to force a periodic broadcast even
			// when there hasn't been an update and/or move away from using channels toward a model where
			// old updates can be purged from the buffer in favor of more recent updates (since only the
			// most recent health state really matters to gates).
			log.Warning("A streaming health buffer is full. Closing the channel")
			close(ch)
			delete(hs.clients, ch)
		}
	}
}

func (hs *healthStreamer) AppendDetails(details []*kv) []*kv {
	hs.mu.Lock()
	defer hs.mu.Unlock()
	if hs.state.Target.TabletType == topodatapb.TabletType_MASTER {
		return details
	}
	sbm := time.Duration(hs.state.RealtimeStats.SecondsBehindMaster) * time.Second
	class := healthyClass
	switch {
	case sbm > hs.unhealthyThreshold.Get():
		class = unhealthyClass
	case sbm > hs.degradedThreshold:
		class = unhappyClass
	}
	details = append(details, &kv{
		Key:   "Replication Lag",
		Class: class,
		Value: fmt.Sprintf("%ds", hs.state.RealtimeStats.SecondsBehindMaster),
	})
	if hs.state.RealtimeStats.HealthError != "" {
		details = append(details, &kv{
			Key:   "Replication Error",
			Class: unhappyClass,
			Value: hs.state.RealtimeStats.HealthError,
		})
	}

	return details
}

func (hs *healthStreamer) SetUnhealthyThreshold(v time.Duration) {
	hs.unhealthyThreshold.Set(v)
	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	for ch := range hs.clients {
		select {
		case ch <- shr:
		default:
			log.Info("Resetting health streamer clients due to unhealthy threshold change")
			close(ch)
			delete(hs.clients, ch)
		}
	}
}

// reload reloads the schema from the underlying mysql
func (hs *healthStreamer) reload() error {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	// Schema Reload to happen only on master.
	if hs.state.Target.TabletType != topodatapb.TabletType_MASTER {
		return nil
	}

	ctx := hs.ctx
	conn, err := hs.conns.Get(ctx)
	if err != nil {
		return err
	}
	defer conn.Recycle()

	if !hs.initSuccess {
		hs.initSuccess, err = hs.InitSchemaLocked(conn)
		if err != nil {
			return err
		}
	}

	maxrows := 10000
	qr, err := conn.Exec(ctx, mysql.DetectSchemaChange, maxrows, false)
	if err != nil {
		return err
	}

	// If no change detected, then return
	if len(qr.Rows) == 0 {
		return nil
	}

	var tables []string
	var tablePredicates []string
	for _, row := range qr.Rows {
		table := row[0].ToString()
		tables = append(tables, table)

		tableName := sqlparser.NewStrLiteral(table)
		tablePredicates = append(tablePredicates, "table_name = "+sqlparser.String(tableName))
	}
	tableNamePredicates := strings.Join(tablePredicates, " OR ")
	del := fmt.Sprintf("%s WHERE %s", mysql.ClearSchemaCopy, tableNamePredicates)
	upd := fmt.Sprintf("%s AND %s", mysql.InsertIntoSchemaCopy, tableNamePredicates)

	// Reload the schema in a transaction.
	_, err = conn.Exec(ctx, "begin", 1, false)
	if err != nil {
		return err
	}
	defer conn.Exec(ctx, "rollback", 1, false)

	_, err = conn.Exec(ctx, del, 1, false)
	if err != nil {
		return err
	}

	_, err = conn.Exec(ctx, upd, 1, false)
	if err != nil {
		return err
	}

	_, err = conn.Exec(ctx, "commit", 1, false)
	if err != nil {
		return err
	}

	hs.state.RealtimeStats.TableSchemaChanged = tables
	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	hs.broadCastToClients(shr)
	hs.state.RealtimeStats.TableSchemaChanged = nil

	return nil
}

func (hs *healthStreamer) InitSchemaLocked(conn *connpool.DBConn) (bool, error) {
	for _, query := range mysql.VTDatabaseInit {
		_, err := conn.Exec(hs.ctx, query, 1, false)
		if err != nil {
			return false, err
		}
	}

	return true, nil
}
<commit_msg>use IN instead of ORing together a bunch of equality comparisons<commit_after>/*
Copyright 2020 The Vitess Authors.

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

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

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

package tabletserver

import (
	"flag"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/sqlparser"

	"vitess.io/vitess/go/vt/dbconfigs"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/timer"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/connpool"

	"context"

	"github.com/golang/protobuf/proto"

	"vitess.io/vitess/go/history"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager/vreplication"
	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
)

var (
	// blpFunc is a legaacy feature.
	// TODO(sougou): remove after legacy resharding worflows are removed.
	blpFunc = vreplication.StatusSummary

	errUnintialized = "tabletserver uninitialized"

	streamHealthBufferSize = flag.Uint("stream_health_buffer_size", 20, "max streaming health entries to buffer per streaming health client")
)

// healthStreamer streams health information to callers.
type healthStreamer struct {
	stats              *tabletenv.Stats
	degradedThreshold  time.Duration
	unhealthyThreshold sync2.AtomicDuration

	mu      sync.Mutex
	ctx     context.Context
	cancel  context.CancelFunc
	clients map[chan *querypb.StreamHealthResponse]struct{}
	state   *querypb.StreamHealthResponse

	history *history.History

	ticks       *timer.Timer
	dbConfig    dbconfigs.Connector
	conns       *connpool.Pool
	initSuccess bool
}

func newHealthStreamer(env tabletenv.Env, alias topodatapb.TabletAlias) *healthStreamer {
	var newTimer *timer.Timer
	var pool *connpool.Pool
	if env.Config().SignalWhenSchemaChange {
		reloadTime := env.Config().SchemaReloadIntervalSeconds.Get()
		newTimer = timer.NewTimer(reloadTime)
		// We need one connection for the reloader.
		pool = connpool.NewPool(env, "", tabletenv.ConnPoolConfig{
			Size:               1,
			IdleTimeoutSeconds: env.Config().OltpReadPool.IdleTimeoutSeconds,
		})
	}
	return &healthStreamer{
		stats:              env.Stats(),
		degradedThreshold:  env.Config().Healthcheck.DegradedThresholdSeconds.Get(),
		unhealthyThreshold: sync2.NewAtomicDuration(env.Config().Healthcheck.UnhealthyThresholdSeconds.Get()),
		clients:            make(map[chan *querypb.StreamHealthResponse]struct{}),

		state: &querypb.StreamHealthResponse{
			Target:      &querypb.Target{},
			TabletAlias: &alias,
			RealtimeStats: &querypb.RealtimeStats{
				HealthError: errUnintialized,
			},
		},

		history: history.New(5),
		ticks:   newTimer,
		conns:   pool,
	}
}

func (hs *healthStreamer) InitDBConfig(target querypb.Target, cp dbconfigs.Connector) {
	// Weird test failures happen if we don't instantiate
	// a separate variable.
	inner := target
	hs.state.Target = &inner
	hs.dbConfig = cp
}

func (hs *healthStreamer) Open() {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel != nil {
		return
	}
	hs.ctx, hs.cancel = context.WithCancel(context.Background())
	if hs.conns != nil {
		// if we don't have a live conns object, it means we are not configured to signal when the schema changes
		hs.conns.Open(hs.dbConfig, hs.dbConfig, hs.dbConfig)
		hs.ticks.Start(func() {
			if err := hs.reload(); err != nil {
				log.Errorf("periodic schema reload failed in health stream: %v", err)
			}
		})
	}

}

func (hs *healthStreamer) Close() {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel != nil {
		if hs.ticks != nil {
			hs.ticks.Stop()
			hs.conns.Close()
		}
		hs.cancel()
		hs.cancel = nil
	}
}

func (hs *healthStreamer) Stream(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
	ch, hsCtx := hs.register()
	if hsCtx == nil {
		return vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "tabletserver is shutdown")
	}
	defer hs.unregister(ch)

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-hsCtx.Done():
			return vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "tabletserver is shutdown")
		case shr, ok := <-ch:
			if !ok {
				return vterrors.Errorf(vtrpcpb.Code_RESOURCE_EXHAUSTED, "stream health buffer overflowed. client should reconnect for up-to-date status")
			}
			if err := callback(shr); err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
	}
}

func (hs *healthStreamer) register() (chan *querypb.StreamHealthResponse, context.Context) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	if hs.cancel == nil {
		return nil, nil
	}

	ch := make(chan *querypb.StreamHealthResponse, *streamHealthBufferSize)
	hs.clients[ch] = struct{}{}

	// Send the current state immediately.
	ch <- proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	return ch, hs.ctx
}

func (hs *healthStreamer) unregister(ch chan *querypb.StreamHealthResponse) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	delete(hs.clients, ch)
}

func (hs *healthStreamer) ChangeState(tabletType topodatapb.TabletType, terTimestamp time.Time, lag time.Duration, err error, serving bool) {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	hs.state.Target.TabletType = tabletType
	if tabletType == topodatapb.TabletType_MASTER {
		hs.state.TabletExternallyReparentedTimestamp = terTimestamp.Unix()
	} else {
		hs.state.TabletExternallyReparentedTimestamp = 0
	}
	if err != nil {
		hs.state.RealtimeStats.HealthError = err.Error()
	} else {
		hs.state.RealtimeStats.HealthError = ""
	}
	hs.state.RealtimeStats.SecondsBehindMaster = uint32(lag.Seconds())
	hs.state.Serving = serving

	hs.state.RealtimeStats.SecondsBehindMasterFilteredReplication, hs.state.RealtimeStats.BinlogPlayersCount = blpFunc()
	hs.state.RealtimeStats.Qps = hs.stats.QPSRates.TotalRate()

	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)

	hs.broadCastToClients(shr)
	hs.history.Add(&historyRecord{
		Time:       time.Now(),
		serving:    shr.Serving,
		tabletType: shr.Target.TabletType,
		lag:        lag,
		err:        err,
	})
}

func (hs *healthStreamer) broadCastToClients(shr *querypb.StreamHealthResponse) {
	for ch := range hs.clients {
		select {
		case ch <- shr:
		default:
			// We can't block this state change on broadcasting to a streaming health client, but we
			// also don't want to silently fail to inform a streaming health client of a state change
			// because it can allow a vtgate to get wedged in a state where it's wrong about whether
			// a tablet is healthy and can't automatically recover (see
			//  https://github.com/vitessio/vitess/issues/5445). If we can't send a health update
			// to this client we'll close() the channel which will ultimate fail the streaming health
			// RPC and cause vtgates to reconnect.
			//
			// An alternative approach for streaming health would be to force a periodic broadcast even
			// when there hasn't been an update and/or move away from using channels toward a model where
			// old updates can be purged from the buffer in favor of more recent updates (since only the
			// most recent health state really matters to gates).
			log.Warning("A streaming health buffer is full. Closing the channel")
			close(ch)
			delete(hs.clients, ch)
		}
	}
}

func (hs *healthStreamer) AppendDetails(details []*kv) []*kv {
	hs.mu.Lock()
	defer hs.mu.Unlock()
	if hs.state.Target.TabletType == topodatapb.TabletType_MASTER {
		return details
	}
	sbm := time.Duration(hs.state.RealtimeStats.SecondsBehindMaster) * time.Second
	class := healthyClass
	switch {
	case sbm > hs.unhealthyThreshold.Get():
		class = unhealthyClass
	case sbm > hs.degradedThreshold:
		class = unhappyClass
	}
	details = append(details, &kv{
		Key:   "Replication Lag",
		Class: class,
		Value: fmt.Sprintf("%ds", hs.state.RealtimeStats.SecondsBehindMaster),
	})
	if hs.state.RealtimeStats.HealthError != "" {
		details = append(details, &kv{
			Key:   "Replication Error",
			Class: unhappyClass,
			Value: hs.state.RealtimeStats.HealthError,
		})
	}

	return details
}

func (hs *healthStreamer) SetUnhealthyThreshold(v time.Duration) {
	hs.unhealthyThreshold.Set(v)
	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	for ch := range hs.clients {
		select {
		case ch <- shr:
		default:
			log.Info("Resetting health streamer clients due to unhealthy threshold change")
			close(ch)
			delete(hs.clients, ch)
		}
	}
}

// reload reloads the schema from the underlying mysql
func (hs *healthStreamer) reload() error {
	hs.mu.Lock()
	defer hs.mu.Unlock()

	// Schema Reload to happen only on master.
	if hs.state.Target.TabletType != topodatapb.TabletType_MASTER {
		return nil
	}

	ctx := hs.ctx
	conn, err := hs.conns.Get(ctx)
	if err != nil {
		return err
	}
	defer conn.Recycle()

	if !hs.initSuccess {
		hs.initSuccess, err = hs.InitSchemaLocked(conn)
		if err != nil {
			return err
		}
	}

	maxrows := 10000
	qr, err := conn.Exec(ctx, mysql.DetectSchemaChange, maxrows, false)
	if err != nil {
		return err
	}

	// If no change detected, then return
	if len(qr.Rows) == 0 {
		return nil
	}

	var tables []string
	var tableNames []string
	for _, row := range qr.Rows {
		table := row[0].ToString()
		tables = append(tables, table)

		escapedTblName := sqlparser.String(sqlparser.NewStrLiteral(table))
		tableNames = append(tableNames, escapedTblName)
	}
	tableNamePredicate := fmt.Sprintf("table_name IN (%s)", strings.Join(tableNames, ", "))
	del := fmt.Sprintf("%s WHERE %s", mysql.ClearSchemaCopy, tableNamePredicate)
	upd := fmt.Sprintf("%s AND %s", mysql.InsertIntoSchemaCopy, tableNamePredicate)

	// Reload the schema in a transaction.
	_, err = conn.Exec(ctx, "begin", 1, false)
	if err != nil {
		return err
	}
	defer conn.Exec(ctx, "rollback", 1, false)

	_, err = conn.Exec(ctx, del, 1, false)
	if err != nil {
		return err
	}

	_, err = conn.Exec(ctx, upd, 1, false)
	if err != nil {
		return err
	}

	_, err = conn.Exec(ctx, "commit", 1, false)
	if err != nil {
		return err
	}

	hs.state.RealtimeStats.TableSchemaChanged = tables
	shr := proto.Clone(hs.state).(*querypb.StreamHealthResponse)
	hs.broadCastToClients(shr)
	hs.state.RealtimeStats.TableSchemaChanged = nil

	return nil
}

func (hs *healthStreamer) InitSchemaLocked(conn *connpool.DBConn) (bool, error) {
	for _, query := range mysql.VTDatabaseInit {
		_, err := conn.Exec(hs.ctx, query, 1, false)
		if err != nil {
			return false, err
		}
	}

	return true, nil
}
<|endoftext|>


<commit_before>package engine

import (
	"github.com/keybase/client/go/libkb"
	keybase_1 "github.com/keybase/client/protocol/go"
	"testing"
)

func runTrack(fu *FakeUser, username string) (idUI *FakeIdentifyUI, res *IdentifyRes, err error) {
	idUI = &FakeIdentifyUI{
		Proofs: make(map[string]string),
		Fapr:   keybase_1.FinishAndPromptRes{TrackRemote: true},
	}
	arg := TrackEngineArg{
		TheirName: username,
	}
	ctx := Context{
		LogUI:    G.UI.GetLogUI(),
		TrackUI:  idUI,
		SecretUI: fu.NewSecretUI(),
	}
	eng := NewTrackEngine(&arg)
	err = RunEngine(eng, &ctx, nil, nil)
	res = eng.Result()
	return
}

func assertTracked(t *testing.T, fu *FakeUser, theirName string) {
	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		t.Fatal(err)
	}
	them, err := libkb.LoadUser(libkb.LoadUserArg{Name: theirName})
	if err != nil {
		t.Fatal(err)
	}
	s, err := me.GetTrackingStatementFor(them.GetName(), them.GetUid())
	if err != nil {
		t.Fatal(err)
	}
	if s == nil {
		t.Fatal("expeted a tracking statement; but didn't see one")
	}

}

func trackAlice(t *testing.T, fu *FakeUser) {
	idUI, res, err := runTrack(fu, "t_alice")
	if err != nil {
		t.Fatal(err)
	}
	checkAliceProofs(t, idUI, res)
	assertTracked(t, fu, "t_alice")
	return
}

func trackBob(t *testing.T, fu *FakeUser) {
	idUI, res, err := runTrack(fu, "t_bob")
	if err != nil {
		t.Fatal(err)
	}
	checkBobProofs(t, idUI, res)
	assertTracked(t, fu, "t_bob")
	return
}

func TestTrack(t *testing.T) {
	tc := libkb.SetupTest(t, "track")
	defer tc.Cleanup()
	fu := CreateAndSignupFakeUser(t, "track")

	trackAlice(t, fu)

	// Assert that we gracefully handle the case of no login
	G.LoginState.Logout()
	_, _, err := runTrack(fu, "t_bob")
	if err == nil {
		t.Fatal("expected logout error; got no error")
	} else if _, ok := err.(libkb.LoginRequiredError); !ok {
		t.Fatalf("expected a LoginRequireError; got %s", err.Error())
	}
	fu.LoginOrBust(t)
	trackBob(t, fu)
	return
}
<commit_msg>libkb.SetupTest -> SetupEngineTest<commit_after>package engine

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase_1 "github.com/keybase/client/protocol/go"
)

func runTrack(fu *FakeUser, username string) (idUI *FakeIdentifyUI, res *IdentifyRes, err error) {
	idUI = &FakeIdentifyUI{
		Proofs: make(map[string]string),
		Fapr:   keybase_1.FinishAndPromptRes{TrackRemote: true},
	}
	arg := TrackEngineArg{
		TheirName: username,
	}
	ctx := Context{
		LogUI:    G.UI.GetLogUI(),
		TrackUI:  idUI,
		SecretUI: fu.NewSecretUI(),
	}
	eng := NewTrackEngine(&arg)
	err = RunEngine(eng, &ctx, nil, nil)
	res = eng.Result()
	return
}

func assertTracked(t *testing.T, fu *FakeUser, theirName string) {
	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		t.Fatal(err)
	}
	them, err := libkb.LoadUser(libkb.LoadUserArg{Name: theirName})
	if err != nil {
		t.Fatal(err)
	}
	s, err := me.GetTrackingStatementFor(them.GetName(), them.GetUid())
	if err != nil {
		t.Fatal(err)
	}
	if s == nil {
		t.Fatal("expeted a tracking statement; but didn't see one")
	}

}

func trackAlice(t *testing.T, fu *FakeUser) {
	idUI, res, err := runTrack(fu, "t_alice")
	if err != nil {
		t.Fatal(err)
	}
	checkAliceProofs(t, idUI, res)
	assertTracked(t, fu, "t_alice")
	return
}

func trackBob(t *testing.T, fu *FakeUser) {
	idUI, res, err := runTrack(fu, "t_bob")
	if err != nil {
		t.Fatal(err)
	}
	checkBobProofs(t, idUI, res)
	assertTracked(t, fu, "t_bob")
	return
}

func TestTrack(t *testing.T) {
	tc := SetupEngineTest(t, "track")
	defer tc.Cleanup()
	fu := CreateAndSignupFakeUser(t, "track")

	trackAlice(t, fu)

	// Assert that we gracefully handle the case of no login
	G.LoginState.Logout()
	_, _, err := runTrack(fu, "t_bob")
	if err == nil {
		t.Fatal("expected logout error; got no error")
	} else if _, ok := err.(libkb.LoginRequiredError); !ok {
		t.Fatalf("expected a LoginRequireError; got %s", err.Error())
	}
	fu.LoginOrBust(t)
	trackBob(t, fu)
	return
}
<|endoftext|>


<commit_before>/*
	Utilities for isolating and swarming. See swarming_test.go for usage examples.
*/
package swarming

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"time"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/sklog"
)

const (
	SWARMING_SERVER          = "chromium-swarm.appspot.com"
	SWARMING_SERVER_PRIVATE  = "chrome-swarming.appspot.com"
	LUCI_CLIENT_REPO         = "https://chromium.googlesource.com/infra/luci/client-py"
	RECOMMENDED_IO_TIMEOUT   = 20 * time.Minute
	RECOMMENDED_HARD_TIMEOUT = 1 * time.Hour
	RECOMMENDED_PRIORITY     = 90
	RECOMMENDED_EXPIRATION   = 4 * time.Hour
	// "priority 0 can only be used for terminate request"
	HIGHEST_PRIORITY = 1
	LOWEST_PRIORITY  = 255
)

type SwarmingClient struct {
	WorkDir            string
	isolateClient      *isolate.Client
	isolateServer      string
	SwarmingPy         string
	SwarmingServer     string
	ServiceAccountJSON string
}

type SwarmingTask struct {
	Title          string
	IsolatedHash   string
	OutputDir      string
	Dimensions     map[string]string
	Tags           map[string]string
	CipdPackages   []string
	Priority       int
	Expiration     time.Duration
	Idempotent     bool
	ServiceAccount string
	TaskID         string // Populated after the task is triggered.
}

type ShardOutputFormat struct {
	Output string `json:"output"`
	State  string `json:"state"`
}

type TaskOutputFormat struct {
	Shards []ShardOutputFormat `json:"shards"`
}

func (t *SwarmingTask) Trigger(ctx context.Context, s *SwarmingClient, hardTimeout, ioTimeout time.Duration) error {
	if err := _VerifyBinaryExists(ctx, s.SwarmingPy); err != nil {
		return fmt.Errorf("Could not find swarming binary: %s", err)
	}

	// Run swarming trigger.
	dumpJSON := path.Join(t.OutputDir, fmt.Sprintf("%s-trigger-output.json", t.Title))
	triggerArgs := []string{
		"trigger",
		"--swarming", s.SwarmingServer,
		"--isolate-server", s.isolateServer,
		"--priority", strconv.Itoa(t.Priority),
		"--shards", strconv.Itoa(1),
		"--task-name", t.Title,
		"--dump-json", dumpJSON,
		"--expiration", strconv.FormatFloat(t.Expiration.Seconds(), 'f', 0, 64),
		"--io-timeout", strconv.FormatFloat(ioTimeout.Seconds(), 'f', 0, 64),
		"--hard-timeout", strconv.FormatFloat(hardTimeout.Seconds(), 'f', 0, 64),
		"--verbose",
	}
	if t.ServiceAccount != "" {
		triggerArgs = append(triggerArgs, "--service-account", t.ServiceAccount)
	}
	for k, v := range t.Dimensions {
		triggerArgs = append(triggerArgs, "--dimension", k, v)
	}
	for k, v := range t.Tags {
		triggerArgs = append(triggerArgs, "--tag", fmt.Sprintf("%s:%s", k, v))
	}
	for _, c := range t.CipdPackages {
		triggerArgs = append(triggerArgs, "--cipd-package", c)
	}
	if t.Idempotent {
		triggerArgs = append(triggerArgs, "--idempotent")
	}
	if s.ServiceAccountJSON != "" {
		triggerArgs = append(triggerArgs, "--auth-service-account-json", s.ServiceAccountJSON)
	}
	triggerArgs = append(triggerArgs, "--isolated", t.IsolatedHash)

	err := exec.Run(ctx, &exec.Command{
		Name: s.SwarmingPy,
		Args: triggerArgs,
		// Triggering a task should be immediate. Setting a 15m timeout incase
		// something goes wrong.
		Timeout:   15 * time.Minute,
		LogStdout: true,
		LogStderr: true,
	})
	if err != nil {
		return fmt.Errorf("Swarming trigger for %s failed with: %s", t.Title, err)
	}

	// Read the taskID from the dumpJSON and set it to the task object.
	type Task struct {
		TaskID string `json:"task_id"`
	}
	type Tasks struct {
		Tasks map[string]Task `json:"tasks"`
	}
	var tasks Tasks
	f, err := os.Open(dumpJSON)
	if err != nil {
		return err
	}
	if err := json.NewDecoder(f).Decode(&tasks); err != nil {
		return fmt.Errorf("Could not decode %s: %s", dumpJSON, err)
	}
	t.TaskID = tasks.Tasks[t.Title].TaskID

	return nil
}

// Collect collects the swarming task. It is a blocking call that returns only after the task
// completes. It returns the following:
// * Output of the task.
// * Location of the ${ISOLATED_OUTDIR}.
// * State of the task. Eg: COMPLETED/KILLED.
// * Error is non-nil if something goes wrong. If the command to collect returns a non-zero exit
//   code then error is non-nil but all of the above (output, outdir, state) are also returned if
//   known. This is useful for checking if a task failed because it was cancelled.
func (t *SwarmingTask) Collect(ctx context.Context, s *SwarmingClient, logStdout, logStderr bool) (string, string, string, error) {
	if verifyErr := _VerifyBinaryExists(ctx, s.SwarmingPy); verifyErr != nil {
		return "", "", "", fmt.Errorf("Could not find swarming binary: %s", verifyErr)
	}
	dumpJSON := path.Join(t.OutputDir, fmt.Sprintf("%s-trigger-output.json", t.Title))

	// Run swarming collect.
	collectArgs := []string{
		"collect",
		"--json", dumpJSON,
		"--swarming", s.SwarmingServer,
		"--task-output-dir", t.OutputDir,
		"--verbose",
	}
	if s.ServiceAccountJSON != "" {
		collectArgs = append(collectArgs, "--auth-service-account-json", s.ServiceAccountJSON)
	}
	collectCmdErr := exec.Run(ctx, &exec.Command{
		Name:      s.SwarmingPy,
		Args:      collectArgs,
		Timeout:   t.Expiration,
		LogStdout: logStdout,
		LogStderr: logStderr,
	})

	// Read and parse the summary file if it exists before checking for the error.
	outputSummaryFile := path.Join(t.OutputDir, "summary.json")
	output := ""
	state := ""
	if _, statErr := os.Stat(outputSummaryFile); statErr == nil {

		outputSummary, readErr := ioutil.ReadFile(outputSummaryFile)
		if readErr != nil {
			return "", "", "", fmt.Errorf("Could not read output summary %s: %s", outputSummaryFile, readErr)
		}
		var summaryOutput TaskOutputFormat
		if decodeErr := json.NewDecoder(bytes.NewReader(outputSummary)).Decode(&summaryOutput); decodeErr != nil {
			return "", "", "", fmt.Errorf("Could not decode %s: %s", outputSummaryFile, decodeErr)
		}
		output = summaryOutput.Shards[0].Output
		state = summaryOutput.Shards[0].State
	}

	// Directory that will contain output written to ${ISOLATED_OUTDIR}.
	outputDir := path.Join(t.OutputDir, "0")
	if collectCmdErr != nil {
		return output, outputDir, state, fmt.Errorf("Swarming collect for %s failed with: %s", t.Title, collectCmdErr)
	}
	return output, outputDir, state, nil
}

// NewSwarmingClient returns an instance of Swarming populated with default
// values.
func NewSwarmingClient(ctx context.Context, workDir, swarmingServer, isolateServer, serviceAccountJSON string) (*SwarmingClient, error) {
	if _, err := os.Stat(workDir); os.IsNotExist(err) {
		if err := os.MkdirAll(workDir, 0700); err != nil {
			return nil, fmt.Errorf("Could not create %s: %s", workDir, err)
		}
	}
	// Checkout luci client-py to get access to swarming.py for triggering and
	// collecting tasks.
	luciClient, err := git.NewCheckout(ctx, LUCI_CLIENT_REPO, workDir)
	if err != nil {
		return nil, fmt.Errorf("Could not checkout %s: %s", LUCI_CLIENT_REPO, err)
	}
	if err := luciClient.Update(ctx); err != nil {
		return nil, err
	}
	swarmingPy := path.Join(luciClient.Dir(), "swarming.py")

	// Create an isolate client.
	isolateClient, err := isolate.NewClientWithServiceAccount(workDir, isolateServer, serviceAccountJSON)
	if err != nil {
		return nil, fmt.Errorf("Failed to create isolate client: %s", err)
	}

	return &SwarmingClient{
		WorkDir:            workDir,
		isolateClient:      isolateClient,
		isolateServer:      isolateServer,
		SwarmingPy:         swarmingPy,
		SwarmingServer:     swarmingServer,
		ServiceAccountJSON: serviceAccountJSON,
	}, nil
}

func (s *SwarmingClient) GetIsolateClient() *isolate.Client {
	return s.isolateClient
}

// CreateIsolatedGenJSON creates isolated.gen.json files in the work dir. They then
// can be passed on to BatchArchiveTargets.
func (s *SwarmingClient) CreateIsolatedGenJSON(isolatePath, baseDir, osType, taskName string, extraVars map[string]string, blackList []string) (string, error) {
	// Verify that isolatePath is an absolute path.
	if !path.IsAbs(isolatePath) {
		return "", fmt.Errorf("isolate path %s must be an absolute path", isolatePath)
	}

	isolatedPath := path.Join(s.WorkDir, fmt.Sprintf("%s.isolated", taskName))
	isolatedGenJSONPath := path.Join(s.WorkDir, fmt.Sprintf("%s.isolated.gen.json", taskName))
	t := &isolate.Task{
		BaseDir:     baseDir,
		Blacklist:   blackList,
		ExtraVars:   extraVars,
		IsolateFile: isolatePath,
		OsType:      osType,
	}
	if err := isolate.WriteIsolatedGenJson(t, isolatedGenJSONPath, isolatedPath); err != nil {
		return "", err
	}
	return isolatedGenJSONPath, nil
}

// BatchArchiveTargets batcharchives the specified isolated.gen.json files.
func (s *SwarmingClient) BatchArchiveTargets(ctx context.Context, isolatedGenJSONs []string, d time.Duration) (map[string]string, error) {
	// Run isolate batcharchive.
	dumpJSON := path.Join(s.WorkDir, "isolate-output.json")
	if err := s.isolateClient.BatchArchiveTasks(ctx, isolatedGenJSONs, dumpJSON); err != nil {
		return nil, err
	}

	// Read the isolate hashes from the dump JSON.
	dumpFile, err := ioutil.ReadFile(dumpJSON)
	if err != nil {
		return nil, fmt.Errorf("Could not read JSON output %s: %s", dumpJSON, err)
	}
	var tasksToHashes map[string]string
	if err := json.NewDecoder(bytes.NewReader(dumpFile)).Decode(&tasksToHashes); err != nil {
		return nil, fmt.Errorf("Could not decode %s: %s", dumpJSON, err)
	}

	return tasksToHashes, nil
}

// Trigger swarming using the specified hashes and dimensions.
func (s *SwarmingClient) TriggerSwarmingTasks(ctx context.Context, tasksToHashes, dimensions, tags map[string]string, cipdPackages []string, priority int, expiration, hardTimeout, ioTimeout time.Duration, idempotent, addTaskNameAsTag bool, serviceAccount string) ([]*SwarmingTask, error) {
	tasks := []*SwarmingTask{}

	for taskName, hash := range tasksToHashes {
		taskOutputDir := path.Join(s.WorkDir, taskName)
		if err := os.MkdirAll(taskOutputDir, 0700); err != nil {
			return nil, fmt.Errorf("Could not create %s: %s", taskOutputDir, err)
		}
		taskTags := map[string]string{}
		for k, v := range tags {
			taskTags[k] = v
		}
		if addTaskNameAsTag {
			taskTags["name"] = taskName
		}
		task := &SwarmingTask{
			Title:          taskName,
			IsolatedHash:   hash,
			OutputDir:      taskOutputDir,
			Dimensions:     dimensions,
			Tags:           taskTags,
			CipdPackages:   cipdPackages,
			Priority:       priority,
			Expiration:     expiration,
			Idempotent:     idempotent,
			ServiceAccount: serviceAccount,
		}
		if err := task.Trigger(ctx, s, hardTimeout, ioTimeout); err != nil {
			return nil, fmt.Errorf("Could not trigger task %s: %s", taskName, err)
		}
		sklog.Infof("Triggered the task: %v", task)
		tasks = append(tasks, task)
	}

	return tasks, nil
}

func (s *SwarmingClient) Cleanup() {
	if err := os.RemoveAll(s.WorkDir); err != nil {
		sklog.Errorf("Could not cleanup swarming work dir: %s", err)
	}
}

func _VerifyBinaryExists(ctx context.Context, binary string) error {
	err := exec.Run(ctx, &exec.Command{
		Name:      binary,
		Args:      []string{"help"},
		Timeout:   60 * time.Second,
		LogStdout: false,
		LogStderr: true,
	})
	if err != nil {
		return fmt.Errorf("Error finding the binary %s: %s", binary, err)
	}
	return nil
}
<commit_msg>[CT] Close file descriptor in swarming.go<commit_after>/*
	Utilities for isolating and swarming. See swarming_test.go for usage examples.
*/
package swarming

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"time"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	SWARMING_SERVER          = "chromium-swarm.appspot.com"
	SWARMING_SERVER_PRIVATE  = "chrome-swarming.appspot.com"
	LUCI_CLIENT_REPO         = "https://chromium.googlesource.com/infra/luci/client-py"
	RECOMMENDED_IO_TIMEOUT   = 20 * time.Minute
	RECOMMENDED_HARD_TIMEOUT = 1 * time.Hour
	RECOMMENDED_PRIORITY     = 90
	RECOMMENDED_EXPIRATION   = 4 * time.Hour
	// "priority 0 can only be used for terminate request"
	HIGHEST_PRIORITY = 1
	LOWEST_PRIORITY  = 255
)

type SwarmingClient struct {
	WorkDir            string
	isolateClient      *isolate.Client
	isolateServer      string
	SwarmingPy         string
	SwarmingServer     string
	ServiceAccountJSON string
}

type SwarmingTask struct {
	Title          string
	IsolatedHash   string
	OutputDir      string
	Dimensions     map[string]string
	Tags           map[string]string
	CipdPackages   []string
	Priority       int
	Expiration     time.Duration
	Idempotent     bool
	ServiceAccount string
	TaskID         string // Populated after the task is triggered.
}

type ShardOutputFormat struct {
	Output string `json:"output"`
	State  string `json:"state"`
}

type TaskOutputFormat struct {
	Shards []ShardOutputFormat `json:"shards"`
}

func (t *SwarmingTask) Trigger(ctx context.Context, s *SwarmingClient, hardTimeout, ioTimeout time.Duration) error {
	if err := _VerifyBinaryExists(ctx, s.SwarmingPy); err != nil {
		return fmt.Errorf("Could not find swarming binary: %s", err)
	}

	// Run swarming trigger.
	dumpJSON := path.Join(t.OutputDir, fmt.Sprintf("%s-trigger-output.json", t.Title))
	triggerArgs := []string{
		"trigger",
		"--swarming", s.SwarmingServer,
		"--isolate-server", s.isolateServer,
		"--priority", strconv.Itoa(t.Priority),
		"--shards", strconv.Itoa(1),
		"--task-name", t.Title,
		"--dump-json", dumpJSON,
		"--expiration", strconv.FormatFloat(t.Expiration.Seconds(), 'f', 0, 64),
		"--io-timeout", strconv.FormatFloat(ioTimeout.Seconds(), 'f', 0, 64),
		"--hard-timeout", strconv.FormatFloat(hardTimeout.Seconds(), 'f', 0, 64),
		"--verbose",
	}
	if t.ServiceAccount != "" {
		triggerArgs = append(triggerArgs, "--service-account", t.ServiceAccount)
	}
	for k, v := range t.Dimensions {
		triggerArgs = append(triggerArgs, "--dimension", k, v)
	}
	for k, v := range t.Tags {
		triggerArgs = append(triggerArgs, "--tag", fmt.Sprintf("%s:%s", k, v))
	}
	for _, c := range t.CipdPackages {
		triggerArgs = append(triggerArgs, "--cipd-package", c)
	}
	if t.Idempotent {
		triggerArgs = append(triggerArgs, "--idempotent")
	}
	if s.ServiceAccountJSON != "" {
		triggerArgs = append(triggerArgs, "--auth-service-account-json", s.ServiceAccountJSON)
	}
	triggerArgs = append(triggerArgs, "--isolated", t.IsolatedHash)

	err := exec.Run(ctx, &exec.Command{
		Name: s.SwarmingPy,
		Args: triggerArgs,
		// Triggering a task should be immediate. Setting a 15m timeout incase
		// something goes wrong.
		Timeout:   15 * time.Minute,
		LogStdout: true,
		LogStderr: true,
	})
	if err != nil {
		return fmt.Errorf("Swarming trigger for %s failed with: %s", t.Title, err)
	}

	// Read the taskID from the dumpJSON and set it to the task object.
	type Task struct {
		TaskID string `json:"task_id"`
	}
	type Tasks struct {
		Tasks map[string]Task `json:"tasks"`
	}
	var tasks Tasks
	f, err := os.Open(dumpJSON)
	if err != nil {
		return err
	}
	defer util.Close(f)
	if err := json.NewDecoder(f).Decode(&tasks); err != nil {
		return fmt.Errorf("Could not decode %s: %s", dumpJSON, err)
	}
	t.TaskID = tasks.Tasks[t.Title].TaskID

	return nil
}

// Collect collects the swarming task. It is a blocking call that returns only after the task
// completes. It returns the following:
// * Output of the task.
// * Location of the ${ISOLATED_OUTDIR}.
// * State of the task. Eg: COMPLETED/KILLED.
// * Error is non-nil if something goes wrong. If the command to collect returns a non-zero exit
//   code then error is non-nil but all of the above (output, outdir, state) are also returned if
//   known. This is useful for checking if a task failed because it was cancelled.
func (t *SwarmingTask) Collect(ctx context.Context, s *SwarmingClient, logStdout, logStderr bool) (string, string, string, error) {
	if verifyErr := _VerifyBinaryExists(ctx, s.SwarmingPy); verifyErr != nil {
		return "", "", "", fmt.Errorf("Could not find swarming binary: %s", verifyErr)
	}
	dumpJSON := path.Join(t.OutputDir, fmt.Sprintf("%s-trigger-output.json", t.Title))

	// Run swarming collect.
	collectArgs := []string{
		"collect",
		"--json", dumpJSON,
		"--swarming", s.SwarmingServer,
		"--task-output-dir", t.OutputDir,
		"--verbose",
	}
	if s.ServiceAccountJSON != "" {
		collectArgs = append(collectArgs, "--auth-service-account-json", s.ServiceAccountJSON)
	}
	collectCmdErr := exec.Run(ctx, &exec.Command{
		Name:      s.SwarmingPy,
		Args:      collectArgs,
		Timeout:   t.Expiration,
		LogStdout: logStdout,
		LogStderr: logStderr,
	})

	// Read and parse the summary file if it exists before checking for the error.
	outputSummaryFile := path.Join(t.OutputDir, "summary.json")
	output := ""
	state := ""
	if _, statErr := os.Stat(outputSummaryFile); statErr == nil {

		outputSummary, readErr := ioutil.ReadFile(outputSummaryFile)
		if readErr != nil {
			return "", "", "", fmt.Errorf("Could not read output summary %s: %s", outputSummaryFile, readErr)
		}
		var summaryOutput TaskOutputFormat
		if decodeErr := json.NewDecoder(bytes.NewReader(outputSummary)).Decode(&summaryOutput); decodeErr != nil {
			return "", "", "", fmt.Errorf("Could not decode %s: %s", outputSummaryFile, decodeErr)
		}
		output = summaryOutput.Shards[0].Output
		state = summaryOutput.Shards[0].State
	}

	// Directory that will contain output written to ${ISOLATED_OUTDIR}.
	outputDir := path.Join(t.OutputDir, "0")
	if collectCmdErr != nil {
		return output, outputDir, state, fmt.Errorf("Swarming collect for %s failed with: %s", t.Title, collectCmdErr)
	}
	return output, outputDir, state, nil
}

// NewSwarmingClient returns an instance of Swarming populated with default
// values.
func NewSwarmingClient(ctx context.Context, workDir, swarmingServer, isolateServer, serviceAccountJSON string) (*SwarmingClient, error) {
	if _, err := os.Stat(workDir); os.IsNotExist(err) {
		if err := os.MkdirAll(workDir, 0700); err != nil {
			return nil, fmt.Errorf("Could not create %s: %s", workDir, err)
		}
	}
	// Checkout luci client-py to get access to swarming.py for triggering and
	// collecting tasks.
	luciClient, err := git.NewCheckout(ctx, LUCI_CLIENT_REPO, workDir)
	if err != nil {
		return nil, fmt.Errorf("Could not checkout %s: %s", LUCI_CLIENT_REPO, err)
	}
	if err := luciClient.Update(ctx); err != nil {
		return nil, err
	}
	swarmingPy := path.Join(luciClient.Dir(), "swarming.py")

	// Create an isolate client.
	isolateClient, err := isolate.NewClientWithServiceAccount(workDir, isolateServer, serviceAccountJSON)
	if err != nil {
		return nil, fmt.Errorf("Failed to create isolate client: %s", err)
	}

	return &SwarmingClient{
		WorkDir:            workDir,
		isolateClient:      isolateClient,
		isolateServer:      isolateServer,
		SwarmingPy:         swarmingPy,
		SwarmingServer:     swarmingServer,
		ServiceAccountJSON: serviceAccountJSON,
	}, nil
}

func (s *SwarmingClient) GetIsolateClient() *isolate.Client {
	return s.isolateClient
}

// CreateIsolatedGenJSON creates isolated.gen.json files in the work dir. They then
// can be passed on to BatchArchiveTargets.
func (s *SwarmingClient) CreateIsolatedGenJSON(isolatePath, baseDir, osType, taskName string, extraVars map[string]string, blackList []string) (string, error) {
	// Verify that isolatePath is an absolute path.
	if !path.IsAbs(isolatePath) {
		return "", fmt.Errorf("isolate path %s must be an absolute path", isolatePath)
	}

	isolatedPath := path.Join(s.WorkDir, fmt.Sprintf("%s.isolated", taskName))
	isolatedGenJSONPath := path.Join(s.WorkDir, fmt.Sprintf("%s.isolated.gen.json", taskName))
	t := &isolate.Task{
		BaseDir:     baseDir,
		Blacklist:   blackList,
		ExtraVars:   extraVars,
		IsolateFile: isolatePath,
		OsType:      osType,
	}
	if err := isolate.WriteIsolatedGenJson(t, isolatedGenJSONPath, isolatedPath); err != nil {
		return "", err
	}
	return isolatedGenJSONPath, nil
}

// BatchArchiveTargets batcharchives the specified isolated.gen.json files.
func (s *SwarmingClient) BatchArchiveTargets(ctx context.Context, isolatedGenJSONs []string, d time.Duration) (map[string]string, error) {
	// Run isolate batcharchive.
	dumpJSON := path.Join(s.WorkDir, "isolate-output.json")
	if err := s.isolateClient.BatchArchiveTasks(ctx, isolatedGenJSONs, dumpJSON); err != nil {
		return nil, err
	}

	// Read the isolate hashes from the dump JSON.
	dumpFile, err := ioutil.ReadFile(dumpJSON)
	if err != nil {
		return nil, fmt.Errorf("Could not read JSON output %s: %s", dumpJSON, err)
	}
	var tasksToHashes map[string]string
	if err := json.NewDecoder(bytes.NewReader(dumpFile)).Decode(&tasksToHashes); err != nil {
		return nil, fmt.Errorf("Could not decode %s: %s", dumpJSON, err)
	}

	return tasksToHashes, nil
}

// Trigger swarming using the specified hashes and dimensions.
func (s *SwarmingClient) TriggerSwarmingTasks(ctx context.Context, tasksToHashes, dimensions, tags map[string]string, cipdPackages []string, priority int, expiration, hardTimeout, ioTimeout time.Duration, idempotent, addTaskNameAsTag bool, serviceAccount string) ([]*SwarmingTask, error) {
	tasks := []*SwarmingTask{}

	for taskName, hash := range tasksToHashes {
		taskOutputDir := path.Join(s.WorkDir, taskName)
		if err := os.MkdirAll(taskOutputDir, 0700); err != nil {
			return nil, fmt.Errorf("Could not create %s: %s", taskOutputDir, err)
		}
		taskTags := map[string]string{}
		for k, v := range tags {
			taskTags[k] = v
		}
		if addTaskNameAsTag {
			taskTags["name"] = taskName
		}
		task := &SwarmingTask{
			Title:          taskName,
			IsolatedHash:   hash,
			OutputDir:      taskOutputDir,
			Dimensions:     dimensions,
			Tags:           taskTags,
			CipdPackages:   cipdPackages,
			Priority:       priority,
			Expiration:     expiration,
			Idempotent:     idempotent,
			ServiceAccount: serviceAccount,
		}
		if err := task.Trigger(ctx, s, hardTimeout, ioTimeout); err != nil {
			return nil, fmt.Errorf("Could not trigger task %s: %s", taskName, err)
		}
		sklog.Infof("Triggered the task: %v", task)
		tasks = append(tasks, task)
	}

	return tasks, nil
}

func (s *SwarmingClient) Cleanup() {
	if err := os.RemoveAll(s.WorkDir); err != nil {
		sklog.Errorf("Could not cleanup swarming work dir: %s", err)
	}
}

func _VerifyBinaryExists(ctx context.Context, binary string) error {
	err := exec.Run(ctx, &exec.Command{
		Name:      binary,
		Args:      []string{"help"},
		Timeout:   60 * time.Second,
		LogStdout: false,
		LogStderr: true,
	})
	if err != nil {
		return fmt.Errorf("Error finding the binary %s: %s", binary, err)
	}
	return nil
}
<|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.
*/

/*
  Generate my.cnf files from templates.
*/

package mysqlctl

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"strconv"
)

// Mycnf is a memory structure that contains a bunch of interesting
// parameters to start mysqld. It can be used to generate standard
// my.cnf files from a server id and mysql port. It can also be
// populated from an existing my.cnf, or by command line parameters.
type Mycnf struct {
	// ServerID is the unique id for this server.
	// Used to create a bunch of named directories.
	ServerID uint32

	// MysqlPort is the port for the MySQL server running on this machine.
	// It is mainly used to communicate with topology server.
	MysqlPort int32

	// DataDir is where the table files are
	// (used by vt software for Clone)
	DataDir string

	// InnodbDataHomeDir is the data directory for innodb.
	// (used by vt software for Clone)
	InnodbDataHomeDir string

	// InnodbLogGroupHomeDir is the logs directory for innodb.
	// (used by vt software for Clone)
	InnodbLogGroupHomeDir string

	// SocketFile is the path to the local mysql.sock file.
	// (used by vt software to check server is running)
	SocketFile string

	// GeneralLogPath is the path to store general logs at,
	// if general-log is enabled.
	// (unused by vt software for now)
	GeneralLogPath string

	// ErrorLogPath is the path to store error logs at.
	// (unused by vt software for now)
	ErrorLogPath string

	// SlowLogPath is the slow query log path
	// (unused by vt software for now)
	SlowLogPath string

	// RelayLogPath is the path of the relay logs
	// (unused by vt software for now)
	RelayLogPath string

	// RelayLogIndexPath is the file name for the relay log index
	// (unused by vt software for now)
	RelayLogIndexPath string

	// RelayLogInfoPath is the file name for the relay log info file
	// (unused by vt software for now)
	RelayLogInfoPath string

	// BinLogPath is the base path for binlogs
	// (used by vt software for binlog streaming)
	BinLogPath string

	// MasterInfoFile is the master.info file location.
	// (unused by vt software for now)
	MasterInfoFile string

	// PidFile is the mysql.pid file location
	// (used by vt software to check server is running)
	PidFile string

	// TmpDir is where to create temporary tables
	// (unused by vt software for now)
	TmpDir string

	// SlaveLoadTmpDir is where to create tmp files for replication
	// (unused by vt software for now)
	SlaveLoadTmpDir string

	mycnfMap map[string]string
	path     string // the actual path that represents this mycnf
}

func (cnf *Mycnf) lookup(key string) string {
	key = normKey([]byte(key))
	return cnf.mycnfMap[key]
}

func (cnf *Mycnf) lookupWithDefault(key, defaultVal string) string {
	val := cnf.lookup(key)
	if val == "" {
		if defaultVal == "" {
			panic(fmt.Errorf("Value for key '%v' not set and no default value set", key))
		}
		return defaultVal
	}
	return val
}

func (cnf *Mycnf) lookupAndCheck(key string) string {
	return cnf.lookupWithDefault(key, "")
}

func normKey(bkey []byte) string {
	// FIXME(msolomon) People are careless about hyphen vs underscore - we should normalize.
	// But you have to normalize to hyphen, or mysqld_safe can fail.
	return string(bytes.Replace(bytes.TrimSpace(bkey), []byte("_"), []byte("-"), -1))
}

// ReadMycnf will read an existing my.cnf from disk, and update the passed in Mycnf object
// with values from the my.cnf on disk.
func ReadMycnf(mycnf *Mycnf) (*Mycnf, error) {
	var err error
	defer func(err *error) {
		if x := recover(); x != nil {
			*err = x.(error)
		}
	}(&err)

	f, err := os.Open(mycnf.path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buf := bufio.NewReader(f)
	if err != nil {
		return nil, err
	}
	mycnf.mycnfMap = make(map[string]string)
	var lval, rval string
	var parts [][]byte

	for {
		line, _, err := buf.ReadLine()
		if err == io.EOF {
			break
		}
		line = bytes.TrimSpace(line)

		parts = bytes.Split(line, []byte("="))
		if len(parts) < 2 {
			continue
		}
		lval = normKey(parts[0])
		rval = string(bytes.TrimSpace(parts[1]))
		mycnf.mycnfMap[lval] = rval
	}

	serverIDStr := mycnf.lookupAndCheck("server-id")
	serverID, err := strconv.Atoi(serverIDStr)
	if err != nil {
		return nil, fmt.Errorf("Failed to convert server-id %v", err)
	}
	mycnf.ServerID = uint32(serverID)

	portStr := mycnf.lookupAndCheck("port")
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("Failed: failed to convert port %v", err)
	}

	mycnf.MysqlPort = int32(port)
	mycnf.DataDir = mycnf.lookupWithDefault("datadir", mycnf.DataDir)
	mycnf.InnodbDataHomeDir = mycnf.lookupWithDefault("innodb_data_home_dir", mycnf.InnodbDataHomeDir)
	mycnf.InnodbLogGroupHomeDir = mycnf.lookupWithDefault("innodb_log_group_home_dir", mycnf.InnodbLogGroupHomeDir)
	mycnf.SocketFile = mycnf.lookupWithDefault("socket", mycnf.SocketFile)
	mycnf.GeneralLogPath = mycnf.lookupWithDefault("general_log_file", mycnf.GeneralLogPath)
	mycnf.ErrorLogPath = mycnf.lookupWithDefault("log-error", mycnf.ErrorLogPath)
	mycnf.SlowLogPath = mycnf.lookupWithDefault("slow-query-log-file", mycnf.SlowLogPath)
	mycnf.RelayLogPath = mycnf.lookupWithDefault("relay-log", mycnf.RelayLogPath)
	mycnf.RelayLogIndexPath = mycnf.lookupWithDefault("relay-log-index", mycnf.RelayLogIndexPath)
	mycnf.RelayLogInfoPath = mycnf.lookupWithDefault("relay-log-info-file", mycnf.RelayLogInfoPath)
	mycnf.BinLogPath = mycnf.lookupWithDefault("log-bin", mycnf.BinLogPath)
	mycnf.MasterInfoFile = mycnf.lookupWithDefault("master-info-file", mycnf.MasterInfoFile)
	mycnf.PidFile = mycnf.lookupWithDefault("pid-file", mycnf.PidFile)
	mycnf.TmpDir = mycnf.lookupWithDefault("tmpdir", mycnf.TmpDir)
	mycnf.SlaveLoadTmpDir = mycnf.lookupWithDefault("slave_load_tmpdir", mycnf.SlaveLoadTmpDir)

	return mycnf, nil
}
<commit_msg>bug: fix mycnf error propagation<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.
*/

/*
  Generate my.cnf files from templates.
*/

package mysqlctl

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"strconv"
)

// Mycnf is a memory structure that contains a bunch of interesting
// parameters to start mysqld. It can be used to generate standard
// my.cnf files from a server id and mysql port. It can also be
// populated from an existing my.cnf, or by command line parameters.
type Mycnf struct {
	// ServerID is the unique id for this server.
	// Used to create a bunch of named directories.
	ServerID uint32

	// MysqlPort is the port for the MySQL server running on this machine.
	// It is mainly used to communicate with topology server.
	MysqlPort int32

	// DataDir is where the table files are
	// (used by vt software for Clone)
	DataDir string

	// InnodbDataHomeDir is the data directory for innodb.
	// (used by vt software for Clone)
	InnodbDataHomeDir string

	// InnodbLogGroupHomeDir is the logs directory for innodb.
	// (used by vt software for Clone)
	InnodbLogGroupHomeDir string

	// SocketFile is the path to the local mysql.sock file.
	// (used by vt software to check server is running)
	SocketFile string

	// GeneralLogPath is the path to store general logs at,
	// if general-log is enabled.
	// (unused by vt software for now)
	GeneralLogPath string

	// ErrorLogPath is the path to store error logs at.
	// (unused by vt software for now)
	ErrorLogPath string

	// SlowLogPath is the slow query log path
	// (unused by vt software for now)
	SlowLogPath string

	// RelayLogPath is the path of the relay logs
	// (unused by vt software for now)
	RelayLogPath string

	// RelayLogIndexPath is the file name for the relay log index
	// (unused by vt software for now)
	RelayLogIndexPath string

	// RelayLogInfoPath is the file name for the relay log info file
	// (unused by vt software for now)
	RelayLogInfoPath string

	// BinLogPath is the base path for binlogs
	// (used by vt software for binlog streaming)
	BinLogPath string

	// MasterInfoFile is the master.info file location.
	// (unused by vt software for now)
	MasterInfoFile string

	// PidFile is the mysql.pid file location
	// (used by vt software to check server is running)
	PidFile string

	// TmpDir is where to create temporary tables
	// (unused by vt software for now)
	TmpDir string

	// SlaveLoadTmpDir is where to create tmp files for replication
	// (unused by vt software for now)
	SlaveLoadTmpDir string

	mycnfMap map[string]string
	path     string // the actual path that represents this mycnf
}

func (cnf *Mycnf) lookup(key string) string {
	key = normKey([]byte(key))
	return cnf.mycnfMap[key]
}

func (cnf *Mycnf) lookupWithDefault(key, defaultVal string) (string, error) {
	val := cnf.lookup(key)
	if val == "" {
		if defaultVal == "" {
			return "", fmt.Errorf("value for key '%v' not set and no default value set", key)
		}
		return defaultVal, nil
	}
	return val, nil
}

func (cnf *Mycnf) lookupInt(key string) (int, error) {
	val, err := cnf.lookupWithDefault(key, "")
	if err != nil {
		return 0, err
	}
	ival, err := strconv.Atoi(val)
	if err != nil {
		return 0, fmt.Errorf("failed to convert %s: %v", key, err)
	}
	return ival, nil
}

func normKey(bkey []byte) string {
	// FIXME(msolomon) People are careless about hyphen vs underscore - we should normalize.
	// But you have to normalize to hyphen, or mysqld_safe can fail.
	return string(bytes.Replace(bytes.TrimSpace(bkey), []byte("_"), []byte("-"), -1))
}

// ReadMycnf will read an existing my.cnf from disk, and update the passed in Mycnf object
// with values from the my.cnf on disk.
func ReadMycnf(mycnf *Mycnf) (*Mycnf, error) {
	f, err := os.Open(mycnf.path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buf := bufio.NewReader(f)
	if err != nil {
		return nil, err
	}
	mycnf.mycnfMap = make(map[string]string)
	var lval, rval string
	var parts [][]byte

	for {
		line, _, err := buf.ReadLine()
		if err == io.EOF {
			break
		}
		line = bytes.TrimSpace(line)

		parts = bytes.Split(line, []byte("="))
		if len(parts) < 2 {
			continue
		}
		lval = normKey(parts[0])
		rval = string(bytes.TrimSpace(parts[1]))
		mycnf.mycnfMap[lval] = rval
	}

	serverID, err := mycnf.lookupInt("server-id")
	if err != nil {
		return nil, err
	}
	mycnf.ServerID = uint32(serverID)

	port, err := mycnf.lookupInt("port")
	if err != nil {
		return nil, err
	}
	mycnf.MysqlPort = int32(port)

	mapping := map[string]*string{
		"datadir":                   &mycnf.DataDir,
		"innodb_data_home_dir":      &mycnf.InnodbDataHomeDir,
		"innodb_log_group_home_dir": &mycnf.InnodbLogGroupHomeDir,
		"socket":                    &mycnf.SocketFile,
		"general_log_file":          &mycnf.GeneralLogPath,
		"log-error":                 &mycnf.ErrorLogPath,
		"slow-query-log-file":       &mycnf.SlowLogPath,
		"relay-log":                 &mycnf.RelayLogPath,
		"relay-log-index":           &mycnf.RelayLogIndexPath,
		"relay-log-info-file":       &mycnf.RelayLogInfoPath,
		"log-bin":                   &mycnf.BinLogPath,
		"master-info-file":          &mycnf.MasterInfoFile,
		"pid-file":                  &mycnf.PidFile,
		"tmpdir":                    &mycnf.TmpDir,
		"slave_load_tmpdir":         &mycnf.SlaveLoadTmpDir,
	}
	for key, member := range mapping {
		val, err := mycnf.lookupWithDefault(key, *member)
		if err != nil {
			return nil, err
		}
		*member = val
	}

	return mycnf, nil
}
<|endoftext|>


<commit_before>package utils

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"text/template"
	"time"

	"github.com/ThomasRooney/gexpect"
)

// Deis points to the CLI used to run tests.
var Deis = os.Getenv("DEIS_BINARY") + " "

func init() {
	if Deis == " " {
		Deis = "deis "
	}
}

// DeisTestConfig allows tests to be repeated against different
// targets, with different example apps, using specific credentials, and so on.
type DeisTestConfig struct {
	AuthKey            string
	Hosts              string
	Domain             string
	SSHKey             string
	ClusterName        string
	UserName           string
	Password           string
	Email              string
	ExampleApp         string
	AppDomain          string
	AppName            string
	ProcessNum         string
	ImageID            string
	Version            string
	AppUser            string
	SSLCertificatePath string
	SSLKeyPath         string
}

// randomApp is used for the test run if DEIS_TEST_APP isn't set
var randomApp = GetRandomApp()

// GetGlobalConfig returns a test configuration object.
func GetGlobalConfig() *DeisTestConfig {
	authKey := os.Getenv("DEIS_TEST_AUTH_KEY")
	if authKey == "" {
		authKey = "deis"
	}
	hosts := os.Getenv("DEIS_TEST_HOSTS")
	if hosts == "" {
		hosts = "172.17.8.100"
	}
	domain := os.Getenv("DEIS_TEST_DOMAIN")
	if domain == "" {
		domain = "local3.deisapp.com"
	}
	sshKey := os.Getenv("DEIS_TEST_SSH_KEY")
	if sshKey == "" {
		sshKey = "~/.vagrant.d/insecure_private_key"
	}
	exampleApp := os.Getenv("DEIS_TEST_APP")
	if exampleApp == "" {
		exampleApp = randomApp
	}
	appDomain := os.Getenv("DEIS_TEST_APP_DOMAIN")
	if appDomain == "" {
		appDomain = fmt.Sprintf("test.%s", domain)
	}

	// generate a self-signed certifcate for the app domain
	keyOut, err := filepath.Abs(appDomain + ".key")
	if err != nil {
		log.Fatal(err)
	}
	certOut, err := filepath.Abs(appDomain + ".cert")
	if err != nil {
		log.Fatal(err)
	}
	cmd := exec.Command("openssl", "req", "-new", "-newkey", "rsa:4096", "-nodes", "-x509",
		"-days", "1",
		"-subj", fmt.Sprintf("/C=US/ST=Colorado/L=Boulder/CN=%s", appDomain),
		"-keyout", keyOut,
		"-out", certOut)
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatal(err)
	}

	var envCfg = DeisTestConfig{
		AuthKey:            authKey,
		Hosts:              hosts,
		Domain:             domain,
		SSHKey:             sshKey,
		ClusterName:        "dev",
		UserName:           "test",
		Password:           "asdf1234",
		Email:              "test@test.co.nz",
		ExampleApp:         exampleApp,
		AppDomain:          appDomain,
		AppName:            "sample",
		ProcessNum:         "2",
		ImageID:            "buildtest",
		Version:            "2",
		AppUser:            "test1",
		SSLCertificatePath: certOut,
		SSLKeyPath:         keyOut,
	}
	return &envCfg
}

// HTTPClient returns a client for use with the integration tests.
func HTTPClient() *http.Client {
	// disable security check for self-signed certificates
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	return &http.Client{Transport: tr}
}

func doCurl(url string) ([]byte, error) {
	client := HTTPClient()
	response, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)

	if !strings.Contains(string(body), "Powered by") {
		return nil, fmt.Errorf("App not started (%d)\nBody: (%s)", response.StatusCode, string(body))
	}

	return body, nil
}

// Curl connects to an endpoint to see if the endpoint is responding.
func Curl(t *testing.T, url string) {
	CurlWithFail(t, url, false, "")
}

// CurlApp is a convenience function to see if the example app is running.
func CurlApp(t *testing.T, cfg DeisTestConfig) {
	CurlWithFail(t, fmt.Sprintf("http://%s.%s", cfg.AppName, cfg.Domain), false, "")
}

// CurlWithFail connects to a Deis endpoint to see if the example app is running.
func CurlWithFail(t *testing.T, url string, failFlag bool, expect string) {
	// FIXME: try the curl a few times
	for i := 0; i < 20; i++ {
		body, err := doCurl(url)
		if err == nil {
			fmt.Println(string(body))
			return
		}
		time.Sleep(1 * time.Second)
	}

	// once more to fail with an error
	body, err := doCurl(url)

	switch failFlag {
	case true:
		if err != nil {
			if strings.Contains(string(err.Error()), expect) {
				fmt.Println("(Error expected...ok) " + expect)
			} else {
				t.Fatal(err)
			}
		} else {
			if strings.Contains(string(body), expect) {
				fmt.Println("(Error expected...ok) " + expect)
			} else {
				t.Fatal(err)
			}
		}
	case false:
		if err != nil {
			t.Fatal(err)
		} else {
			fmt.Println(string(body))
		}
	}
}

// AuthPasswd tests whether `deis auth:passwd` updates a user's password.
func AuthPasswd(t *testing.T, params *DeisTestConfig, password string) {
	fmt.Println("deis auth:passwd")
	child, err := gexpect.Spawn(Deis + " auth:passwd")
	if err != nil {
		t.Fatalf("command not started\n%v", err)
	}
	fmt.Println("current password:")
	err = child.Expect("current password: ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(params.Password)
	fmt.Println("new password:")
	err = child.Expect("new password: ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(password)
	fmt.Println("new password (confirm):")
	err = child.Expect("new password (confirm): ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(password)
	err = child.Expect("Password change succeeded")
	if err != nil {
		t.Fatalf("command executiuon failed\n%v", err)
	}
	child.Close()
}

// CheckList executes a command and optionally tests whether its output does
// or does not contain a given string.
func CheckList(
	t *testing.T, cmd string, params interface{}, contain string, notflag bool) {
	var cmdBuf bytes.Buffer
	tmpl := template.Must(template.New("cmd").Parse(cmd))
	if err := tmpl.Execute(&cmdBuf, params); err != nil {
		t.Fatal(err)
	}
	cmdString := cmdBuf.String()
	fmt.Println(cmdString)
	var cmdl *exec.Cmd
	if strings.Contains(cmd, "cat") {
		cmdl = exec.Command("sh", "-c", cmdString)
	} else {
		cmdl = exec.Command("sh", "-c", Deis+cmdString)
	}
	stdout, _, err := RunCommandWithStdoutStderr(cmdl)
	if err != nil {
		t.Fatal(err)
	}
	if notflag && strings.Contains(stdout.String(), contain) {
		t.Fatalf("Didn't expect '%s' in command output:\n%v", contain, stdout)
	}
	if !notflag && !strings.Contains(stdout.String(), contain) {
		t.Fatalf("Expected '%s' in command output:\n%v", contain, stdout)
	}
}

// Execute takes command string and parameters required to execute the command,
// a failflag to check whether the command is expected to fail, and an expect
// string to check whether the command has failed according to failflag.
//
// If failflag is true and the command failed, check the stdout and stderr for
// the expect string.
func Execute(t *testing.T, cmd string, params interface{}, failFlag bool, expect string) {
	var cmdBuf bytes.Buffer
	tmpl := template.Must(template.New("cmd").Parse(cmd))
	if err := tmpl.Execute(&cmdBuf, params); err != nil {
		t.Fatal(err)
	}
	cmdString := cmdBuf.String()
	fmt.Println(cmdString)
	var cmdl *exec.Cmd
	if strings.Contains(cmd, "git ") {
		cmdl = exec.Command("sh", "-c", cmdString)
	} else {
		cmdl = exec.Command("sh", "-c", Deis+cmdString)
	}

	switch failFlag {
	case true:
		if stdout, stderr, err := RunCommandWithStdoutStderr(cmdl); err != nil {
			if strings.Contains(stdout.String(), expect) || strings.Contains(stderr.String(), expect) {
				fmt.Println("(Error expected...ok)")
			} else {
				t.Fatal(err)
			}
		} else {
			if strings.Contains(stdout.String(), expect) || strings.Contains(stderr.String(), expect) {
				fmt.Println("(Error expected...ok)" + expect)
			} else {
				t.Fatal(err)
			}
		}
	case false:
		if _, _, err := RunCommandWithStdoutStderr(cmdl); err != nil {
			t.Fatal(err)
		} else {
			fmt.Println("ok")
		}
	}
}

// AppsDestroyTest destroys a Deis app and checks that it was successful.
func AppsDestroyTest(t *testing.T, params *DeisTestConfig) {
	fmt.Printf("destroying app %s...\n", params.ExampleApp)
	cmd := "apps:destroy --app={{.AppName}} --confirm={{.AppName}}"
	if err := Chdir(params.ExampleApp); err != nil {
		t.Fatal(err)
	}
	Execute(t, cmd, params, false, "")
	if err := Chdir(".."); err != nil {
		t.Fatal(err)
	}
	if err := Rmdir(params.ExampleApp); err != nil {
		t.Fatal(err)
	}
}

// GetRandomApp returns a known working example app at random for testing.
func GetRandomApp() string {
	rand.Seed(int64(time.Now().Unix()))
	apps := []string{
		"example-clojure-ring",
		// "example-dart",
		"example-dockerfile-python",
		"example-go",
		"example-java-jetty",
		"example-nodejs-express",
		// "example-php",
		"example-play",
		"example-python-django",
		"example-python-flask",
		"example-ruby-sinatra",
		"example-scala",
		"example-dockerfile-http",
	}
	return apps[rand.Intn(len(apps))]
}
<commit_msg>fix(tests): abort if warning is found in output<commit_after>package utils

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"text/template"
	"time"

	"github.com/ThomasRooney/gexpect"
)

// Deis points to the CLI used to run tests.
var Deis = os.Getenv("DEIS_BINARY") + " "

func init() {
	if Deis == " " {
		Deis = "deis "
	}
}

// DeisTestConfig allows tests to be repeated against different
// targets, with different example apps, using specific credentials, and so on.
type DeisTestConfig struct {
	AuthKey            string
	Hosts              string
	Domain             string
	SSHKey             string
	ClusterName        string
	UserName           string
	Password           string
	Email              string
	ExampleApp         string
	AppDomain          string
	AppName            string
	ProcessNum         string
	ImageID            string
	Version            string
	AppUser            string
	SSLCertificatePath string
	SSLKeyPath         string
}

// randomApp is used for the test run if DEIS_TEST_APP isn't set
var randomApp = GetRandomApp()

// GetGlobalConfig returns a test configuration object.
func GetGlobalConfig() *DeisTestConfig {
	authKey := os.Getenv("DEIS_TEST_AUTH_KEY")
	if authKey == "" {
		authKey = "deis"
	}
	hosts := os.Getenv("DEIS_TEST_HOSTS")
	if hosts == "" {
		hosts = "172.17.8.100"
	}
	domain := os.Getenv("DEIS_TEST_DOMAIN")
	if domain == "" {
		domain = "local3.deisapp.com"
	}
	sshKey := os.Getenv("DEIS_TEST_SSH_KEY")
	if sshKey == "" {
		sshKey = "~/.vagrant.d/insecure_private_key"
	}
	exampleApp := os.Getenv("DEIS_TEST_APP")
	if exampleApp == "" {
		exampleApp = randomApp
	}
	appDomain := os.Getenv("DEIS_TEST_APP_DOMAIN")
	if appDomain == "" {
		appDomain = fmt.Sprintf("test.%s", domain)
	}

	// generate a self-signed certifcate for the app domain
	keyOut, err := filepath.Abs(appDomain + ".key")
	if err != nil {
		log.Fatal(err)
	}
	certOut, err := filepath.Abs(appDomain + ".cert")
	if err != nil {
		log.Fatal(err)
	}
	cmd := exec.Command("openssl", "req", "-new", "-newkey", "rsa:4096", "-nodes", "-x509",
		"-days", "1",
		"-subj", fmt.Sprintf("/C=US/ST=Colorado/L=Boulder/CN=%s", appDomain),
		"-keyout", keyOut,
		"-out", certOut)
	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}
	if err := cmd.Wait(); err != nil {
		log.Fatal(err)
	}

	var envCfg = DeisTestConfig{
		AuthKey:            authKey,
		Hosts:              hosts,
		Domain:             domain,
		SSHKey:             sshKey,
		ClusterName:        "dev",
		UserName:           "test",
		Password:           "asdf1234",
		Email:              "test@test.co.nz",
		ExampleApp:         exampleApp,
		AppDomain:          appDomain,
		AppName:            "sample",
		ProcessNum:         "2",
		ImageID:            "buildtest",
		Version:            "2",
		AppUser:            "test1",
		SSLCertificatePath: certOut,
		SSLKeyPath:         keyOut,
	}
	return &envCfg
}

// HTTPClient returns a client for use with the integration tests.
func HTTPClient() *http.Client {
	// disable security check for self-signed certificates
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	return &http.Client{Transport: tr}
}

func doCurl(url string) ([]byte, error) {
	client := HTTPClient()
	response, err := client.Get(url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)

	if !strings.Contains(string(body), "Powered by") {
		return nil, fmt.Errorf("App not started (%d)\nBody: (%s)", response.StatusCode, string(body))
	}

	return body, nil
}

// Curl connects to an endpoint to see if the endpoint is responding.
func Curl(t *testing.T, url string) {
	CurlWithFail(t, url, false, "")
}

// CurlApp is a convenience function to see if the example app is running.
func CurlApp(t *testing.T, cfg DeisTestConfig) {
	CurlWithFail(t, fmt.Sprintf("http://%s.%s", cfg.AppName, cfg.Domain), false, "")
}

// CurlWithFail connects to a Deis endpoint to see if the example app is running.
func CurlWithFail(t *testing.T, url string, failFlag bool, expect string) {
	// FIXME: try the curl a few times
	for i := 0; i < 20; i++ {
		body, err := doCurl(url)
		if err == nil {
			fmt.Println(string(body))
			return
		}
		time.Sleep(1 * time.Second)
	}

	// once more to fail with an error
	body, err := doCurl(url)

	switch failFlag {
	case true:
		if err != nil {
			if strings.Contains(string(err.Error()), expect) {
				fmt.Println("(Error expected...ok) " + expect)
			} else {
				t.Fatal(err)
			}
		} else {
			if strings.Contains(string(body), expect) {
				fmt.Println("(Error expected...ok) " + expect)
			} else {
				t.Fatal(err)
			}
		}
	case false:
		if err != nil {
			t.Fatal(err)
		} else {
			fmt.Println(string(body))
		}
	}
}

// AuthPasswd tests whether `deis auth:passwd` updates a user's password.
func AuthPasswd(t *testing.T, params *DeisTestConfig, password string) {
	fmt.Println("deis auth:passwd")
	child, err := gexpect.Spawn(Deis + " auth:passwd")
	if err != nil {
		t.Fatalf("command not started\n%v", err)
	}
	fmt.Println("current password:")
	err = child.Expect("current password: ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(params.Password)
	fmt.Println("new password:")
	err = child.Expect("new password: ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(password)
	fmt.Println("new password (confirm):")
	err = child.Expect("new password (confirm): ")
	if err != nil {
		t.Fatalf("expect password failed\n%v", err)
	}
	child.SendLine(password)
	err = child.Expect("Password change succeeded")
	if err != nil {
		t.Fatalf("command executiuon failed\n%v", err)
	}
	child.Close()
}

// CheckList executes a command and optionally tests whether its output does
// or does not contain a given string.
func CheckList(
	t *testing.T, cmd string, params interface{}, contain string, notflag bool) {
	var cmdBuf bytes.Buffer
	tmpl := template.Must(template.New("cmd").Parse(cmd))
	if err := tmpl.Execute(&cmdBuf, params); err != nil {
		t.Fatal(err)
	}
	cmdString := cmdBuf.String()
	fmt.Println(cmdString)
	var cmdl *exec.Cmd
	if strings.Contains(cmd, "cat") {
		cmdl = exec.Command("sh", "-c", cmdString)
	} else {
		cmdl = exec.Command("sh", "-c", Deis+cmdString)
	}
	stdout, _, err := RunCommandWithStdoutStderr(cmdl)
	if err != nil {
		t.Fatal(err)
	}
	if notflag && strings.Contains(stdout.String(), contain) {
		t.Fatalf("Didn't expect '%s' in command output:\n%v", contain, stdout)
	}
	if !notflag && !strings.Contains(stdout.String(), contain) {
		t.Fatalf("Expected '%s' in command output:\n%v", contain, stdout)
	}
}

// Execute takes command string and parameters required to execute the command,
// a failflag to check whether the command is expected to fail, and an expect
// string to check whether the command has failed according to failflag.
//
// If failflag is true and the command failed, check the stdout and stderr for
// the expect string.
func Execute(t *testing.T, cmd string, params interface{}, failFlag bool, expect string) {
	var cmdBuf bytes.Buffer
	tmpl := template.Must(template.New("cmd").Parse(cmd))
	if err := tmpl.Execute(&cmdBuf, params); err != nil {
		t.Fatal(err)
	}
	cmdString := cmdBuf.String()
	fmt.Println(cmdString)
	var cmdl *exec.Cmd
	if strings.Contains(cmd, "git ") {
		cmdl = exec.Command("sh", "-c", cmdString)
	} else {
		cmdl = exec.Command("sh", "-c", Deis+cmdString)
	}

	switch failFlag {
	case true:
		if stdout, stderr, err := RunCommandWithStdoutStderr(cmdl); err != nil {
			if strings.Contains(stdout.String(), expect) || strings.Contains(stderr.String(), expect) {
				fmt.Println("(Error expected...ok)")
			} else {
				t.Fatal(err)
			}
		} else {
			if strings.Contains(stdout.String(), expect) || strings.Contains(stderr.String(), expect) {
				fmt.Println("(Error expected...ok)" + expect)
			} else {
				t.Fatal(err)
			}
		}
	case false:
		stdout, stderr, err := RunCommandWithStdoutStderr(cmdl)
		if err != nil {
			t.Fatal(err)
		}

		if containsWarning(stdout.String()) || containsWarning(stderr.String()) {
			t.Fatal("Warning found in output, aborting")
		}

		fmt.Println("ok")
	}
}

// AppsDestroyTest destroys a Deis app and checks that it was successful.
func AppsDestroyTest(t *testing.T, params *DeisTestConfig) {
	fmt.Printf("destroying app %s...\n", params.ExampleApp)
	cmd := "apps:destroy --app={{.AppName}} --confirm={{.AppName}}"
	if err := Chdir(params.ExampleApp); err != nil {
		t.Fatal(err)
	}
	Execute(t, cmd, params, false, "")
	if err := Chdir(".."); err != nil {
		t.Fatal(err)
	}
	if err := Rmdir(params.ExampleApp); err != nil {
		t.Fatal(err)
	}
}

// GetRandomApp returns a known working example app at random for testing.
func GetRandomApp() string {
	rand.Seed(int64(time.Now().Unix()))
	apps := []string{
		"example-clojure-ring",
		// "example-dart",
		"example-dockerfile-python",
		"example-go",
		"example-java-jetty",
		"example-nodejs-express",
		// "example-php",
		"example-play",
		"example-python-django",
		"example-python-flask",
		"example-ruby-sinatra",
		"example-scala",
		"example-dockerfile-http",
	}
	return apps[rand.Intn(len(apps))]
}

func containsWarning(out string) bool {
	if strings.Contains(out, "WARNING") {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package toolbox

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
)

/*
FeatureSet contains an instance of every available toolbox feature. Together they may initialise, run self tests,
and receive configuration from JSON.
*/
type FeatureSet struct {
	AESDecrypt         AESDecrypt          `json:"AESDecrypt"`
	BrowserPhantomJS   BrowserPhantomJS    `json:"BrowserPhantomJS"`
	BrowserSlimerJS    BrowserSlimerJS     `json:"BrowserSlimerJS"`
	PublicContact      PublicContact       `json:"PublicContact"`
	EnvControl         EnvControl          `json:"EnvControl"`
	IMAPAccounts       IMAPAccounts        `json:"IMAPAccounts"`
	Joke               Joke                `json:"Joke"`
	RSS                RSS                 `json:"RSS"`
	SendMail           SendMail            `json:"SendMail"`
	Shell              Shell               `json:"Shell"`
	TextSearch         TextSearch          `json:"TextSearch"`
	Twilio             Twilio              `json:"Twilio"`
	Twitter            Twitter             `json:"Twitter"`
	TwoFACodeGenerator TwoFACodeGenerator  `json:"TwoFACodeGenerator"`
	WolframAlpha       WolframAlpha        `json:"WolframAlpha"`
	LookupByTrigger    map[Trigger]Feature `json:"-"`
}

//var TestFeatureSet = FeatureSet{} // Features are assigned by init_test.go

// Run initialisation routine on all features, and then populate lookup table for all configured features.
func (fs *FeatureSet) Initialise() error {
	fs.LookupByTrigger = map[Trigger]Feature{}
	triggers := map[Trigger]Feature{
		fs.AESDecrypt.Trigger():         &fs.AESDecrypt,         // a
		fs.BrowserPhantomJS.Trigger():   &fs.BrowserPhantomJS,   // bp
		fs.BrowserSlimerJS.Trigger():    &fs.BrowserSlimerJS,    // bs
		fs.PublicContact.Trigger():      &fs.PublicContact,      // c
		fs.EnvControl.Trigger():         &fs.EnvControl,         // e
		fs.TextSearch.Trigger():         &fs.TextSearch,         // g
		fs.IMAPAccounts.Trigger():       &fs.IMAPAccounts,       // i
		fs.Joke.Trigger():               &fs.Joke,               // j
		fs.RSS.Trigger():                &fs.RSS,                // r
		fs.SendMail.Trigger():           &fs.SendMail,           // m
		fs.Shell.Trigger():              &fs.Shell,              // s
		fs.Twilio.Trigger():             &fs.Twilio,             // p
		fs.Twitter.Trigger():            &fs.Twitter,            // t
		fs.TwoFACodeGenerator.Trigger(): &fs.TwoFACodeGenerator, // 2
		fs.WolframAlpha.Trigger():       &fs.WolframAlpha,       // w
	}
	errs := make([]string, 0, 0)
	for trigger, featureRef := range triggers {
		/*
			Collect initialisation errors (if any) from all failed features so that all mistakes can be presented to
			caller at once.
		*/
		if featureRef.IsConfigured() {
			if err := featureRef.Initialise(); err != nil {
				errs = append(errs, err.Error())
			}
			fs.LookupByTrigger[trigger] = featureRef
		}
	}
	if len(errs) != 0 {
		return errors.New(strings.Join(errs, " | "))
	}
	return nil
}

// Run self test of all configured features in parallel. Return test errors if any.
func (fs *FeatureSet) SelfTest() error {
	ret := make([]string, 0, 0)
	retMutex := &sync.Mutex{}
	wait := &sync.WaitGroup{}
	wait.Add(len(fs.LookupByTrigger))
	for a, featureRef := range fs.LookupByTrigger {
		fmt.Println("Self test is testing", a)
		go func(ref Feature) {
			err := ref.SelfTest()
			if err != nil {
				retMutex.Lock()
				ret = append(ret, fmt.Sprintf("%s: %s", ref.Trigger(), err.Error()))
				retMutex.Unlock()
			}
			wait.Done()
		}(featureRef)
	}
	wait.Wait()
	if len(ret) == 0 {
		return nil
	}
	return errors.New(strings.Join(ret, " | "))
}

// Deserialise feature configuration from JSON configuration. The function does not initialise features automatically.
func (fs *FeatureSet) DeserialiseFromJSON(configJSON json.RawMessage) error {
	// Turn input JSON into map[string]json.RawMessage, map key is the feature key in JSON.
	var configMap map[string]json.RawMessage
	if err := json.Unmarshal(configJSON, &configMap); err != nil {
		return fmt.Errorf("FeatureSet.DeserialiseFromJSON: failed to retrieve config map - %v", err)
	}
	// Here are the feature keys
	features := map[string]Feature{
		"AESDecrypt":         &fs.AESDecrypt,
		"BrowserPhantomJS":   &fs.BrowserPhantomJS,
		"BrowserSlimerJS":    &fs.BrowserSlimerJS,
		"EnvControl":         &fs.EnvControl,
		"IMAPAccounts":       &fs.IMAPAccounts,
		"Joke":               &fs.Joke,
		"RSS":                &fs.RSS,
		"SendMail":           &fs.SendMail,
		"Shell":              &fs.Shell,
		"Twilio":             &fs.Twilio,
		"Twitter":            &fs.Twitter,
		"TwoFACodeGenerator": &fs.TwoFACodeGenerator,
		"WolframAlpha":       &fs.WolframAlpha,
	}
	for featureKey, featureRef := range features {
		if featureJSON, exists := configMap[featureKey]; exists {
			if err := json.Unmarshal(featureJSON, &featureRef); err != nil {
				return fmt.Errorf("FeatureSet.DeserialiseFromJSON: failed to deserialise JSON key %s - %v", featureKey, err)
			}
		}
	}
	return nil
}

// Return all configured & initialised triggers, sorted in alphabetical order.
func (fs *FeatureSet) GetTriggers() []string {
	ret := make([]string, 0, 8)
	if fs.LookupByTrigger == nil {
		return ret
	}
	for trigger := range fs.LookupByTrigger {
		ret = append(ret, string(trigger))
	}
	sort.Strings(ret)
	return ret
}
<commit_msg>remove debug output from FeatureSet.SelfTest<commit_after>package toolbox

import (
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
)

/*
FeatureSet contains an instance of every available toolbox feature. Together they may initialise, run self tests,
and receive configuration from JSON.
*/
type FeatureSet struct {
	AESDecrypt         AESDecrypt          `json:"AESDecrypt"`
	BrowserPhantomJS   BrowserPhantomJS    `json:"BrowserPhantomJS"`
	BrowserSlimerJS    BrowserSlimerJS     `json:"BrowserSlimerJS"`
	PublicContact      PublicContact       `json:"PublicContact"`
	EnvControl         EnvControl          `json:"EnvControl"`
	IMAPAccounts       IMAPAccounts        `json:"IMAPAccounts"`
	Joke               Joke                `json:"Joke"`
	RSS                RSS                 `json:"RSS"`
	SendMail           SendMail            `json:"SendMail"`
	Shell              Shell               `json:"Shell"`
	TextSearch         TextSearch          `json:"TextSearch"`
	Twilio             Twilio              `json:"Twilio"`
	Twitter            Twitter             `json:"Twitter"`
	TwoFACodeGenerator TwoFACodeGenerator  `json:"TwoFACodeGenerator"`
	WolframAlpha       WolframAlpha        `json:"WolframAlpha"`
	LookupByTrigger    map[Trigger]Feature `json:"-"`
}

//var TestFeatureSet = FeatureSet{} // Features are assigned by init_test.go

// Run initialisation routine on all features, and then populate lookup table for all configured features.
func (fs *FeatureSet) Initialise() error {
	fs.LookupByTrigger = map[Trigger]Feature{}
	triggers := map[Trigger]Feature{
		fs.AESDecrypt.Trigger():         &fs.AESDecrypt,         // a
		fs.BrowserPhantomJS.Trigger():   &fs.BrowserPhantomJS,   // bp
		fs.BrowserSlimerJS.Trigger():    &fs.BrowserSlimerJS,    // bs
		fs.PublicContact.Trigger():      &fs.PublicContact,      // c
		fs.EnvControl.Trigger():         &fs.EnvControl,         // e
		fs.TextSearch.Trigger():         &fs.TextSearch,         // g
		fs.IMAPAccounts.Trigger():       &fs.IMAPAccounts,       // i
		fs.Joke.Trigger():               &fs.Joke,               // j
		fs.RSS.Trigger():                &fs.RSS,                // r
		fs.SendMail.Trigger():           &fs.SendMail,           // m
		fs.Shell.Trigger():              &fs.Shell,              // s
		fs.Twilio.Trigger():             &fs.Twilio,             // p
		fs.Twitter.Trigger():            &fs.Twitter,            // t
		fs.TwoFACodeGenerator.Trigger(): &fs.TwoFACodeGenerator, // 2
		fs.WolframAlpha.Trigger():       &fs.WolframAlpha,       // w
	}
	errs := make([]string, 0, 0)
	for trigger, featureRef := range triggers {
		/*
			Collect initialisation errors (if any) from all failed features so that all mistakes can be presented to
			caller at once.
		*/
		if featureRef.IsConfigured() {
			if err := featureRef.Initialise(); err != nil {
				errs = append(errs, err.Error())
			}
			fs.LookupByTrigger[trigger] = featureRef
		}
	}
	if len(errs) != 0 {
		return errors.New(strings.Join(errs, " | "))
	}
	return nil
}

// Run self test of all configured features in parallel. Return test errors if any.
func (fs *FeatureSet) SelfTest() error {
	ret := make([]string, 0, 0)
	retMutex := &sync.Mutex{}
	wait := &sync.WaitGroup{}
	wait.Add(len(fs.LookupByTrigger))
	for _, featureRef := range fs.LookupByTrigger {
		go func(ref Feature) {
			err := ref.SelfTest()
			if err != nil {
				retMutex.Lock()
				ret = append(ret, fmt.Sprintf("%s: %s", ref.Trigger(), err.Error()))
				retMutex.Unlock()
			}
			wait.Done()
		}(featureRef)
	}
	wait.Wait()
	if len(ret) == 0 {
		return nil
	}
	return errors.New(strings.Join(ret, " | "))
}

// Deserialise feature configuration from JSON configuration. The function does not initialise features automatically.
func (fs *FeatureSet) DeserialiseFromJSON(configJSON json.RawMessage) error {
	// Turn input JSON into map[string]json.RawMessage, map key is the feature key in JSON.
	var configMap map[string]json.RawMessage
	if err := json.Unmarshal(configJSON, &configMap); err != nil {
		return fmt.Errorf("FeatureSet.DeserialiseFromJSON: failed to retrieve config map - %v", err)
	}
	// Here are the feature keys
	features := map[string]Feature{
		"AESDecrypt":         &fs.AESDecrypt,
		"BrowserPhantomJS":   &fs.BrowserPhantomJS,
		"BrowserSlimerJS":    &fs.BrowserSlimerJS,
		"EnvControl":         &fs.EnvControl,
		"IMAPAccounts":       &fs.IMAPAccounts,
		"Joke":               &fs.Joke,
		"RSS":                &fs.RSS,
		"SendMail":           &fs.SendMail,
		"Shell":              &fs.Shell,
		"Twilio":             &fs.Twilio,
		"Twitter":            &fs.Twitter,
		"TwoFACodeGenerator": &fs.TwoFACodeGenerator,
		"WolframAlpha":       &fs.WolframAlpha,
	}
	for featureKey, featureRef := range features {
		if featureJSON, exists := configMap[featureKey]; exists {
			if err := json.Unmarshal(featureJSON, &featureRef); err != nil {
				return fmt.Errorf("FeatureSet.DeserialiseFromJSON: failed to deserialise JSON key %s - %v", featureKey, err)
			}
		}
	}
	return nil
}

// Return all configured & initialised triggers, sorted in alphabetical order.
func (fs *FeatureSet) GetTriggers() []string {
	ret := make([]string, 0, 8)
	if fs.LookupByTrigger == nil {
		return ret
	}
	for trigger := range fs.LookupByTrigger {
		ret = append(ret, string(trigger))
	}
	sort.Strings(ret)
	return ret
}
<|endoftext|>


<commit_before>package postgres

import (
	"fmt"
	"time"

	"gitgud.io/softashell/comfy-translator/translator"
	lru "github.com/hashicorp/golang-lru"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

type translationError int

const (
	errorNone           translationError = iota // Everything is fine
	errorMinor                                  // Connection timed out or something like that
	errorBadTranslation                         // Returned really bad translation
)

type Cache struct {
	db       *gorm.DB
	lrustore map[string]*lru.TwoQueueCache
}

type Item struct {
	Translation string
	ErrorCode   translationError
	ErrorText   string
	Timestamp   int64
}

type PgItem struct {
	Text string `gorm:"primarykey"`
	Bucket string `gorm:"primaryKey"`
	Translation string
	ErrorCode   translationError
	ErrorText   string
	Timestamp   time.Time
}

func NewCache(connStr string, cacheSize int, translators []string) (*Cache, error) {
	db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{})
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	err = db.AutoMigrate(&PgItem{})
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	store := make(map[string]*lru.TwoQueueCache)
	for _, t := range translators {
		s, err := lru.New2Q(5000)
		if err != nil {
			log.Fatalf("Can't start memory cache for %s", t)
		}

		store[t] = s
	}

	cache := &Cache{db: db, lrustore: store}

	return cache, nil
}

func (c *Cache) Close() error {
	return nil
}

func (c *Cache) Put(bucketName, text, translation string, cerr error) error {
	errorCode := errorNone
	errorText := ""

	if cerr != nil {
		errorText = cerr.Error()

		switch cerr.(type) {
		case translator.BadTranslationError:
			errorCode = errorBadTranslation
		default:
			errorCode = errorMinor
		}
	}

	pgItem := PgItem{
		Text: 			text,
		Bucket: 		bucketName,
		Translation: 	translation,
		ErrorCode:   	errorCode,
		ErrorText:   	errorText,
		Timestamp:   	time.Now().UTC(),
	}

	result := c.db.Create(&pgItem)
	if result.Error != nil {
		log.Fatal(errors.Wrap(result.Error, "failed to execute insert"))
	}

	// Add to memory cache
	i := Item{
		Translation: translation,
		ErrorCode:   errorCode,
		ErrorText:   errorText,
		Timestamp:   time.Now().UTC().Unix(),
	}

	c.lrustore[bucketName].Add(text, i)

	return result.Error
}

func (c *Cache) Get(bucketName, text string) (string, bool, error) {
	var found bool
	var translation string
	var errorCode translationError
	var errorText string
	var timestamp int64

	item, ok := c.lrustore[bucketName].Get(text)
	if ok {
		i := item.(Item)

		translation = i.Translation
		errorCode = i.ErrorCode
		errorText = i.ErrorText
	} else {
		i := PgItem{}
		result := c.db.Limit(1).Find(&i, PgItem{Text: text,	Bucket: bucketName})
		if result.Error != nil {
			return translation, found, nil
		}

		translation = i.Translation
		errorCode = i.ErrorCode
		errorText = i.ErrorText
	}

	if translation != "" {
		found = true
	}

	if errorCode != errorNone {
		errorTime := time.Unix(timestamp, 0)

		if time.Since(errorTime) > getCacheExpiration(errorCode) {
			i := PgItem{}
			result := c.db.Delete(&i, PgItem{Text: text,	Bucket: bucketName})
			if result.Error != nil {
				log.Warn("unable to delete item: ", result.Error)
			}

			c.lrustore[bucketName].Remove(text)

			// Act as if nothing was found
			return "", false, nil
		}

		return "", found, fmt.Errorf("%s", errorText)
	}

	return translation, found, nil
}
<commit_msg>Change table name<commit_after>package postgres

import (
	"fmt"
	"time"

	"gitgud.io/softashell/comfy-translator/translator"
	lru "github.com/hashicorp/golang-lru"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

type translationError int

const (
	errorNone           translationError = iota // Everything is fine
	errorMinor                                  // Connection timed out or something like that
	errorBadTranslation                         // Returned really bad translation
)

type Cache struct {
	db       *gorm.DB
	lrustore map[string]*lru.TwoQueueCache
}

type Item struct {
	Translation string
	ErrorCode   translationError
	ErrorText   string
	Timestamp   int64
}

type Translation struct {
	Text string `gorm:"primarykey"`
	Bucket string `gorm:"primaryKey"`
	Translation string
	ErrorCode   translationError
	ErrorText   string
	Timestamp   time.Time
}

func NewCache(connStr string, cacheSize int, translators []string) (*Cache, error) {
	db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{})
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	err = db.AutoMigrate(&Translation{})
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	store := make(map[string]*lru.TwoQueueCache)
	for _, t := range translators {
		s, err := lru.New2Q(5000)
		if err != nil {
			log.Fatalf("Can't start memory cache for %s", t)
		}

		store[t] = s
	}

	cache := &Cache{db: db, lrustore: store}

	return cache, nil
}

func (c *Cache) Close() error {
	return nil
}

func (c *Cache) Put(bucketName, text, translation string, cerr error) error {
	errorCode := errorNone
	errorText := ""

	if cerr != nil {
		errorText = cerr.Error()

		switch cerr.(type) {
		case translator.BadTranslationError:
			errorCode = errorBadTranslation
		default:
			errorCode = errorMinor
		}
	}

	pgItem := Translation{
		Text: 			text,
		Bucket: 		bucketName,
		Translation: 	translation,
		ErrorCode:   	errorCode,
		ErrorText:   	errorText,
		Timestamp:   	time.Now().UTC(),
	}

	result := c.db.Create(&pgItem)
	if result.Error != nil {
		log.Fatal(errors.Wrap(result.Error, "failed to execute insert"))
	}

	// Add to memory cache
	i := Item{
		Translation: translation,
		ErrorCode:   errorCode,
		ErrorText:   errorText,
		Timestamp:   time.Now().UTC().Unix(),
	}

	c.lrustore[bucketName].Add(text, i)

	return result.Error
}

func (c *Cache) Get(bucketName, text string) (string, bool, error) {
	var found bool
	var translation string
	var errorCode translationError
	var errorText string
	var timestamp int64

	item, ok := c.lrustore[bucketName].Get(text)
	if ok {
		i := item.(Item)

		translation = i.Translation
		errorCode = i.ErrorCode
		errorText = i.ErrorText
	} else {
		i := Translation{}
		result := c.db.Limit(1).Find(&i, Translation{Text: text,	Bucket: bucketName})
		if result.Error != nil {
			return translation, found, nil
		}

		translation = i.Translation
		errorCode = i.ErrorCode
		errorText = i.ErrorText
	}

	if translation != "" {
		found = true
	}

	if errorCode != errorNone {
		errorTime := time.Unix(timestamp, 0)

		if time.Since(errorTime) > getCacheExpiration(errorCode) {
			i := Translation{}
			result := c.db.Delete(&i, Translation{Text: text,	Bucket: bucketName})
			if result.Error != nil {
				log.Warn("unable to delete item: ", result.Error)
			}

			c.lrustore[bucketName].Remove(text)

			// Act as if nothing was found
			return "", false, nil
		}

		return "", found, fmt.Errorf("%s", errorText)
	}

	return translation, found, nil
}
<|endoftext|>