<commit_before>package spql

import (
	"fmt"
	"reflect"

	"github.com/SparrowDb/sparrowdb/backup"
	"github.com/SparrowDb/sparrowdb/db"
	"github.com/SparrowDb/sparrowdb/errors"
	"github.com/SparrowDb/sparrowdb/model"
	"github.com/SparrowDb/sparrowdb/slog"
)

// QueryExecutor holds query executor data
type QueryExecutor struct {
	dbManager *db.DBManager
}

// ExecuteQuery execute query in channel
func (qe *QueryExecutor) ExecuteQuery(query *Query) <-chan *QueryResult {
	results := make(chan *QueryResult)
	go qe.executeQuery(query, results)
	return results
}

func (qe *QueryExecutor) executeQuery(query *Query, results chan *QueryResult) {
	defer close(results)
	defer func() {
		if r := recover(); r != nil {
			slog.Errorf("%s", r)
		}
	}()

	inputs := make([]reflect.Value, 2)
	inputs[0] = reflect.ValueOf(query)
	inputs[1] = reflect.ValueOf(results)

	reflect.ValueOf(qe).MethodByName(query.Method).Call(inputs)
}

// CreateDatabase process create database from query string
func (qe *QueryExecutor) CreateDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*CreateDatabaseStmt)

	databaseCfg := db.DatabaseDescriptor{
		Name:           qp.Name,
		MaxDataLogSize: qp.MaxDataLogSize,
		MaxCacheSize:   qp.MaxCacheSize,
		BloomFilterFp:  qp.BloomFilterFp,
		CronExp:        qp.CronExp,
		Path:           qp.Path,
		SnapshotPath:   qp.SnapshotPath,
	}

	err := qe.dbManager.CreateDatabase(databaseCfg)
	qr := QueryResult{Database: databaseCfg.Name}

	if err != nil {
		qr.AddErrorStr(err.Error())
	}

	results <- &qr
}

// DropDatabase process drop database from results <- qrquery string
func (qe *QueryExecutor) DropDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*DropDatabaseStmt)

	err := qe.dbManager.DropDatabase(qp.Name)
	qr := QueryResult{Database: qp.Name}

	if err != nil {
		qr.AddErrorStr(err.Error())
	}

	results <- &qr
}

// ShowDatabases process show databases from query string
func (qe *QueryExecutor) ShowDatabases(query *Query, results chan *QueryResult) {
	n := qe.dbManager.GetDatabasesNames()
	qr := QueryResult{}

	for _, v := range n {
		qr.AddValue(v)
	}

	results <- &qr
}

// InfoDatabase process show database info
func (qe *QueryExecutor) InfoDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*InfoDatabaseStmt)
	qr := QueryResult{}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		qr.AddValue(db.Descriptor)
		qr.AddValue(db.Info())
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// Delete delets entry from database with tombstone
func (qe *QueryExecutor) Delete(query *Query, results chan *QueryResult) {
	qp := query.Params.(*DeleteStmt)
	qr := QueryResult{}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		result := <-qe.dbManager.GetData(qp.Name, qp.Key)

		// Check if found requested data or DataDefinition is deleted
		if result == nil || result.Status == model.DataDefinitionRemoved {
			qr.AddErrorStr(fmt.Sprintf(errDataNotFound, qp.Key, qp.Name))
		} else {
			tbs := model.NewTombstone(result)
			db.InsertData(tbs)
		}
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// Select do query in database
func (qe *QueryExecutor) Select(query *Query, results chan *QueryResult) {
	qp := query.Params.(*SelectStmt)
	qr := QueryResult{Database: qp.Name}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok {
		qe.doSelect(qp, &qr, db, results)
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
		results <- &qr
	}
}

func (qe *QueryExecutor) doSelect(qp *SelectStmt, qr *QueryResult, db *db.Database, result chan *QueryResult) {
	// empty means query all
	if qp.Key == "" {

	} else {
		if d, _, ok := db.GetDataByKey(qp.Key); ok {
			qr.AddValue(d.QueryResult())
			result <- qr
		}
	}
}

// CreateSnapshot process to create snapshot of database
func (qe *QueryExecutor) CreateSnapshot(query *Query, results chan *QueryResult) {
	qp := query.Params.(*CreateSnapshotStmt)
	qr := QueryResult{Database: qp.Name}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		err := backup.CreateSnapshot(db.Descriptor.Path, db.Descriptor.SnapshotPath)
		if err != nil {
			qr.AddErrorStr(errors.ErrCreateDatabase.Error())
		}
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// NewQueryExecutor returns new QueryExecutor
func NewQueryExecutor(dbm *db.DBManager) *QueryExecutor {
	return &QueryExecutor{
		dbManager: dbm,
	}
}
<commit_msg>tombstone keeps version<commit_after>package spql

import (
	"fmt"
	"reflect"

	"github.com/SparrowDb/sparrowdb/backup"
	"github.com/SparrowDb/sparrowdb/db"
	"github.com/SparrowDb/sparrowdb/errors"
	"github.com/SparrowDb/sparrowdb/model"
	"github.com/SparrowDb/sparrowdb/slog"
)

// QueryExecutor holds query executor data
type QueryExecutor struct {
	dbManager *db.DBManager
}

// ExecuteQuery execute query in channel
func (qe *QueryExecutor) ExecuteQuery(query *Query) <-chan *QueryResult {
	results := make(chan *QueryResult)
	go qe.executeQuery(query, results)
	return results
}

func (qe *QueryExecutor) executeQuery(query *Query, results chan *QueryResult) {
	defer close(results)
	defer func() {
		if r := recover(); r != nil {
			slog.Errorf("%s", r)
		}
	}()

	inputs := make([]reflect.Value, 2)
	inputs[0] = reflect.ValueOf(query)
	inputs[1] = reflect.ValueOf(results)

	reflect.ValueOf(qe).MethodByName(query.Method).Call(inputs)
}

// CreateDatabase process create database from query string
func (qe *QueryExecutor) CreateDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*CreateDatabaseStmt)

	databaseCfg := db.DatabaseDescriptor{
		Name:           qp.Name,
		MaxDataLogSize: qp.MaxDataLogSize,
		MaxCacheSize:   qp.MaxCacheSize,
		BloomFilterFp:  qp.BloomFilterFp,
		CronExp:        qp.CronExp,
		Path:           qp.Path,
		SnapshotPath:   qp.SnapshotPath,
	}

	err := qe.dbManager.CreateDatabase(databaseCfg)
	qr := QueryResult{Database: databaseCfg.Name}

	if err != nil {
		qr.AddErrorStr(err.Error())
	}

	results <- &qr
}

// DropDatabase process drop database from results <- qrquery string
func (qe *QueryExecutor) DropDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*DropDatabaseStmt)

	err := qe.dbManager.DropDatabase(qp.Name)
	qr := QueryResult{Database: qp.Name}

	if err != nil {
		qr.AddErrorStr(err.Error())
	}

	results <- &qr
}

// ShowDatabases process show databases from query string
func (qe *QueryExecutor) ShowDatabases(query *Query, results chan *QueryResult) {
	n := qe.dbManager.GetDatabasesNames()
	qr := QueryResult{}

	for _, v := range n {
		qr.AddValue(v)
	}

	results <- &qr
}

// InfoDatabase process show database info
func (qe *QueryExecutor) InfoDatabase(query *Query, results chan *QueryResult) {
	qp := query.Params.(*InfoDatabaseStmt)
	qr := QueryResult{}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		qr.AddValue(db.Descriptor)
		qr.AddValue(db.Info())
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// Delete delets entry from database with tombstone
func (qe *QueryExecutor) Delete(query *Query, results chan *QueryResult) {
	qp := query.Params.(*DeleteStmt)
	qr := QueryResult{}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		storedDf, idx, exists := db.GetDataByKey(qp.Key)

		// Check if found requested data or DataDefinition is deleted
		if exists == false || storedDf.Status == model.DataDefinitionRemoved {
			qr.AddErrorStr(fmt.Sprintf(errDataNotFound, qp.Key, qp.Name))
		} else {
			tbs := model.NewTombstone(storedDf)

			tbs.Revision = 0
			tbs.Version = append(tbs.Version, storedDf.Version...)
			tbs.Version = append(tbs.Version, idx)
			db.InsertData(tbs)
		}
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// Select do query in database
func (qe *QueryExecutor) Select(query *Query, results chan *QueryResult) {
	qp := query.Params.(*SelectStmt)
	qr := QueryResult{Database: qp.Name}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok {
		qe.doSelect(qp, &qr, db, results)
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
		results <- &qr
	}
}

func (qe *QueryExecutor) doSelect(qp *SelectStmt, qr *QueryResult, db *db.Database, result chan *QueryResult) {
	// empty means query all
	if qp.Key == "" {

	} else {
		if d, _, ok := db.GetDataByKey(qp.Key); ok {
			qr.AddValue(d.QueryResult())
			result <- qr
		}
	}
}

// CreateSnapshot process to create snapshot of database
func (qe *QueryExecutor) CreateSnapshot(query *Query, results chan *QueryResult) {
	qp := query.Params.(*CreateSnapshotStmt)
	qr := QueryResult{Database: qp.Name}

	if db, ok := qe.dbManager.GetDatabase(qp.Name); ok == true {
		err := backup.CreateSnapshot(db.Descriptor.Path, db.Descriptor.SnapshotPath)
		if err != nil {
			qr.AddErrorStr(errors.ErrCreateDatabase.Error())
		}
	} else {
		qr.AddErrorStr(errors.ErrDatabaseNotFound.Error())
	}

	results <- &qr
}

// NewQueryExecutor returns new QueryExecutor
func NewQueryExecutor(dbm *db.DBManager) *QueryExecutor {
	return &QueryExecutor{
		dbManager: dbm,
	}
}
<|endoftext|>


<commit_before>// +build linux
package libipvs

import (
	"fmt"
	"net"
	"os/exec"
	"strings"
	"syscall"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	schedMethods = []string{
		RoundRobin,
		LeastConnection,
		DestinationHashing,
		SourceHashing,
	}

	protocols = []string{
		"TCP",
		"UDP",
		"FWM",
	}

	fwdMethods = []FwdMethod{
		IP_VS_CONN_F_MASQ,
		IP_VS_CONN_F_TUNNEL,
		IP_VS_CONN_F_DROUTE,
	}

	fwdMethodStrings = []string{
		"Masq",
		"Tunnel",
		"Route",
	}
)

func checkDestination(t *testing.T, checkPresent bool, protocol, serviceAddress, realAddress, fwdMethod string) {
	var (
		realServerStart bool
		realServers     []string
	)

	out, err := exec.Command("ipvsadm", "-Ln").CombinedOutput()
	require.NoError(t, err)

	for _, o := range strings.Split(string(out), "\n") {
		cmpStr := serviceAddress
		if protocol == "FWM" {
			cmpStr = " " + cmpStr
		}

		if strings.Contains(o, cmpStr) {
			realServerStart = true
			continue
		}

		if realServerStart {
			if !strings.Contains(o, "->") {
				break
			}

			realServers = append(realServers, o)
		}
	}

	for _, r := range realServers {
		if strings.Contains(r, realAddress) {
			parts := strings.Fields(r)
			assert.Equal(t, fwdMethod, parts[2])
			return
		}
	}

	if checkPresent {
		t.Fatalf("Did not find the destination %s fwdMethod %s in ipvs output", realAddress, fwdMethod)
	}
}

func checkService(t *testing.T, checkPresent bool, protocol, schedMethod, serviceAddress string) {
	out, err := exec.Command("ipvsadm", "-Ln").CombinedOutput()
	require.NoError(t, err)

	for _, o := range strings.Split(string(out), "\n") {
		cmpStr := serviceAddress
		if protocol == "FWM" {
			cmpStr = " " + cmpStr
		}

		if strings.Contains(o, cmpStr) {
			parts := strings.Split(o, " ")
			assert.Equal(t, protocol, parts[0])
			assert.Equal(t, serviceAddress, parts[2])
			assert.Equal(t, schedMethod, parts[3])

			if !checkPresent {
				t.Fatalf("Did not expect the service %s in ipvs output", serviceAddress)
			}

			return
		}
	}

	if checkPresent {
		t.Fatalf("Did not find the service %s in ipvs output", serviceAddress)
	}
}

func clearIpvsState(t *testing.T) {
	_, err := exec.Command("ipvsadm", "--clear").CombinedOutput()
	require.NoError(t, err)
}

func TestService(t *testing.T) {
	defer clearIpvsState(t)

	i, err := New()
	require.NoError(t, err)

	for _, protocol := range protocols {
		for _, schedMethod := range schedMethods {
			var serviceAddress string

			s := Service{
				AddressFamily: syscall.AF_INET,
				SchedName:     schedMethod,
			}

			switch protocol {
			case "FWM":
				s.FWMark = 1234
				serviceAddress = fmt.Sprintf("%d", 1234)
			case "TCP":
				s.Protocol = syscall.IPPROTO_TCP
				s.Port = 80
				s.Address = net.ParseIP("1.2.3.4")
				s.Netmask = 0xFFFFFFFF
				serviceAddress = "1.2.3.4:80"
			case "UDP":
				s.Protocol = syscall.IPPROTO_UDP
				s.Port = 53
				s.Address = net.ParseIP("2.3.4.5")
				serviceAddress = "2.3.4.5:53"
			}

			err := i.NewService(&s)
			assert.NoError(t, err)
			checkService(t, true, protocol, schedMethod, serviceAddress)
			var lastMethod string
			for _, updateSchedMethod := range schedMethods {
				if updateSchedMethod == schedMethod {
					continue
				}

				s.SchedName = updateSchedMethod
				err = i.UpdateService(&s)
				assert.NoError(t, err)
				checkService(t, true, protocol, updateSchedMethod, serviceAddress)
				lastMethod = updateSchedMethod
			}

			err = i.DelService(&s)
			checkService(t, false, protocol, lastMethod, serviceAddress)
		}
	}
}

func TestDestination(t *testing.T) {
	defer clearIpvsState(t)

	i, err := New()
	require.NoError(t, err)

	for _, protocol := range protocols {
		var serviceAddress string

		s := Service{
			AddressFamily: syscall.AF_INET,
			SchedName:     RoundRobin,
		}

		switch protocol {
		case "FWM":
			s.FWMark = 1234
			serviceAddress = fmt.Sprintf("%d", 1234)
		case "TCP":
			s.Protocol = syscall.IPPROTO_TCP
			s.Port = 80
			s.Address = net.ParseIP("1.2.3.4")
			s.Netmask = 0xFFFFFFFF
			serviceAddress = "1.2.3.4:80"
		case "UDP":
			s.Protocol = syscall.IPPROTO_UDP
			s.Port = 53
			s.Address = net.ParseIP("2.3.4.5")
			serviceAddress = "2.3.4.5:53"
		}

		err := i.NewService(&s)
		assert.NoError(t, err)
		checkService(t, true, protocol, RoundRobin, serviceAddress)

		s.SchedName = ""
		for j, fwdMethod := range fwdMethods {
			d1 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.2"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress := "10.1.1.2:5000"
			err := i.NewDestination(&s, &d1)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])
			d2 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.3"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress = "10.1.1.3:5000"
			err = i.NewDestination(&s, &d2)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])

			d3 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.4"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress = "10.1.1.4:5000"
			err = i.NewDestination(&s, &d3)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])

			for m, updateFwdMethod := range fwdMethods {
				if updateFwdMethod == fwdMethod {
					continue
				}
				d1.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.2:5000"
				err = i.UpdateDestination(&s, &d1)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])

				d2.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.3:5000"
				err = i.UpdateDestination(&s, &d2)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])

				d3.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.4:5000"
				err = i.UpdateDestination(&s, &d3)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])
			}

			err = i.DelDestination(&s, &d1)
			assert.NoError(t, err)
			err = i.DelDestination(&s, &d2)
			assert.NoError(t, err)
			err = i.DelDestination(&s, &d3)
			assert.NoError(t, err)
		}
	}
	// clearIpvs(t)
}
<commit_msg>Remove commented line<commit_after>// +build linux
package libipvs

import (
	"fmt"
	"net"
	"os/exec"
	"strings"
	"syscall"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	schedMethods = []string{
		RoundRobin,
		LeastConnection,
		DestinationHashing,
		SourceHashing,
	}

	protocols = []string{
		"TCP",
		"UDP",
		"FWM",
	}

	fwdMethods = []FwdMethod{
		IP_VS_CONN_F_MASQ,
		IP_VS_CONN_F_TUNNEL,
		IP_VS_CONN_F_DROUTE,
	}

	fwdMethodStrings = []string{
		"Masq",
		"Tunnel",
		"Route",
	}
)

func checkDestination(t *testing.T, checkPresent bool, protocol, serviceAddress, realAddress, fwdMethod string) {
	var (
		realServerStart bool
		realServers     []string
	)

	out, err := exec.Command("ipvsadm", "-Ln").CombinedOutput()
	require.NoError(t, err)

	for _, o := range strings.Split(string(out), "\n") {
		cmpStr := serviceAddress
		if protocol == "FWM" {
			cmpStr = " " + cmpStr
		}

		if strings.Contains(o, cmpStr) {
			realServerStart = true
			continue
		}

		if realServerStart {
			if !strings.Contains(o, "->") {
				break
			}

			realServers = append(realServers, o)
		}
	}

	for _, r := range realServers {
		if strings.Contains(r, realAddress) {
			parts := strings.Fields(r)
			assert.Equal(t, fwdMethod, parts[2])
			return
		}
	}

	if checkPresent {
		t.Fatalf("Did not find the destination %s fwdMethod %s in ipvs output", realAddress, fwdMethod)
	}
}

func checkService(t *testing.T, checkPresent bool, protocol, schedMethod, serviceAddress string) {
	out, err := exec.Command("ipvsadm", "-Ln").CombinedOutput()
	require.NoError(t, err)

	for _, o := range strings.Split(string(out), "\n") {
		cmpStr := serviceAddress
		if protocol == "FWM" {
			cmpStr = " " + cmpStr
		}

		if strings.Contains(o, cmpStr) {
			parts := strings.Split(o, " ")
			assert.Equal(t, protocol, parts[0])
			assert.Equal(t, serviceAddress, parts[2])
			assert.Equal(t, schedMethod, parts[3])

			if !checkPresent {
				t.Fatalf("Did not expect the service %s in ipvs output", serviceAddress)
			}

			return
		}
	}

	if checkPresent {
		t.Fatalf("Did not find the service %s in ipvs output", serviceAddress)
	}
}

func clearIpvsState(t *testing.T) {
	_, err := exec.Command("ipvsadm", "--clear").CombinedOutput()
	require.NoError(t, err)
}

func TestService(t *testing.T) {
	defer clearIpvsState(t)

	i, err := New()
	require.NoError(t, err)

	for _, protocol := range protocols {
		for _, schedMethod := range schedMethods {
			var serviceAddress string

			s := Service{
				AddressFamily: syscall.AF_INET,
				SchedName:     schedMethod,
			}

			switch protocol {
			case "FWM":
				s.FWMark = 1234
				serviceAddress = fmt.Sprintf("%d", 1234)
			case "TCP":
				s.Protocol = syscall.IPPROTO_TCP
				s.Port = 80
				s.Address = net.ParseIP("1.2.3.4")
				s.Netmask = 0xFFFFFFFF
				serviceAddress = "1.2.3.4:80"
			case "UDP":
				s.Protocol = syscall.IPPROTO_UDP
				s.Port = 53
				s.Address = net.ParseIP("2.3.4.5")
				serviceAddress = "2.3.4.5:53"
			}

			err := i.NewService(&s)
			assert.NoError(t, err)
			checkService(t, true, protocol, schedMethod, serviceAddress)
			var lastMethod string
			for _, updateSchedMethod := range schedMethods {
				if updateSchedMethod == schedMethod {
					continue
				}

				s.SchedName = updateSchedMethod
				err = i.UpdateService(&s)
				assert.NoError(t, err)
				checkService(t, true, protocol, updateSchedMethod, serviceAddress)
				lastMethod = updateSchedMethod
			}

			err = i.DelService(&s)
			checkService(t, false, protocol, lastMethod, serviceAddress)
		}
	}
}

func TestDestination(t *testing.T) {
	defer clearIpvsState(t)

	i, err := New()
	require.NoError(t, err)

	for _, protocol := range protocols {
		var serviceAddress string

		s := Service{
			AddressFamily: syscall.AF_INET,
			SchedName:     RoundRobin,
		}

		switch protocol {
		case "FWM":
			s.FWMark = 1234
			serviceAddress = fmt.Sprintf("%d", 1234)
		case "TCP":
			s.Protocol = syscall.IPPROTO_TCP
			s.Port = 80
			s.Address = net.ParseIP("1.2.3.4")
			s.Netmask = 0xFFFFFFFF
			serviceAddress = "1.2.3.4:80"
		case "UDP":
			s.Protocol = syscall.IPPROTO_UDP
			s.Port = 53
			s.Address = net.ParseIP("2.3.4.5")
			serviceAddress = "2.3.4.5:53"
		}

		err := i.NewService(&s)
		assert.NoError(t, err)
		checkService(t, true, protocol, RoundRobin, serviceAddress)

		s.SchedName = ""
		for j, fwdMethod := range fwdMethods {
			d1 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.2"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress := "10.1.1.2:5000"
			err := i.NewDestination(&s, &d1)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])
			d2 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.3"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress = "10.1.1.3:5000"
			err = i.NewDestination(&s, &d2)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])

			d3 := Destination{
				AddressFamily: syscall.AF_INET,
				Address:       net.ParseIP("10.1.1.4"),
				Port:          5000,
				Weight:        1,
				FwdMethod:     fwdMethod,
			}

			realAddress = "10.1.1.4:5000"
			err = i.NewDestination(&s, &d3)
			assert.NoError(t, err)
			checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[j])

			for m, updateFwdMethod := range fwdMethods {
				if updateFwdMethod == fwdMethod {
					continue
				}
				d1.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.2:5000"
				err = i.UpdateDestination(&s, &d1)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])

				d2.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.3:5000"
				err = i.UpdateDestination(&s, &d2)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])

				d3.FwdMethod = updateFwdMethod
				realAddress = "10.1.1.4:5000"
				err = i.UpdateDestination(&s, &d3)
				assert.NoError(t, err)
				checkDestination(t, true, protocol, serviceAddress, realAddress, fwdMethodStrings[m])
			}

			err = i.DelDestination(&s, &d1)
			assert.NoError(t, err)
			err = i.DelDestination(&s, &d2)
			assert.NoError(t, err)
			err = i.DelDestination(&s, &d3)
			assert.NoError(t, err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	. "github.com/logrusorgru/aurora"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

var infile string
var subfile string
var addsubs bool
var toplevel string
var webvtt bool
var jasonfile string
var cmdtemplate string
var completed string
var batch string

// Variant struct for HLS variants
type Variant struct {
	Name      string `json:"name"`
	Aspect    string `json:"aspect"`
	Rate      string `json:"framerate"`
	Vbr       string `json:"vbitrate"`
	Abr       string `json:"abitrate"`
	Bandwidth string
}

// Create variant's destination directory
func (v *Variant) mkDest() string {
	dest := fmt.Sprintf("%s/%s", toplevel, v.Name)
	os.MkdirAll(dest, 0755)
	return dest
}

func (v *Variant) mkInputs() string {
	inputs := fmt.Sprintf(" -i %s", infile)
	if addsubs {
		inputs = fmt.Sprintf(" -i %s -i %s/%s  ", infile, toplevel, subfile)
	}
	return inputs
}

// This Variant method assembles the ffmpeg command
func (v *Variant) mkCmd(cmdtemplate string) string {
	data, err := ioutil.ReadFile(cmdtemplate)
	chk(err, "Error reading template file")
	inputs := v.mkInputs()
	r := strings.NewReplacer("INFILE", inputs, "ASPECT", v.Aspect,
		"VBITRATE", v.Vbr, "FRAMERATE", v.Rate, "ABITRATE", v.Abr,
		"TOPLEVEL", toplevel, "NAME", v.Name, "\n", " ")
	cmd := fmt.Sprintf("%s\n", r.Replace(string(data)))
	//fmt.Println(cmd)
	return cmd
}

// Read actual bitrate from first segment to set bandwidth in master.m3u8
func (v *Variant) readRate() {
	cmd := fmt.Sprintf("ffprobe -i %s/%s/index0.ts", toplevel, v.Name)
	data := chkExec(cmd)
	two := strings.Split(data, "bitrate: ")[1]
	rate := strings.Split(two, " kb/s")[0]
	v.Bandwidth = fmt.Sprintf("%v000", rate)
}

// Start transcoding the variant
func (v *Variant) start() {
	v.mkDest()
	fmt.Printf("* Variants: %s %s \r", Cyan(completed), v.Aspect)
	completed += fmt.Sprintf("%s ", v.Aspect)
	cmd := v.mkCmd(cmdtemplate)
	chkExec(cmd)
	v.readRate()
	fmt.Printf("* Variants: %s  \r", Cyan(completed))
}

// #EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=7483000,RESOLUTION=1920:1080,
// hd1920/index.m3u8
func (v *Variant) mkStanza() string {
	stanza := fmt.Sprintf("#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=%v,RESOLUTION=%v", v.Bandwidth, v.Aspect)
	if addsubs {
		stanza = fmt.Sprintf("%s,SUBTITLES=\"webvtt\"", stanza)
	}
	return stanza
}

func chkExec(cmd string) string {
	// Executes external commands and checks for runtime errors
	parts := strings.Fields(cmd)
	data, err := exec.Command(parts[0], parts[1:]...).CombinedOutput()
	chk(err, fmt.Sprintf("Error running \n %s \n %v", cmd, string(data)))
	return string(data)
}

// probes for Closed Captions in video file.
func hasCaptions() bool {
	cmd := fmt.Sprintf("ffprobe -i %s", infile)
	data := chkExec(cmd)
	if strings.Contains(data, "Captions") {
		return true
	}
	return false
}

// Captions are segmented along with the first variant and then moved to toplevel/subs
func mvCaptions(vardir string) {
	srcdir := fmt.Sprintf("%s/%s", toplevel, vardir)
	destdir := fmt.Sprintf("%s/subs", toplevel)
	os.MkdirAll(destdir, 0755)
	files, err := ioutil.ReadDir(srcdir)
	chk(err, "Error moving Captions")
	for _, f := range files {
		if strings.Contains(f.Name(), "vtt") {
			os.Rename(fmt.Sprintf("%s/%s", srcdir, f.Name()), fmt.Sprintf("%s/%s", destdir, f.Name()))
		}
	}
}

// return a subtitle stanza for use in the  master.m3u8
func mkSubStanza() string {
	return "#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID=\"webvtt\",NAME=\"English\",DEFAULT=YES,AUTOSELECT=YES,FORCED=NO,LANGUAGE=\"en\",URI=\"subs/vtt_index.m3u8\"\n"
}

// Read json file for variants
func dataToVariants() []Variant {
	var variants []Variant
	data, err := ioutil.ReadFile(jasonfile)
	chk(err, "Error reading JSON file")
	json.Unmarshal(data, &variants)
	return variants
}

// Set the toplevel dir for variants by splitting video file name at the "."
func mkTopLevel() {
	if toplevel == "" {
		toplevel = strings.Split(infile, `.`)[0]
	}
	os.MkdirAll(toplevel, 0755)
}

func extractCaptions() string {
	fmt.Println("Extracting captions")
	assfile := fmt.Sprintf("%s.ass", toplevel)
	cmd := fmt.Sprintf("ffmpeg -y -f lavfi -fix_sub_duration -i movie=%s[out0+subcc] %s/%s", infile, toplevel, assfile)
	chkExec(cmd)
	return assfile
}

func mkSubfile() {
	if webvtt {
		return
	} else {
		addsubs = false
		if subfile == "" {
			if hasCaptions() {
				subfile = extractCaptions()
			}
		}
		if subfile != "" {
			addsubs = true
		}
	}
}

// Generic catchall error checking
func chk(err error, mesg string) {
	if err != nil {
		fmt.Printf("%s\n", mesg)
		//panic(err)
	}
}

// Make all variants and write master.m3u8
func mkAll(variants []Variant) {
	mkTopLevel()
	mkSubfile()
	var m3u8Master = fmt.Sprintf("%s/master.m3u8", toplevel)
	fp, err := os.Create(m3u8Master)
	chk(err, "in mkAll")
	defer fp.Close()
	w := bufio.NewWriter(fp)
	w.WriteString("#EXTM3U\n")
	fmt.Println("\n* Video file:", Cyan(infile), "\n* Toplevel:", Cyan(toplevel), "\n* Subtitle file:", Cyan(subfile))
	for _, v := range variants {
		v.start()
		if addsubs && !(webvtt) {
			mvCaptions(v.Name)
			w.WriteString(mkSubStanza())
			addsubs = false
			subfile = ""
			webvtt = true
		}
		w.WriteString(fmt.Sprintf("%s\n", v.mkStanza()))
		w.WriteString(fmt.Sprintf("%s/index.m3u8\n", v.Name))
	}
	fmt.Println("\n\n")
	w.Flush()
}

func main() {
	flag.StringVar(&infile, "i", "", "Video file to segment (either -i or -b is required)")
	flag.StringVar(&subfile, "s", "", "subtitle file to segment (optional)")
	flag.StringVar(&toplevel, "d", "", "override top level directory for hls files (optional)")
	flag.StringVar(&jasonfile, "j", `./hls.json`, "JSON file of variants (optional)")
	flag.StringVar(&cmdtemplate, "t", `./cmd.template`, "command template file (optional)")
	flag.StringVar(&batch, "b", "", "batch mode, list multiple input files (either -i or -b is required)")

	flag.Parse()
	fmt.Println(batch)
	variants := dataToVariants()
	webvtt = false
	if batch != "" {
		batch = strings.Replace(batch, " ", ",", -1)
		for _, b := range strings.Split(batch, ",") {
			fmt.Println(b)
			webvtt = false
			infile = b
			completed = ""
			toplevel = ""
			mkTopLevel()
			variants := dataToVariants()
			mkAll(variants)
		}
	} else {
		if infile != "" {
			mkAll(variants)
		} else {
			flag.PrintDefaults()
		}

	}
}
<commit_msg>fixed inputs<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	. "github.com/logrusorgru/aurora"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

var infile string
var subfile string
var addsubs bool
var toplevel string
var webvtt bool
var jasonfile string
var cmdtemplate string
var completed string
var batch string

// Variant struct for HLS variants
type Variant struct {
	Name      string `json:"name"`
	Aspect    string `json:"aspect"`
	Rate      string `json:"framerate"`
	Vbr       string `json:"vbitrate"`
	Abr       string `json:"abitrate"`
	Bandwidth string
}

// Create variant's destination directory
func (v *Variant) mkDest() string {
	dest := fmt.Sprintf("%s/%s", toplevel, v.Name)
	os.MkdirAll(dest, 0755)
	return dest
}

func (v *Variant) mkInputs() string {
	inputs := fmt.Sprintf(" -i %s", infile)
	if addsubs {
		inputs = fmt.Sprintf(" -i %s -i %s  ", infile, subfile)
	}
	return inputs
}

// This Variant method assembles the ffmpeg command
func (v *Variant) mkCmd(cmdtemplate string) string {
	data, err := ioutil.ReadFile(cmdtemplate)
	chk(err, "Error reading template file")
	inputs := v.mkInputs()
	r := strings.NewReplacer("INPUTS", inputs, "ASPECT", v.Aspect,
		"VBITRATE", v.Vbr, "FRAMERATE", v.Rate, "ABITRATE", v.Abr,
		"TOPLEVEL", toplevel, "NAME", v.Name, "\n", " ")
	cmd := fmt.Sprintf("%s\n", r.Replace(string(data)))
	//fmt.Println(cmd)
	return cmd
}

// Read actual bitrate from first segment to set bandwidth in master.m3u8
func (v *Variant) readRate() {
	cmd := fmt.Sprintf("ffprobe -i %s/%s/index0.ts", toplevel, v.Name)
	data := chkExec(cmd)
	two := strings.Split(data, "bitrate: ")[1]
	rate := strings.Split(two, " kb/s")[0]
	v.Bandwidth = fmt.Sprintf("%v000", rate)
}

// Start transcoding the variant
func (v *Variant) start() {
	v.mkDest()
	fmt.Printf("* Variants: %s %s \r", Cyan(completed), v.Aspect)
	completed += fmt.Sprintf("%s ", v.Aspect)
	cmd := v.mkCmd(cmdtemplate)
	chkExec(cmd)
	v.readRate()
	fmt.Printf("* Variants: %s  \r", Cyan(completed))
}

// #EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=7483000,RESOLUTION=1920:1080,
// hd1920/index.m3u8
func (v *Variant) mkStanza() string {
	stanza := fmt.Sprintf("#EXT-X-STREAM-INF:PROGRAM-ID=1,BANDWIDTH=%v,RESOLUTION=%v", v.Bandwidth, v.Aspect)
	if addsubs {
		stanza = fmt.Sprintf("%s,SUBTITLES=\"webvtt\"", stanza)
	}
	return stanza
}

func chkExec(cmd string) string {
	// Executes external commands and checks for runtime errors
	parts := strings.Fields(cmd)
	data, err := exec.Command(parts[0], parts[1:]...).CombinedOutput()
	chk(err, fmt.Sprintf("Error running \n %s \n %v", cmd, string(data)))
	return string(data)
}

// probes for Closed Captions in video file.
func hasCaptions() bool {
	cmd := fmt.Sprintf("ffprobe -i %s", infile)
	data := chkExec(cmd)
	if strings.Contains(data, "Captions") {
		return true
	}
	return false
}

// Captions are segmented along with the first variant and then moved to toplevel/subs
func mvCaptions(vardir string) {
	srcdir := fmt.Sprintf("%s/%s", toplevel, vardir)
	destdir := fmt.Sprintf("%s/subs", toplevel)
	os.MkdirAll(destdir, 0755)
	files, err := ioutil.ReadDir(srcdir)
	chk(err, "Error moving Captions")
	for _, f := range files {
		if strings.Contains(f.Name(), "vtt") {
			os.Rename(fmt.Sprintf("%s/%s", srcdir, f.Name()), fmt.Sprintf("%s/%s", destdir, f.Name()))
		}
	}
}

// return a subtitle stanza for use in the  master.m3u8
func mkSubStanza() string {
	return "#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID=\"webvtt\",NAME=\"English\",DEFAULT=YES,AUTOSELECT=YES,FORCED=NO,LANGUAGE=\"en\",URI=\"subs/vtt_index.m3u8\"\n"
}

// Read json file for variants
func dataToVariants() []Variant {
	var variants []Variant
	data, err := ioutil.ReadFile(jasonfile)
	chk(err, "Error reading JSON file")
	json.Unmarshal(data, &variants)
	return variants
}

// Set the toplevel dir for variants by splitting video file name at the "."
func mkTopLevel() {
	if toplevel == "" {
		toplevel = strings.Split(infile, `.`)[0]
	}
	os.MkdirAll(toplevel, 0755)
}

func extractCaptions() string {
	fmt.Println("Extracting captions")
	assfile := fmt.Sprintf("%s/%s.ass",toplevel, toplevel)
	cmd := fmt.Sprintf("ffmpeg -y -f lavfi -fix_sub_duration -i movie=%s[out0+subcc] %s", infile, assfile)
	chkExec(cmd)
	return assfile
}

func mkSubfile() {
	if webvtt {
		return
	} else {
		addsubs = false
		if subfile == "" {
			if hasCaptions() {
				subfile = extractCaptions()
			}
		}
		if subfile != "" {

			addsubs = true
		}
	}
}

// Generic catchall error checking
func chk(err error, mesg string) {
	if err != nil {
		fmt.Printf("%s\n", mesg)
		//panic(err)
	}
}

// Make all variants and write master.m3u8
func mkAll(variants []Variant) {
	mkTopLevel()
	mkSubfile()
	var m3u8Master = fmt.Sprintf("%s/master.m3u8", toplevel)
	fp, err := os.Create(m3u8Master)
	chk(err, "in mkAll")
	defer fp.Close()
	w := bufio.NewWriter(fp)
	w.WriteString("#EXTM3U\n")
	fmt.Println("\n* Video file:", Cyan(infile), "\n* Toplevel:", Cyan(toplevel), "\n* Subtitle file:", Cyan(subfile))
	for _, v := range variants {
		v.start()
		if addsubs && !(webvtt) {
			mvCaptions(v.Name)
			w.WriteString(mkSubStanza())
			addsubs = false
			subfile = ""
			webvtt = true
		}
		w.WriteString(fmt.Sprintf("%s\n", v.mkStanza()))
		w.WriteString(fmt.Sprintf("%s/index.m3u8\n", v.Name))
	}
	fmt.Println("\n\n")
	w.Flush()
}

func main() {
	flag.StringVar(&infile, "i", "", "Video file to segment (either -i or -b is required)")
	flag.StringVar(&subfile, "s", "", "subtitle file to segment (optional)")
	flag.StringVar(&toplevel, "d", "", "override top level directory for hls files (optional)")
	flag.StringVar(&jasonfile, "j", `./hls.json`, "JSON file of variants (optional)")
	flag.StringVar(&cmdtemplate, "t", `./cmd.template`, "command template file (optional)")
	flag.StringVar(&batch, "b", "", "batch mode, list multiple input files (either -i or -b is required)")

	flag.Parse()
	fmt.Println(batch)
	variants := dataToVariants()

	if batch != "" {
		batch = strings.Replace(batch, " ", ",", -1)
		for _, b := range strings.Split(batch, ",") {
			fmt.Println(b)
			webvtt = false
			infile = b
			completed = ""
			toplevel = ""
			mkTopLevel()
			variants := dataToVariants()
			mkAll(variants)
		}
	} else {
		if infile != "" {
			mkAll(variants)
		} else {
			flag.PrintDefaults()
		}

	}
}
<|endoftext|>


<commit_before>package irc

import (
	"bytes"
	"github.com/Lukasa/GoBot/struc"
	"sort"
	"strings"
)

// ParseIRCMessage takes a byte array of an IRC message and parses it into an IRCMessage structure.
// This function isn't necessarily cheap, so expects to be run as a goroutine.
// There are a ton of casual assumptions in here that make this function less-than-resilient. I'm going to call it
// "low-hanging fruit for future improvement" and move on with my life.
func ParseIRCMessage(msg []byte, out chan *struc.IRCMessage) {
	parsedMsg := struc.NewIRCMessage()
	response := false
	prefix := ""
	responseCode := ""
	command := -1
	arguments := make([]string, 0)
	newMsg := []byte{}
	trailer := ""

	// First, split the message on whitespace.
	components := bytes.Fields(msg)

	// Check whether the first component begins with a ':'. If it does, it's an IRC prefix.
	if components[0][0] == ':' {
		prefix = string(components[0][1:])
		components = components[1:]
	}

	// The new first component is either a command or a response code. Easiest way to tell is to work out
	// if the first character is actually a digit.
	if (components[0][0] <= '9') && (components[0][0] >= '0') {
		responseCode = string(components[0])
		response = true
	} else {
		// Find the index of the command string in the Commands array.
		commandStr := string(bytes.ToUpper(components[0]))
		command = sort.SearchStrings(struc.Commands, commandStr) // This isn't quite right if we get a command that we don't know.
	}

	components = components[1:]

	// It's possible the message ends there. If it doesn't, the remaining values are either arguments or the 'trailing'
	// section. The trailing section begins with a colon and the arguments shouldn't (I don't think), so use that as the
	// delimiter.
	for _, arg := range components {
		if arg[0] == ':' {
			break
		}

		arguments = append(arguments, string(arg))
	}

	// Anything left after this stage is part of the 'trailing' section. Join it up.
	components = components[len(arguments):]
	newMsg = bytes.Join(components, []byte{' '})

	// Strip the leading colon.
	if len(newMsg) > 0 {
		newMsg = newMsg[1:]
	}

	trailer = string(newMsg)

	// Append everything to the IRCMessage.
	parsedMsg.Prefix = prefix
	parsedMsg.Response = response
	parsedMsg.Command = command
	parsedMsg.ResponseCode = responseCode
	parsedMsg.Arguments = arguments
	parsedMsg.Trailing = trailer

	// Send it on the channel and exit.
	out <- parsedMsg
	return
}

// UnparseIRCMessage builds an IRC message into a byte array suitable for sending on the wire. This method is not necessarily
// cheap, so it expects to be run as a goroutine.
func UnparseIRCMessage(msg *struc.IRCMessage, out chan []byte) {
	components := make([]string, 0, 5)

	if msg.Prefix != "" {
		components = append(components, ":"+msg.Prefix)
	}

	if msg.Response {
		components = append(components, msg.ResponseCode)
	} else {
		components = append(components, struc.Commands[msg.Command])
	}

	for _, arg := range msg.Arguments {
		components = append(components, arg)
	}

	if msg.Trailing != "" {
		components = append(components, ":"+msg.Trailing)
	}

	strMsg := strings.Join(components, " ")
	out <- []byte(strMsg)
}
<commit_msg>We may not have other components.<commit_after>package irc

import (
	"bytes"
	"github.com/Lukasa/GoBot/struc"
	"sort"
	"strings"
)

// ParseIRCMessage takes a byte array of an IRC message and parses it into an IRCMessage structure.
// This function isn't necessarily cheap, so expects to be run as a goroutine.
// There are a ton of casual assumptions in here that make this function less-than-resilient. I'm going to call it
// "low-hanging fruit for future improvement" and move on with my life.
func ParseIRCMessage(msg []byte, out chan *struc.IRCMessage) {
	parsedMsg := struc.NewIRCMessage()
	response := false
	prefix := ""
	responseCode := ""
	command := -1
	arguments := make([]string, 0)
	newMsg := []byte{}
	trailer := ""

	// First, split the message on whitespace.
	components := bytes.Fields(msg)

	// Check whether the first component begins with a ':'. If it does, it's an IRC prefix.
	if components[0][0] == ':' {
		prefix = string(components[0][1:])
		components = components[1:]
	}

	if len(components) != 0 {
		// The new first component is either a command or a response code. Easiest way to tell is to work out
		// if the first character is actually a digit.
		if (components[0][0] <= '9') && (components[0][0] >= '0') {
			responseCode = string(components[0])
			response = true
		} else {
			// Find the index of the command string in the Commands array.
			commandStr := string(bytes.ToUpper(components[0]))
			command = sort.SearchStrings(struc.Commands, commandStr) // This isn't quite right if we get a command that we don't know.
		}

		components = components[1:]

		// It's possible the message ends there. If it doesn't, the remaining values are either arguments or the 'trailing'
		// section. The trailing section begins with a colon and the arguments shouldn't (I don't think), so use that as the
		// delimiter.
		for _, arg := range components {
			if arg[0] == ':' {
				break
			}

			arguments = append(arguments, string(arg))
		}

		// Anything left after this stage is part of the 'trailing' section. Join it up.
		components = components[len(arguments):]
		newMsg = bytes.Join(components, []byte{' '})

		// Strip the leading colon.
		if len(newMsg) > 0 {
			newMsg = newMsg[1:]
		}

		trailer = string(newMsg)
	}

	// Append everything to the IRCMessage.
	parsedMsg.Prefix = prefix
	parsedMsg.Response = response
	parsedMsg.Command = command
	parsedMsg.ResponseCode = responseCode
	parsedMsg.Arguments = arguments
	parsedMsg.Trailing = trailer

	// Send it on the channel and exit.
	out <- parsedMsg
	return
}

// UnparseIRCMessage builds an IRC message into a byte array suitable for sending on the wire. This method is not necessarily
// cheap, so it expects to be run as a goroutine.
func UnparseIRCMessage(msg *struc.IRCMessage, out chan []byte) {
	components := make([]string, 0, 5)

	if msg.Prefix != "" {
		components = append(components, ":"+msg.Prefix)
	}

	if msg.Response {
		components = append(components, msg.ResponseCode)
	} else {
		components = append(components, struc.Commands[msg.Command])
	}

	for _, arg := range msg.Arguments {
		components = append(components, arg)
	}

	if msg.Trailing != "" {
		components = append(components, ":"+msg.Trailing)
	}

	strMsg := strings.Join(components, " ")
	out <- []byte(strMsg)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"summercat.com/irc"
)

// Channel holds everything to do with a channel.
type Channel struct {
	// Canonicalized.
	Name string

	// Client id to Client.
	Members map[uint64]*Client

	Topic string
}

// Server holds the state for a server.
// I put everything global to a server in an instance of struct rather than
// have global variables.
type Server struct {
	Config Config

	// Client id to Client.
	Clients map[uint64]*Client

	// Canonicalized nickname to Client.
	// The reason I have this as well as clients is to track unregistered
	// clients.
	Nicks map[string]*Client

	// Channel name (canonicalized) to Channel.
	Channels map[string]*Channel

	// When we close this channel, this indicates that we're shutting down.
	// Other goroutines can check if this channel is closed.
	ShutdownChan chan struct{}

	// Tell the server something on this channel.
	ToServerChan chan Event

	// TCP listener.
	Listener net.Listener

	// WaitGroup to ensure all goroutines clean up before we end.
	WG sync.WaitGroup
}

// Event holds a message containing something to tell the server.
type Event struct {
	Type    EventType
	Client  *Client
	Message irc.Message
}

// EventType is a type of event we can tell the server about.
type EventType int

const (
	// NullEvent is a default event. This means the event was not populated.
	NullEvent EventType = iota

	// NewClientEvent means a new client connected.
	NewClientEvent

	// DeadClientEvent means client died for some reason. Clean it up.
	// It's useful to be able to know immediately and inform the client if we're
	// going to decide they are getting cut off (e.g., malformed message).
	DeadClientEvent

	// MessageFromClientEvent means a client sent a message.
	MessageFromClientEvent

	// WakeUpEvent means the server should wake up and do bookkeeping.
	WakeUpEvent
)

func main() {
	log.SetFlags(0)

	args, err := getArgs()
	if err != nil {
		log.Fatal(err)
	}

	server, err := newServer(args.ConfigFile)
	if err != nil {
		log.Fatal(err)
	}

	err = server.start()
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Server shutdown cleanly.")
}

func newServer(configFile string) (*Server, error) {
	s := Server{
		Clients:  make(map[uint64]*Client),
		Nicks:    make(map[string]*Client),
		Channels: make(map[string]*Channel),

		// shutdown() closes this channel.
		ShutdownChan: make(chan struct{}),

		// We never manually close this channel.
		ToServerChan: make(chan Event),
	}

	err := s.checkAndParseConfig(configFile)
	if err != nil {
		return nil, fmt.Errorf("Configuration problem: %s", err)
	}

	return &s, nil
}

// start starts up the server.
//
// We open the TCP port, start goroutines, and then receive messages on our
// channels.
func (s *Server) start() error {
	// TODO: TLS
	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%s", s.Config.ListenHost,
		s.Config.ListenPort))
	if err != nil {
		return fmt.Errorf("Unable to listen: %s", err)
	}
	s.Listener = ln

	// acceptConnections accepts connections on the TCP listener.
	s.WG.Add(1)
	go s.acceptConnections()

	// Alarm is a goroutine to wake up this one periodically so we can do things
	// like ping clients.
	s.WG.Add(1)
	go s.alarm()

MessageLoop:
	for {
		select {
		case evt := <-s.ToServerChan:
			if evt.Type == NewClientEvent {
				log.Printf("New client connection: %s", evt.Client)
				s.Clients[evt.Client.ID] = evt.Client
				continue
			}

			if evt.Type == DeadClientEvent {
				_, exists := s.Clients[evt.Client.ID]
				if exists {
					log.Printf("Client %s died.", evt.Client)
					evt.Client.quit("I/O error")
				}
				continue
			}

			if evt.Type == MessageFromClientEvent {
				_, exists := s.Clients[evt.Client.ID]
				if exists {
					log.Printf("Client %s: Message: %s", evt.Client, evt.Message)
					s.handleMessage(evt.Client, evt.Message)
				}
				continue
			}

			if evt.Type == WakeUpEvent {
				s.checkAndPingClients()
				continue
			}

			log.Fatalf("Unexpected event: %d", evt.Type)

		case <-s.ShutdownChan:
			break MessageLoop
		}
	}

	// We don't need to drain any channels. None close that will have any
	// goroutines blocked on them.

	s.WG.Wait()

	return nil
}

// shutdown starts server shutdown.
func (s *Server) shutdown() {
	log.Printf("Server shutdown initiated.")

	// Closing ShutdownChan indicates to other goroutines that we're shutting
	// down.
	close(s.ShutdownChan)

	err := s.Listener.Close()
	if err != nil {
		log.Printf("Problem closing TCP listener: %s", err)
	}

	// All clients need to be told. This also closes their write channels.
	for _, client := range s.Clients {
		client.quit("Server shutting down")
	}
}

// acceptConnections accepts TCP connections and tells the main server loop
// through a channel. It sets up separate goroutines for reading/writing to
// and from the client.
func (s *Server) acceptConnections() {
	defer s.WG.Done()

	id := uint64(0)

	for {
		if s.isShuttingDown() {
			break
		}

		conn, err := s.Listener.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %s", err)
			continue
		}

		client := NewClient(s, id, conn)

		// Handle rollover of uint64. Unlikely to happen (outside abuse) but.
		if id+1 == 0 {
			log.Fatalf("Unique ids rolled over!")
		}
		id++

		// ToServerChan is synchronous. We want to make sure server knows about the
		// client before it starts hearing anything from its other channels about
		// the client.
		s.newEvent(Event{Type: NewClientEvent, Client: client})

		s.WG.Add(1)
		go client.readLoop()
		s.WG.Add(1)
		go client.writeLoop()
	}

	log.Printf("Connection accepter shutting down.")
}

// Return true if the server is shutting down.
func (s *Server) isShuttingDown() bool {
	// No messages get sent to this channel, so if we receive a message on it,
	// then we know the channel was closed.
	select {
	case <-s.ShutdownChan:
		return true
	default:
		return false
	}
}

// Alarm sends a message to the server goroutine to wake it up.
// It sleeps and then repeats.
func (s *Server) alarm() {
	defer s.WG.Done()

	for {
		if s.isShuttingDown() {
			break
		}

		time.Sleep(s.Config.WakeupTime)

		s.newEvent(Event{Type: WakeUpEvent})
	}

	log.Printf("Alarm shutting down.")
}

// checkAndPingClients looks at each connected client.
//
// If they've been idle a short time, we send them a PING (if they're
// registered).
//
// If they've been idle a long time, we kill their connection.
func (s *Server) checkAndPingClients() {
	now := time.Now()

	for _, client := range s.Clients {
		timeIdle := now.Sub(client.LastActivityTime)
		timeSincePing := now.Sub(client.LastPingTime)

		if client.Registered {
			// Was it active recently enough that we don't need to do anything?
			if timeIdle < s.Config.PingTime {
				continue
			}

			// It's been idle a while.

			// Has it been idle long enough that we consider it dead?
			if timeIdle > s.Config.DeadTime {
				client.quit(fmt.Sprintf("Ping timeout: %d seconds",
					int(timeIdle.Seconds())))
				continue
			}

			// Should we ping it? We might have pinged it recently.
			if timeSincePing < s.Config.PingTime {
				continue
			}

			s.messageClient(client, "PING", []string{s.Config.ServerName})
			client.LastPingTime = now
			continue
		}

		if timeIdle > s.Config.DeadTime {
			client.quit("Idle too long.")
		}
	}
}

// Send an IRC message to a client. Appears to be from the server.
// This works by writing to a client's channel.
//
// Note: Only the server goroutine should call this (due to channel use).
func (s *Server) messageClient(c *Client, command string, params []string) {
	// For numeric messages, we need to prepend the nick.
	// Use * for the nick in cases where the client doesn't have one yet.
	// This is what ircd-ratbox does. Maybe not RFC...
	isNumeric := true
	for _, c := range command {
		if c < 48 || c > 57 {
			isNumeric = false
		}
	}

	if isNumeric {
		nick := "*"
		if len(c.DisplayNick) > 0 {
			nick = c.DisplayNick
		}
		newParams := []string{nick}

		newParams = append(newParams, params...)
		params = newParams
	}

	c.WriteChan <- irc.Message{
		Prefix:  s.Config.ServerName,
		Command: command,
		Params:  params,
	}
}

// newEvent tells the server something happens.
//
// Any goroutine can call this function.
//
// It sends the server a message on its ToServerChan.
//
// It will not block on shutdown as we select on the shutdown channel which we
// close when shutting down the server. This means receive on the shutdown
// channel will proceed at that point.
func (s *Server) newEvent(evt Event) {
	select {
	case s.ToServerChan <- evt:
	case <-s.ShutdownChan:
	}
}
<commit_msg>ircd: Move event loop to separate function<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"summercat.com/irc"
)

// Channel holds everything to do with a channel.
type Channel struct {
	// Canonicalized.
	Name string

	// Client id to Client.
	Members map[uint64]*Client

	Topic string
}

// Server holds the state for a server.
// I put everything global to a server in an instance of struct rather than
// have global variables.
type Server struct {
	Config Config

	// Client id to Client.
	Clients map[uint64]*Client

	// Canonicalized nickname to Client.
	// The reason I have this as well as clients is to track unregistered
	// clients.
	Nicks map[string]*Client

	// Channel name (canonicalized) to Channel.
	Channels map[string]*Channel

	// When we close this channel, this indicates that we're shutting down.
	// Other goroutines can check if this channel is closed.
	ShutdownChan chan struct{}

	// Tell the server something on this channel.
	ToServerChan chan Event

	// TCP listener.
	Listener net.Listener

	// WaitGroup to ensure all goroutines clean up before we end.
	WG sync.WaitGroup
}

// Event holds a message containing something to tell the server.
type Event struct {
	Type    EventType
	Client  *Client
	Message irc.Message
}

// EventType is a type of event we can tell the server about.
type EventType int

const (
	// NullEvent is a default event. This means the event was not populated.
	NullEvent EventType = iota

	// NewClientEvent means a new client connected.
	NewClientEvent

	// DeadClientEvent means client died for some reason. Clean it up.
	// It's useful to be able to know immediately and inform the client if we're
	// going to decide they are getting cut off (e.g., malformed message).
	DeadClientEvent

	// MessageFromClientEvent means a client sent a message.
	MessageFromClientEvent

	// WakeUpEvent means the server should wake up and do bookkeeping.
	WakeUpEvent
)

func main() {
	log.SetFlags(0)

	args, err := getArgs()
	if err != nil {
		log.Fatal(err)
	}

	server, err := newServer(args.ConfigFile)
	if err != nil {
		log.Fatal(err)
	}

	err = server.start()
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Server shutdown cleanly.")
}

func newServer(configFile string) (*Server, error) {
	s := Server{
		Clients:  make(map[uint64]*Client),
		Nicks:    make(map[string]*Client),
		Channels: make(map[string]*Channel),

		// shutdown() closes this channel.
		ShutdownChan: make(chan struct{}),

		// We never manually close this channel.
		ToServerChan: make(chan Event),
	}

	err := s.checkAndParseConfig(configFile)
	if err != nil {
		return nil, fmt.Errorf("Configuration problem: %s", err)
	}

	return &s, nil
}

// start starts up the server.
//
// We open the TCP port, start goroutines, and then receive messages on our
// channels.
func (s *Server) start() error {
	// TODO: TLS
	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%s", s.Config.ListenHost,
		s.Config.ListenPort))
	if err != nil {
		return fmt.Errorf("Unable to listen: %s", err)
	}
	s.Listener = ln

	// acceptConnections accepts connections on the TCP listener.
	s.WG.Add(1)
	go s.acceptConnections()

	// Alarm is a goroutine to wake up this one periodically so we can do things
	// like ping clients.
	s.WG.Add(1)
	go s.alarm()

	s.eventLoop()

	// We don't need to drain any channels. None close that will have any
	// goroutines blocked on them.

	s.WG.Wait()

	return nil
}

// eventLoop processes events on the server's channel.
//
// It continues until the shutdown channel closes, indicating shutdown.
func (s *Server) eventLoop() {
	for {
		select {
		case evt := <-s.ToServerChan:
			if evt.Type == NewClientEvent {
				log.Printf("New client connection: %s", evt.Client)
				s.Clients[evt.Client.ID] = evt.Client
				continue
			}

			if evt.Type == DeadClientEvent {
				_, exists := s.Clients[evt.Client.ID]
				if exists {
					log.Printf("Client %s died.", evt.Client)
					evt.Client.quit("I/O error")
				}
				continue
			}

			if evt.Type == MessageFromClientEvent {
				_, exists := s.Clients[evt.Client.ID]
				if exists {
					log.Printf("Client %s: Message: %s", evt.Client, evt.Message)
					s.handleMessage(evt.Client, evt.Message)
				}
				continue
			}

			if evt.Type == WakeUpEvent {
				s.checkAndPingClients()
				continue
			}

			log.Fatalf("Unexpected event: %d", evt.Type)

		case <-s.ShutdownChan:
			return
		}
	}
}

// shutdown starts server shutdown.
func (s *Server) shutdown() {
	log.Printf("Server shutdown initiated.")

	// Closing ShutdownChan indicates to other goroutines that we're shutting
	// down.
	close(s.ShutdownChan)

	err := s.Listener.Close()
	if err != nil {
		log.Printf("Problem closing TCP listener: %s", err)
	}

	// All clients need to be told. This also closes their write channels.
	for _, client := range s.Clients {
		client.quit("Server shutting down")
	}
}

// acceptConnections accepts TCP connections and tells the main server loop
// through a channel. It sets up separate goroutines for reading/writing to
// and from the client.
func (s *Server) acceptConnections() {
	defer s.WG.Done()

	id := uint64(0)

	for {
		if s.isShuttingDown() {
			break
		}

		conn, err := s.Listener.Accept()
		if err != nil {
			log.Printf("Failed to accept connection: %s", err)
			continue
		}

		client := NewClient(s, id, conn)

		// Handle rollover of uint64. Unlikely to happen (outside abuse) but.
		if id+1 == 0 {
			log.Fatalf("Unique ids rolled over!")
		}
		id++

		// ToServerChan is synchronous. We want to make sure server knows about the
		// client before it starts hearing anything from its other channels about
		// the client.
		s.newEvent(Event{Type: NewClientEvent, Client: client})

		s.WG.Add(1)
		go client.readLoop()
		s.WG.Add(1)
		go client.writeLoop()
	}

	log.Printf("Connection accepter shutting down.")
}

// Return true if the server is shutting down.
func (s *Server) isShuttingDown() bool {
	// No messages get sent to this channel, so if we receive a message on it,
	// then we know the channel was closed.
	select {
	case <-s.ShutdownChan:
		return true
	default:
		return false
	}
}

// Alarm sends a message to the server goroutine to wake it up.
// It sleeps and then repeats.
func (s *Server) alarm() {
	defer s.WG.Done()

	for {
		if s.isShuttingDown() {
			break
		}

		time.Sleep(s.Config.WakeupTime)

		s.newEvent(Event{Type: WakeUpEvent})
	}

	log.Printf("Alarm shutting down.")
}

// checkAndPingClients looks at each connected client.
//
// If they've been idle a short time, we send them a PING (if they're
// registered).
//
// If they've been idle a long time, we kill their connection.
func (s *Server) checkAndPingClients() {
	now := time.Now()

	for _, client := range s.Clients {
		timeIdle := now.Sub(client.LastActivityTime)
		timeSincePing := now.Sub(client.LastPingTime)

		if client.Registered {
			// Was it active recently enough that we don't need to do anything?
			if timeIdle < s.Config.PingTime {
				continue
			}

			// It's been idle a while.

			// Has it been idle long enough that we consider it dead?
			if timeIdle > s.Config.DeadTime {
				client.quit(fmt.Sprintf("Ping timeout: %d seconds",
					int(timeIdle.Seconds())))
				continue
			}

			// Should we ping it? We might have pinged it recently.
			if timeSincePing < s.Config.PingTime {
				continue
			}

			s.messageClient(client, "PING", []string{s.Config.ServerName})
			client.LastPingTime = now
			continue
		}

		if timeIdle > s.Config.DeadTime {
			client.quit("Idle too long.")
		}
	}
}

// Send an IRC message to a client. Appears to be from the server.
// This works by writing to a client's channel.
//
// Note: Only the server goroutine should call this (due to channel use).
func (s *Server) messageClient(c *Client, command string, params []string) {
	// For numeric messages, we need to prepend the nick.
	// Use * for the nick in cases where the client doesn't have one yet.
	// This is what ircd-ratbox does. Maybe not RFC...
	isNumeric := true
	for _, c := range command {
		if c < 48 || c > 57 {
			isNumeric = false
		}
	}

	if isNumeric {
		nick := "*"
		if len(c.DisplayNick) > 0 {
			nick = c.DisplayNick
		}
		newParams := []string{nick}

		newParams = append(newParams, params...)
		params = newParams
	}

	c.WriteChan <- irc.Message{
		Prefix:  s.Config.ServerName,
		Command: command,
		Params:  params,
	}
}

// newEvent tells the server something happens.
//
// Any goroutine can call this function.
//
// It sends the server a message on its ToServerChan.
//
// It will not block on shutdown as we select on the shutdown channel which we
// close when shutting down the server. This means receive on the shutdown
// channel will proceed at that point.
func (s *Server) newEvent(evt Event) {
	select {
	case s.ToServerChan <- evt:
	case <-s.ShutdownChan:
	}
}
<|endoftext|>


<commit_before>// Package wsync computes a list of operations needed to mutate one file
// into another file, re-using as much of the former as possible.
//
// Base on code from: https://bitbucket.org/kardianos/rsync/
// Original rsync algorithm: http://www.samba.org/~tridge/phd_thesis.pdf
//
// The main change in our fork is supporting blocks of sizes less than
// the context's block size (instead of just passing them as OpData),
// and being able to pick from a hash library that can span multiple
// files, and not just the 'old version' of a file (at the same path).
// This allows us to handle renames gracefully (incl. partial rewrites)
//
// Definitions
//   Source: The final content.
//   Target: The content to be made into final content.
//   Signature: The sequence of hashes used to identify the content.
package wsync

import (
	"crypto/md5"
	"io"
	"os"

	"github.com/itchio/lake"
	"github.com/pkg/errors"
)

// MaxDataOp is the maximum number of 'fresh bytes' that can be contained
// in a single Data operation
const MaxDataOp = (4 * 1024 * 1024)

// NewContext creates a new Context, given a blocksize.
// It uses MD5 as a 'strong hash' (in the sense of an RSync paper,
// and compared to the very weak rolling hash)
func NewContext(BlockSize int) *Context {
	return &Context{
		blockSize:    BlockSize,
		uniqueHasher: md5.New(),
	}
}

type devNullReader struct{}

var _ io.Reader = (*devNullReader)(nil)

func (dvr *devNullReader) Read(buf []byte) (int, error) {
	for i := range buf {
		buf[i] = 0
	}
	return len(buf), nil
}

// ApplyPatch applies the difference to the target.
func (ctx *Context) ApplyPatch(output io.Writer, pool lake.Pool, ops chan Operation) error {
	return ctx.ApplyPatchFull(output, pool, ops, true)
}

// ApplyPatchFull is like ApplyPatch but accepts an ApplyWound channel
func (ctx *Context) ApplyPatchFull(output io.Writer, pool lake.Pool, ops chan Operation, failFast bool) error {
	for op := range ops {
		err := ctx.ApplySingleFull(output, pool, op, failFast)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

func (ctx *Context) ApplySingle(output io.Writer, pool lake.Pool, op Operation) error {
	return ctx.ApplySingleFull(output, pool, op, true)
}

func (ctx *Context) ApplySingleFull(output io.Writer, pool lake.Pool, op Operation, failFast bool) error {
	const minBufferSize = 32 * 1024 // golang's io.Copy default szie
	if len(ctx.buffer) < minBufferSize {
		ctx.buffer = make([]byte, minBufferSize)
	}
	buffer := ctx.buffer

	blockSize := int64(ctx.blockSize)

	switch op.Type {
	case OpBlockRange:
		fileSize := pool.GetSize(op.FileIndex)
		fixedSize := (op.BlockSpan - 1) * blockSize
		lastIndex := op.BlockIndex + (op.BlockSpan - 1)
		lastSize := blockSize
		if blockSize*(lastIndex+1) > fileSize {
			lastSize = fileSize % blockSize
		}
		opSize := (fixedSize + lastSize)

		target, err := pool.GetReadSeeker(op.FileIndex)
		if err != nil {
			if failFast {
				return errors.WithStack(err)
			}
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, opSize), buffer)
			return nil
		}

		_, err = target.Seek(blockSize*op.BlockIndex, os.SEEK_SET)
		if err != nil {
			if failFast {
				return errors.WithStack(err)
			}
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, opSize), buffer)
			return nil
		}

		copied, err := io.CopyBuffer(output, io.LimitReader(target, opSize), buffer)
		if err != nil {
			if failFast {
				return errors.Wrapf(err, "While copying %d bytes", blockSize*op.BlockSpan)
			}

			remaining := opSize - copied
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, remaining), buffer)
			return nil
		}
	case OpData:
		_, err := output.Write(op.Data)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

// ComputeDiff creates the operation list to mutate the target signature into the source.
// Any data operation from the OperationWriter must have the data copied out
// within the span of the function; the data buffer underlying the operation
// data is reused.
func (ctx *Context) ComputeDiff(source io.Reader, library *BlockLibrary, ops OperationWriter, preferredFileIndex int64) (err error) {
	minBufferSize := (ctx.blockSize * 2) + MaxDataOp
	if len(ctx.buffer) < minBufferSize {
		ctx.buffer = make([]byte, minBufferSize)
	}
	buffer := ctx.buffer

	type section struct {
		tail int
		head int
	}

	var data, sum section
	var n, validTo int
	var αPop, αPush, β, β1, β2 uint32
	var rolling, lastRun bool
	var shortSize int32

	// Store the previous non-data operation for combining.
	var prevOp *Operation

	// Send the last operation if there is one waiting.
	defer func() {
		if prevOp == nil {
			return
		}

		err = ops(*prevOp)
		prevOp = nil
	}()

	// Combine OpBlockRanges together. To achieve this, we store the previous
	// non-data operation and determine if it can be extended.
	enqueue := func(op Operation) error {
		switch op.Type {
		case OpBlockRange:
			if prevOp != nil {
				if prevOp.Type == OpBlockRange && prevOp.FileIndex == op.FileIndex && prevOp.BlockIndex+prevOp.BlockSpan == op.BlockIndex {
					// combine [prevOp][op] into [ prevOp ]
					prevOp.BlockSpan += op.BlockSpan
					return nil
				}

				opErr := ops(*prevOp)
				if opErr != nil {
					return errors.WithStack(opErr)
				}
				// prevOp has been completely sent off, can no longer be combined with anything
				prevOp = nil
			}
			prevOp = &op
		case OpData:
			// Never save a data operation, as it would corrupt the buffer.
			if prevOp != nil {
				opErr := ops(*prevOp)
				if opErr != nil {
					return errors.WithStack(opErr)
				}
			}
			opErr := ops(op)
			if opErr != nil {
				return errors.WithStack(opErr)
			}
			prevOp = nil
		}
		return nil
	}

	for !lastRun {
		// Determine if the buffer should be extended.
		if sum.tail+ctx.blockSize > validTo {
			// Determine if the buffer should be wrapped.
			if validTo+ctx.blockSize > len(buffer) {
				// Before wrapping the buffer, send any trailing data off.
				if data.tail < data.head {
					err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:data.head]})
					if err != nil {
						return errors.WithStack(err)
					}
				}
				// Wrap the buffer.
				l := validTo - sum.tail
				copy(buffer[:l], buffer[sum.tail:validTo])

				// Reset indexes.
				validTo = l
				sum.tail = 0
				data.head = 0
				data.tail = 0
			}

			n, err = io.ReadAtLeast(source, buffer[validTo:validTo+ctx.blockSize], ctx.blockSize)
			validTo += n
			if err != nil {
				if errors.Cause(err) != io.EOF && errors.Cause(err) != io.ErrUnexpectedEOF {
					return errors.WithStack(err)
				}
				lastRun = true

				shortSize = int32(n)
			}
		}

		// Set the hash sum window head. Must either be a block size
		// or be at the end of the buffer.
		sum.head = min(sum.tail+ctx.blockSize, validTo)

		skip := false

		// Compute the rolling hash.
		if rolling {
			βold := β

			αPush = uint32(buffer[sum.head-1])
			β1 = (β1 - αPop + αPush) % _M
			β2 = (β2 - uint32(sum.head-sum.tail)*αPop + β1) % _M
			β = β1 + _M*β2

			if β == βold {
				skip = true
			}
		} else {
			β, β1, β2 = βhash(buffer[sum.tail:sum.head])
			rolling = true
		}

		var blockHash *BlockHash

		if !skip {
			// Determine if there is a hash match.
			if hh, ok := library.hashLookup[β]; ok {
				blockHash = ctx.findUniqueHash(hh, buffer[sum.tail:sum.head], shortSize, preferredFileIndex)
			}
		}

		// Send data off if there is data available and a hash is found (so the buffer before it
		// must be flushed first), or the data chunk size has reached it's maximum size (for buffer
		// allocation purposes) or to flush the end of the data.
		if data.tail < data.head && (blockHash != nil || data.head-data.tail >= MaxDataOp) {
			err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:data.head]})
			if err != nil {
				return errors.WithStack(err)
			}
			data.tail = data.head
		}

		if blockHash != nil {
			err = enqueue(Operation{Type: OpBlockRange, FileIndex: blockHash.FileIndex, BlockIndex: blockHash.BlockIndex, BlockSpan: 1})
			if err != nil {
				return errors.WithStack(err)
			}
			rolling = false
			sum.tail += ctx.blockSize

			// There is prior knowledge that any available data
			// buffered will have already been sent. Thus we can
			// assume data.head and data.tail are the same.
			// May trigger "data wrap".
			data.head = sum.tail
			data.tail = sum.tail
		} else {
			if lastRun {
				err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:validTo]})
				if err != nil {
					return errors.WithStack(err)
				}
			} else {
				// The following is for the next loop iteration, so don't try to calculate if last.
				if rolling {
					αPop = uint32(buffer[sum.tail])
				}
				sum.tail++

				// May trigger "data wrap".
				data.head = sum.tail
			}
		}
	}
	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<commit_msg>rsync: Clean operations - don't send trailing zero-length DATA ops<commit_after>// Package wsync computes a list of operations needed to mutate one file
// into another file, re-using as much of the former as possible.
//
// Base on code from: https://bitbucket.org/kardianos/rsync/
// Original rsync algorithm: http://www.samba.org/~tridge/phd_thesis.pdf
//
// The main change in our fork is supporting blocks of sizes less than
// the context's block size (instead of just passing them as OpData),
// and being able to pick from a hash library that can span multiple
// files, and not just the 'old version' of a file (at the same path).
// This allows us to handle renames gracefully (incl. partial rewrites)
//
// Definitions
//   Source: The final content.
//   Target: The content to be made into final content.
//   Signature: The sequence of hashes used to identify the content.
package wsync

import (
	"crypto/md5"
	"io"
	"os"

	"github.com/itchio/lake"
	"github.com/pkg/errors"
)

// MaxDataOp is the maximum number of 'fresh bytes' that can be contained
// in a single Data operation
const MaxDataOp = (4 * 1024 * 1024)

// NewContext creates a new Context, given a blocksize.
// It uses MD5 as a 'strong hash' (in the sense of an RSync paper,
// and compared to the very weak rolling hash)
func NewContext(BlockSize int) *Context {
	return &Context{
		blockSize:    BlockSize,
		uniqueHasher: md5.New(),
	}
}

type devNullReader struct{}

var _ io.Reader = (*devNullReader)(nil)

func (dvr *devNullReader) Read(buf []byte) (int, error) {
	for i := range buf {
		buf[i] = 0
	}
	return len(buf), nil
}

// ApplyPatch applies the difference to the target.
func (ctx *Context) ApplyPatch(output io.Writer, pool lake.Pool, ops chan Operation) error {
	return ctx.ApplyPatchFull(output, pool, ops, true)
}

// ApplyPatchFull is like ApplyPatch but accepts an ApplyWound channel
func (ctx *Context) ApplyPatchFull(output io.Writer, pool lake.Pool, ops chan Operation, failFast bool) error {
	for op := range ops {
		err := ctx.ApplySingleFull(output, pool, op, failFast)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

func (ctx *Context) ApplySingle(output io.Writer, pool lake.Pool, op Operation) error {
	return ctx.ApplySingleFull(output, pool, op, true)
}

func (ctx *Context) ApplySingleFull(output io.Writer, pool lake.Pool, op Operation, failFast bool) error {
	const minBufferSize = 32 * 1024 // golang's io.Copy default szie
	if len(ctx.buffer) < minBufferSize {
		ctx.buffer = make([]byte, minBufferSize)
	}
	buffer := ctx.buffer

	blockSize := int64(ctx.blockSize)

	switch op.Type {
	case OpBlockRange:
		fileSize := pool.GetSize(op.FileIndex)
		fixedSize := (op.BlockSpan - 1) * blockSize
		lastIndex := op.BlockIndex + (op.BlockSpan - 1)
		lastSize := blockSize
		if blockSize*(lastIndex+1) > fileSize {
			lastSize = fileSize % blockSize
		}
		opSize := (fixedSize + lastSize)

		target, err := pool.GetReadSeeker(op.FileIndex)
		if err != nil {
			if failFast {
				return errors.WithStack(err)
			}
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, opSize), buffer)
			return nil
		}

		_, err = target.Seek(blockSize*op.BlockIndex, os.SEEK_SET)
		if err != nil {
			if failFast {
				return errors.WithStack(err)
			}
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, opSize), buffer)
			return nil
		}

		copied, err := io.CopyBuffer(output, io.LimitReader(target, opSize), buffer)
		if err != nil {
			if failFast {
				return errors.Wrapf(err, "While copying %d bytes", blockSize*op.BlockSpan)
			}

			remaining := opSize - copied
			io.CopyBuffer(output, io.LimitReader(&devNullReader{}, remaining), buffer)
			return nil
		}
	case OpData:
		_, err := output.Write(op.Data)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

// makeOperationCleaner filters out empty DATA ops that are *not*
// the first operation. This makes patches slightly cleaner.
func makeOperationCleaner(ops OperationWriter) OperationWriter {
	sendCount := 0

	return func(op Operation) error {
		if sendCount > 0 && op.Type == OpData && len(op.Data) == 0 {
			// discard
			return nil
		}

		sendCount++
		return ops(op)
	}
}

// ComputeDiff creates the operation list to mutate the target signature into the source.
// Any data operation from the OperationWriter must have the data copied out
// within the span of the function; the data buffer underlying the operation
// data is reused.
func (ctx *Context) ComputeDiff(source io.Reader, library *BlockLibrary, ops OperationWriter, preferredFileIndex int64) (err error) {
	ops = makeOperationCleaner(ops)

	minBufferSize := (ctx.blockSize * 2) + MaxDataOp
	if len(ctx.buffer) < minBufferSize {
		ctx.buffer = make([]byte, minBufferSize)
	}
	buffer := ctx.buffer

	type section struct {
		tail int
		head int
	}

	var data, sum section
	var n, validTo int
	var αPop, αPush, β, β1, β2 uint32
	var rolling, lastRun bool
	var shortSize int32

	// Store the previous non-data operation for combining.
	var prevOp *Operation

	// Send the last operation if there is one waiting.
	defer func() {
		if prevOp != nil {
			err = ops(*prevOp)
			prevOp = nil
		}
	}()

	// Combine OpBlockRanges together. To achieve this, we store the previous
	// non-data operation and determine if it can be extended.
	enqueue := func(op Operation) error {
		switch op.Type {
		case OpBlockRange:
			if prevOp != nil {
				if prevOp.Type == OpBlockRange && prevOp.FileIndex == op.FileIndex && prevOp.BlockIndex+prevOp.BlockSpan == op.BlockIndex {
					// combine [prevOp][op] into [ prevOp ]
					prevOp.BlockSpan += op.BlockSpan
					return nil
				}

				err := ops(*prevOp)
				if err != nil {
					return errors.WithStack(err)
				}
				// prevOp has been completely sent off, can no longer be combined with anything
				prevOp = nil
			}
			prevOp = &op
		case OpData:
			// Never save a data operation, as it would corrupt the buffer.
			if prevOp != nil {
				err := ops(*prevOp)
				prevOp = nil
				if err != nil {
					return errors.WithStack(err)
				}
			}

			err := ops(op)
			if err != nil {
				return errors.WithStack(err)
			}
		}
		return nil
	}

	for !lastRun {
		// Determine if the buffer should be extended.
		if sum.tail+ctx.blockSize > validTo {
			// Determine if the buffer should be wrapped.
			if validTo+ctx.blockSize > len(buffer) {
				// Before wrapping the buffer, send any trailing data off.
				if data.tail < data.head {
					err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:data.head]})
					if err != nil {
						return errors.WithStack(err)
					}
				}
				// Wrap the buffer.
				l := validTo - sum.tail
				copy(buffer[:l], buffer[sum.tail:validTo])

				// Reset indexes.
				validTo = l
				sum.tail = 0
				data.head = 0
				data.tail = 0
			}

			n, err = io.ReadAtLeast(source, buffer[validTo:validTo+ctx.blockSize], ctx.blockSize)
			validTo += n
			if err != nil {
				if errors.Cause(err) != io.EOF && errors.Cause(err) != io.ErrUnexpectedEOF {
					return errors.WithStack(err)
				}
				lastRun = true

				shortSize = int32(n)
			}
		}

		// Set the hash sum window head. Must either be a block size
		// or be at the end of the buffer.
		sum.head = min(sum.tail+ctx.blockSize, validTo)

		skip := false

		// Compute the rolling hash.
		if rolling {
			βold := β

			αPush = uint32(buffer[sum.head-1])
			β1 = (β1 - αPop + αPush) % _M
			β2 = (β2 - uint32(sum.head-sum.tail)*αPop + β1) % _M
			β = β1 + _M*β2

			if β == βold {
				skip = true
			}
		} else {
			β, β1, β2 = βhash(buffer[sum.tail:sum.head])
			rolling = true
		}

		var blockHash *BlockHash

		if !skip {
			// Determine if there is a hash match.
			if hh, ok := library.hashLookup[β]; ok {
				blockHash = ctx.findUniqueHash(hh, buffer[sum.tail:sum.head], shortSize, preferredFileIndex)
			}
		}

		// Send data off if there is data available and a hash is found (so the buffer before it
		// must be flushed first), or the data chunk size has reached it's maximum size (for buffer
		// allocation purposes) or to flush the end of the data.
		if data.tail < data.head && (blockHash != nil || data.head-data.tail >= MaxDataOp) {
			err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:data.head]})
			if err != nil {
				return errors.WithStack(err)
			}
			data.tail = data.head
		}

		if blockHash != nil {
			err = enqueue(Operation{Type: OpBlockRange, FileIndex: blockHash.FileIndex, BlockIndex: blockHash.BlockIndex, BlockSpan: 1})
			if err != nil {
				return errors.WithStack(err)
			}
			rolling = false
			sum.tail += ctx.blockSize

			// There is prior knowledge that any available data
			// buffered will have already been sent. Thus we can
			// assume data.head and data.tail are the same.
			// May trigger "data wrap".
			data.head = sum.tail
			data.tail = sum.tail
		} else {
			if lastRun {
				err = enqueue(Operation{Type: OpData, Data: buffer[data.tail:validTo]})
				if err != nil {
					return errors.WithStack(err)
				}
			} else {
				// The following is for the next loop iteration, so don't try to calculate if last.
				if rolling {
					αPop = uint32(buffer[sum.tail])
				}
				sum.tail++

				// May trigger "data wrap".
				data.head = sum.tail
			}
		}
	}
	return nil
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	serial "github.com/tarm/serial"
)

type Scale struct {
	port io.ReadWriteCloser
}

func (scale Scale) Open() error {
	c := serial.Config{Name: "/dev/ttyUSB0", Baud: 9600, ReadTimeout: time.Second * 1}
	port, err := serial.OpenPort(&c)
	scale.port = port
	return scale, err
}

func (scale Scale) Read() {
	port.Write("IP\r\n")
	buf := make([]byte, 128)
	n, err = port.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("%q", buf[:n])

}

func main() {
	scale = scale.Scale{}
	err = scale.Open()
	if err != nil {
		log.Fatal(err)
	}
	value, err = scale.Read()
	fmt.Println(value)
}
<commit_msg>basic reading from the scale<commit_after>package main

import (
	"bufio"
	"fmt"
	serial "github.com/tarm/serial"
	"log"
	"os"
	"time"
)

type Scale struct {
}

func (scale Scale) Open() (port *serial.Port, err error) {
	c := &serial.Config{Name: "/dev/ttyUSB0", Baud: 9600}
	port, err = serial.OpenPort(c)
	return
}

func (scale Scale) Read(port *serial.Port) (value string) {
	port.Write([]byte("IP\r\n"))
	scanner := bufio.NewScanner(port)
	scanner.Scan()

	value = scanner.Text()
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}
	return
}

func main() {
	scale := Scale{}
	port, err := scale.Open()
	if err != nil {
		log.Fatal(err)
	}
	for {
		value := scale.Read(port)
		fmt.Println(time.Now(), value)
	}
}
<|endoftext|>


<commit_before>package main

import "testing"
import "os/user"

func Test_getCurrentUser_1(t *testing.T) {
    _, _, _, _, err := getCurrentUser()
    if err != nil {
        t.Error("Error from getCurrentUser")
    }
}

func Test_getUser_1(t *testing.T) {
	mockuser := &user.User{Username: "vagrant", HomeDir: "/home/vagrant", Uid: "1000", Gid: "1000"}
	username, homedir, uid, gid, err := getUser(mockuser)
	if err != nil {
		t.Error("Got error from getUser " + err.Error())
	}
	if username == "vagrant" {
		t.Log("username passed.")
	} else {
		t.Errorf("Username failed: %s", username)
	}
	if homedir == "/home/vagrant" {
		t.Log("homedir passed.")
	} else {
		t.Errorf("homedir failed: %s", homedir)
	}
	if uid == 1000 {
		t.Log("uid passed.")
	} else {
		t.Errorf("uid failed: %i", uid)
	}
	if gid == 1000 {
		t.Log("git passed.")
	} else {
		t.Errorf("gid failed: %i", gid)
	}
}

func Test_getUser_2(t *testing.T) {
	mockuser := &user.User{Username: "", HomeDir: "/home/vagrant", Uid: "1000", Gid: "1000"}
	_, _, _, _, err := getUser(mockuser)
	if err == nil {
		t.Error("No error from getUser")
	}
}

func Test_getUser_3(t *testing.T) {
	mockuser := &user.User{Username: "Foo", HomeDir: "", Uid: "1000", Gid: "1000"}
	_, _, _, _, err := getUser(mockuser)
	if err == nil {
		t.Error("No error from getUser")
	}
}

<commit_msg>go fmt<commit_after>package main

import "testing"
import "os/user"

func Test_getCurrentUser_1(t *testing.T) {
	_, _, _, _, err := getCurrentUser()
	if err != nil {
		t.Error("Error from getCurrentUser")
	}
}

func Test_getUser_1(t *testing.T) {
	mockuser := &user.User{Username: "vagrant", HomeDir: "/home/vagrant", Uid: "1000", Gid: "1000"}
	username, homedir, uid, gid, err := getUser(mockuser)
	if err != nil {
		t.Error("Got error from getUser " + err.Error())
	}
	if username == "vagrant" {
		t.Log("username passed.")
	} else {
		t.Errorf("Username failed: %s", username)
	}
	if homedir == "/home/vagrant" {
		t.Log("homedir passed.")
	} else {
		t.Errorf("homedir failed: %s", homedir)
	}
	if uid == 1000 {
		t.Log("uid passed.")
	} else {
		t.Errorf("uid failed: %i", uid)
	}
	if gid == 1000 {
		t.Log("git passed.")
	} else {
		t.Errorf("gid failed: %i", gid)
	}
}

func Test_getUser_2(t *testing.T) {
	mockuser := &user.User{Username: "", HomeDir: "/home/vagrant", Uid: "1000", Gid: "1000"}
	_, _, _, _, err := getUser(mockuser)
	if err == nil {
		t.Error("No error from getUser")
	}
}

func Test_getUser_3(t *testing.T) {
	mockuser := &user.User{Username: "Foo", HomeDir: "", Uid: "1000", Gid: "1000"}
	_, _, _, _, err := getUser(mockuser)
	if err == nil {
		t.Error("No error from getUser")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
	pbg "github.com/brotherlogic/goserver/proto"
	"github.com/brotherlogic/goserver/utils"
)

const (
	intentWait = time.Second
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	config     *pb.Config
	serving    bool
	LastIntent time.Time
	LastMaster time.Time
	worldMutex *sync.Mutex
	world      map[string]map[string]struct{}
	getter     getter
}

type prodGetter struct{}

func (g *prodGetter) getJobs(server *pbd.RegistryEntry) (*pbs.JobList, error) {
	list := &pbs.JobList{}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	conn, err := grpc.Dial(server.GetIp()+":"+strconv.Itoa(int(server.GetPort())), grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		return list, err
	}

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(ctx, &pbs.Empty{}, grpc.FailFast(false))
	return r, err
}

func (s *Server) checkerThread(i *pb.Intent) {
	for true {
		time.Sleep(time.Minute)

		if len(s.world[i.GetSpec().GetName()]) != int(i.Count) {
			s.Log(fmt.Sprintf("MISMATCH: %v, %v", i, s.world[i.GetSpec().GetName()]))
		}
	}
}

func (g *prodGetter) getSlaves() (*pbd.ServiceList, error) {
	ret := &pbd.ServiceList{}

	conn, err := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	if err != nil {
		return ret, err
	}
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err != nil {
		return ret, err
	}

	for _, s := range r.Services {
		if s.GetName() == "gobuildslave" {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret, nil
}

type mainChecker struct {
	prev   []string
	logger func(string)
}

func getIP(servertype, servername string) (string, int) {
	conn, _ := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err != nil {
		return "", -1
	}
	for _, s := range r.Services {
		if s.Name == servertype && s.Identifier == servername {
			return s.Ip, int(s.Port)
		}
	}

	return "", -1
}

func (t *mainChecker) getprev() []string {
	return t.prev
}
func (t *mainChecker) setprev(v []string) {
	t.prev = v
}

func (t *mainChecker) assess(server string) (*pbs.JobList, *pbs.Config) {
	list := &pbs.JobList{}
	conf := &pbs.Config{}

	ip, port := getIP("gobuildslave", server)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	conn, err := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		return list, conf
	}

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(ctx, &pbs.Empty{}, grpc.FailFast(false))
	if err != nil {
		return list, conf
	}

	r2, err := slave.GetConfig(ctx, &pbs.Empty{}, grpc.FailFast(false))
	if err != nil {
		return list, conf
	}

	return r, r2
}

func (t *mainChecker) master(entry *pbd.RegistryEntry, master bool) bool {
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	conn, _ := grpc.Dial(entry.GetIp()+":"+strconv.Itoa(int(entry.GetPort())), grpc.WithInsecure())
	defer conn.Close()

	server := pbg.NewGoserverServiceClient(conn)
	_, err := server.Mote(ctx, &pbg.MoteRequest{Master: master}, grpc.FailFast(false))
	if err != nil {
		t.logger(fmt.Sprintf("Master REJECT(%v): %v", entry, err))
	}

	return err == nil
}

func runJob(job *pbs.JobSpec, server string) {
	if server != "" {
		ip, port := getIP("gobuildslave", server)
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
		defer conn.Close()

		slave := pbs.NewGoBuildSlaveClient(conn)
		job.Server = server
		slave.Run(ctx, job, grpc.FailFast(false))
	}
}

func (t *mainChecker) discover() *pbd.ServiceList {
	ret := &pbd.ServiceList{}

	conn, _ := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err == nil {
		for _, s := range r.Services {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildMasterServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//GetState gets the state of the server
func (s Server) GetState() []*pbg.State {
	return []*pbg.State{&pbg.State{Key: "last_intent", TimeValue: s.LastIntent.Unix()},
		&pbg.State{Key: "last_master", TimeValue: s.LastMaster.Unix()},
		&pbg.State{Key: "world", Text: fmt.Sprintf("%v", s.world)}}
}

//Compare compares current state to desired state
func (s Server) Compare(ctx context.Context, in *pb.Empty) (*pb.CompareResponse, error) {
	resp := &pb.CompareResponse{}
	list, _ := getFleetStatus(&mainChecker{logger: s.Log})
	cc := &pb.Config{}
	for _, jlist := range list {
		for _, job := range jlist.GetDetails() {
			cc.Intents = append(cc.Intents, &pb.Intent{Spec: job.GetSpec()})
		}
	}
	resp.Current = cc
	resp.Desired = s.config

	return resp, nil
}

func getConfig(c checker) *pb.Config {
	list, _ := getFleetStatus(c)
	config := &pb.Config{}

	for _, jlist := range list {
		for _, job := range jlist.Details {
			found := false
			for _, ij := range config.Intents {
				if job.Spec.Name == ij.Spec.Name {
					ij.Count++
					found = true
				}
			}

			if !found {
				config.Intents = append(config.Intents, &pb.Intent{Spec: &pbs.JobSpec{Name: job.Spec.Name}, Count: 1})
			}
		}
	}

	return config
}

// MatchIntent tries to match the intent with the state of production
func (s *Server) MatchIntent() {
	checker := &mainChecker{logger: s.Log}
	for s.serving {
		time.Sleep(intentWait)
		s.LastIntent = time.Now()

		state := getConfig(checker)
		diff := configDiff(s.config, state)
		joblist := runJobs(diff)
		for _, job := range joblist {
			runJob(job, chooseServer(job, checker))
		}
	}
}

// SetMaster sets up the master settings
func (s *Server) SetMaster() {
	t := time.Now()
	checker := &mainChecker{logger: s.Log}
	s.LastMaster = time.Now()

	fleet := checker.discover()
	matcher := make(map[string][]*pbd.RegistryEntry)
	hasMaster := make(map[string]int)
	for _, entry := range fleet.GetServices() {
		if !entry.GetIgnoresMaster() {
			if _, ok := matcher[entry.GetName()]; !ok {
				if entry.GetMaster() {
					hasMaster[entry.GetName()]++
				}
				matcher[entry.GetName()] = []*pbd.RegistryEntry{entry}
			} else {
				if entry.GetMaster() {
					hasMaster[entry.GetName()] = 1
					matcher[entry.GetName()] = append(matcher[entry.GetName()], entry)
				}
			}
		}
	}

	for key, entries := range matcher {
		if hasMaster[key] > 1 {
			hasMaster[key] = 1
			seen := false
			for _, entry := range entries {
				if seen && entry.GetMaster() {
					checker.master(entry, false)
				} else if entry.GetMaster() {
					seen = true
				}
			}
		}

		if hasMaster[key] == 0 {
			for _, entry := range entries {
				if checker.master(entry, true) {
					entry.Master = true
					break
				}
			}
		}
	}
	s.LogFunction("SetMasterRun", t)
}

//Init builds up the server
func Init(config *pb.Config) *Server {
	s := &Server{
		&goserver.GoServer{},
		config,
		true,
		time.Now(),
		time.Now(),
		&sync.Mutex{},
		make(map[string]map[string]struct{}),
		&prodGetter{},
	}
	return s
}

func (s *Server) becomeMaster() {
	for true {
		time.Sleep(time.Second * 5)
		_, _, err := utils.Resolve("gobuildmaster")
		if err != nil {
			s.Registry.Master = true
		}
	}
}

func main() {
	config, err := loadConfig("config.pb")
	if err != nil {
		log.Fatalf("Fatal loading of config: %v", err)
	}

	s := Init(config)

	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServer("gobuildmaster", false)
	s.RegisterServingTask(s.MatchIntent)
	s.RegisterServingTask(s.becomeMaster)
	s.RegisterRepeatingTask(s.SetMaster, time.Second)
	s.RegisterRepeatingTask(s.buildWorld, time.Minute)

	go s.checkerThread(s.config.GetIntents()[0])

	err = s.Serve()
	if err != nil {
		log.Fatalf("Serve error: %v", err)
	}
}
<commit_msg>Runs an up. This closes #96<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"strconv"
	"sync"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	pbs "github.com/brotherlogic/gobuildslave/proto"
	pbg "github.com/brotherlogic/goserver/proto"
	"github.com/brotherlogic/goserver/utils"
)

const (
	intentWait = time.Second
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	config     *pb.Config
	serving    bool
	LastIntent time.Time
	LastMaster time.Time
	worldMutex *sync.Mutex
	world      map[string]map[string]struct{}
	getter     getter
}

type prodGetter struct{}

func (g *prodGetter) getJobs(server *pbd.RegistryEntry) (*pbs.JobList, error) {
	list := &pbs.JobList{}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	conn, err := grpc.Dial(server.GetIp()+":"+strconv.Itoa(int(server.GetPort())), grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		return list, err
	}

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(ctx, &pbs.Empty{}, grpc.FailFast(false))
	return r, err
}

func (s *Server) checkerThread(i *pb.Intent) {
	for true {
		time.Sleep(time.Minute)

		if len(s.world[i.GetSpec().GetName()]) != int(i.Count) {
			s.Log(fmt.Sprintf("MISMATCH: %v, %v", i, s.world[i.GetSpec().GetName()]))
			if len(s.world[i.GetSpec().GetName()]) < int(i.Count) {
				s.Log(fmt.Sprintf("UP : %v", i.GetSpec().GetName()))
			}
		}
	}
}

func (g *prodGetter) getSlaves() (*pbd.ServiceList, error) {
	ret := &pbd.ServiceList{}

	conn, err := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	if err != nil {
		return ret, err
	}
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err != nil {
		return ret, err
	}

	for _, s := range r.Services {
		if s.GetName() == "gobuildslave" {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret, nil
}

type mainChecker struct {
	prev   []string
	logger func(string)
}

func getIP(servertype, servername string) (string, int) {
	conn, _ := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err != nil {
		return "", -1
	}
	for _, s := range r.Services {
		if s.Name == servertype && s.Identifier == servername {
			return s.Ip, int(s.Port)
		}
	}

	return "", -1
}

func (t *mainChecker) getprev() []string {
	return t.prev
}
func (t *mainChecker) setprev(v []string) {
	t.prev = v
}

func (t *mainChecker) assess(server string) (*pbs.JobList, *pbs.Config) {
	list := &pbs.JobList{}
	conf := &pbs.Config{}

	ip, port := getIP("gobuildslave", server)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	conn, err := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
	defer conn.Close()
	if err != nil {
		return list, conf
	}

	slave := pbs.NewGoBuildSlaveClient(conn)
	r, err := slave.List(ctx, &pbs.Empty{}, grpc.FailFast(false))
	if err != nil {
		return list, conf
	}

	r2, err := slave.GetConfig(ctx, &pbs.Empty{}, grpc.FailFast(false))
	if err != nil {
		return list, conf
	}

	return r, r2
}

func (t *mainChecker) master(entry *pbd.RegistryEntry, master bool) bool {
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()
	conn, _ := grpc.Dial(entry.GetIp()+":"+strconv.Itoa(int(entry.GetPort())), grpc.WithInsecure())
	defer conn.Close()

	server := pbg.NewGoserverServiceClient(conn)
	_, err := server.Mote(ctx, &pbg.MoteRequest{Master: master}, grpc.FailFast(false))
	if err != nil {
		t.logger(fmt.Sprintf("Master REJECT(%v): %v", entry, err))
	}

	return err == nil
}

func runJob(job *pbs.JobSpec, server string) {
	if server != "" {
		ip, port := getIP("gobuildslave", server)
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		conn, _ := grpc.Dial(ip+":"+strconv.Itoa(port), grpc.WithInsecure())
		defer conn.Close()

		slave := pbs.NewGoBuildSlaveClient(conn)
		job.Server = server
		slave.Run(ctx, job, grpc.FailFast(false))
	}
}

func (t *mainChecker) discover() *pbd.ServiceList {
	ret := &pbd.ServiceList{}

	conn, _ := grpc.Dial(utils.RegistryIP+":"+strconv.Itoa(utils.RegistryPort), grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	r, err := registry.ListAllServices(ctx, &pbd.Empty{}, grpc.FailFast(false))
	if err == nil {
		for _, s := range r.Services {
			ret.Services = append(ret.Services, s)
		}
	}

	return ret
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildMasterServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//GetState gets the state of the server
func (s Server) GetState() []*pbg.State {
	return []*pbg.State{&pbg.State{Key: "last_intent", TimeValue: s.LastIntent.Unix()},
		&pbg.State{Key: "last_master", TimeValue: s.LastMaster.Unix()},
		&pbg.State{Key: "world", Text: fmt.Sprintf("%v", s.world)}}
}

//Compare compares current state to desired state
func (s Server) Compare(ctx context.Context, in *pb.Empty) (*pb.CompareResponse, error) {
	resp := &pb.CompareResponse{}
	list, _ := getFleetStatus(&mainChecker{logger: s.Log})
	cc := &pb.Config{}
	for _, jlist := range list {
		for _, job := range jlist.GetDetails() {
			cc.Intents = append(cc.Intents, &pb.Intent{Spec: job.GetSpec()})
		}
	}
	resp.Current = cc
	resp.Desired = s.config

	return resp, nil
}

func getConfig(c checker) *pb.Config {
	list, _ := getFleetStatus(c)
	config := &pb.Config{}

	for _, jlist := range list {
		for _, job := range jlist.Details {
			found := false
			for _, ij := range config.Intents {
				if job.Spec.Name == ij.Spec.Name {
					ij.Count++
					found = true
				}
			}

			if !found {
				config.Intents = append(config.Intents, &pb.Intent{Spec: &pbs.JobSpec{Name: job.Spec.Name}, Count: 1})
			}
		}
	}

	return config
}

// MatchIntent tries to match the intent with the state of production
func (s *Server) MatchIntent() {
	checker := &mainChecker{logger: s.Log}
	for s.serving {
		time.Sleep(intentWait)
		s.LastIntent = time.Now()

		state := getConfig(checker)
		diff := configDiff(s.config, state)
		joblist := runJobs(diff)
		for _, job := range joblist {
			runJob(job, chooseServer(job, checker))
		}
	}
}

// SetMaster sets up the master settings
func (s *Server) SetMaster() {
	t := time.Now()
	checker := &mainChecker{logger: s.Log}
	s.LastMaster = time.Now()

	fleet := checker.discover()
	matcher := make(map[string][]*pbd.RegistryEntry)
	hasMaster := make(map[string]int)
	for _, entry := range fleet.GetServices() {
		if !entry.GetIgnoresMaster() {
			if _, ok := matcher[entry.GetName()]; !ok {
				if entry.GetMaster() {
					hasMaster[entry.GetName()]++
				}
				matcher[entry.GetName()] = []*pbd.RegistryEntry{entry}
			} else {
				if entry.GetMaster() {
					hasMaster[entry.GetName()] = 1
					matcher[entry.GetName()] = append(matcher[entry.GetName()], entry)
				}
			}
		}
	}

	for key, entries := range matcher {
		if hasMaster[key] > 1 {
			hasMaster[key] = 1
			seen := false
			for _, entry := range entries {
				if seen && entry.GetMaster() {
					checker.master(entry, false)
				} else if entry.GetMaster() {
					seen = true
				}
			}
		}

		if hasMaster[key] == 0 {
			for _, entry := range entries {
				if checker.master(entry, true) {
					entry.Master = true
					break
				}
			}
		}
	}
	s.LogFunction("SetMasterRun", t)
}

//Init builds up the server
func Init(config *pb.Config) *Server {
	s := &Server{
		&goserver.GoServer{},
		config,
		true,
		time.Now(),
		time.Now(),
		&sync.Mutex{},
		make(map[string]map[string]struct{}),
		&prodGetter{},
	}
	return s
}

func (s *Server) becomeMaster() {
	for true {
		time.Sleep(time.Second * 5)
		_, _, err := utils.Resolve("gobuildmaster")
		if err != nil {
			s.Registry.Master = true
		}
	}
}

func main() {
	config, err := loadConfig("config.pb")
	if err != nil {
		log.Fatalf("Fatal loading of config: %v", err)
	}

	s := Init(config)

	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServer("gobuildmaster", false)
	s.RegisterServingTask(s.MatchIntent)
	s.RegisterServingTask(s.becomeMaster)
	s.RegisterRepeatingTask(s.SetMaster, time.Second)
	s.RegisterRepeatingTask(s.buildWorld, time.Minute)

	for i := 0; i < len(s.config.GetIntents()); i++ {
		go s.checkerThread(s.config.GetIntents()[i])
	}

	err = s.Serve()
	if err != nil {
		log.Fatalf("Serve error: %v", err)
	}
}
<|endoftext|>