<commit_before>/*
Copyright 2019 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 vreplication

import (
	"flag"
	"fmt"
	"strings"
	"time"

	querypb "vitess.io/vitess/go/vt/proto/query"

	"vitess.io/vitess/go/vt/vtgate/evalengine"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/binlog/binlogplayer"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"

	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
)

var (
	// idleTimeout is set to slightly above 1s, compared to heartbeatTime
	// set by VStreamer at slightly below 1s. This minimizes conflicts
	// between the two timeouts.
	idleTimeout         = 1100 * time.Millisecond
	dbLockRetryDelay    = 1 * time.Second
	relayLogMaxSize     = flag.Int("relay_log_max_size", 250000, "Maximum buffer size (in bytes) for VReplication target buffering. If single rows are larger than this, a single row is buffered at a time.")
	relayLogMaxItems    = flag.Int("relay_log_max_items", 1000, "Maximum number of rows for VReplication target buffering.")
	copyTimeout         = 1 * time.Hour
	replicaLagTolerance = 10 * time.Second
)

// vreplicator provides the core logic to start vreplication streams
type vreplicator struct {
	vre      *Engine
	id       uint32
	dbClient *vdbClient
	// source
	source          *binlogdatapb.BinlogSource
	sourceVStreamer VStreamerClient

	stats *binlogplayer.Stats
	// mysqld is used to fetch the local schema.
	mysqld    mysqlctl.MysqlDaemon
	pkInfoMap map[string][]*PrimaryKeyInfo

	originalFKCheckSetting int64
}

// newVReplicator creates a new vreplicator. The valid fields from the source are:
// Keyspce, Shard, Filter, OnDdl, ExternalMySql and StopAfterCopy.
// The Filter consists of Rules. Each Rule has a Match and an (inner) Filter field.
// The Match can be a table name or, if it begins with a "/", a wildcard.
// The Filter can be empty: get all rows and columns.
// The Filter can be a keyrange, like "-80": get all rows that are within the keyrange.
// The Filter can be a select expression. Examples.
//   "select * from t", same as an empty Filter,
//   "select * from t where in_keyrange('-80')", same as "-80",
//   "select * from t where in_keyrange(col1, 'hash', '-80')",
//   "select col1, col2 from t where...",
//   "select col1, keyspace_id() as ksid from t where...",
//   "select id, count(*), sum(price) from t group by id".
//   Only "in_keyrange" expressions are supported in the where clause.
//   The select expressions can be any valid non-aggregate expressions,
//   or count(*), or sum(col).
//   If the target column name does not match the source expression, an
//   alias like "a+b as targetcol" must be used.
//   More advanced constructs can be used. Please see the table plan builder
//   documentation for more info.
func newVReplicator(id uint32, source *binlogdatapb.BinlogSource, sourceVStreamer VStreamerClient, stats *binlogplayer.Stats, dbClient binlogplayer.DBClient, mysqld mysqlctl.MysqlDaemon, vre *Engine) *vreplicator {
	return &vreplicator{
		vre:             vre,
		id:              id,
		source:          source,
		sourceVStreamer: sourceVStreamer,
		stats:           stats,
		dbClient:        newVDBClient(dbClient, stats),
		mysqld:          mysqld,
	}
}

// Replicate starts a vreplication stream. It can be in one of three phases:
// 1. Init: If a request is issued with no starting position, we assume that the
// contents of the tables must be copied first. During this phase, the list of
// tables to be copied is inserted into the copy_state table. A successful insert
// gets us out of this phase.
// 2. Copy: If the copy_state table has rows, then we are in this phase. During this
// phase, we repeatedly invoke copyNext until all the tables are copied. After each
// table is successfully copied, it's removed from the copy_state table. We exit this
// phase when there are no rows left in copy_state.
// 3. Replicate: In this phase, we replicate binlog events indefinitely, unless
// a stop position was requested. This phase differs from the Init phase because
// there is a replication position.
// If a request had a starting position, then we go directly into phase 3.
// During these phases, the state of vreplication is reported as 'Init', 'Copying',
// or 'Running'. They all mean the same thing. The difference in the phases depends
// on the criteria defined above. The different states reported are mainly
// informational. The 'Stopped' state is, however, honored.
// All phases share the same plan building framework. We leverage the fact the
// row representation of a read (during copy) and a binlog event are identical.
// However, there are some subtle differences, explained in the plan builder
// code.
func (vr *vreplicator) Replicate(ctx context.Context) error {
	err := vr.replicate(ctx)
	if err != nil {
		log.Errorf("Replicate error: %s", err.Error())
		if err := vr.setMessage(err.Error()); err != nil {
			log.Errorf("Failed to set error state: %v", err)
		}
	}
	return err
}

func (vr *vreplicator) replicate(ctx context.Context) error {
	pkInfo, err := vr.buildPkInfoMap(ctx)
	if err != nil {
		return err
	}
	vr.pkInfoMap = pkInfo
	if err := vr.getSettingFKCheck(); err != nil {
		return err
	}
	//defensive guard, should be a no-op since it should happen after copy is done
	defer vr.resetFKCheckAfterCopy()
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}
		// This rollback is a no-op. It's here for safety
		// in case the functions below leave transactions open.
		vr.dbClient.Rollback()

		settings, numTablesToCopy, err := vr.readSettings(ctx)
		if err != nil {
			return err
		}
		// If any of the operations below changed state to Stopped, we should return.
		if settings.State == binlogplayer.BlpStopped {
			return nil
		}

		switch {
		case numTablesToCopy != 0:
			if err := vr.clearFKCheck(); err != nil {
				log.Warningf("Unable to clear FK check %v", err)
				return err
			}
			if err := newVCopier(vr).copyNext(ctx, settings); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Copy"}, 1)
				return err
			}
		case settings.StartPos.IsZero():
			if err := newVCopier(vr).initTablesForCopy(ctx); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Copy"}, 1)
				return err
			}
		default:
			if err := vr.resetFKCheckAfterCopy(); err != nil {
				log.Warningf("Unable to reset FK check %v", err)
				return err
			}
			if vr.source.StopAfterCopy {
				return vr.setState(binlogplayer.BlpStopped, "Stopped after copy.")
			}
			if err := vr.setState(binlogplayer.BlpRunning, ""); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Replicate"}, 1)
				return err
			}
			return newVPlayer(vr, settings, nil, mysql.Position{}, "replicate").play(ctx)
		}
	}
}

// PrimaryKeyInfo is used to store charset and collation for primary keys where applicable
type PrimaryKeyInfo struct {
	Name      string
	CharSet   string
	Collation string
}

func (vr *vreplicator) buildPkInfoMap(ctx context.Context) (map[string][]*PrimaryKeyInfo, error) {
	schema, err := vr.mysqld.GetSchema(ctx, vr.dbClient.DBName(), []string{"/.*/"}, nil, false)
	if err != nil {
		return nil, err
	}
	queryTemplate := "select character_set_name, collation_name,column_name,data_type from information_schema.columns where table_schema=%s and table_name=%s;"
	pkInfoMap := make(map[string][]*PrimaryKeyInfo)
	for _, td := range schema.TableDefinitions {

		query := fmt.Sprintf(queryTemplate, encodeString(vr.dbClient.DBName()), encodeString(td.Name))
		qr, err := vr.mysqld.FetchSuperQuery(ctx, query)
		if err != nil {
			return nil, err
		}
		if len(qr.Rows) > 0 && len(qr.Fields) != 4 {
			return nil, fmt.Errorf("incorrect result returned for collation query")
		}

		var pks []string
		if len(td.PrimaryKeyColumns) != 0 {
			pks = td.PrimaryKeyColumns
		} else {
			pks = td.Columns
		}
		var pkInfos []*PrimaryKeyInfo
		for _, pk := range pks {
			charSet := ""
			collation := ""
			for _, row := range qr.Rows {
				columnName := row[2].ToString()
				if strings.EqualFold(columnName, pk) {
					var currentField *querypb.Field
					for _, field := range td.Fields {
						if field.Name == pk {
							currentField = field
							break
						}
					}
					if currentField != nil && sqltypes.IsText(currentField.Type) {
						charSet = row[0].ToString()
						collation = row[1].ToString()
					}
					break
				}
			}
			pkInfos = append(pkInfos, &PrimaryKeyInfo{
				Name:      pk,
				CharSet:   charSet,
				Collation: collation,
			})
		}
		pkInfoMap[td.Name] = pkInfos
	}
	return pkInfoMap, nil
}

func (vr *vreplicator) readSettings(ctx context.Context) (settings binlogplayer.VRSettings, numTablesToCopy int64, err error) {
	settings, err = binlogplayer.ReadVRSettings(vr.dbClient, vr.id)
	if err != nil {
		return settings, numTablesToCopy, fmt.Errorf("error reading VReplication settings: %v", err)
	}

	query := fmt.Sprintf("select count(*) from _vt.copy_state where vrepl_id=%d", vr.id)
	qr, err := withDDL.Exec(ctx, query, vr.dbClient.ExecuteFetch)
	if err != nil {
		return settings, numTablesToCopy, err
	}
	if len(qr.Rows) == 0 || len(qr.Rows[0]) == 0 {
		return settings, numTablesToCopy, fmt.Errorf("unexpected result from %s: %v", query, qr)
	}
	numTablesToCopy, err = evalengine.ToInt64(qr.Rows[0][0])
	if err != nil {
		return settings, numTablesToCopy, err
	}
	return settings, numTablesToCopy, nil
}

func (vr *vreplicator) setMessage(message string) error {
	vr.stats.History.Add(&binlogplayer.StatsHistoryRecord{
		Time:    time.Now(),
		Message: message,
	})
	query := fmt.Sprintf("update _vt.vreplication set message=%v where id=%v", encodeString(binlogplayer.MessageTruncate(message)), vr.id)
	if _, err := vr.dbClient.Execute(query); err != nil {
		return fmt.Errorf("could not set message: %v: %v", query, err)
	}
	return nil
}

func (vr *vreplicator) setState(state, message string) error {
	if message != "" {
		vr.stats.History.Add(&binlogplayer.StatsHistoryRecord{
			Time:    time.Now(),
			Message: message,
		})
	}
	vr.stats.State.Set(state)
	query := fmt.Sprintf("update _vt.vreplication set state='%v', message=%v where id=%v", state, encodeString(binlogplayer.MessageTruncate(message)), vr.id)
	if _, err := vr.dbClient.ExecuteFetch(query, 1); err != nil {
		return fmt.Errorf("could not set state: %v: %v", query, err)
	}
	return nil
}

func encodeString(in string) string {
	var buf strings.Builder
	sqltypes.NewVarChar(in).EncodeSQL(&buf)
	return buf.String()
}

func (vr *vreplicator) getSettingFKCheck() error {
	qr, err := vr.dbClient.Execute("select @@foreign_key_checks;")
	if err != nil {
		return err
	}
	if qr.RowsAffected != 1 || len(qr.Fields) != 1 {
		return fmt.Errorf("unable to select @@foreign_key_checks")
	}
	vr.originalFKCheckSetting, err = evalengine.ToInt64(qr.Rows[0][0])
	if err != nil {
		return err
	}
	return nil
}

func (vr *vreplicator) resetFKCheckAfterCopy() error {
	_, err := vr.dbClient.Execute(fmt.Sprintf("set foreign_key_checks=%d;", vr.originalFKCheckSetting))
	return err
}

func (vr *vreplicator) clearFKCheck() error {
	_, err := vr.dbClient.Execute("set foreign_key_checks=0;")
	return err
}
<commit_msg>Bump default max vreplication row buffering to 5k<commit_after>/*
Copyright 2019 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 vreplication

import (
	"flag"
	"fmt"
	"strings"
	"time"

	querypb "vitess.io/vitess/go/vt/proto/query"

	"vitess.io/vitess/go/vt/vtgate/evalengine"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/binlog/binlogplayer"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl"

	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
)

var (
	// idleTimeout is set to slightly above 1s, compared to heartbeatTime
	// set by VStreamer at slightly below 1s. This minimizes conflicts
	// between the two timeouts.
	idleTimeout         = 1100 * time.Millisecond
	dbLockRetryDelay    = 1 * time.Second
	relayLogMaxSize     = flag.Int("relay_log_max_size", 250000, "Maximum buffer size (in bytes) for VReplication target buffering. If single rows are larger than this, a single row is buffered at a time.")
	relayLogMaxItems    = flag.Int("relay_log_max_items", 5000, "Maximum number of rows for VReplication target buffering.")
	copyTimeout         = 1 * time.Hour
	replicaLagTolerance = 10 * time.Second
)

// vreplicator provides the core logic to start vreplication streams
type vreplicator struct {
	vre      *Engine
	id       uint32
	dbClient *vdbClient
	// source
	source          *binlogdatapb.BinlogSource
	sourceVStreamer VStreamerClient

	stats *binlogplayer.Stats
	// mysqld is used to fetch the local schema.
	mysqld    mysqlctl.MysqlDaemon
	pkInfoMap map[string][]*PrimaryKeyInfo

	originalFKCheckSetting int64
}

// newVReplicator creates a new vreplicator. The valid fields from the source are:
// Keyspce, Shard, Filter, OnDdl, ExternalMySql and StopAfterCopy.
// The Filter consists of Rules. Each Rule has a Match and an (inner) Filter field.
// The Match can be a table name or, if it begins with a "/", a wildcard.
// The Filter can be empty: get all rows and columns.
// The Filter can be a keyrange, like "-80": get all rows that are within the keyrange.
// The Filter can be a select expression. Examples.
//   "select * from t", same as an empty Filter,
//   "select * from t where in_keyrange('-80')", same as "-80",
//   "select * from t where in_keyrange(col1, 'hash', '-80')",
//   "select col1, col2 from t where...",
//   "select col1, keyspace_id() as ksid from t where...",
//   "select id, count(*), sum(price) from t group by id".
//   Only "in_keyrange" expressions are supported in the where clause.
//   The select expressions can be any valid non-aggregate expressions,
//   or count(*), or sum(col).
//   If the target column name does not match the source expression, an
//   alias like "a+b as targetcol" must be used.
//   More advanced constructs can be used. Please see the table plan builder
//   documentation for more info.
func newVReplicator(id uint32, source *binlogdatapb.BinlogSource, sourceVStreamer VStreamerClient, stats *binlogplayer.Stats, dbClient binlogplayer.DBClient, mysqld mysqlctl.MysqlDaemon, vre *Engine) *vreplicator {
	return &vreplicator{
		vre:             vre,
		id:              id,
		source:          source,
		sourceVStreamer: sourceVStreamer,
		stats:           stats,
		dbClient:        newVDBClient(dbClient, stats),
		mysqld:          mysqld,
	}
}

// Replicate starts a vreplication stream. It can be in one of three phases:
// 1. Init: If a request is issued with no starting position, we assume that the
// contents of the tables must be copied first. During this phase, the list of
// tables to be copied is inserted into the copy_state table. A successful insert
// gets us out of this phase.
// 2. Copy: If the copy_state table has rows, then we are in this phase. During this
// phase, we repeatedly invoke copyNext until all the tables are copied. After each
// table is successfully copied, it's removed from the copy_state table. We exit this
// phase when there are no rows left in copy_state.
// 3. Replicate: In this phase, we replicate binlog events indefinitely, unless
// a stop position was requested. This phase differs from the Init phase because
// there is a replication position.
// If a request had a starting position, then we go directly into phase 3.
// During these phases, the state of vreplication is reported as 'Init', 'Copying',
// or 'Running'. They all mean the same thing. The difference in the phases depends
// on the criteria defined above. The different states reported are mainly
// informational. The 'Stopped' state is, however, honored.
// All phases share the same plan building framework. We leverage the fact the
// row representation of a read (during copy) and a binlog event are identical.
// However, there are some subtle differences, explained in the plan builder
// code.
func (vr *vreplicator) Replicate(ctx context.Context) error {
	err := vr.replicate(ctx)
	if err != nil {
		log.Errorf("Replicate error: %s", err.Error())
		if err := vr.setMessage(err.Error()); err != nil {
			log.Errorf("Failed to set error state: %v", err)
		}
	}
	return err
}

func (vr *vreplicator) replicate(ctx context.Context) error {
	pkInfo, err := vr.buildPkInfoMap(ctx)
	if err != nil {
		return err
	}
	vr.pkInfoMap = pkInfo
	if err := vr.getSettingFKCheck(); err != nil {
		return err
	}
	//defensive guard, should be a no-op since it should happen after copy is done
	defer vr.resetFKCheckAfterCopy()
	for {
		select {
		case <-ctx.Done():
			return nil
		default:
		}
		// This rollback is a no-op. It's here for safety
		// in case the functions below leave transactions open.
		vr.dbClient.Rollback()

		settings, numTablesToCopy, err := vr.readSettings(ctx)
		if err != nil {
			return err
		}
		// If any of the operations below changed state to Stopped, we should return.
		if settings.State == binlogplayer.BlpStopped {
			return nil
		}

		switch {
		case numTablesToCopy != 0:
			if err := vr.clearFKCheck(); err != nil {
				log.Warningf("Unable to clear FK check %v", err)
				return err
			}
			if err := newVCopier(vr).copyNext(ctx, settings); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Copy"}, 1)
				return err
			}
		case settings.StartPos.IsZero():
			if err := newVCopier(vr).initTablesForCopy(ctx); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Copy"}, 1)
				return err
			}
		default:
			if err := vr.resetFKCheckAfterCopy(); err != nil {
				log.Warningf("Unable to reset FK check %v", err)
				return err
			}
			if vr.source.StopAfterCopy {
				return vr.setState(binlogplayer.BlpStopped, "Stopped after copy.")
			}
			if err := vr.setState(binlogplayer.BlpRunning, ""); err != nil {
				vr.stats.ErrorCounts.Add([]string{"Replicate"}, 1)
				return err
			}
			return newVPlayer(vr, settings, nil, mysql.Position{}, "replicate").play(ctx)
		}
	}
}

// PrimaryKeyInfo is used to store charset and collation for primary keys where applicable
type PrimaryKeyInfo struct {
	Name      string
	CharSet   string
	Collation string
}

func (vr *vreplicator) buildPkInfoMap(ctx context.Context) (map[string][]*PrimaryKeyInfo, error) {
	schema, err := vr.mysqld.GetSchema(ctx, vr.dbClient.DBName(), []string{"/.*/"}, nil, false)
	if err != nil {
		return nil, err
	}
	queryTemplate := "select character_set_name, collation_name,column_name,data_type from information_schema.columns where table_schema=%s and table_name=%s;"
	pkInfoMap := make(map[string][]*PrimaryKeyInfo)
	for _, td := range schema.TableDefinitions {

		query := fmt.Sprintf(queryTemplate, encodeString(vr.dbClient.DBName()), encodeString(td.Name))
		qr, err := vr.mysqld.FetchSuperQuery(ctx, query)
		if err != nil {
			return nil, err
		}
		if len(qr.Rows) > 0 && len(qr.Fields) != 4 {
			return nil, fmt.Errorf("incorrect result returned for collation query")
		}

		var pks []string
		if len(td.PrimaryKeyColumns) != 0 {
			pks = td.PrimaryKeyColumns
		} else {
			pks = td.Columns
		}
		var pkInfos []*PrimaryKeyInfo
		for _, pk := range pks {
			charSet := ""
			collation := ""
			for _, row := range qr.Rows {
				columnName := row[2].ToString()
				if strings.EqualFold(columnName, pk) {
					var currentField *querypb.Field
					for _, field := range td.Fields {
						if field.Name == pk {
							currentField = field
							break
						}
					}
					if currentField != nil && sqltypes.IsText(currentField.Type) {
						charSet = row[0].ToString()
						collation = row[1].ToString()
					}
					break
				}
			}
			pkInfos = append(pkInfos, &PrimaryKeyInfo{
				Name:      pk,
				CharSet:   charSet,
				Collation: collation,
			})
		}
		pkInfoMap[td.Name] = pkInfos
	}
	return pkInfoMap, nil
}

func (vr *vreplicator) readSettings(ctx context.Context) (settings binlogplayer.VRSettings, numTablesToCopy int64, err error) {
	settings, err = binlogplayer.ReadVRSettings(vr.dbClient, vr.id)
	if err != nil {
		return settings, numTablesToCopy, fmt.Errorf("error reading VReplication settings: %v", err)
	}

	query := fmt.Sprintf("select count(*) from _vt.copy_state where vrepl_id=%d", vr.id)
	qr, err := withDDL.Exec(ctx, query, vr.dbClient.ExecuteFetch)
	if err != nil {
		return settings, numTablesToCopy, err
	}
	if len(qr.Rows) == 0 || len(qr.Rows[0]) == 0 {
		return settings, numTablesToCopy, fmt.Errorf("unexpected result from %s: %v", query, qr)
	}
	numTablesToCopy, err = evalengine.ToInt64(qr.Rows[0][0])
	if err != nil {
		return settings, numTablesToCopy, err
	}
	return settings, numTablesToCopy, nil
}

func (vr *vreplicator) setMessage(message string) error {
	vr.stats.History.Add(&binlogplayer.StatsHistoryRecord{
		Time:    time.Now(),
		Message: message,
	})
	query := fmt.Sprintf("update _vt.vreplication set message=%v where id=%v", encodeString(binlogplayer.MessageTruncate(message)), vr.id)
	if _, err := vr.dbClient.Execute(query); err != nil {
		return fmt.Errorf("could not set message: %v: %v", query, err)
	}
	return nil
}

func (vr *vreplicator) setState(state, message string) error {
	if message != "" {
		vr.stats.History.Add(&binlogplayer.StatsHistoryRecord{
			Time:    time.Now(),
			Message: message,
		})
	}
	vr.stats.State.Set(state)
	query := fmt.Sprintf("update _vt.vreplication set state='%v', message=%v where id=%v", state, encodeString(binlogplayer.MessageTruncate(message)), vr.id)
	if _, err := vr.dbClient.ExecuteFetch(query, 1); err != nil {
		return fmt.Errorf("could not set state: %v: %v", query, err)
	}
	return nil
}

func encodeString(in string) string {
	var buf strings.Builder
	sqltypes.NewVarChar(in).EncodeSQL(&buf)
	return buf.String()
}

func (vr *vreplicator) getSettingFKCheck() error {
	qr, err := vr.dbClient.Execute("select @@foreign_key_checks;")
	if err != nil {
		return err
	}
	if qr.RowsAffected != 1 || len(qr.Fields) != 1 {
		return fmt.Errorf("unable to select @@foreign_key_checks")
	}
	vr.originalFKCheckSetting, err = evalengine.ToInt64(qr.Rows[0][0])
	if err != nil {
		return err
	}
	return nil
}

func (vr *vreplicator) resetFKCheckAfterCopy() error {
	_, err := vr.dbClient.Execute(fmt.Sprintf("set foreign_key_checks=%d;", vr.originalFKCheckSetting))
	return err
}

func (vr *vreplicator) clearFKCheck() error {
	_, err := vr.dbClient.Execute("set foreign_key_checks=0;")
	return err
}
<|endoftext|>


<commit_before>package daemon

import (
	"encoding/json"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"

	"gopkg.in/sorcix/irc.v1"
)

type Backlog struct {
	directory string
	targets   map[string]io.ReadWriter
	metadata  *Metadata
}

type Metadata struct {
	channels    map[string]ChannelMetadata
	currentNick string
}

type ChannelMetadata struct {
	nicks map[string]bool
	topic string
}

// Sync on disk and in memory representation of metadata
func (m *Metadata) sync(directory string) {
	filename := filepath.Join(directory, "metadata.json")
	ticker := time.NewTicker(15 * time.Second)

	for {
		<-ticker.C

		marshal, err := json.Marshal(m)
		if err != nil {
			log.Println("Couldn't serialize json", err)
		}

		fp, err := os.Open(filename)
		if err != nil {
			log.Fatal(err)
		}

		fp.Write(marshal)
	}
}

func NewBacklog(logdir string, address string) *Backlog {
	directory := filepath.Join(logdir, address)

	if err := os.MkdirAll(directory, 0700); err != nil {
		log.Fatal("Couldn't create backlog directory:", err)
	}

	meta := &Metadata{
		channels:    make(map[string]ChannelMetadata),
		currentNick: "",
	}

	go meta.sync(directory)

	return &Backlog{
		directory: directory,
		targets:   make(map[string]io.ReadWriter, 10),
		metadata:  meta,
	}
}

func (b *Backlog) AddMessage(msg *irc.Message) error {

	switch msg.Command {
	case irc.PRIVMSG, irc.NOTICE:
		break
	case irc.JOIN:
		break
	case irc.PART:
		break

	}
	return nil
}
<commit_msg>os.Open is readonly<commit_after>package daemon

import (
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"time"

	"gopkg.in/sorcix/irc.v1"
)

type Backlog struct {
	directory string
	targets   map[string]io.ReadWriter
	metadata  *Metadata
}

type Metadata struct {
	channels    map[string]ChannelMetadata
	currentNick string
}

type ChannelMetadata struct {
	nicks map[string]bool
	topic string
}

// Sync on disk and in memory representation of metadata
func (m *Metadata) sync(directory string) {
	filename := filepath.Join(directory, "metadata.json")
	ticker := time.NewTicker(15 * time.Second)

	for {
		<-ticker.C

		marshal, err := json.Marshal(m)
		if err != nil {
			log.Println("Couldn't serialize json", err)
		}

		if err := ioutil.WriteFile(filename, marshal, 0700); err != nil {
			log.Fatal(err)
		}
	}
}

func NewBacklog(logdir string, address string) *Backlog {
	directory := filepath.Join(logdir, address)

	if err := os.MkdirAll(directory, 0700); err != nil {
		log.Fatal("Couldn't create backlog directory:", err)
	}

	meta := &Metadata{
		channels:    make(map[string]ChannelMetadata),
		currentNick: "",
	}

	go meta.sync(directory)

	return &Backlog{
		directory: directory,
		targets:   make(map[string]io.ReadWriter, 10),
		metadata:  meta,
	}
}

func (b *Backlog) AddMessage(msg *irc.Message) error {

	switch msg.Command {
	case irc.PRIVMSG, irc.NOTICE:
		break
	case irc.JOIN:
		break
	case irc.PART:
		break

	}
	return nil
}
<|endoftext|>


<commit_before>package daemon

import (
	"io"
	"os/exec"
	"sync"
	"time"

	"github.com/docker/docker/daemon/execdriver"
	"github.com/docker/docker/pkg/log"
	"github.com/docker/docker/runconfig"
)

const defaultTimeIncrement = 100

// containerMonitor monitors the execution of a container's main process.
// If a restart policy is specified for the cotnainer the monitor will ensure that the
// process is restarted based on the rules of the policy.  When the container is finally stopped
// the monitor will reset and cleanup any of the container resources such as networking allocations
// and the rootfs
type containerMonitor struct {
	mux sync.Mutex

	// container is the container being monitored
	container *Container

	// restartPolicy is the current policy being applied to the container monitor
	restartPolicy runconfig.RestartPolicy

	// failureCount is the number of times the container has failed to
	// start in a row
	failureCount int

	// shouldStop signals the monitor that the next time the container exits it is
	// either because docker or the user asked for the container to be stopped
	shouldStop bool

	// startSignal is a channel that is closes after the container initially starts
	startSignal chan struct{}

	// stopChan is used to signal to the monitor whenever there is a wait for the
	// next restart so that the timeIncrement is not honored and the user is not
	// left waiting for nothing to happen during this time
	stopChan chan struct{}

	// timeIncrement is the amount of time to wait between restarts
	// this is in milliseconds
	timeIncrement int

	// lastStartTime is the time which the monitor last exec'd the container's process
	lastStartTime time.Time
}

// newContainerMonitor returns an initialized containerMonitor for the provided container
// honoring the provided restart policy
func newContainerMonitor(container *Container, policy runconfig.RestartPolicy) *containerMonitor {
	return &containerMonitor{
		container:     container,
		restartPolicy: policy,
		timeIncrement: defaultTimeIncrement,
		stopChan:      make(chan struct{}, 1),
		startSignal:   make(chan struct{}, 1),
	}
}

// Stop signals to the container monitor that it should stop monitoring the container
// for exits the next time the process dies
func (m *containerMonitor) ExitOnNext() {
	m.mux.Lock()

	// we need to protect having a double close of the channel when stop is called
	// twice or else we will get a panic
	if !m.shouldStop {
		m.shouldStop = true
		close(m.stopChan)
	}

	m.mux.Unlock()
}

// Close closes the container's resources such as networking allocations and
// unmounts the contatiner's root filesystem
func (m *containerMonitor) Close() error {
	// Cleanup networking and mounts
	m.container.cleanup()

	// FIXME: here is race condition between two RUN instructions in Dockerfile
	// because they share same runconfig and change image. Must be fixed
	// in builder/builder.go
	if err := m.container.toDisk(); err != nil {
		log.Errorf("Error dumping container %s state to disk: %s", m.container.ID, err)

		return err
	}

	return nil
}

// Start starts the containers process and monitors it according to the restart policy
func (m *containerMonitor) Start() error {
	var (
		err        error
		exitStatus int
	)

	// this variable indicates that we under container.Lock
	underLock := true

	// ensure that when the monitor finally exits we release the networking and unmount the rootfs
	defer func() {
		if !underLock {
			m.container.Lock()
			defer m.container.Unlock()
		}
		m.Close()
	}()

	// reset the restart count
	m.container.RestartCount = -1

	for {
		m.container.RestartCount++

		if err := m.container.startLoggingToDisk(); err != nil {
			m.resetContainer()

			return err
		}

		pipes := execdriver.NewPipes(m.container.stdin, m.container.stdout, m.container.stderr, m.container.Config.OpenStdin)

		m.container.LogEvent("start")

		m.lastStartTime = time.Now()

		if exitStatus, err = m.container.daemon.Run(m.container, pipes, m.callback); err != nil {
			// if we receive an internal error from the initial start of a container then lets
			// return it instead of entering the restart loop
			if m.container.RestartCount == 0 {
				m.resetContainer()

				return err
			}

			log.Errorf("Error running container: %s", err)
		}

		// here container.Lock is already lost
		underLock = false

		m.resetMonitor(err == nil && exitStatus == 0)

		if m.shouldRestart(exitStatus) {
			m.container.State.SetRestarting(exitStatus)

			m.container.LogEvent("die")

			m.resetContainer()

			// sleep with a small time increment between each restart to help avoid issues cased by quickly
			// restarting the container because of some types of errors ( networking cut out, etc... )
			m.waitForNextRestart()

			// we need to check this before reentering the loop because the waitForNextRestart could have
			// been terminated by a request from a user
			if m.shouldStop {
				m.container.State.SetStopped(exitStatus)

				return err
			}

			continue
		}

		m.container.State.SetStopped(exitStatus)

		m.container.LogEvent("die")

		m.resetContainer()

		break
	}

	return err
}

// resetMonitor resets the stateful fields on the containerMonitor based on the
// previous runs success or failure.  Reguardless of success, if the container had
// an execution time of more than 10s then reset the timer back to the default
func (m *containerMonitor) resetMonitor(successful bool) {
	executionTime := time.Now().Sub(m.lastStartTime).Seconds()

	if executionTime > 10 {
		m.timeIncrement = defaultTimeIncrement
	} else {
		// otherwise we need to increment the amount of time we wait before restarting
		// the process.  We will build up by multiplying the increment by 2
		m.timeIncrement *= 2
	}

	// the container exited successfully so we need to reset the failure counter
	if successful {
		m.failureCount = 0
	} else {
		m.failureCount++
	}
}

// waitForNextRestart waits with the default time increment to restart the container unless
// a user or docker asks for the container to be stopped
func (m *containerMonitor) waitForNextRestart() {
	select {
	case <-time.After(time.Duration(m.timeIncrement) * time.Millisecond):
	case <-m.stopChan:
	}
}

// shouldRestart checks the restart policy and applies the rules to determine if
// the container's process should be restarted
func (m *containerMonitor) shouldRestart(exitStatus int) bool {
	m.mux.Lock()
	defer m.mux.Unlock()

	// do not restart if the user or docker has requested that this container be stopped
	if m.shouldStop {
		return false
	}

	switch m.restartPolicy.Name {
	case "always":
		return true
	case "on-failure":
		// the default value of 0 for MaximumRetryCount means that we will not enforce a maximum count
		if max := m.restartPolicy.MaximumRetryCount; max != 0 && m.failureCount >= max {
			log.Debugf("stopping restart of container %s because maximum failure could of %d has been reached", max)
			return false
		}

		return exitStatus != 0
	}

	return false
}

// callback ensures that the container's state is properly updated after we
// received ack from the execution drivers
func (m *containerMonitor) callback(command *execdriver.Command) {
	if command.Tty {
		// The callback is called after the process Start()
		// so we are in the parent process. In TTY mode, stdin/out/err is the PtySlace
		// which we close here.
		if c, ok := command.Stdout.(io.Closer); ok {
			c.Close()
		}
	}

	m.container.State.SetRunning(command.Pid())

	// signal that the process has started
	// close channel only if not closed
	select {
	case <-m.startSignal:
	default:
		close(m.startSignal)
	}

	if err := m.container.ToDisk(); err != nil {
		log.Debugf("%s", err)
	}
}

// resetContainer resets the container's IO and ensures that the command is able to be executed again
// by copying the data into a new struct
func (m *containerMonitor) resetContainer() {
	container := m.container

	if container.Config.OpenStdin {
		if err := container.stdin.Close(); err != nil {
			log.Errorf("%s: Error close stdin: %s", container.ID, err)
		}
	}

	if err := container.stdout.Clean(); err != nil {
		log.Errorf("%s: Error close stdout: %s", container.ID, err)
	}

	if err := container.stderr.Clean(); err != nil {
		log.Errorf("%s: Error close stderr: %s", container.ID, err)
	}

	if container.command != nil && container.command.Terminal != nil {
		if err := container.command.Terminal.Close(); err != nil {
			log.Errorf("%s: Error closing terminal: %s", container.ID, err)
		}
	}

	// Re-create a brand new stdin pipe once the container exited
	if container.Config.OpenStdin {
		container.stdin, container.stdinPipe = io.Pipe()
	}

	c := container.command.Cmd

	container.command.Cmd = exec.Cmd{
		Stdin:       c.Stdin,
		Stdout:      c.Stdout,
		Stderr:      c.Stderr,
		Path:        c.Path,
		Env:         c.Env,
		ExtraFiles:  c.ExtraFiles,
		Args:        c.Args,
		Dir:         c.Dir,
		SysProcAttr: c.SysProcAttr,
	}
}
<commit_msg>Cleanup: Make channels unbuffered<commit_after>package daemon

import (
	"io"
	"os/exec"
	"sync"
	"time"

	"github.com/docker/docker/daemon/execdriver"
	"github.com/docker/docker/pkg/log"
	"github.com/docker/docker/runconfig"
)

const defaultTimeIncrement = 100

// containerMonitor monitors the execution of a container's main process.
// If a restart policy is specified for the cotnainer the monitor will ensure that the
// process is restarted based on the rules of the policy.  When the container is finally stopped
// the monitor will reset and cleanup any of the container resources such as networking allocations
// and the rootfs
type containerMonitor struct {
	mux sync.Mutex

	// container is the container being monitored
	container *Container

	// restartPolicy is the current policy being applied to the container monitor
	restartPolicy runconfig.RestartPolicy

	// failureCount is the number of times the container has failed to
	// start in a row
	failureCount int

	// shouldStop signals the monitor that the next time the container exits it is
	// either because docker or the user asked for the container to be stopped
	shouldStop bool

	// startSignal is a channel that is closes after the container initially starts
	startSignal chan struct{}

	// stopChan is used to signal to the monitor whenever there is a wait for the
	// next restart so that the timeIncrement is not honored and the user is not
	// left waiting for nothing to happen during this time
	stopChan chan struct{}

	// timeIncrement is the amount of time to wait between restarts
	// this is in milliseconds
	timeIncrement int

	// lastStartTime is the time which the monitor last exec'd the container's process
	lastStartTime time.Time
}

// newContainerMonitor returns an initialized containerMonitor for the provided container
// honoring the provided restart policy
func newContainerMonitor(container *Container, policy runconfig.RestartPolicy) *containerMonitor {
	return &containerMonitor{
		container:     container,
		restartPolicy: policy,
		timeIncrement: defaultTimeIncrement,
		stopChan:      make(chan struct{}),
		startSignal:   make(chan struct{}),
	}
}

// Stop signals to the container monitor that it should stop monitoring the container
// for exits the next time the process dies
func (m *containerMonitor) ExitOnNext() {
	m.mux.Lock()

	// we need to protect having a double close of the channel when stop is called
	// twice or else we will get a panic
	if !m.shouldStop {
		m.shouldStop = true
		close(m.stopChan)
	}

	m.mux.Unlock()
}

// Close closes the container's resources such as networking allocations and
// unmounts the contatiner's root filesystem
func (m *containerMonitor) Close() error {
	// Cleanup networking and mounts
	m.container.cleanup()

	// FIXME: here is race condition between two RUN instructions in Dockerfile
	// because they share same runconfig and change image. Must be fixed
	// in builder/builder.go
	if err := m.container.toDisk(); err != nil {
		log.Errorf("Error dumping container %s state to disk: %s", m.container.ID, err)

		return err
	}

	return nil
}

// Start starts the containers process and monitors it according to the restart policy
func (m *containerMonitor) Start() error {
	var (
		err        error
		exitStatus int
	)

	// this variable indicates that we under container.Lock
	underLock := true

	// ensure that when the monitor finally exits we release the networking and unmount the rootfs
	defer func() {
		if !underLock {
			m.container.Lock()
			defer m.container.Unlock()
		}
		m.Close()
	}()

	// reset the restart count
	m.container.RestartCount = -1

	for {
		m.container.RestartCount++

		if err := m.container.startLoggingToDisk(); err != nil {
			m.resetContainer()

			return err
		}

		pipes := execdriver.NewPipes(m.container.stdin, m.container.stdout, m.container.stderr, m.container.Config.OpenStdin)

		m.container.LogEvent("start")

		m.lastStartTime = time.Now()

		if exitStatus, err = m.container.daemon.Run(m.container, pipes, m.callback); err != nil {
			// if we receive an internal error from the initial start of a container then lets
			// return it instead of entering the restart loop
			if m.container.RestartCount == 0 {
				m.resetContainer()

				return err
			}

			log.Errorf("Error running container: %s", err)
		}

		// here container.Lock is already lost
		underLock = false

		m.resetMonitor(err == nil && exitStatus == 0)

		if m.shouldRestart(exitStatus) {
			m.container.State.SetRestarting(exitStatus)

			m.container.LogEvent("die")

			m.resetContainer()

			// sleep with a small time increment between each restart to help avoid issues cased by quickly
			// restarting the container because of some types of errors ( networking cut out, etc... )
			m.waitForNextRestart()

			// we need to check this before reentering the loop because the waitForNextRestart could have
			// been terminated by a request from a user
			if m.shouldStop {
				m.container.State.SetStopped(exitStatus)

				return err
			}

			continue
		}

		m.container.State.SetStopped(exitStatus)

		m.container.LogEvent("die")

		m.resetContainer()

		break
	}

	return err
}

// resetMonitor resets the stateful fields on the containerMonitor based on the
// previous runs success or failure.  Reguardless of success, if the container had
// an execution time of more than 10s then reset the timer back to the default
func (m *containerMonitor) resetMonitor(successful bool) {
	executionTime := time.Now().Sub(m.lastStartTime).Seconds()

	if executionTime > 10 {
		m.timeIncrement = defaultTimeIncrement
	} else {
		// otherwise we need to increment the amount of time we wait before restarting
		// the process.  We will build up by multiplying the increment by 2
		m.timeIncrement *= 2
	}

	// the container exited successfully so we need to reset the failure counter
	if successful {
		m.failureCount = 0
	} else {
		m.failureCount++
	}
}

// waitForNextRestart waits with the default time increment to restart the container unless
// a user or docker asks for the container to be stopped
func (m *containerMonitor) waitForNextRestart() {
	select {
	case <-time.After(time.Duration(m.timeIncrement) * time.Millisecond):
	case <-m.stopChan:
	}
}

// shouldRestart checks the restart policy and applies the rules to determine if
// the container's process should be restarted
func (m *containerMonitor) shouldRestart(exitStatus int) bool {
	m.mux.Lock()
	defer m.mux.Unlock()

	// do not restart if the user or docker has requested that this container be stopped
	if m.shouldStop {
		return false
	}

	switch m.restartPolicy.Name {
	case "always":
		return true
	case "on-failure":
		// the default value of 0 for MaximumRetryCount means that we will not enforce a maximum count
		if max := m.restartPolicy.MaximumRetryCount; max != 0 && m.failureCount >= max {
			log.Debugf("stopping restart of container %s because maximum failure could of %d has been reached", max)
			return false
		}

		return exitStatus != 0
	}

	return false
}

// callback ensures that the container's state is properly updated after we
// received ack from the execution drivers
func (m *containerMonitor) callback(command *execdriver.Command) {
	if command.Tty {
		// The callback is called after the process Start()
		// so we are in the parent process. In TTY mode, stdin/out/err is the PtySlace
		// which we close here.
		if c, ok := command.Stdout.(io.Closer); ok {
			c.Close()
		}
	}

	m.container.State.SetRunning(command.Pid())

	// signal that the process has started
	// close channel only if not closed
	select {
	case <-m.startSignal:
	default:
		close(m.startSignal)
	}

	if err := m.container.ToDisk(); err != nil {
		log.Debugf("%s", err)
	}
}

// resetContainer resets the container's IO and ensures that the command is able to be executed again
// by copying the data into a new struct
func (m *containerMonitor) resetContainer() {
	container := m.container

	if container.Config.OpenStdin {
		if err := container.stdin.Close(); err != nil {
			log.Errorf("%s: Error close stdin: %s", container.ID, err)
		}
	}

	if err := container.stdout.Clean(); err != nil {
		log.Errorf("%s: Error close stdout: %s", container.ID, err)
	}

	if err := container.stderr.Clean(); err != nil {
		log.Errorf("%s: Error close stderr: %s", container.ID, err)
	}

	if container.command != nil && container.command.Terminal != nil {
		if err := container.command.Terminal.Close(); err != nil {
			log.Errorf("%s: Error closing terminal: %s", container.ID, err)
		}
	}

	// Re-create a brand new stdin pipe once the container exited
	if container.Config.OpenStdin {
		container.stdin, container.stdinPipe = io.Pipe()
	}

	c := container.command.Cmd

	container.command.Cmd = exec.Cmd{
		Stdin:       c.Stdin,
		Stdout:      c.Stdout,
		Stderr:      c.Stderr,
		Path:        c.Path,
		Env:         c.Env,
		ExtraFiles:  c.ExtraFiles,
		Args:        c.Args,
		Dir:         c.Dir,
		SysProcAttr: c.SysProcAttr,
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2015 Ludovic Fauvet
// Licensed under the MIT license

package daemon

import (
	"errors"
	"fmt"
	"github.com/etix/mirrorbits/cli"
	. "github.com/etix/mirrorbits/config"
	"github.com/etix/mirrorbits/core"
	"github.com/etix/mirrorbits/database"
	"github.com/etix/mirrorbits/mirrors"
	"github.com/etix/mirrorbits/scan"
	"github.com/etix/mirrorbits/utils"
	"github.com/garyburd/redigo/redis"
	"github.com/op/go-logging"
	"math/rand"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	healthCheckThreads = 10
	userAgent          = "Mirrorbits/" + core.VERSION + " PING CHECK"
	clientTimeout      = time.Duration(20 * time.Second)
	clientDeadline     = time.Duration(40 * time.Second)
	redirectError      = errors.New("Redirect not allowed")
	mirrorNotScanned   = errors.New("Mirror has not yet been scanned")

	log = logging.MustGetLogger("main")
)

type Monitor struct {
	redis           *database.Redis
	cache           *mirrors.Cache
	mirrors         map[string]*Mirror
	mapLock         sync.Mutex
	httpClient      http.Client
	healthCheckChan chan string
	syncChan        chan string
	stop            chan bool
	wg              sync.WaitGroup
	formatLongestID int

	cluster *cluster
}

type Mirror struct {
	mirrors.Mirror
	checking  bool
	scanning  bool
	lastCheck int64
}

func (m *Mirror) NeedHealthCheck() bool {
	return utils.ElapsedSec(m.lastCheck, int64(60*GetConfig().CheckInterval))
}

func (m *Mirror) NeedSync() bool {
	return utils.ElapsedSec(m.LastSync, int64(60*GetConfig().ScanInterval))
}

func (m *Mirror) IsScanning() bool {
	return m.scanning
}

func (m *Mirror) IsChecking() bool {
	return m.checking
}

func NewMonitor(r *database.Redis, c *mirrors.Cache) *Monitor {
	monitor := new(Monitor)
	monitor.redis = r
	monitor.cache = c
	monitor.cluster = NewCluster(r)
	monitor.mirrors = make(map[string]*Mirror)
	monitor.healthCheckChan = make(chan string, healthCheckThreads*5)
	monitor.syncChan = make(chan string)
	monitor.stop = make(chan bool)

	rand.Seed(time.Now().UnixNano())

	transport := http.Transport{
		DisableKeepAlives:   true,
		MaxIdleConnsPerHost: 0,
		Dial: func(network, addr string) (net.Conn, error) {
			deadline := time.Now().Add(clientDeadline)
			c, err := net.DialTimeout(network, addr, clientTimeout)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	monitor.httpClient = http.Client{
		CheckRedirect: checkRedirect,
		Transport:     &transport,
	}
	return monitor
}

func (m *Monitor) Stop() {
	select {
	case _, _ = <-m.stop:
		return
	default:
		m.cluster.Stop()
		close(m.stop)
	}
}

func (m *Monitor) Wait() {
	m.wg.Wait()
}

// Return an error if the endpoint is an unauthorized redirect
func checkRedirect(req *http.Request, via []*http.Request) error {
	if GetConfig().DisallowRedirects {
		return redirectError
	}
	return nil
}

// Main monitor loop
func (m *Monitor) MonitorLoop() {
	m.wg.Add(1)

	mirrorUpdateEvent := make(chan string, 10)
	m.redis.Pubsub.SubscribeEvent(database.MIRROR_UPDATE, mirrorUpdateEvent)

	// Scan the local repository
	m.retry(func() error {
		return m.scanRepository()
	}, 1*time.Second)

	// Synchronize the list of all known mirrors
	m.retry(func() error {
		ids, err := m.mirrorsID()
		if err != nil {
			return err
		}
		m.syncMirrorList(ids...)
		return nil
	}, 500*time.Millisecond)

	// Start the cluster manager
	m.cluster.Start()

	// Start the health check routines
	for i := 0; i < healthCheckThreads; i++ {
		go m.healthCheckLoop()
	}

	// Start the mirror sync routines
	for i := 0; i < GetConfig().ConcurrentSync; i++ {
		go m.syncLoop()
	}

	// Setup recurrent tasks
	repositoryScanInterval := -1
	repositoryScanTicker := time.NewTicker(1 * time.Minute)
	mirrorCheckTicker := time.NewTicker(1 * time.Second)
	configUpdateTicker := time.NewTicker(1 * time.Second)

	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case id := <-mirrorUpdateEvent:
			m.syncMirrorList(id)
		case <-configUpdateTicker.C:
			if repositoryScanInterval != GetConfig().RepositoryScanInterval {
				repositoryScanInterval = GetConfig().RepositoryScanInterval

				if repositoryScanInterval == 0 {
					repositoryScanTicker.Stop()
				} else {
					repositoryScanTicker = time.NewTicker(time.Duration(repositoryScanInterval) * time.Minute)
				}
			}
		case <-repositoryScanTicker.C:
			m.scanRepository()
		case <-mirrorCheckTicker.C:
			m.mapLock.Lock()
			for k, v := range m.mirrors {
				if !v.Enabled {
					// Ignore disabled mirrors
					continue
				}
				if v.NeedHealthCheck() && !v.IsChecking() && m.cluster.IsHandled(k) {
					select {
					case m.healthCheckChan <- k:
						m.mirrors[k].checking = true
					default:
					}
				}
				if v.NeedSync() && !v.IsScanning() && m.cluster.IsHandled(k) {
					select {
					case m.syncChan <- k:
						m.mirrors[k].scanning = true
					default:
					}
				}
			}
			m.mapLock.Unlock()
		}
	}
}

// Returns a list of all mirrors ID
func (m *Monitor) mirrorsID() ([]string, error) {
	rconn := m.redis.Get()
	defer rconn.Close()

	return redis.Strings(rconn.Do("LRANGE", "MIRRORS", "0", "-1"))
}

// Sync the remote mirror struct with the local dataset
// TODO needs improvements
func (m *Monitor) syncMirrorList(mirrorsIDs ...string) ([]mirrors.Mirror, error) {

	mlist := make([]mirrors.Mirror, 0, len(mirrorsIDs))

	for _, id := range mirrorsIDs {
		if len(id) > m.formatLongestID {
			m.formatLongestID = len(id)
		}
		mirror, err := m.cache.GetMirror(id)
		if err != nil && err != redis.ErrNil {
			log.Error("Fetching mirror %s failed: %s", id, err.Error())
			continue
		} else if err == redis.ErrNil {
			// Mirror has been deleted
			m.mapLock.Lock()
			delete(m.mirrors, id)
			m.cluster.RemoveMirror(&mirror)
			m.mapLock.Unlock()
			continue
		}
		mlist = append(mlist, mirror)

		m.mapLock.Lock()
		m.cluster.AddMirror(&mirror)
		m.mapLock.Unlock()
	}

	m.mapLock.Lock()

	// Prepare the list of mirrors
	for _, e := range mlist {
		var isChecking bool = false
		var isScanning bool = false
		var lastCheck int64 = 0
		if _, ok := m.mirrors[e.ID]; ok {
			isChecking = m.mirrors[e.ID].checking
			isScanning = m.mirrors[e.ID].scanning
			lastCheck = m.mirrors[e.ID].lastCheck
		}
		m.mirrors[e.ID] = &Mirror{
			Mirror:    e,
			checking:  isChecking,
			scanning:  isScanning,
			lastCheck: lastCheck,
		}
	}
	m.mapLock.Unlock()

	log.Debug("%d mirror%s updated", len(mirrorsIDs), utils.Plural(len(mirrorsIDs)))
	return mlist, nil
}

// Main health check loop
// TODO merge with the monitorLoop?
func (m *Monitor) healthCheckLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.healthCheckChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			if m.healthCheck(mirror.Mirror) == mirrorNotScanned {
				// Not removing the 'checking' lock is intended here so the mirror won't
				// be checked again until the rsync/ftp scan is finished.
				continue
			}

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastCheck = time.Now().UTC().Unix()
				m.mirrors[k].checking = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Main sync loop
// TODO merge with the monitorLoop?
func (m *Monitor) syncLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.syncChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			conn := m.redis.Get()
			scanning, err := scan.IsScanning(conn, k)
			if err != nil {
				log.Error("syncloop: ", err.Error())
				conn.Close()
				goto unlock
			} else if scanning {
				// A scan is already in progress on another node
				conn.Close()
				goto unlock
			}
			conn.Close()

			log.Debug("Scanning %s", k)

			err = cli.NoSyncMethod

			// First try to scan with rsync
			if mirror.RsyncURL != "" {
				err = scan.Scan(scan.RSYNC, m.redis, mirror.RsyncURL, k, m.stop)
			}
			// If it failed or rsync wasn't supported
			// fallback to FTP
			if err != nil && mirror.FtpURL != "" {
				err = scan.Scan(scan.FTP, m.redis, mirror.FtpURL, k, m.stop)
			}

			if err == scan.ScanInProgress {
				log.Warning("%-30.30s Scan already in progress", k)
				goto unlock
			}

			if mirror.Up == false {
				select {
				case m.healthCheckChan <- k:
				default:
				}
			}

		unlock:
			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].scanning = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Do an actual health check against a given mirror
func (m *Monitor) healthCheck(mirror mirrors.Mirror) error {
	format := "%-" + fmt.Sprintf("%d.%ds", m.formatLongestID+4, m.formatLongestID+4)

	file, size, err := m.getRandomFile(mirror.ID)
	if err != nil {
		if err == redis.ErrNil {
			return mirrorNotScanned
		} else {
			log.Warning(format+"Error: Cannot obtain a random file: %s", mirror.ID, err)
		}
		return err
	}

	req, err := http.NewRequest("HEAD", strings.TrimRight(mirror.HttpURL, "/")+file, nil)
	req.Header.Set("User-Agent", userAgent)
	req.Close = true

	start := time.Now()
	resp, err := m.httpClient.Do(req)
	elapsed := time.Since(start)

	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			log.Debug("Op: %s | Net: %s | Addr: %s | Err: %s | Temporary: %t", opErr.Op, opErr.Net, opErr.Addr, opErr.Error(), opErr.Temporary())
		}
		mirrors.MarkMirrorDown(m.redis, mirror.ID, "Unreachable")
		log.Error(format+"Error: %s (%dms)", mirror.ID, err.Error(), elapsed/time.Millisecond)
		return err
	}
	resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")

	if resp.StatusCode == 404 {
		mirrors.MarkMirrorDown(m.redis, mirror.ID, fmt.Sprintf("File not found %s (error 404)", file))
		if GetConfig().DisableOnMissingFile {
			mirrors.DisableMirror(m.redis, mirror.ID)
		}
		log.Error(format+"Error: File %s not found (error 404)", mirror.ID, file)
	} else if resp.StatusCode != 200 {
		mirrors.MarkMirrorDown(m.redis, mirror.ID, fmt.Sprintf("Got status code %d", resp.StatusCode))
		log.Warning(format+"Down! Status: %d", mirror.ID, resp.StatusCode)
	} else {
		mirrors.MarkMirrorUp(m.redis, mirror.ID)
		rsize, err := strconv.ParseInt(contentLength, 10, 64)
		if err == nil && rsize != size {
			log.Warning(format+"File size mismatch! [%s] (%dms)", mirror.ID, file, elapsed/time.Millisecond)
		} else {
			log.Notice(format+"Up! (%dms)", mirror.ID, elapsed/time.Millisecond)
		}
	}
	return nil
}

// Get a random filename known to be served by the given mirror
func (m *Monitor) getRandomFile(identifier string) (file string, size int64, err error) {
	sinterKey := fmt.Sprintf("HANDLEDFILES_%s", identifier)

	rconn := m.redis.Get()
	defer rconn.Close()

	file, err = redis.String(rconn.Do("SRANDMEMBER", sinterKey))
	if err != nil {
		return
	}

	size, err = redis.Int64(rconn.Do("HGET", fmt.Sprintf("FILE_%s", file), "size"))
	if err != nil {
		return
	}

	return
}

// Trigger a sync of the local repository
func (m *Monitor) scanRepository() error {
	err := scan.ScanSource(m.redis, m.stop)
	if err != nil {
		log.Error("Scanning source failed: %s", err.Error())
	}
	return err
}

// Retry a function until no errors is returned while still allowing
// the process to be stopped.
func (m *Monitor) retry(fn func() error, delay time.Duration) {
	for {
		err := fn()
		if err == nil {
			break
		}
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case <-time.After(delay):
		}
	}
}
<commit_msg>monitor: remove the ugly timer checking for configuration change<commit_after>// Copyright (c) 2014-2015 Ludovic Fauvet
// Licensed under the MIT license

package daemon

import (
	"errors"
	"fmt"
	"github.com/etix/mirrorbits/cli"
	. "github.com/etix/mirrorbits/config"
	"github.com/etix/mirrorbits/core"
	"github.com/etix/mirrorbits/database"
	"github.com/etix/mirrorbits/mirrors"
	"github.com/etix/mirrorbits/scan"
	"github.com/etix/mirrorbits/utils"
	"github.com/garyburd/redigo/redis"
	"github.com/op/go-logging"
	"math/rand"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	healthCheckThreads = 10
	userAgent          = "Mirrorbits/" + core.VERSION + " PING CHECK"
	clientTimeout      = time.Duration(20 * time.Second)
	clientDeadline     = time.Duration(40 * time.Second)
	redirectError      = errors.New("Redirect not allowed")
	mirrorNotScanned   = errors.New("Mirror has not yet been scanned")

	log = logging.MustGetLogger("main")
)

type Monitor struct {
	redis           *database.Redis
	cache           *mirrors.Cache
	mirrors         map[string]*Mirror
	mapLock         sync.Mutex
	httpClient      http.Client
	healthCheckChan chan string
	syncChan        chan string
	stop            chan bool
	configNotifier  chan bool
	wg              sync.WaitGroup
	formatLongestID int

	cluster *cluster
}

type Mirror struct {
	mirrors.Mirror
	checking  bool
	scanning  bool
	lastCheck int64
}

func (m *Mirror) NeedHealthCheck() bool {
	return utils.ElapsedSec(m.lastCheck, int64(60*GetConfig().CheckInterval))
}

func (m *Mirror) NeedSync() bool {
	return utils.ElapsedSec(m.LastSync, int64(60*GetConfig().ScanInterval))
}

func (m *Mirror) IsScanning() bool {
	return m.scanning
}

func (m *Mirror) IsChecking() bool {
	return m.checking
}

func NewMonitor(r *database.Redis, c *mirrors.Cache) *Monitor {
	monitor := new(Monitor)
	monitor.redis = r
	monitor.cache = c
	monitor.cluster = NewCluster(r)
	monitor.mirrors = make(map[string]*Mirror)
	monitor.healthCheckChan = make(chan string, healthCheckThreads*5)
	monitor.syncChan = make(chan string)
	monitor.stop = make(chan bool)
	monitor.configNotifier = make(chan bool, 1)

	SubscribeConfig(monitor.configNotifier)

	rand.Seed(time.Now().UnixNano())

	transport := http.Transport{
		DisableKeepAlives:   true,
		MaxIdleConnsPerHost: 0,
		Dial: func(network, addr string) (net.Conn, error) {
			deadline := time.Now().Add(clientDeadline)
			c, err := net.DialTimeout(network, addr, clientTimeout)
			if err != nil {
				return nil, err
			}
			c.SetDeadline(deadline)
			return c, nil
		},
	}

	monitor.httpClient = http.Client{
		CheckRedirect: checkRedirect,
		Transport:     &transport,
	}
	return monitor
}

func (m *Monitor) Stop() {
	select {
	case _, _ = <-m.stop:
		return
	default:
		m.cluster.Stop()
		close(m.stop)
	}
}

func (m *Monitor) Wait() {
	m.wg.Wait()
}

// Return an error if the endpoint is an unauthorized redirect
func checkRedirect(req *http.Request, via []*http.Request) error {
	if GetConfig().DisallowRedirects {
		return redirectError
	}
	return nil
}

// Main monitor loop
func (m *Monitor) MonitorLoop() {
	m.wg.Add(1)

	mirrorUpdateEvent := make(chan string, 10)
	m.redis.Pubsub.SubscribeEvent(database.MIRROR_UPDATE, mirrorUpdateEvent)

	// Scan the local repository
	m.retry(func() error {
		return m.scanRepository()
	}, 1*time.Second)

	// Synchronize the list of all known mirrors
	m.retry(func() error {
		ids, err := m.mirrorsID()
		if err != nil {
			return err
		}
		m.syncMirrorList(ids...)
		return nil
	}, 500*time.Millisecond)

	// Start the cluster manager
	m.cluster.Start()

	// Start the health check routines
	for i := 0; i < healthCheckThreads; i++ {
		go m.healthCheckLoop()
	}

	// Start the mirror sync routines
	for i := 0; i < GetConfig().ConcurrentSync; i++ {
		go m.syncLoop()
	}

	// Setup recurrent tasks
	var repositoryScanTicker <-chan time.Time
	repositoryScanInterval := -1
	mirrorCheckTicker := time.NewTicker(1 * time.Second)

	select {
	case m.configNotifier <- true:
	default:
	}

	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case id := <-mirrorUpdateEvent:
			m.syncMirrorList(id)
		case <-m.configNotifier:
			if repositoryScanInterval != GetConfig().RepositoryScanInterval {
				repositoryScanInterval = GetConfig().RepositoryScanInterval

				if repositoryScanInterval == 0 {
					repositoryScanTicker = nil
				} else {
					repositoryScanTicker = time.Tick(time.Duration(repositoryScanInterval) * time.Minute)
				}
			}
		case <-repositoryScanTicker:
			m.scanRepository()
		case <-mirrorCheckTicker.C:
			m.mapLock.Lock()
			for k, v := range m.mirrors {
				if !v.Enabled {
					// Ignore disabled mirrors
					continue
				}
				if v.NeedHealthCheck() && !v.IsChecking() && m.cluster.IsHandled(k) {
					select {
					case m.healthCheckChan <- k:
						m.mirrors[k].checking = true
					default:
					}
				}
				if v.NeedSync() && !v.IsScanning() && m.cluster.IsHandled(k) {
					select {
					case m.syncChan <- k:
						m.mirrors[k].scanning = true
					default:
					}
				}
			}
			m.mapLock.Unlock()
		}
	}
}

// Returns a list of all mirrors ID
func (m *Monitor) mirrorsID() ([]string, error) {
	rconn := m.redis.Get()
	defer rconn.Close()

	return redis.Strings(rconn.Do("LRANGE", "MIRRORS", "0", "-1"))
}

// Sync the remote mirror struct with the local dataset
// TODO needs improvements
func (m *Monitor) syncMirrorList(mirrorsIDs ...string) ([]mirrors.Mirror, error) {

	mlist := make([]mirrors.Mirror, 0, len(mirrorsIDs))

	for _, id := range mirrorsIDs {
		if len(id) > m.formatLongestID {
			m.formatLongestID = len(id)
		}
		mirror, err := m.cache.GetMirror(id)
		if err != nil && err != redis.ErrNil {
			log.Error("Fetching mirror %s failed: %s", id, err.Error())
			continue
		} else if err == redis.ErrNil {
			// Mirror has been deleted
			m.mapLock.Lock()
			delete(m.mirrors, id)
			m.cluster.RemoveMirror(&mirror)
			m.mapLock.Unlock()
			continue
		}
		mlist = append(mlist, mirror)

		m.mapLock.Lock()
		m.cluster.AddMirror(&mirror)
		m.mapLock.Unlock()
	}

	m.mapLock.Lock()

	// Prepare the list of mirrors
	for _, e := range mlist {
		var isChecking bool = false
		var isScanning bool = false
		var lastCheck int64 = 0
		if _, ok := m.mirrors[e.ID]; ok {
			isChecking = m.mirrors[e.ID].checking
			isScanning = m.mirrors[e.ID].scanning
			lastCheck = m.mirrors[e.ID].lastCheck
		}
		m.mirrors[e.ID] = &Mirror{
			Mirror:    e,
			checking:  isChecking,
			scanning:  isScanning,
			lastCheck: lastCheck,
		}
	}
	m.mapLock.Unlock()

	log.Debug("%d mirror%s updated", len(mirrorsIDs), utils.Plural(len(mirrorsIDs)))
	return mlist, nil
}

// Main health check loop
// TODO merge with the monitorLoop?
func (m *Monitor) healthCheckLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.healthCheckChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			if m.healthCheck(mirror.Mirror) == mirrorNotScanned {
				// Not removing the 'checking' lock is intended here so the mirror won't
				// be checked again until the rsync/ftp scan is finished.
				continue
			}

			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].lastCheck = time.Now().UTC().Unix()
				m.mirrors[k].checking = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Main sync loop
// TODO merge with the monitorLoop?
func (m *Monitor) syncLoop() {
	m.wg.Add(1)
	for {
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case k := <-m.syncChan:
			m.mapLock.Lock()
			mirror := m.mirrors[k]
			m.mapLock.Unlock()

			conn := m.redis.Get()
			scanning, err := scan.IsScanning(conn, k)
			if err != nil {
				log.Error("syncloop: ", err.Error())
				conn.Close()
				goto unlock
			} else if scanning {
				// A scan is already in progress on another node
				conn.Close()
				goto unlock
			}
			conn.Close()

			log.Debug("Scanning %s", k)

			err = cli.NoSyncMethod

			// First try to scan with rsync
			if mirror.RsyncURL != "" {
				err = scan.Scan(scan.RSYNC, m.redis, mirror.RsyncURL, k, m.stop)
			}
			// If it failed or rsync wasn't supported
			// fallback to FTP
			if err != nil && mirror.FtpURL != "" {
				err = scan.Scan(scan.FTP, m.redis, mirror.FtpURL, k, m.stop)
			}

			if err == scan.ScanInProgress {
				log.Warning("%-30.30s Scan already in progress", k)
				goto unlock
			}

			if mirror.Up == false {
				select {
				case m.healthCheckChan <- k:
				default:
				}
			}

		unlock:
			m.mapLock.Lock()
			if _, ok := m.mirrors[k]; ok {
				m.mirrors[k].scanning = false
			}
			m.mapLock.Unlock()
		}
	}
}

// Do an actual health check against a given mirror
func (m *Monitor) healthCheck(mirror mirrors.Mirror) error {
	format := "%-" + fmt.Sprintf("%d.%ds", m.formatLongestID+4, m.formatLongestID+4)

	file, size, err := m.getRandomFile(mirror.ID)
	if err != nil {
		if err == redis.ErrNil {
			return mirrorNotScanned
		} else {
			log.Warning(format+"Error: Cannot obtain a random file: %s", mirror.ID, err)
		}
		return err
	}

	req, err := http.NewRequest("HEAD", strings.TrimRight(mirror.HttpURL, "/")+file, nil)
	req.Header.Set("User-Agent", userAgent)
	req.Close = true

	start := time.Now()
	resp, err := m.httpClient.Do(req)
	elapsed := time.Since(start)

	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			log.Debug("Op: %s | Net: %s | Addr: %s | Err: %s | Temporary: %t", opErr.Op, opErr.Net, opErr.Addr, opErr.Error(), opErr.Temporary())
		}
		mirrors.MarkMirrorDown(m.redis, mirror.ID, "Unreachable")
		log.Error(format+"Error: %s (%dms)", mirror.ID, err.Error(), elapsed/time.Millisecond)
		return err
	}
	resp.Body.Close()

	contentLength := resp.Header.Get("Content-Length")

	if resp.StatusCode == 404 {
		mirrors.MarkMirrorDown(m.redis, mirror.ID, fmt.Sprintf("File not found %s (error 404)", file))
		if GetConfig().DisableOnMissingFile {
			mirrors.DisableMirror(m.redis, mirror.ID)
		}
		log.Error(format+"Error: File %s not found (error 404)", mirror.ID, file)
	} else if resp.StatusCode != 200 {
		mirrors.MarkMirrorDown(m.redis, mirror.ID, fmt.Sprintf("Got status code %d", resp.StatusCode))
		log.Warning(format+"Down! Status: %d", mirror.ID, resp.StatusCode)
	} else {
		mirrors.MarkMirrorUp(m.redis, mirror.ID)
		rsize, err := strconv.ParseInt(contentLength, 10, 64)
		if err == nil && rsize != size {
			log.Warning(format+"File size mismatch! [%s] (%dms)", mirror.ID, file, elapsed/time.Millisecond)
		} else {
			log.Notice(format+"Up! (%dms)", mirror.ID, elapsed/time.Millisecond)
		}
	}
	return nil
}

// Get a random filename known to be served by the given mirror
func (m *Monitor) getRandomFile(identifier string) (file string, size int64, err error) {
	sinterKey := fmt.Sprintf("HANDLEDFILES_%s", identifier)

	rconn := m.redis.Get()
	defer rconn.Close()

	file, err = redis.String(rconn.Do("SRANDMEMBER", sinterKey))
	if err != nil {
		return
	}

	size, err = redis.Int64(rconn.Do("HGET", fmt.Sprintf("FILE_%s", file), "size"))
	if err != nil {
		return
	}

	return
}

// Trigger a sync of the local repository
func (m *Monitor) scanRepository() error {
	err := scan.ScanSource(m.redis, m.stop)
	if err != nil {
		log.Error("Scanning source failed: %s", err.Error())
	}
	return err
}

// Retry a function until no errors is returned while still allowing
// the process to be stopped.
func (m *Monitor) retry(fn func() error, delay time.Duration) {
	for {
		err := fn()
		if err == nil {
			break
		}
		select {
		case <-m.stop:
			m.wg.Done()
			return
		case <-time.After(delay):
		}
	}
}
<|endoftext|>


<commit_before>package libchunk_test

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/advanderveer/libchunk"
	"github.com/boltdb/bolt"
	"github.com/restic/chunker"
)

const KeySize = 32

type KeyFunc func([]byte) [KeySize]byte

type Chunker interface {
	Next(data []byte) (chunker.Chunk, error)
}

type quiter interface {
	Fatalf(format string, args ...interface{})
}

var secret = libchunk.Secret{
	0x3D, 0xA3, 0x35, 0x8B, 0x4D, 0xC1, 0x73, 0x00, //polynomial
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, //random bytes
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
}

func randr(size int64) io.Reader {
	return io.LimitReader(rand.New(rand.NewSource(time.Now().UnixNano())), size)
}

func randb(size int64) []byte {
	b, err := ioutil.ReadAll(randr(size))
	if err != nil {
		panic(err)
	}

	return b
}

var boltSize int64
var boltPath string
var boltKeys [][32]byte
var boltInput []byte

func BenchmarkConcurrentChunkingSha2Bolt(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		db    *bolt.DB
		hash  KeyFunc
		chnkr Chunker

		//@TODO add a (local)store interface: fs/bolt
	)

	boltSize = int64(1024 * 1024 * 1024)
	boltInput = randb(boltSize)
	r := bytes.NewReader(boltInput)
	chnkr = chunker.New(r, secret.Pol())
	buf := make([]byte, chunker.MaxSize)
	dir, _ := ioutil.TempDir("", "libchunk")
	boltPath = filepath.Join(dir, "db.bolt")

	hash = func(b []byte) [KeySize]byte {
		return sha256.Sum256(b)
	}

	db, _ = bolt.Open(boltPath, 0777, nil)
	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	defer db.Close()

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		r.Seek(0, 0)
		chnkr = chunker.New(r, secret.Pol())
		wg := &sync.WaitGroup{}
		for {
			chunk, err := chnkr.Next(buf)
			if err == io.EOF {
				break
			}

			wg.Add(1)
			go func() {
				defer wg.Done()
				k := hash(chunk.Data)

				wg.Add(1)
				go func() {
					defer wg.Done()
					encrypted := aead.Seal(nil, k[:], chunk.Data, nil)
					if bytes.Equal(encrypted, chunk.Data) {
						b.Fatal("huh")
					}

					wg.Add(1)
					go func() {
						defer wg.Done()

						err = db.Batch(func(tx *bolt.Tx) error {
							b, _ := tx.CreateBucketIfNotExists([]byte("a"))
							return b.Put(k[:], encrypted)
						})

						if err != nil {
							b.Fatal(err)
						}

						boltKeys = append(boltKeys, k)
					}()
				}()
			}()
		}

		wg.Wait()
	}
}

func BenchmarkConcurrentChunkingSha2FromBolt(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		db    *bolt.DB
	)

	db, _ = bolt.Open(boltPath, 0777, nil)
	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	defer db.Close()
	defer os.Remove(boltPath)

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		outf, _ := ioutil.TempFile("", "libchunk_")
		defer outf.Close()
		defer os.Remove(outf.Name())

		for _, k := range boltKeys {
			var chunk []byte
			err = db.View(func(tx *bolt.Tx) error {
				b := tx.Bucket([]byte("a"))
				v := b.Get(k[:])
				if v == nil || len(v) < 1 {
					return fmt.Errorf("not found")
				}

				chunk = make([]byte, len(v))
				copy(chunk, v)
				return nil
			})

			if err != nil {
				b.Fatal(err)
			}

			plaintext, err := aead.Open(nil, k[:], chunk, nil)
			if err != nil {
				b.Fatal(err)
			}

			//@TODO benchmark output to file: benchmarked at 700Mb/s
			//@TODO benchmark push to s3, goal: 1GB/s
			_, err = outf.Write(plaintext)
			if err != nil {
				b.Fatal(err)
			}
		}
	}
}

var diskDir string
var diskKeys [][32]byte
var diskSize int64

func BenchmarkConcurrentChunkingSha2Disk(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		hash  KeyFunc
		chnkr Chunker
	)

	diskSize = int64(1024 * 1024 * 1024)
	r := bytes.NewReader(randb(diskSize))
	chnkr = chunker.New(r, secret.Pol())
	buf := make([]byte, chunker.MaxSize)
	diskDir, _ = ioutil.TempDir("", "libchunk")

	hash = func(b []byte) [KeySize]byte { return sha256.Sum256(b) }

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	b.ResetTimer()
	b.SetBytes(int64(diskSize))
	for i := 0; i < b.N; i++ {
		r.Seek(0, 0)
		chnkr = chunker.New(r, secret.Pol())
		wg := &sync.WaitGroup{}
		for {
			chunk, err := chnkr.Next(buf)
			if err == io.EOF {
				break
			}

			wg.Add(1)
			go func() {
				defer wg.Done()
				k := hash(chunk.Data)

				wg.Add(1)
				go func() {
					defer wg.Done()
					encrypted := aead.Seal(nil, k[:], chunk.Data, nil)

					wg.Add(1)
					go func() {
						defer wg.Done()

						f, err := os.OpenFile(filepath.Join(diskDir, fmt.Sprintf("%x", k)), os.O_CREATE|os.O_RDWR|os.O_EXCL, 0777)
						if err != nil {
							if os.IsExist(err) {
								return
							}

							b.Fatal(err)
						}

						defer f.Close()
						_, err = f.Write(encrypted)
						if err != nil {
							b.Fatal(err)
						}

						diskKeys = append(diskKeys, k)
					}()
				}()
			}()
		}

		wg.Wait()
	}
}

func BenchmarkConcurrentChunkingSha2FromFS(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
	)

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	defer os.RemoveAll(diskDir)

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		outf, _ := ioutil.TempFile("", "libchunk_")
		defer outf.Close()
		defer os.Remove(outf.Name())

		for _, k := range diskKeys {
			kpath := filepath.Join(diskDir, fmt.Sprintf("%x", k))

			func() {
				f, err := os.OpenFile(kpath, os.O_RDONLY, 0777)
				if err != nil {
					b.Fatal(err)
				}

				defer f.Close()
				chunk, err := ioutil.ReadAll(f)
				if err != nil {
					b.Fatal(err)
				}

				plaintext, err := aead.Open(nil, k[:], chunk, nil)
				if err != nil {
					b.Fatal(err)
				}

				_, err = outf.Write(plaintext)
				if err != nil {
					b.Fatal(err)
				}
			}()
		}
	}
}

func tempDB(t quiter) (p string, db *bolt.DB) {
	dbdir, err := ioutil.TempDir("", "libchunk_")
	if err != nil {
		t.Fatalf("failed to create temp dir for db: %v", err)
	}

	dbpath := filepath.Join(dbdir, "db.bolt")
	db, err = bolt.Open(dbpath, 0666, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		t.Fatalf("failed to open chunks database '%s': %v", dbpath, err)
	}

	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(libchunk.BucketNameChunks)
		if err != nil {
			return fmt.Errorf("failed to create bucket: %s", err)
		}
		return nil
	})

	if err != nil {
		t.Fatalf("failed to create buckets: %v", err)
	}

	return dbpath, db
}

func TestSplitJoinSmallNonRandom(t *testing.T) {
	_, db := tempDB(t)
	input := []byte("foo bar") //@TODO find somall content that generates multiple chunks
	output := bytes.NewBuffer(nil)

	pr, pw := io.Pipe()
	go func() {
		defer pw.Close()
		err := libchunk.Split(db, secret, bytes.NewBuffer(input), pw)
		if err != nil {
			t.Fatalf("failed to split: %v", err)
		}
	}()

	err := libchunk.Join(db, secret, pr, output)
	if err != nil {
		t.Fatalf("failed to join: %v", err)
	}

	if !bytes.Equal(input, output.Bytes()) {
		t.Errorf("expected joined output (len %d) to be the same as input (len %d)", output.Len(), len(input))
	}

	//assert encryption at rest
	//assert throughput

	//assert:
	//  - assert writer output
	//  - different polynomials generate different
	//    - key
	//    - encrypted content
	//  - different secret parts generate random noise
	//  -assert multiple chunks being outputed

	//assert db file size
	//assert buffer content
	//assert chunks in db
	//assert deduplication
	//assert encryption

}
<commit_msg>added a store interface<commit_after>package libchunk_test

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/advanderveer/libchunk"
	"github.com/boltdb/bolt"
	"github.com/restic/chunker"
)

const KeySize = 32

type KeyFunc func([]byte) [KeySize]byte

type Chunker interface {
	Next(data []byte) (chunker.Chunk, error)
}

//Store holds chunks
type Store interface {
	Put(k [KeySize]byte, chunk []byte) error
	Get(k [KeySize]byte) (chunk []byte, err error)
}

//RemoteStore holds chunks at a remote location such
//that it becomes efficient to keep an up-to-date local
//index to prevent unnesseary interactions
type RemoteStore interface {
	//@TODO think of a remote store interface that uses
	//A local store as an index
}

//
// Test types
//

type fsStore string

func (s fsStore) Put(k [KeySize]byte, chunk []byte) error {
	f, err := os.OpenFile(filepath.Join(string(s), fmt.Sprintf("%x", k)), os.O_CREATE|os.O_RDWR|os.O_EXCL, 0777)
	if err != nil {
		if os.IsExist(err) {
			return nil
		}

		return err
	}

	defer f.Close()
	_, err = f.Write(chunk)
	if err != nil {
		return err
	}

	return nil
}

func (s fsStore) Get(k [KeySize]byte) ([]byte, error) {
	kpath := filepath.Join(string(s), fmt.Sprintf("%x", k))
	f, err := os.OpenFile(kpath, os.O_RDONLY, 0777)
	if err != nil {
		return nil, err
	}

	defer f.Close()
	chunk, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}

	return chunk, nil
}

type boltStore struct {
	db         *bolt.DB
	bucketName []byte
}

func (s *boltStore) Put(k [KeySize]byte, chunk []byte) (err error) {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b, _ := tx.CreateBucketIfNotExists(s.bucketName)
		existing := b.Get(k[:])
		if existing != nil {
			return nil
		}

		return b.Put(k[:], chunk)
	})
}

func (s *boltStore) Get(k [KeySize]byte) (chunk []byte, err error) {
	return chunk, s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(s.bucketName)
		v := b.Get(k[:])
		if v == nil || len(v) < 1 {
			return fmt.Errorf("not found")
		}

		chunk = make([]byte, len(v))
		copy(chunk, v)
		return nil
	})
}

type quiter interface {
	Fatalf(format string, args ...interface{})
}

var secret = libchunk.Secret{
	0x3D, 0xA3, 0x35, 0x8B, 0x4D, 0xC1, 0x73, 0x00, //polynomial
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, //random bytes
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
}

func randr(size int64) io.Reader {
	return io.LimitReader(rand.New(rand.NewSource(time.Now().UnixNano())), size)
}

func randb(size int64) []byte {
	b, err := ioutil.ReadAll(randr(size))
	if err != nil {
		panic(err)
	}

	return b
}

var boltSize int64
var boltPath string
var boltKeys [][32]byte
var boltInput []byte

func BenchmarkConcurrentChunkingSha2Bolt(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		hash  KeyFunc
		chnkr Chunker
		store Store
	)

	boltSize = int64(1024 * 1024 * 1024)
	boltInput = randb(boltSize)
	r := bytes.NewReader(boltInput)
	chnkr = chunker.New(r, secret.Pol())
	buf := make([]byte, chunker.MaxSize)
	dir, _ := ioutil.TempDir("", "libchunk")
	boltPath = filepath.Join(dir, "db.bolt")

	hash = func(b []byte) [KeySize]byte {
		return sha256.Sum256(b)
	}

	db, _ := bolt.Open(boltPath, 0777, nil)
	store = &boltStore{db, []byte("a")}

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	defer db.Close()

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		r.Seek(0, 0)
		chnkr = chunker.New(r, secret.Pol())
		wg := &sync.WaitGroup{}
		for {
			chunk, err := chnkr.Next(buf)
			if err == io.EOF {
				break
			}

			wg.Add(1)
			go func() {
				defer wg.Done()
				k := hash(chunk.Data)

				wg.Add(1)
				go func() {
					defer wg.Done()
					encrypted := aead.Seal(nil, k[:], chunk.Data, nil)
					if bytes.Equal(encrypted, chunk.Data) {
						b.Fatal("huh")
					}

					wg.Add(1)
					go func() {
						defer wg.Done()

						err = store.Put(k, encrypted)
						if err != nil {
							b.Fatal(err)
						}

						boltKeys = append(boltKeys, k)
					}()
				}()
			}()
		}

		wg.Wait()
	}
}

func BenchmarkConcurrentChunkingSha2FromBolt(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		store Store
	)

	db, _ := bolt.Open(boltPath, 0777, nil)

	store = &boltStore{db, []byte("a")}

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	defer db.Close()
	defer os.Remove(boltPath)

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		outf, _ := ioutil.TempFile("", "libchunk_")
		defer outf.Close()
		defer os.Remove(outf.Name())

		for _, k := range boltKeys {
			chunk, err := store.Get(k)
			if err != nil {
				b.Fatal(err)
			}

			plaintext, err := aead.Open(nil, k[:], chunk, nil)
			if err != nil {
				b.Fatal(err)
			}

			//@TODO benchmark output to file: benchmarked at 700Mb/s
			//@TODO benchmark push to s3, goal: 1GB/s
			_, err = outf.Write(plaintext)
			if err != nil {
				b.Fatal(err)
			}
		}
	}
}

var diskDir string
var diskKeys [][32]byte
var diskSize int64

func BenchmarkConcurrentChunkingSha2Disk(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		hash  KeyFunc
		chnkr Chunker
		store Store
	)

	diskSize = int64(1024 * 1024 * 1024)
	r := bytes.NewReader(randb(diskSize))
	chnkr = chunker.New(r, secret.Pol())
	buf := make([]byte, chunker.MaxSize)
	diskDir, _ = ioutil.TempDir("", "libchunk")

	store = fsStore(diskDir)

	hash = func(b []byte) [KeySize]byte { return sha256.Sum256(b) }

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	b.ResetTimer()
	b.SetBytes(int64(diskSize))
	for i := 0; i < b.N; i++ {
		r.Seek(0, 0)
		chnkr = chunker.New(r, secret.Pol())
		wg := &sync.WaitGroup{}
		for {
			chunk, err := chnkr.Next(buf)
			if err == io.EOF {
				break
			}

			wg.Add(1)
			go func() {
				defer wg.Done()
				k := hash(chunk.Data)

				wg.Add(1)
				go func() {
					defer wg.Done()
					encrypted := aead.Seal(nil, k[:], chunk.Data, nil)

					wg.Add(1)
					go func() {
						defer wg.Done()

						err = store.Put(k, encrypted)
						if err != nil {
							b.Fatal(err)
						}

						diskKeys = append(diskKeys, k)
					}()
				}()
			}()
		}

		wg.Wait()
	}
}

func BenchmarkConcurrentChunkingSha2FromFS(b *testing.B) {
	var (
		err   error
		block cipher.Block
		aead  cipher.AEAD
		store Store
	)

	block, err = aes.NewCipher(secret[:])
	if err != nil {
		b.Fatalf("failed to create AES block cipher: %v", err)
	}

	aead, err = cipher.NewGCMWithNonceSize(block, sha256.Size)
	if err != nil {
		b.Fatalf("failed to setup GCM cipher mode: %v", err)
	}

	store = fsStore(diskDir)
	defer os.RemoveAll(diskDir)

	b.ResetTimer()
	b.SetBytes(int64(boltSize))
	for i := 0; i < b.N; i++ {
		outf, _ := ioutil.TempFile("", "libchunk_")
		defer outf.Close()
		defer os.Remove(outf.Name())

		for _, k := range diskKeys {

			func() {
				chunk, err := store.Get(k)
				if err != nil {
					b.Fatal(err)
				}

				plaintext, err := aead.Open(nil, k[:], chunk, nil)
				if err != nil {
					b.Fatal(err)
				}

				_, err = outf.Write(plaintext)
				if err != nil {
					b.Fatal(err)
				}
			}()
		}
	}
}

func tempDB(t quiter) (p string, db *bolt.DB) {
	dbdir, err := ioutil.TempDir("", "libchunk_")
	if err != nil {
		t.Fatalf("failed to create temp dir for db: %v", err)
	}

	dbpath := filepath.Join(dbdir, "db.bolt")
	db, err = bolt.Open(dbpath, 0666, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		t.Fatalf("failed to open chunks database '%s': %v", dbpath, err)
	}

	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(libchunk.BucketNameChunks)
		if err != nil {
			return fmt.Errorf("failed to create bucket: %s", err)
		}
		return nil
	})

	if err != nil {
		t.Fatalf("failed to create buckets: %v", err)
	}

	return dbpath, db
}

func TestSplitJoinSmallNonRandom(t *testing.T) {
	_, db := tempDB(t)
	input := []byte("foo bar") //@TODO find somall content that generates multiple chunks
	output := bytes.NewBuffer(nil)

	pr, pw := io.Pipe()
	go func() {
		defer pw.Close()
		err := libchunk.Split(db, secret, bytes.NewBuffer(input), pw)
		if err != nil {
			t.Fatalf("failed to split: %v", err)
		}
	}()

	err := libchunk.Join(db, secret, pr, output)
	if err != nil {
		t.Fatalf("failed to join: %v", err)
	}

	if !bytes.Equal(input, output.Bytes()) {
		t.Errorf("expected joined output (len %d) to be the same as input (len %d)", output.Len(), len(input))
	}

	//assert encryption at rest
	//assert throughput

	//assert:
	//  - assert writer output
	//  - different polynomials generate different
	//    - key
	//    - encrypted content
	//  - different secret parts generate random noise
	//  -assert multiple chunks being outputed

	//assert db file size
	//assert buffer content
	//assert chunks in db
	//assert deduplication
	//assert encryption

}
<|endoftext|>


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

// Package backup is a simple library for backing up a config file and restoring
// it using a temporary file.
//
// Example usage:
//
//    tmp, err := backup.Save(someFile)
//    if err != nil {
//       return fmt.Errorf("backing up %q: %v", somePath, err)
//    }
//    defer tmp.Release()
//    // ... do real work ...
//    tmp.Restore()
package backup

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
)

// A File records the locations of an original file and a backup copy of that
// file.  The Restore method replaces the contents of the original with the
// backup, overwriting any changes that were made since the backup was created.
type File struct {
	orig, tmp string // file paths
}

// New creates a backup copy of the specified file, located in the same
// directory. The caller should ensure the Release method is called when the
// backup is no longer needed, to clean up.
func New(orig string) (*File, error) {
	f, err := os.Open(orig)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	dir, base := filepath.Split(orig)
	tf, err := ioutil.TempFile(dir, base+".bkup")
	if err != nil {
		return nil, err
	}
	_, err = io.Copy(tf, f)
	cerr := tf.Close()
	if err != nil {
		return nil, err
	} else if cerr != nil {
		return nil, cerr
	}
	return &File{orig, tf.Name()}, nil
}

// Restore puts the original version of the backed up file back in place.
func (f *File) Restore() error {
	if err := os.Rename(f.tmp, f.orig); err != nil {
		return err
	}
	f.tmp = ""
	return nil
}

// Release removes the temporary file copy if it hasn't already been moved.
func (f *File) Release() {
	if f.tmp != "" {
		if err := os.Remove(f.tmp); err != nil {
			log.Printf("Warning: removing backup of %q failed: %v", f.orig, err)
		}
	}
}
<commit_msg>Use filepath.Dir & filepath.Base instead of filepath.Split. (#3053)<commit_after>/*
 * Copyright 2018 The Kythe Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package backup is a simple library for backing up a config file and restoring
// it using a temporary file.
//
// Example usage:
//
//    tmp, err := backup.Save(someFile)
//    if err != nil {
//       return fmt.Errorf("backing up %q: %v", somePath, err)
//    }
//    defer tmp.Release()
//    // ... do real work ...
//    tmp.Restore()
package backup

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
)

// A File records the locations of an original file and a backup copy of that
// file.  The Restore method replaces the contents of the original with the
// backup, overwriting any changes that were made since the backup was created.
type File struct {
	orig, tmp string // file paths
}

// New creates a backup copy of the specified file, located in the same
// directory. The caller should ensure the Release method is called when the
// backup is no longer needed, to clean up.
func New(orig string) (*File, error) {
	f, err := os.Open(orig)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	dir := filepath.Dir(orig)
	base := filepath.Base(orig)
	tf, err := ioutil.TempFile(dir, base+".bkup")
	if err != nil {
		return nil, err
	}
	_, err = io.Copy(tf, f)
	cerr := tf.Close()
	if err != nil {
		return nil, err
	} else if cerr != nil {
		return nil, cerr
	}
	return &File{orig, tf.Name()}, nil
}

// Restore puts the original version of the backed up file back in place.
func (f *File) Restore() error {
	if err := os.Rename(f.tmp, f.orig); err != nil {
		return err
	}
	f.tmp = ""
	return nil
}

// Release removes the temporary file copy if it hasn't already been moved.
func (f *File) Release() {
	if f.tmp != "" {
		if err := os.Remove(f.tmp); err != nil {
			log.Printf("Warning: removing backup of %q failed: %v", f.orig, err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris

package syscall_test

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"unsafe"
)

func TestDirent(t *testing.T) {
	const (
		direntBufSize   = 2048
		filenameMinSize = 11
	)

	d := t.TempDir()
	t.Logf("tmpdir: %s", d)

	for i, c := range []byte("0123456789") {
		name := string(bytes.Repeat([]byte{c}, filenameMinSize+i))
		err := os.WriteFile(filepath.Join(d, name), nil, 0644)
		if err != nil {
			t.Fatalf("writefile: %v", err)
		}
	}

	buf := bytes.Repeat([]byte("DEADBEAF"), direntBufSize/8)
	fd, err := syscall.Open(d, syscall.O_RDONLY, 0)
	if err != nil {
		t.Fatalf("syscall.open: %v", err)
	}
	defer syscall.Close(fd)
	n, err := syscall.ReadDirent(fd, buf)
	if err != nil {
		t.Fatalf("syscall.readdir: %v", err)
	}
	buf = buf[:n]

	names := make([]string, 0, 10)
	for len(buf) > 0 {
		var bc int
		bc, _, names = syscall.ParseDirent(buf, -1, names)
		buf = buf[bc:]
	}

	sort.Strings(names)
	t.Logf("names: %q", names)

	if len(names) != 10 {
		t.Errorf("got %d names; expected 10", len(names))
	}
	for i, name := range names {
		ord, err := strconv.Atoi(name[:1])
		if err != nil {
			t.Fatalf("names[%d] is non-integer %q: %v", i, names[i], err)
		}
		if expected := string(strings.Repeat(name[:1], filenameMinSize+ord)); name != expected {
			t.Errorf("names[%d] is %q (len %d); expected %q (len %d)", i, name, len(name), expected, len(expected))
		}
	}
}

func TestDirentRepeat(t *testing.T) {
	const N = 100
	// Note: the size of the buffer is small enough that the loop
	// below will need to execute multiple times. See issue #31368.
	size := N * unsafe.Offsetof(syscall.Dirent{}.Name) / 4
	if runtime.GOOS == "freebsd" || runtime.GOOS == "netbsd" {
		if size < 1024 {
			size = 1024 // DIRBLKSIZ, see issue 31403.
		}
		if runtime.GOOS == "freebsd" {
			t.Skip("need to fix issue 31416 first")
		}
	}

	// Make a directory containing N files
	d := t.TempDir()

	var files []string
	for i := 0; i < N; i++ {
		files = append(files, fmt.Sprintf("file%d", i))
	}
	for _, file := range files {
		err := os.WriteFile(filepath.Join(d, file), []byte("contents"), 0644)
		if err != nil {
			t.Fatalf("writefile: %v", err)
		}
	}

	// Read the directory entries using ReadDirent.
	fd, err := syscall.Open(d, syscall.O_RDONLY, 0)
	if err != nil {
		t.Fatalf("syscall.open: %v", err)
	}
	defer syscall.Close(fd)
	var files2 []string
	for {
		buf := make([]byte, size)
		n, err := syscall.ReadDirent(fd, buf)
		if err != nil {
			t.Fatalf("syscall.readdir: %v", err)
		}
		if n == 0 {
			break
		}
		buf = buf[:n]
		for len(buf) > 0 {
			var consumed int
			consumed, _, files2 = syscall.ParseDirent(buf, -1, files2)
			buf = buf[consumed:]
		}
	}

	// Check results
	sort.Strings(files)
	sort.Strings(files2)
	if strings.Join(files, "|") != strings.Join(files2, "|") {
		t.Errorf("bad file list: want\n%q\ngot\n%q", files, files2)
	}
}
<commit_msg>syscall: in TestDirent, make as many ReadDirent calls as are needed<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris

package syscall_test

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"unsafe"
)

func TestDirent(t *testing.T) {
	const (
		direntBufSize   = 2048 // arbitrary? See https://go.dev/issue/37323.
		filenameMinSize = 11
	)

	d := t.TempDir()
	t.Logf("tmpdir: %s", d)

	for i, c := range []byte("0123456789") {
		name := string(bytes.Repeat([]byte{c}, filenameMinSize+i))
		err := os.WriteFile(filepath.Join(d, name), nil, 0644)
		if err != nil {
			t.Fatalf("writefile: %v", err)
		}
	}

	names := make([]string, 0, 10)

	fd, err := syscall.Open(d, syscall.O_RDONLY, 0)
	if err != nil {
		t.Fatalf("syscall.open: %v", err)
	}
	defer syscall.Close(fd)

	buf := bytes.Repeat([]byte{0xCD}, direntBufSize)
	for {
		n, err := syscall.ReadDirent(fd, buf)
		if err == syscall.EINVAL {
			// On linux, 'man getdents64' says that EINVAL indicates “result buffer is too small”.
			// Try a bigger buffer.
			t.Logf("ReadDirent: %v; retrying with larger buffer", err)
			buf = bytes.Repeat([]byte{0xCD}, len(buf)*2)
			continue
		}
		if err != nil {
			t.Fatalf("syscall.readdir: %v", err)
		}
		t.Logf("ReadDirent: read %d bytes", n)
		if n == 0 {
			break
		}

		var consumed, count int
		consumed, count, names = syscall.ParseDirent(buf[:n], -1, names)
		t.Logf("ParseDirent: %d new name(s)", count)
		if consumed != n {
			t.Fatalf("ParseDirent: consumed %d bytes; expected %d", consumed, n)
		}
	}

	sort.Strings(names)
	t.Logf("names: %q", names)

	if len(names) != 10 {
		t.Errorf("got %d names; expected 10", len(names))
	}
	for i, name := range names {
		ord, err := strconv.Atoi(name[:1])
		if err != nil {
			t.Fatalf("names[%d] is non-integer %q: %v", i, names[i], err)
		}
		if expected := string(strings.Repeat(name[:1], filenameMinSize+ord)); name != expected {
			t.Errorf("names[%d] is %q (len %d); expected %q (len %d)", i, name, len(name), expected, len(expected))
		}
	}
}

func TestDirentRepeat(t *testing.T) {
	const N = 100
	// Note: the size of the buffer is small enough that the loop
	// below will need to execute multiple times. See issue #31368.
	size := N * unsafe.Offsetof(syscall.Dirent{}.Name) / 4
	if runtime.GOOS == "freebsd" || runtime.GOOS == "netbsd" {
		if size < 1024 {
			size = 1024 // DIRBLKSIZ, see issue 31403.
		}
		if runtime.GOOS == "freebsd" {
			t.Skip("need to fix issue 31416 first")
		}
	}

	// Make a directory containing N files
	d := t.TempDir()

	var files []string
	for i := 0; i < N; i++ {
		files = append(files, fmt.Sprintf("file%d", i))
	}
	for _, file := range files {
		err := os.WriteFile(filepath.Join(d, file), []byte("contents"), 0644)
		if err != nil {
			t.Fatalf("writefile: %v", err)
		}
	}

	// Read the directory entries using ReadDirent.
	fd, err := syscall.Open(d, syscall.O_RDONLY, 0)
	if err != nil {
		t.Fatalf("syscall.open: %v", err)
	}
	defer syscall.Close(fd)
	var files2 []string
	for {
		buf := make([]byte, size)
		n, err := syscall.ReadDirent(fd, buf)
		if err != nil {
			t.Fatalf("syscall.readdir: %v", err)
		}
		if n == 0 {
			break
		}
		buf = buf[:n]
		for len(buf) > 0 {
			var consumed int
			consumed, _, files2 = syscall.ParseDirent(buf, -1, files2)
			buf = buf[consumed:]
		}
	}

	// Check results
	sort.Strings(files)
	sort.Strings(files2)
	if strings.Join(files, "|") != strings.Join(files2, "|") {
		t.Errorf("bad file list: want\n%q\ngot\n%q", files, files2)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ppc64 ppc64le

package unix

// Linux getrandom system call number.
// See GetRandom in getrandom_linux.go.
const randomTrap uintptr = 359
<commit_msg>libgo: fix getrandom build for 32-bit ppc<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ppc ppc64 ppc64le

package unix

// Linux getrandom system call number.
// See GetRandom in getrandom_linux.go.
const randomTrap uintptr = 359
<|endoftext|>