<commit_before>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2017

package instana

import (
	"encoding/json"
	"time"

	"github.com/opentracing/opentracing-go/ext"
)

type typedSpanData interface {
	Type() RegisteredSpanType
	Kind() SpanKind
}

// SpanKind represents values of field `k` in OpenTracing span representation. It represents
// the direction of the call associated with a span.
type SpanKind uint8

// Valid span kinds
const (
	// The kind of a span associated with an inbound call, this must be the first span in the trace.
	EntrySpanKind SpanKind = iota + 1
	// The kind of a span associated with an outbound call, e.g. an HTTP client request, posting to a message bus, etc.
	ExitSpanKind
	// The default kind for a span that is associated with a call within the same service.
	IntermediateSpanKind
)

// String returns string representation of a span kind suitable for use as a value for `data.sdk.type`
// tag of an SDK span. By default all spans are intermediate unless they are explicitly set to be "entry" or "exit"
func (k SpanKind) String() string {
	switch k {
	case EntrySpanKind:
		return "entry"
	case ExitSpanKind:
		return "exit"
	default:
		return "intermediate"
	}
}

// Span represents the OpenTracing span document to be sent to the agent
type Span struct {
	TraceID         int64
	TraceIDHi       int64
	ParentID        int64
	SpanID          int64
	Ancestor        *TraceReference
	Timestamp       uint64
	Duration        uint64
	Name            string
	From            *fromS
	Batch           *batchInfo
	Kind            int
	Ec              int
	Data            typedSpanData
	Synthetic       bool
	CorrelationType string
	CorrelationID   string
	ForeignTrace    bool
}

func newSpan(span *spanS) Span {
	data := RegisteredSpanType(span.Operation).ExtractData(span)
	sp := Span{
		TraceID:         span.context.TraceID,
		TraceIDHi:       span.context.TraceIDHi,
		ParentID:        span.context.ParentID,
		SpanID:          span.context.SpanID,
		Timestamp:       uint64(span.Start.UnixNano()) / uint64(time.Millisecond),
		Duration:        uint64(span.Duration) / uint64(time.Millisecond),
		Name:            string(data.Type()),
		Ec:              span.ErrorCount,
		CorrelationType: span.Correlation.Type,
		CorrelationID:   span.Correlation.ID,
		ForeignTrace:    span.context.ForeignTrace,
		Kind:            int(data.Kind()),
		Data:            data,
	}

	if bs, ok := span.Tags[batchSizeTag].(int); ok {
		if bs > 1 {
			sp.Batch = &batchInfo{Size: bs}
		}
		delete(span.Tags, batchSizeTag)
	}

	if syn, ok := span.Tags[syntheticCallTag].(bool); ok {
		sp.Synthetic = syn
		delete(span.Tags, syntheticCallTag)
	}

	if len(span.context.Links) > 0 {
		ancestor := span.context.Links[0]
		sp.Ancestor = &TraceReference{
			TraceID:  ancestor.TraceID,
			ParentID: ancestor.SpanID,
		}
	}

	return sp
}

type TraceReference struct {
	TraceID  string `json:"t"`
	ParentID string `json:"p,omitempty"`
}

// MarshalJSON serializes span to JSON for sending it to Instana
func (sp Span) MarshalJSON() ([]byte, error) {
	var parentID string
	if sp.ParentID != 0 {
		parentID = FormatID(sp.ParentID)
	}

	var longTraceID string
	if sp.TraceIDHi != 0 && sp.Kind == int(EntrySpanKind) {
		longTraceID = FormatLongID(sp.TraceIDHi, sp.TraceID)
	}

	return json.Marshal(struct {
		TraceReference

		SpanID          string          `json:"s"`
		LongTraceID     string          `json:"lt,omitempty"`
		Timestamp       uint64          `json:"ts"`
		Duration        uint64          `json:"d"`
		Name            string          `json:"n"`
		From            *fromS          `json:"f"`
		Batch           *batchInfo      `json:"b,omitempty"`
		Kind            int             `json:"k"`
		Ec              int             `json:"ec,omitempty"`
		Data            typedSpanData   `json:"data"`
		Synthetic       bool            `json:"sy,omitempty"`
		CorrelationType string          `json:"crtp,omitempty"`
		CorrelationID   string          `json:"crid,omitempty"`
		ForeignTrace    bool            `json:"tp,omitempty"`
		Ancestor        *TraceReference `json:"ia,omitempty"`
	}{
		TraceReference{
			FormatID(sp.TraceID),
			parentID,
		},
		FormatID(sp.SpanID),
		longTraceID,
		sp.Timestamp,
		sp.Duration,
		sp.Name,
		sp.From,
		sp.Batch,
		sp.Kind,
		sp.Ec,
		sp.Data,
		sp.Synthetic,
		sp.CorrelationType,
		sp.CorrelationID,
		sp.ForeignTrace,
		sp.Ancestor,
	})
}

type batchInfo struct {
	Size int `json:"s"`
}

// CustomSpanData holds user-defined span tags
type CustomSpanData struct {
	Tags map[string]interface{} `json:"tags,omitempty"`
}

func filterCustomSpanTags(tags map[string]interface{}, st RegisteredSpanType) map[string]interface{} {
	knownTags := st.TagsNames()
	customTags := make(map[string]interface{})

	for k, v := range tags {
		if k == string(ext.SpanKind) {
			continue
		}

		if _, ok := knownTags[k]; ok {
			continue
		}

		customTags[k] = v
	}

	return customTags
}

// SpanData contains fields to be sent in the `data` section of an OT span document. These fields are
// common for all span types.
type SpanData struct {
	Service string          `json:"service,omitempty"`
	Custom  *CustomSpanData `json:"custom,omitempty"`

	st RegisteredSpanType
}

// NewSpanData initializes a new span data from tracer span
func NewSpanData(span *spanS, st RegisteredSpanType) SpanData {
	data := SpanData{
		Service: span.Service,
		st:      st,
	}

	if customTags := filterCustomSpanTags(span.Tags, st); len(customTags) > 0 {
		data.Custom = &CustomSpanData{Tags: customTags}
	}

	return data
}

// Type returns the registered span type suitable for use as the value of `n` field.
func (d SpanData) Type() RegisteredSpanType {
	return d.st
}

// SDKSpanData represents the `data` section of an SDK span sent within an OT span document
type SDKSpanData struct {
	// Deprecated
	SpanData `json:"-"`

	Service string      `json:"service,omitempty"`
	Tags    SDKSpanTags `json:"sdk"`

	sk SpanKind
}

// NewSDKSpanData initializes a new SDK span data from a tracer span
func NewSDKSpanData(span *spanS) SDKSpanData {
	sk := IntermediateSpanKind

	switch span.Tags[string(ext.SpanKind)] {
	case ext.SpanKindRPCServerEnum, string(ext.SpanKindRPCServerEnum),
		ext.SpanKindConsumerEnum, string(ext.SpanKindConsumerEnum),
		"entry":
		sk = EntrySpanKind
	case ext.SpanKindRPCClientEnum, string(ext.SpanKindRPCClientEnum),
		ext.SpanKindProducerEnum, string(ext.SpanKindProducerEnum),
		"exit":
		sk = ExitSpanKind
	}

	return SDKSpanData{
		Service: span.Service,
		Tags:    NewSDKSpanTags(span, sk.String()),
		sk:      sk,
	}
}

// Type returns the registered span type suitable for use as the value of `n` field.
func (d SDKSpanData) Type() RegisteredSpanType {
	return SDKSpanType
}

// Kind returns the kind of the span. It handles the github.com/opentracing/opentracing-go/ext.SpanKindEnum
// values as well as generic "entry" and "exit"
func (d SDKSpanData) Kind() SpanKind {
	return d.sk
}

// KnownTags returns the list of known tags for this span type
// SDKSpanTags contains fields within the `data.sdk` section of an OT span document
type SDKSpanTags struct {
	Name      string                 `json:"name"`
	Type      string                 `json:"type,omitempty"`
	Arguments string                 `json:"arguments,omitempty"`
	Return    string                 `json:"return,omitempty"`
	Custom    map[string]interface{} `json:"custom,omitempty"`
}

// NewSDKSpanTags extracts SDK span tags from a tracer span
func NewSDKSpanTags(span *spanS, spanType string) SDKSpanTags {
	tags := SDKSpanTags{
		Name:   span.Operation,
		Type:   spanType,
		Custom: map[string]interface{}{},
	}

	if len(span.Tags) != 0 {
		tags.Custom["tags"] = span.Tags
	}

	if logs := collectTracerSpanLogs(span); len(logs) > 0 {
		tags.Custom["logs"] = logs
	}

	if len(span.context.Baggage) != 0 {
		tags.Custom["baggage"] = span.context.Baggage
	}

	return tags
}

// readStringTag populates the &dst with the tag value if it's of either string or []byte type
func readStringTag(dst *string, tag interface{}) {
	switch s := tag.(type) {
	case string:
		*dst = s
	case []byte:
		*dst = string(s)
	}
}

// readIntTag populates the &dst with the tag value if it's of any kind of integer type
func readIntTag(dst *int, tag interface{}) {
	switch n := tag.(type) {
	case int:
		*dst = n
	case int8:
		*dst = int(n)
	case int16:
		*dst = int(n)
	case int32:
		*dst = int(n)
	case int64:
		*dst = int(n)
	case uint:
		*dst = int(n)
	case uint8:
		*dst = int(n)
	case uint16:
		*dst = int(n)
	case uint32:
		*dst = int(n)
	case uint64:
		*dst = int(n)
	}
}

func collectTracerSpanLogs(span *spanS) map[uint64]map[string]interface{} {
	logs := make(map[uint64]map[string]interface{})
	for _, l := range span.Logs {
		if _, ok := logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)]; !ok {
			logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)] = make(map[string]interface{})
		}

		for _, f := range l.Fields {
			logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)][f.Key()] = f.Value()
		}
	}

	return logs
}
<commit_msg>Send custom tags for registered spans within `sdk.custom.tags` section<commit_after>// (c) Copyright IBM Corp. 2021
// (c) Copyright Instana Inc. 2017

package instana

import (
	"encoding/json"
	"time"

	"github.com/opentracing/opentracing-go/ext"
)

type typedSpanData interface {
	Type() RegisteredSpanType
	Kind() SpanKind
}

// SpanKind represents values of field `k` in OpenTracing span representation. It represents
// the direction of the call associated with a span.
type SpanKind uint8

// Valid span kinds
const (
	// The kind of a span associated with an inbound call, this must be the first span in the trace.
	EntrySpanKind SpanKind = iota + 1
	// The kind of a span associated with an outbound call, e.g. an HTTP client request, posting to a message bus, etc.
	ExitSpanKind
	// The default kind for a span that is associated with a call within the same service.
	IntermediateSpanKind
)

// String returns string representation of a span kind suitable for use as a value for `data.sdk.type`
// tag of an SDK span. By default all spans are intermediate unless they are explicitly set to be "entry" or "exit"
func (k SpanKind) String() string {
	switch k {
	case EntrySpanKind:
		return "entry"
	case ExitSpanKind:
		return "exit"
	default:
		return "intermediate"
	}
}

// Span represents the OpenTracing span document to be sent to the agent
type Span struct {
	TraceID         int64
	TraceIDHi       int64
	ParentID        int64
	SpanID          int64
	Ancestor        *TraceReference
	Timestamp       uint64
	Duration        uint64
	Name            string
	From            *fromS
	Batch           *batchInfo
	Kind            int
	Ec              int
	Data            typedSpanData
	Synthetic       bool
	CorrelationType string
	CorrelationID   string
	ForeignTrace    bool
}

func newSpan(span *spanS) Span {
	data := RegisteredSpanType(span.Operation).ExtractData(span)
	sp := Span{
		TraceID:         span.context.TraceID,
		TraceIDHi:       span.context.TraceIDHi,
		ParentID:        span.context.ParentID,
		SpanID:          span.context.SpanID,
		Timestamp:       uint64(span.Start.UnixNano()) / uint64(time.Millisecond),
		Duration:        uint64(span.Duration) / uint64(time.Millisecond),
		Name:            string(data.Type()),
		Ec:              span.ErrorCount,
		CorrelationType: span.Correlation.Type,
		CorrelationID:   span.Correlation.ID,
		ForeignTrace:    span.context.ForeignTrace,
		Kind:            int(data.Kind()),
		Data:            data,
	}

	if bs, ok := span.Tags[batchSizeTag].(int); ok {
		if bs > 1 {
			sp.Batch = &batchInfo{Size: bs}
		}
		delete(span.Tags, batchSizeTag)
	}

	if syn, ok := span.Tags[syntheticCallTag].(bool); ok {
		sp.Synthetic = syn
		delete(span.Tags, syntheticCallTag)
	}

	if len(span.context.Links) > 0 {
		ancestor := span.context.Links[0]
		sp.Ancestor = &TraceReference{
			TraceID:  ancestor.TraceID,
			ParentID: ancestor.SpanID,
		}
	}

	return sp
}

type TraceReference struct {
	TraceID  string `json:"t"`
	ParentID string `json:"p,omitempty"`
}

// MarshalJSON serializes span to JSON for sending it to Instana
func (sp Span) MarshalJSON() ([]byte, error) {
	var parentID string
	if sp.ParentID != 0 {
		parentID = FormatID(sp.ParentID)
	}

	var longTraceID string
	if sp.TraceIDHi != 0 && sp.Kind == int(EntrySpanKind) {
		longTraceID = FormatLongID(sp.TraceIDHi, sp.TraceID)
	}

	return json.Marshal(struct {
		TraceReference

		SpanID          string          `json:"s"`
		LongTraceID     string          `json:"lt,omitempty"`
		Timestamp       uint64          `json:"ts"`
		Duration        uint64          `json:"d"`
		Name            string          `json:"n"`
		From            *fromS          `json:"f"`
		Batch           *batchInfo      `json:"b,omitempty"`
		Kind            int             `json:"k"`
		Ec              int             `json:"ec,omitempty"`
		Data            typedSpanData   `json:"data"`
		Synthetic       bool            `json:"sy,omitempty"`
		CorrelationType string          `json:"crtp,omitempty"`
		CorrelationID   string          `json:"crid,omitempty"`
		ForeignTrace    bool            `json:"tp,omitempty"`
		Ancestor        *TraceReference `json:"ia,omitempty"`
	}{
		TraceReference{
			FormatID(sp.TraceID),
			parentID,
		},
		FormatID(sp.SpanID),
		longTraceID,
		sp.Timestamp,
		sp.Duration,
		sp.Name,
		sp.From,
		sp.Batch,
		sp.Kind,
		sp.Ec,
		sp.Data,
		sp.Synthetic,
		sp.CorrelationType,
		sp.CorrelationID,
		sp.ForeignTrace,
		sp.Ancestor,
	})
}

type batchInfo struct {
	Size int `json:"s"`
}

// CustomSpanData holds user-defined span tags
type CustomSpanData struct {
	Tags map[string]interface{} `json:"tags,omitempty"`
}

func filterCustomSpanTags(tags map[string]interface{}, st RegisteredSpanType) map[string]interface{} {
	knownTags := st.TagsNames()
	customTags := make(map[string]interface{})

	for k, v := range tags {
		if k == string(ext.SpanKind) {
			continue
		}

		if _, ok := knownTags[k]; ok {
			continue
		}

		customTags[k] = v
	}

	return customTags
}

// SpanData contains fields to be sent in the `data` section of an OT span document. These fields are
// common for all span types.
type SpanData struct {
	Service string          `json:"service,omitempty"`
	Custom  *CustomSpanData `json:"sdk.custom,omitempty"`

	st RegisteredSpanType
}

// NewSpanData initializes a new span data from tracer span
func NewSpanData(span *spanS, st RegisteredSpanType) SpanData {
	data := SpanData{
		Service: span.Service,
		st:      st,
	}

	if customTags := filterCustomSpanTags(span.Tags, st); len(customTags) > 0 {
		data.Custom = &CustomSpanData{Tags: customTags}
	}

	return data
}

// Type returns the registered span type suitable for use as the value of `n` field.
func (d SpanData) Type() RegisteredSpanType {
	return d.st
}

// SDKSpanData represents the `data` section of an SDK span sent within an OT span document
type SDKSpanData struct {
	// Deprecated
	SpanData `json:"-"`

	Service string      `json:"service,omitempty"`
	Tags    SDKSpanTags `json:"sdk"`

	sk SpanKind
}

// NewSDKSpanData initializes a new SDK span data from a tracer span
func NewSDKSpanData(span *spanS) SDKSpanData {
	sk := IntermediateSpanKind

	switch span.Tags[string(ext.SpanKind)] {
	case ext.SpanKindRPCServerEnum, string(ext.SpanKindRPCServerEnum),
		ext.SpanKindConsumerEnum, string(ext.SpanKindConsumerEnum),
		"entry":
		sk = EntrySpanKind
	case ext.SpanKindRPCClientEnum, string(ext.SpanKindRPCClientEnum),
		ext.SpanKindProducerEnum, string(ext.SpanKindProducerEnum),
		"exit":
		sk = ExitSpanKind
	}

	return SDKSpanData{
		Service: span.Service,
		Tags:    NewSDKSpanTags(span, sk.String()),
		sk:      sk,
	}
}

// Type returns the registered span type suitable for use as the value of `n` field.
func (d SDKSpanData) Type() RegisteredSpanType {
	return SDKSpanType
}

// Kind returns the kind of the span. It handles the github.com/opentracing/opentracing-go/ext.SpanKindEnum
// values as well as generic "entry" and "exit"
func (d SDKSpanData) Kind() SpanKind {
	return d.sk
}

// KnownTags returns the list of known tags for this span type
// SDKSpanTags contains fields within the `data.sdk` section of an OT span document
type SDKSpanTags struct {
	Name      string                 `json:"name"`
	Type      string                 `json:"type,omitempty"`
	Arguments string                 `json:"arguments,omitempty"`
	Return    string                 `json:"return,omitempty"`
	Custom    map[string]interface{} `json:"custom,omitempty"`
}

// NewSDKSpanTags extracts SDK span tags from a tracer span
func NewSDKSpanTags(span *spanS, spanType string) SDKSpanTags {
	tags := SDKSpanTags{
		Name:   span.Operation,
		Type:   spanType,
		Custom: map[string]interface{}{},
	}

	if len(span.Tags) != 0 {
		tags.Custom["tags"] = span.Tags
	}

	if logs := collectTracerSpanLogs(span); len(logs) > 0 {
		tags.Custom["logs"] = logs
	}

	if len(span.context.Baggage) != 0 {
		tags.Custom["baggage"] = span.context.Baggage
	}

	return tags
}

// readStringTag populates the &dst with the tag value if it's of either string or []byte type
func readStringTag(dst *string, tag interface{}) {
	switch s := tag.(type) {
	case string:
		*dst = s
	case []byte:
		*dst = string(s)
	}
}

// readIntTag populates the &dst with the tag value if it's of any kind of integer type
func readIntTag(dst *int, tag interface{}) {
	switch n := tag.(type) {
	case int:
		*dst = n
	case int8:
		*dst = int(n)
	case int16:
		*dst = int(n)
	case int32:
		*dst = int(n)
	case int64:
		*dst = int(n)
	case uint:
		*dst = int(n)
	case uint8:
		*dst = int(n)
	case uint16:
		*dst = int(n)
	case uint32:
		*dst = int(n)
	case uint64:
		*dst = int(n)
	}
}

func collectTracerSpanLogs(span *spanS) map[uint64]map[string]interface{} {
	logs := make(map[uint64]map[string]interface{})
	for _, l := range span.Logs {
		if _, ok := logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)]; !ok {
			logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)] = make(map[string]interface{})
		}

		for _, f := range l.Fields {
			logs[uint64(l.Timestamp.UnixNano())/uint64(time.Millisecond)][f.Key()] = f.Value()
		}
	}

	return logs
}
<|endoftext|>


<commit_before>package git

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	ini "github.com/vaughan0/go-ini"
)

func PrepSubmodules(
	gitDir, checkoutDir, mainRev string,
) error {

	gitModules := filepath.Join(checkoutDir, ".gitmodules")

	submodules, err := ParseSubmodules(gitModules)
	if err != nil {
		if os.IsNotExist(err) {
			// No .gitmodules available.
			return nil
		}
		return err
	}

	log.Printf("Prep %v submodules", len(submodules))

	GetSubmoduleRevs(gitDir, mainRev, submodules)

	errs := make(chan error, len(submodules))

	go func() {
		defer close(errs)

		var wg sync.WaitGroup
		defer wg.Wait()

		// Run only NumCPU in parallel
		semaphore := make(chan struct{}, runtime.NumCPU())

		for _, submodule := range submodules {

			wg.Add(1)
			go func(submodule Submodule) {
				defer wg.Done()
				defer func() { <-semaphore }()
				semaphore <- struct{}{}

				err := prepSubmodule(gitDir, checkoutDir, submodule)
				if err != nil {
					err = fmt.Errorf("processing %v: %v", submodule.Path, err)
				}
				errs <- err
			}(submodule)
		}
	}()

	// errs chan has buffer length len(submodules)
	err = MultipleErrors(errs)
	if err != nil {
		return err
	}
	return nil
}

type ErrMultiple struct {
	errs []error
}

func (em *ErrMultiple) Error() string {
	var s []string
	for _, e := range em.errs {
		s = append(s, e.Error())
	}
	return fmt.Sprint("multiple errors:\n", strings.Join(s, "\n"))
}

// Read errors out of a channel, counting only the non-nil ones.
// If there are zero non-nil errs, nil is returned.
func MultipleErrors(errs <-chan error) error {
	var em ErrMultiple
	for e := range errs {
		if e == nil {
			continue
		}
		em.errs = append(em.errs, e)
	}
	if len(em.errs) == 0 {
		return nil
	}
	return &em
}

// Checkout the working directory of a given submodule.
func prepSubmodule(
	mainGitDir, mainCheckoutDir string,
	submodule Submodule,
) error {

	subGitDir := filepath.Join(mainGitDir, "modules", submodule.Path)

	err := LocalMirror(submodule.URL, subGitDir, submodule.Rev, os.Stderr)
	if err != nil {
		return err
	}

	subCheckoutPath := filepath.Join(mainCheckoutDir, submodule.Path)

	// Note: checkout may recurse onto prepSubmodules.
	err = recursiveCheckout(subGitDir, subCheckoutPath, submodule.Rev)
	if err != nil {
		return err
	}
	return err
}

type Submodule struct {
	Path, URL string
	Rev       string // populated by GetSubmoduleRevs
}

func ParseSubmodules(filename string) (submodules []Submodule, err error) {
	config, err := ini.LoadFile(filename)
	if err != nil {
		return nil, err
	}

	for section := range config {
		if !strings.HasPrefix(section, "submodule") {
			continue
		}
		submodules = append(submodules, Submodule{
			Path: config.Section(section)["path"],
			URL:  config.Section(section)["url"],
		})
	}
	return submodules, nil
}

func GetSubmoduleRevs(gitDir, mainRev string, submodules []Submodule) error {
	for i := range submodules {
		rev, err := GetSubmoduleRev(gitDir, submodules[i].Path, mainRev)
		if err != nil {
			return err
		}
		submodules[i].Rev = rev
	}
	return nil
}

func GetSubmoduleRev(gitDir, submodulePath, mainRev string) (string, error) {
	cmd := Command(gitDir, "git", "ls-tree", mainRev, "--", submodulePath)

	parts, err := cmd.Output()
	if err != nil {
		return "", err
	}

	return strings.Fields(string(parts))[2], nil
}
<commit_msg>Revert "Clean up"<commit_after>package git

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	ini "github.com/vaughan0/go-ini"
)

func PrepSubmodules(
	gitDir, checkoutDir, mainRev string,
) error {

	gitModules := filepath.Join(checkoutDir, ".gitmodules")

	submodules, err := ParseSubmodules(gitModules)
	if err != nil {
		if os.IsNotExist(err) {
			// No .gitmodules available.
			return nil
		}
		return err
	}

	log.Printf("Prep %v submodules", len(submodules))

	GetSubmoduleRevs(gitDir, mainRev, submodules)

	errs := make(chan error, len(submodules))

	go func() {
		defer close(errs)

		var wg sync.WaitGroup
		defer wg.Wait()

		// Run only NumCPU in parallel
		semaphore := make(chan struct{}, runtime.NumCPU())

		for _, submodule := range submodules {

			wg.Add(1)
			go func(submodule Submodule) {
				defer wg.Done()
				defer func() { <-semaphore }()
				semaphore <- struct{}{}

				err := prepSubmodule(gitDir, checkoutDir, submodule)
				if err != nil {
					err = fmt.Errorf("processing %v: %v", submodule.Path, err)
				}
				errs <- err
			}(submodule)
		}
	}()

	// errs chan has buffer length len(submodules)
	err = MultipleErrors(errs)
	if err != nil {
		return err
	}
	return nil
}

type ErrMultiple struct {
	errs []error
}

func (em *ErrMultiple) Error() string {
	var s []string
	for _, e := range em.errs {
		s = append(s, e.Error())
	}
	return fmt.Sprint("multiple errors:\n", strings.Join(s, "\n"))
}

// Read errors out of a channel, counting only the non-nil ones.
// If there are zero non-nil errs, nil is returned.
func MultipleErrors(errs <-chan error) error {
	var em ErrMultiple
	for e := range errs {
		if e == nil {
			continue
		}
		em.errs = append(em.errs, e)
	}
	if len(em.errs) == 0 {
		return nil
	}
	return &em
}

// Checkout the working directory of a given submodule.
func prepSubmodule(
	mainGitDir, mainCheckoutDir string,
	submodule Submodule,
) error {

	subGitDir := filepath.Join(mainGitDir, "modules", submodule.Path)

	err := LocalMirror(submodule.URL, subGitDir, submodule.Rev, os.Stderr)
	if err != nil {
		return err
	}

	subCheckoutPath := filepath.Join(mainCheckoutDir, submodule.Path)

	// Note: checkout may recurse onto prepSubmodules.
	err = recursiveCheckout(subGitDir, subCheckoutPath, submodule.Rev)
	if err != nil {
		return err
	}
	return err
}

type Submodule struct {
	Path, URL string
	Rev       string // populated by GetSubmoduleRevs
}

func ParseSubmodules(filename string) (submodules []Submodule, err error) {
	config, err := ini.LoadFile(filename)
	if err != nil {
		return nil, err
	}

	for section := range config {
		if !strings.HasPrefix(section, "submodule") {
			continue
		}
		submodules = append(submodules, Submodule{
			Path: config.Section(section)["path"],
			URL:  config.Section(section)["url"],
		})
	}
	return submodules, nil
}

func GetSubmoduleRevs(gitDir, mainRev string, submodules []Submodule) error {
	for i := range submodules {
		rev, err := GetSubmoduleRev(gitDir, submodules[i].Path, mainRev)
		if err != nil {
			return err
		}
		submodules[i].Rev = rev
	}
	return nil
}

func GetSubmoduleRev(gitDir, submodulePath, mainRev string) (string, error) {
	cmd := Command(gitDir, "git", "ls-tree", mainRev, "--", submodulePath)
	cmd.Stdout = nil

	parts, err := cmd.Output()
	if err != nil {
		return "", err
	}

	return strings.Fields(string(parts))[2], nil
}
<|endoftext|>


<commit_before>package lidar_test

import (
	"context"
	"errors"
	"time"

	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/dbfakes"
	"github.com/concourse/concourse/atc/lidar"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type Scanner interface {
	Run(ctx context.Context) error
}

var _ = Describe("Scanner", func() {
	var (
		err error

		fakeCheckFactory *dbfakes.FakeCheckFactory
		planFactory      atc.PlanFactory

		scanner Scanner
	)

	BeforeEach(func() {
		planFactory = atc.NewPlanFactory(0)
		fakeCheckFactory = new(dbfakes.FakeCheckFactory)

		scanner = lidar.NewScanner(fakeCheckFactory, planFactory)
	})

	JustBeforeEach(func() {
		err = scanner.Run(context.TODO())
	})

	Describe("Run", func() {
		Context("when fetching resources fails", func() {
			BeforeEach(func() {
				fakeCheckFactory.ResourcesReturns(nil, errors.New("nope"))
			})

			It("errors", func() {
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when fetching resources succeeds", func() {
			var fakeResource *dbfakes.FakeResource

			BeforeEach(func() {
				fakeResource = new(dbfakes.FakeResource)
				fakeResource.NameReturns("some-name")
				fakeResource.TagsReturns([]string{"tag-a", "tag-b"})
				fakeResource.SourceReturns(atc.Source{"some": "source"})

				fakeCheckFactory.ResourcesReturns([]db.Resource{fakeResource}, nil)
			})

			Context("when fetching resource types fails", func() {
				BeforeEach(func() {
					fakeCheckFactory.ResourceTypesByPipelineReturns(nil, errors.New("nope"))
				})

				It("errors", func() {
					Expect(err).To(HaveOccurred())
				})
			})

			Context("when CheckEvery is never", func() {
				BeforeEach(func() {
					fakeResource.CheckEveryReturns(&atc.CheckEvery{Never: true})
					fakeResource.TypeReturns("parent")
					fakeResource.PipelineIDReturns(1)
					fakeResourceType := new(dbfakes.FakeResourceType)
					fakeResourceType.NameReturns("parent")
					fakeResourceType.PipelineIDReturns(1)
					fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{
						1: {fakeResourceType},
					}, nil)
				})

				It("does not check the resource", func() {
					Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(0))
				})
			})

			Context("when fetching resources types succeeds", func() {
				var fakeResourceType *dbfakes.FakeResourceType

				BeforeEach(func() {
					fakeResourceType = new(dbfakes.FakeResourceType)
					fakeResourceType.NameReturns("some-type")
					fakeResourceType.TypeReturns("some-base-type")
					fakeResourceType.TagsReturns([]string{"some-tag"})
					fakeResourceType.SourceReturns(atc.Source{"some": "type-source"})

					fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{1: {fakeResourceType}}, nil)
				})

				Context("when the resource parent type is a base type", func() {
					BeforeEach(func() {
						fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{}, nil)
						fakeResource.TypeReturns("some-type")
					})

					It("creates a check with empty resource types list", func() {
						_, _, resourceTypes, _, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
						var nilResourceTypes db.ResourceTypes
						Expect(resourceTypes).To(Equal(nilResourceTypes))
					})

					Context("when the last check end time is past our interval", func() {
						It("creates a check", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
						})

						Context("when try creating a check panics", func() {
							BeforeEach(func() {
								fakeCheckFactory.TryCreateCheckStub = func(context.Context, db.Checkable, db.ResourceTypes, atc.Version, bool, bool) (db.Build, bool, error) {
									panic("something went wrong")
								}
							})

							It("recovers from the panic", func() {
								Expect(err).ToNot(HaveOccurred())
							})
						})
					})

					Context("when the checkable has a pinned version", func() {
						BeforeEach(func() {
							fakeResource.CurrentPinnedVersionReturns(atc.Version{"some": "version"})
						})

						It("creates a check with that pinned version", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
							_, _, _, fromVersion, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
							Expect(fromVersion).To(Equal(atc.Version{"some": "version"}))
						})
					})

					Context("when the checkable does not have a pinned version", func() {
						BeforeEach(func() {
							fakeResource.CurrentPinnedVersionReturns(nil)
						})

						It("creates a check with a nil pinned version", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
							_, _, _, fromVersion, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
							Expect(fromVersion).To(BeNil())
						})
					})
				})

				Context("when there's a put-only resource", func() {
					BeforeEach(func() {
						By("checkFactory.Resources should not return any put-only resources")
						fakeResourceType.NameReturns("put-only-custom-type")
						fakeResourceType.PipelineIDReturns(1)
					})

					It("does not check the put-only resource", func() {
						Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1),
							"one check created for the unrelated fakeResource")
					})
				})
			})
		})
	})
})
<commit_msg>fix existing unit tests under atc/lidar.<commit_after>package lidar_test

import (
	"context"
	"errors"
	"github.com/concourse/concourse/atc/component"

	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/db/dbfakes"
	"github.com/concourse/concourse/atc/lidar"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

type Scanner interface {
	Run(ctx context.Context, lastRunResult string) (component.RunResult, error)
}

var _ = Describe("Scanner", func() {
	var (
		err error

		fakeCheckFactory *dbfakes.FakeCheckFactory
		planFactory      atc.PlanFactory

		scanner Scanner
	)

	BeforeEach(func() {
		planFactory = atc.NewPlanFactory(0)
		fakeCheckFactory = new(dbfakes.FakeCheckFactory)

		scanner = lidar.NewScanner(fakeCheckFactory, planFactory, 1)
	})

	JustBeforeEach(func() {
		_, err = scanner.Run(context.TODO(), "")
	})

	Describe("Run", func() {
		Context("when fetching resources fails", func() {
			BeforeEach(func() {
				fakeCheckFactory.ResourcesReturns(nil, errors.New("nope"))
			})

			It("errors", func() {
				Expect(err).To(HaveOccurred())
			})
		})

		Context("when fetching resources succeeds", func() {
			var fakeResource *dbfakes.FakeResource

			BeforeEach(func() {
				fakeResource = new(dbfakes.FakeResource)
				fakeResource.NameReturns("some-name")
				fakeResource.TagsReturns([]string{"tag-a", "tag-b"})
				fakeResource.SourceReturns(atc.Source{"some": "source"})

				fakeCheckFactory.ResourcesReturns([]db.Resource{fakeResource}, nil)
			})

			Context("when fetching resource types fails", func() {
				BeforeEach(func() {
					fakeCheckFactory.ResourceTypesByPipelineReturns(nil, errors.New("nope"))
				})

				It("errors", func() {
					Expect(err).To(HaveOccurred())
				})
			})

			Context("when CheckEvery is never", func() {
				BeforeEach(func() {
					fakeResource.CheckEveryReturns(&atc.CheckEvery{Never: true})
					fakeResource.TypeReturns("parent")
					fakeResource.PipelineIDReturns(1)
					fakeResourceType := new(dbfakes.FakeResourceType)
					fakeResourceType.NameReturns("parent")
					fakeResourceType.PipelineIDReturns(1)
					fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{
						1: {fakeResourceType},
					}, nil)
				})

				It("does not check the resource", func() {
					Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(0))
				})
			})

			Context("when fetching resources types succeeds", func() {
				var fakeResourceType *dbfakes.FakeResourceType

				BeforeEach(func() {
					fakeResourceType = new(dbfakes.FakeResourceType)
					fakeResourceType.NameReturns("some-type")
					fakeResourceType.TypeReturns("some-base-type")
					fakeResourceType.TagsReturns([]string{"some-tag"})
					fakeResourceType.SourceReturns(atc.Source{"some": "type-source"})

					fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{1: {fakeResourceType}}, nil)
				})

				Context("when the resource parent type is a base type", func() {
					BeforeEach(func() {
						fakeCheckFactory.ResourceTypesByPipelineReturns(map[int]db.ResourceTypes{}, nil)
						fakeResource.TypeReturns("some-type")
					})

					It("creates a check with empty resource types list", func() {
						_, _, resourceTypes, _, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
						var nilResourceTypes db.ResourceTypes
						Expect(resourceTypes).To(Equal(nilResourceTypes))
					})

					Context("when the last check end time is past our interval", func() {
						It("creates a check", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
						})

						Context("when try creating a check panics", func() {
							BeforeEach(func() {
								fakeCheckFactory.TryCreateCheckStub = func(context.Context, db.Checkable, db.ResourceTypes, atc.Version, bool, bool) (db.Build, bool, error) {
									panic("something went wrong")
								}
							})

							It("recovers from the panic", func() {
								Expect(err).ToNot(HaveOccurred())
							})
						})
					})

					Context("when the checkable has a pinned version", func() {
						BeforeEach(func() {
							fakeResource.CurrentPinnedVersionReturns(atc.Version{"some": "version"})
						})

						It("creates a check with that pinned version", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
							_, _, _, fromVersion, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
							Expect(fromVersion).To(Equal(atc.Version{"some": "version"}))
						})
					})

					Context("when the checkable does not have a pinned version", func() {
						BeforeEach(func() {
							fakeResource.CurrentPinnedVersionReturns(nil)
						})

						It("creates a check with a nil pinned version", func() {
							Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1))
							_, _, _, fromVersion, _, _ := fakeCheckFactory.TryCreateCheckArgsForCall(0)
							Expect(fromVersion).To(BeNil())
						})
					})
				})

				Context("when there's a put-only resource", func() {
					BeforeEach(func() {
						By("checkFactory.Resources should not return any put-only resources")
						fakeResourceType.NameReturns("put-only-custom-type")
						fakeResourceType.PipelineIDReturns(1)
					})

					It("does not check the put-only resource", func() {
						Expect(fakeCheckFactory.TryCreateCheckCallCount()).To(Equal(1),
							"one check created for the unrelated fakeResource")
					})
				})
			})
		})
	})
})
<|endoftext|>


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

package plugin

import (
	"istio.io/istio/pilot/pkg/model"
	istionetworking "istio.io/istio/pilot/pkg/networking"
)

const (
	// AuthzCustom is the name of the authorization plugin (CUSTOM action) passed through the command line
	AuthzCustom = "ext_authz"
	// Authn is the name of the authentication plugin passed through the command line
	Authn = "authn"
	// Authz is the name of the authorization plugin (ALLOW/DENY/AUDIT action) passed through the command line
	Authz = "authz"
	// Health is the name of the health plugin passed through the command line
	Health = "health"
	// MetadataExchange is the name of the telemetry plugin passed through the command line
	MetadataExchange = "metadata_exchange"
)

// InputParams is a set of values passed to Plugin callback methods. Not all fields are guaranteed to
// be set, it's up to the callee to validate required fields are set and emit error if they are not.
// These are for reading only and should not be modified.
type InputParams struct {
	// ListenerProtocol is the protocol/class of listener (TCP, HTTP etc.). Must be set.
	// This is valid only for the inbound listener
	// Outbound listeners could have multiple filter chains, where one filter chain could be
	// a HTTP connection manager with TLS context, while the other could be a tcp proxy with sni
	ListenerProtocol istionetworking.ListenerProtocol
	// Node is the node the response is for.
	Node *model.Proxy
	// ServiceInstance is the service instance colocated with the listener (applies to sidecar).
	ServiceInstance *model.ServiceInstance
	// Service is the service colocated with the listener (applies to sidecar).
	// For outbound TCP listeners, it is the destination service.
	// Push holds stats and other information about the current push.
	Push *model.PushContext
}

// Plugin is called during the construction of a listener.Listener which may alter the Listener in any
// way. Examples include AuthenticationPlugin that sets up mTLS authentication on the inbound Listener
// and outbound Cluster, etc.
type Plugin interface {
	// OnOutboundListener is called whenever a new outbound listener is added to the LDS output for a given service.
	// Can be used to add additional filters on the outbound path.
	OnOutboundListener(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundListener is called whenever a new listener is added to the LDS output for a given service
	// Can be used to add additional filters.
	OnInboundListener(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundFilterChains is called whenever a plugin needs to setup the filter chains, including relevant filter chain
	// configuration, like FilterChainMatch and TLSContext.
	OnInboundFilterChains(in *InputParams) []istionetworking.FilterChain

	// OnInboundPassthrough is called whenever a new passthrough filter chain is added to the LDS output.
	// Can be used to add additional filters.
	OnInboundPassthrough(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundPassthroughFilterChains is called whenever a plugin needs to setup custom pass through filter chain.
	OnInboundPassthroughFilterChains(in *InputParams) []istionetworking.FilterChain
}
<commit_msg>Remove useless comment (#29664)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package plugin

import (
	"istio.io/istio/pilot/pkg/model"
	istionetworking "istio.io/istio/pilot/pkg/networking"
)

const (
	// AuthzCustom is the name of the authorization plugin (CUSTOM action) passed through the command line
	AuthzCustom = "ext_authz"
	// Authn is the name of the authentication plugin passed through the command line
	Authn = "authn"
	// Authz is the name of the authorization plugin (ALLOW/DENY/AUDIT action) passed through the command line
	Authz = "authz"
	// Health is the name of the health plugin passed through the command line
	Health = "health"
	// MetadataExchange is the name of the telemetry plugin passed through the command line
	MetadataExchange = "metadata_exchange"
)

// InputParams is a set of values passed to Plugin callback methods. Not all fields are guaranteed to
// be set, it's up to the callee to validate required fields are set and emit error if they are not.
// These are for reading only and should not be modified.
type InputParams struct {
	// ListenerProtocol is the protocol/class of listener (TCP, HTTP etc.). Must be set.
	// This is valid only for the inbound listener
	// Outbound listeners could have multiple filter chains, where one filter chain could be
	// a HTTP connection manager with TLS context, while the other could be a tcp proxy with sni
	ListenerProtocol istionetworking.ListenerProtocol
	// Node is the node the response is for.
	Node *model.Proxy
	// ServiceInstance is the service instance colocated with the listener (applies to sidecar).
	ServiceInstance *model.ServiceInstance
	// Push holds stats and other information about the current push.
	Push *model.PushContext
}

// Plugin is called during the construction of a listener.Listener which may alter the Listener in any
// way. Examples include AuthenticationPlugin that sets up mTLS authentication on the inbound Listener
// and outbound Cluster, etc.
type Plugin interface {
	// OnOutboundListener is called whenever a new outbound listener is added to the LDS output for a given service.
	// Can be used to add additional filters on the outbound path.
	OnOutboundListener(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundListener is called whenever a new listener is added to the LDS output for a given service
	// Can be used to add additional filters.
	OnInboundListener(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundFilterChains is called whenever a plugin needs to setup the filter chains, including relevant filter chain
	// configuration, like FilterChainMatch and TLSContext.
	OnInboundFilterChains(in *InputParams) []istionetworking.FilterChain

	// OnInboundPassthrough is called whenever a new passthrough filter chain is added to the LDS output.
	// Can be used to add additional filters.
	OnInboundPassthrough(in *InputParams, mutable *istionetworking.MutableObjects) error

	// OnInboundPassthroughFilterChains is called whenever a plugin needs to setup custom pass through filter chain.
	OnInboundPassthroughFilterChains(in *InputParams) []istionetworking.FilterChain
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"github.com/jetblack87/maestro/data"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"path"
	"github.com/samuel/go-zookeeper/zk"
	"time"
	"log"
)

const APP_VERSION = "0.1"

// The flag package provides a default help printer via -h switch
var versionFlag *bool = flag.Bool("v", false, "Print the version number.")
var zookeeper *string = flag.String("zookeeper", "localhost:2181", "The ZooKeeper connection string (defaults to 'localhost:2182').")
var agentName *string = flag.String("name", "", "REQUIRED: The name of the agent.")
var domainName *string = flag.String("domain", "", "REQUIRED: The name of the domain in which this agent lives.")
var agentConfig *string = flag.String("agentConfig", "", "Supply a json file that contains specific configuration for this agent.")
var processesConfig *string = flag.String("processesConfig", "", "Supply a json file that contains specific configuration any processes.")
var logfilePath *string = flag.String("logfile", "stdout", "The path to the logfile.")

var zkdao data.ZkDAO
var request *processStartRequest

const MAX_START_RETRIES = 3

func main() {

	flag.Parse() // Scan the arguments list

	// Setup signal channel
	signalChannel := make(chan os.Signal, 1)
	signal.Notify(signalChannel, os.Interrupt, os.Kill)

    // Setup logging
	if *logfilePath != "stdout" {
		f, err := os.OpenFile(*logfilePath, os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
		if err != nil {
		    log.Printf("error opening file: %v", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	// Check the parameters
	if *versionFlag {
		log.Println("Version:", APP_VERSION)
		os.Exit(0)
	}
	if *agentName == "" {
		panic("-agent is required")
	}
	if *domainName == "" {
		panic("-domain is required")
	}

	log.Printf("maestro agent starting for domain '%s' and agent '%s'\n", *domainName, *domainName)

	zkdao, err := data.NewZkDAO(strings.Split(*zookeeper, ","))
	if err != nil {
		panic(err)
	}

	// Load config into ZK
	err = loadAgentConfig(*agentConfig, *agentName, *domainName)
	if err != nil {
		panic(err)
	}

	// Load processes data into ZK
	err = loadProcessesConfig(*processesConfig, *domainName)
	if err != nil {
		panic(err)
	}

	log.Println("Loading the agent configuration")
	agent, err := zkdao.LoadAgent(data.PathToKey("/maestro/"+*domainName+"/config/agents/"+*agentName), true)
	if err != nil {
		panic(err)
	}

	// Create channel used for watching ZK nodes
	watchChannel := make(chan zk.Event, 1)

	// Remove old runtime config for this agent
	err = zkdao.RemoveRecursive("/maestro/"+*domainName+"/runtime/agents/"+agent.Name)

	// Add processes to the runtime configuration, adding watches to admin_state
	for key := range agent.Processes {
		log.Println("Loading processes from config: " + agent.Processes[key].ProcessClass)
		agent.Processes[key], err = zkdao.LoadProcess(data.PathToKey(agent.Processes[key].ProcessClass), true)
		if err != nil {
			panic(err)
		}
		agent.Processes[key].Key = data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name+"/processes/"+agent.Processes[key].Name)
		if agent.Processes[key].AdminState == "" {
			// Default to on
			log.Println("Defaulting admin_state to 'on'")
			agent.Processes[key].AdminState = "on"
		}
		agent.Processes[key].OperState = "off"
	}

	log.Println("Adding agent to runtime configuration")
	if err != nil {
		log.Printf("Failed to remove agent runtime configuration")
		panic(err)
	}
	err = zkdao.UpdateAgent(data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name), agent, true)
	if err != nil {
		panic(err)
	}

	str, err := zkdao.CreateEphemeral("/maestro/"+*domainName+"/runtime/agents/"+agent.Name+"/eph", []byte("I am alive"))
	if err != nil {
		errMsg := "Error creating ephemeral node: " + str
		log.Println(errMsg)
		panic(err)
	}

	// Add watches for all admin_state nodes
	for key := range agent.Processes {
		adminStatePath := data.KeyToPath(agent.Processes[key].Key)+"/admin_state"
		log.Println("Adding watch to node: " + adminStatePath)
		err := zkdao.Watch(adminStatePath, watchChannel)
		if err != nil {
			log.Println("Failed to add watch to process node:\n" + err.Error())
		}
	}

	log.Println("Starting process monitoring")

	// Create out request (including channels)
	request = &processStartRequest{
		processes : agent.Processes,
		commandChan : make(chan *command, 1),
		resultChan : make(chan *result, 1)}

	go startAndMonitorProcesses(request)

	log.Println("Process monitoring started, waiting on channels")
	for {
		select {
			case w := <-watchChannel:
			if w.Type.String() == "EventNodeDataChanged" {
				adminState, err := zkdao.GetValue(w.Path)
				if err != nil {
					log.Printf("Error getting data for path '%s': %s\n", w.Path, err.Error())
				} else {
					process, err2 := zkdao.LoadProcess(data.PathToKey(path.Dir(w.Path)), true)
					if err2 != nil {
						log.Printf("Error loading process '%s': %s\n", w.Path, err2)
					} else {
						request.commandChan <- &command{process : process, adminState : string(adminState)}
					}
				}
			}
			case r := <-request.resultChan:
			if r.err != nil {
				log.Printf("An error occured running a process:%s\n", r.err.Error())
				// Failed to start, turn off
				r.process.OperState = "off"
				r.process.AdminState = "off"
				zkdao.UpdateProcess(r.process.Key, r.process, false)
			} else {
				var p data.Process
	    		// Update the admin_state and pid in ZK
				if r.process.Key != "" {
					p = r.process
				} else {
					p,err = zkdao.LoadProcess(r.key, false)
					if err != nil {
						log.Printf("Failed to load process: %s\n", err.Error())
					} else {
						p.OperState = r.operState
					}
				}
				log.Printf("Process '%s' oper_state = '%s'\n", p.Key, p.OperState)
	    		zkdao.UpdateProcess(r.process.Key, r.process, false)

	    		// Touch the admin_state node to get process turned back on
	    		if p.AdminState == "on" && p.OperState == "off" {
	    			zkdao.SetValue(data.KeyToPath(p.Key) + "/admin_state", []byte(p.AdminState))
	    		}
			}
			case <-signalChannel: // FIXME this doesn't seem to work (at least not on Windows)
			log.Println("Received signal")
			a,err := zkdao.LoadAgent(data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name),true)
			if err != nil {
				log.Println("Error retrieving agent")
			} else {
				for _, process := range a.Processes {
					request.commandChan <- &command{process : process, adminState : "off"}
				}
			}
			os.Exit(0)
		}
	}
}

func loadAgentConfig(agentConfig, agentName, domainName string) error {
	// Load the file if it was supplied
	if agentConfig != "" {
		log.Println("Loading agent config: " + agentConfig)
		jsonData, err := ioutil.ReadFile(agentConfig)
		if err != nil {
			return err
		}
		var agent data.Agent
		err = json.Unmarshal(jsonData, &agent)
		if err != nil {
			return err
		}
		err = zkdao.UpdateAgent(data.PathToKey("/maestro/"+domainName+"/config/agents/"+agentName), agent, true)
		if err != nil {
			return err
		}
	}
	return nil
}

func loadProcessesConfig(processesConfig, domainName string) error {
	// Load the file if it was supplied
	if processesConfig != "" {
		log.Println("Loading processes config: " + processesConfig)
		jsonData, err := ioutil.ReadFile(processesConfig)
		if err != nil {
			return err
		}
		var processes []data.Process
		err = json.Unmarshal(jsonData, &processes)
		if err != nil {
			return err
		}
		for _, process := range processes {
			err = zkdao.UpdateProcess(data.PathToKey("/maestro/"+domainName+"/config/processes/"+process.Name), process, true)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func startAndMonitorProcesses (startRequest *processStartRequest) {
	// Mapping of process key (string) to command
	processMap := make(map[string]*exec.Cmd)
	// Start all of the processes
	for key := range startRequest.processes {
		if startRequest.processes[key].AdminState == "on" {
        	log.Println("Starting process: " + startRequest.processes[key].Key)
			cmd, err := startProcess(startRequest.processes[key])
			if err != nil {
				log.Printf("Error starting process:\n%s\n", err.Error())
				startRequest.resultChan <- &result{process : startRequest.processes[key], err : err}
			} else {
				processMap[startRequest.processes[key].Key] = cmd
				// Send the result back
				startRequest.processes[key].OperState = "on"
				startRequest.processes[key].Pid = cmd.Process.Pid
				startRequest.resultChan <- &result{process : startRequest.processes[key]}
			}
		}
	}

	// Monitor the processes
	log.Println("Monitoring command channel and processes")
	for {
		select {
			case c := <-startRequest.commandChan:
		    	switch c.adminState {
		    	case "off":
		    		log.Println("Killing process: " + c.process.Key)
		    		if processMap[c.process.Key] != nil {
				        processMap[c.process.Key].Process.Kill()
				        delete(processMap, c.process.Key)
				        c.process.OperState = "off"
       				    startRequest.resultChan <- &result{process : c.process}
				    } else {
				    	log.Println("Process is already stopped: " + c.process.Key)
				    }
		    	case "on":
		        	if processMap[c.process.Key] == nil {
			        	log.Println("Starting process: " + c.process.Key)
		        		cmd, err := startProcess(c.process)
						if err != nil {
							log.Printf("Error starting process:\n%s", err.Error())
							c.process.OperState = "off"
							startRequest.resultChan <- &result{process : c.process,
															   err : err}
						} else {
							processMap[c.process.Key] = cmd
							// Send the result back
							c.process.OperState = "on"
							c.process.Pid = processMap[c.process.Key].Process.Pid
							startRequest.resultChan <- &result{process : c.process}
						}
		        	} else {
		        		log.Println("Process is already running: " + c.process.Key)
		        	}
		    	}
			default:
			   // Check running processes
			   log.Println("Checking processes")
			   for key,process := range processMap {
			   	if process.ProcessState != nil && process.ProcessState.Exited() {
			   		startRequest.resultChan <- &result{key : key,
			   										   operState : "off",
			   										   success : process.ProcessState.Success()}
                    delete(processMap, key)
			   	}
			   }
			   time.Sleep(5 * time.Second)
		}
	}
}

func startProcess (process data.Process) (*exec.Cmd, error) {
	log.Println("Process command: " + process.Command)
	var cmd *exec.Cmd
	if process.Arguments != "" {
		cmd = exec.Command(process.Command, process.Arguments)
	} else {
		cmd = exec.Command(process.Command)
	}
	success := false
	var err error
	for i:=0; i<MAX_START_RETRIES && !success; i++ {
		log.Println("Attempting to start: " + process.Name)
		err = cmd.Start()
		// Create new thread to wait on this process in order to reap it
		go cmd.Wait()
		if err == nil {
			success = true
		}
        if !success {
	        time.Sleep(5 * time.Second)
        }
	}
	return cmd, err
}

// Private structures for communication

type processStartRequest struct {
	processes[] data.Process
	commandChan chan *command
	resultChan chan *result
}

type command struct {
	process data.Process
	adminState string
}

type result struct {
	key string
	operState string
	process data.Process
	success bool
	err error
}
<commit_msg>Moved an error handle<commit_after>package main

import (
	"encoding/json"
	"flag"
	"github.com/jetblack87/maestro/data"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"path"
	"github.com/samuel/go-zookeeper/zk"
	"time"
	"log"
)

const APP_VERSION = "0.1"

// The flag package provides a default help printer via -h switch
var versionFlag *bool = flag.Bool("v", false, "Print the version number.")
var zookeeper *string = flag.String("zookeeper", "localhost:2181", "The ZooKeeper connection string (defaults to 'localhost:2182').")
var agentName *string = flag.String("name", "", "REQUIRED: The name of the agent.")
var domainName *string = flag.String("domain", "", "REQUIRED: The name of the domain in which this agent lives.")
var agentConfig *string = flag.String("agentConfig", "", "Supply a json file that contains specific configuration for this agent.")
var processesConfig *string = flag.String("processesConfig", "", "Supply a json file that contains specific configuration any processes.")
var logfilePath *string = flag.String("logfile", "stdout", "The path to the logfile.")

var zkdao data.ZkDAO
var request *processStartRequest

const MAX_START_RETRIES = 3

func main() {

	flag.Parse() // Scan the arguments list

	// Setup signal channel
	signalChannel := make(chan os.Signal, 1)
	signal.Notify(signalChannel, os.Interrupt, os.Kill)

    // Setup logging
	if *logfilePath != "stdout" {
		f, err := os.OpenFile(*logfilePath, os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
		if err != nil {
		    log.Printf("error opening file: %v", err)
		}
		defer f.Close()
		log.SetOutput(f)
	}

	// Check the parameters
	if *versionFlag {
		log.Println("Version:", APP_VERSION)
		os.Exit(0)
	}
	if *agentName == "" {
		panic("-agent is required")
	}
	if *domainName == "" {
		panic("-domain is required")
	}

	log.Printf("maestro agent starting for domain '%s' and agent '%s'\n", *domainName, *domainName)

	zkdao, err := data.NewZkDAO(strings.Split(*zookeeper, ","))
	if err != nil {
		panic(err)
	}

	// Load config into ZK
	err = loadAgentConfig(*agentConfig, *agentName, *domainName)
	if err != nil {
		panic(err)
	}

	// Load processes data into ZK
	err = loadProcessesConfig(*processesConfig, *domainName)
	if err != nil {
		panic(err)
	}

	log.Println("Loading the agent configuration")
	agent, err := zkdao.LoadAgent(data.PathToKey("/maestro/"+*domainName+"/config/agents/"+*agentName), true)
	if err != nil {
		panic(err)
	}

	// Create channel used for watching ZK nodes
	watchChannel := make(chan zk.Event, 1)

	// Remove old runtime config for this agent
	err = zkdao.RemoveRecursive("/maestro/"+*domainName+"/runtime/agents/"+agent.Name)
	if err != nil {
		log.Printf("Failed to remove agent runtime configuration")
		panic(err)
	}

	// Add processes to the runtime configuration, adding watches to admin_state
	for key := range agent.Processes {
		log.Println("Loading processes from config: " + agent.Processes[key].ProcessClass)
		agent.Processes[key], err = zkdao.LoadProcess(data.PathToKey(agent.Processes[key].ProcessClass), true)
		if err != nil {
			panic(err)
		}
		agent.Processes[key].Key = data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name+"/processes/"+agent.Processes[key].Name)
		if agent.Processes[key].AdminState == "" {
			// Default to on
			log.Println("Defaulting admin_state to 'on'")
			agent.Processes[key].AdminState = "on"
		}
		agent.Processes[key].OperState = "off"
	}

	log.Println("Adding agent to runtime configuration")
	err = zkdao.UpdateAgent(data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name), agent, true)
	if err != nil {
		panic(err)
	}

	str, err := zkdao.CreateEphemeral("/maestro/"+*domainName+"/runtime/agents/"+agent.Name+"/eph", []byte("I am alive"))
	if err != nil {
		errMsg := "Error creating ephemeral node: " + str
		log.Println(errMsg)
		panic(err)
	}

	// Add watches for all admin_state nodes
	for key := range agent.Processes {
		adminStatePath := data.KeyToPath(agent.Processes[key].Key)+"/admin_state"
		log.Println("Adding watch to node: " + adminStatePath)
		err := zkdao.Watch(adminStatePath, watchChannel)
		if err != nil {
			log.Println("Failed to add watch to process node:\n" + err.Error())
		}
	}

	log.Println("Starting process monitoring")

	// Create out request (including channels)
	request = &processStartRequest{
		processes : agent.Processes,
		commandChan : make(chan *command, 1),
		resultChan : make(chan *result, 1)}

	go startAndMonitorProcesses(request)

	log.Println("Process monitoring started, waiting on channels")
	for {
		select {
			case w := <-watchChannel:
			if w.Type.String() == "EventNodeDataChanged" {
				adminState, err := zkdao.GetValue(w.Path)
				if err != nil {
					log.Printf("Error getting data for path '%s': %s\n", w.Path, err.Error())
				} else {
					process, err2 := zkdao.LoadProcess(data.PathToKey(path.Dir(w.Path)), true)
					if err2 != nil {
						log.Printf("Error loading process '%s': %s\n", w.Path, err2)
					} else {
						request.commandChan <- &command{process : process, adminState : string(adminState)}
					}
				}
			}
			case r := <-request.resultChan:
			if r.err != nil {
				log.Printf("An error occured running a process:%s\n", r.err.Error())
				// Failed to start, turn off
				r.process.OperState = "off"
				r.process.AdminState = "off"
				zkdao.UpdateProcess(r.process.Key, r.process, false)
			} else {
				var p data.Process
	    		// Update the admin_state and pid in ZK
				if r.process.Key != "" {
					p = r.process
				} else {
					p,err = zkdao.LoadProcess(r.key, false)
					if err != nil {
						log.Printf("Failed to load process: %s\n", err.Error())
					} else {
						p.OperState = r.operState
					}
				}
				log.Printf("Process '%s' oper_state = '%s'\n", p.Key, p.OperState)
	    		zkdao.UpdateProcess(r.process.Key, r.process, false)

	    		// Touch the admin_state node to get process turned back on
	    		if p.AdminState == "on" && p.OperState == "off" {
	    			zkdao.SetValue(data.KeyToPath(p.Key) + "/admin_state", []byte(p.AdminState))
	    		}
			}
			case <-signalChannel: // FIXME this doesn't seem to work (at least not on Windows)
			log.Println("Received signal")
			a,err := zkdao.LoadAgent(data.PathToKey("/maestro/"+*domainName+"/runtime/agents/"+agent.Name),true)
			if err != nil {
				log.Println("Error retrieving agent")
			} else {
				for _, process := range a.Processes {
					request.commandChan <- &command{process : process, adminState : "off"}
				}
			}
			os.Exit(0)
		}
	}
}

func loadAgentConfig(agentConfig, agentName, domainName string) error {
	// Load the file if it was supplied
	if agentConfig != "" {
		log.Println("Loading agent config: " + agentConfig)
		jsonData, err := ioutil.ReadFile(agentConfig)
		if err != nil {
			return err
		}
		var agent data.Agent
		err = json.Unmarshal(jsonData, &agent)
		if err != nil {
			return err
		}
		err = zkdao.UpdateAgent(data.PathToKey("/maestro/"+domainName+"/config/agents/"+agentName), agent, true)
		if err != nil {
			return err
		}
	}
	return nil
}

func loadProcessesConfig(processesConfig, domainName string) error {
	// Load the file if it was supplied
	if processesConfig != "" {
		log.Println("Loading processes config: " + processesConfig)
		jsonData, err := ioutil.ReadFile(processesConfig)
		if err != nil {
			return err
		}
		var processes []data.Process
		err = json.Unmarshal(jsonData, &processes)
		if err != nil {
			return err
		}
		for _, process := range processes {
			err = zkdao.UpdateProcess(data.PathToKey("/maestro/"+domainName+"/config/processes/"+process.Name), process, true)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func startAndMonitorProcesses (startRequest *processStartRequest) {
	// Mapping of process key (string) to command
	processMap := make(map[string]*exec.Cmd)
	// Start all of the processes
	for key := range startRequest.processes {
		if startRequest.processes[key].AdminState == "on" {
        	log.Println("Starting process: " + startRequest.processes[key].Key)
			cmd, err := startProcess(startRequest.processes[key])
			if err != nil {
				log.Printf("Error starting process:\n%s\n", err.Error())
				startRequest.resultChan <- &result{process : startRequest.processes[key], err : err}
			} else {
				processMap[startRequest.processes[key].Key] = cmd
				// Send the result back
				startRequest.processes[key].OperState = "on"
				startRequest.processes[key].Pid = cmd.Process.Pid
				startRequest.resultChan <- &result{process : startRequest.processes[key]}
			}
		}
	}

	// Monitor the processes
	log.Println("Monitoring command channel and processes")
	for {
		select {
			case c := <-startRequest.commandChan:
		    	switch c.adminState {
		    	case "off":
		    		log.Println("Killing process: " + c.process.Key)
		    		if processMap[c.process.Key] != nil {
				        processMap[c.process.Key].Process.Kill()
				        delete(processMap, c.process.Key)
				        c.process.OperState = "off"
       				    startRequest.resultChan <- &result{process : c.process}
				    } else {
				    	log.Println("Process is already stopped: " + c.process.Key)
				    }
		    	case "on":
		        	if processMap[c.process.Key] == nil {
			        	log.Println("Starting process: " + c.process.Key)
		        		cmd, err := startProcess(c.process)
						if err != nil {
							log.Printf("Error starting process:\n%s", err.Error())
							c.process.OperState = "off"
							startRequest.resultChan <- &result{process : c.process,
															   err : err}
						} else {
							processMap[c.process.Key] = cmd
							// Send the result back
							c.process.OperState = "on"
							c.process.Pid = processMap[c.process.Key].Process.Pid
							startRequest.resultChan <- &result{process : c.process}
						}
		        	} else {
		        		log.Println("Process is already running: " + c.process.Key)
		        	}
		    	}
			default:
			   // Check running processes
			   log.Println("Checking processes")
			   for key,process := range processMap {
			   	if process.ProcessState != nil && process.ProcessState.Exited() {
			   		startRequest.resultChan <- &result{key : key,
			   										   operState : "off",
			   										   success : process.ProcessState.Success()}
                    delete(processMap, key)
			   	}
			   }
			   time.Sleep(5 * time.Second)
		}
	}
}

func startProcess (process data.Process) (*exec.Cmd, error) {
	log.Println("Process command: " + process.Command)
	var cmd *exec.Cmd
	if process.Arguments != "" {
		cmd = exec.Command(process.Command, process.Arguments)
	} else {
		cmd = exec.Command(process.Command)
	}
	success := false
	var err error
	for i:=0; i<MAX_START_RETRIES && !success; i++ {
		log.Println("Attempting to start: " + process.Name)
		err = cmd.Start()
		// Create new thread to wait on this process in order to reap it
		go cmd.Wait()
		if err == nil {
			success = true
		}
        if !success {
	        time.Sleep(5 * time.Second)
        }
	}
	return cmd, err
}

// Private structures for communication

type processStartRequest struct {
	processes[] data.Process
	commandChan chan *command
	resultChan chan *result
}

type command struct {
	process data.Process
	adminState string
}

type result struct {
	key string
	operState string
	process data.Process
	success bool
	err error
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
)

type compare func(interface{}, interface{}) bool

type Tree struct {
	Left   *Tree
	Value  interface{}
	Right  *Tree
	comp   compare
	Levels int
}

func NewTree(f compare) *Tree {
	return &Tree{comp: f}
}

func (t *Tree) Add(v interface{}) bool {
	t.Levels++

	var node *Tree = t
	for node.Value != nil {
		if node.comp(node.Value, v) {
			node = node.Left
		} else {
			node = node.Right
		}
	}

	node.Value = v
	node.Left = NewTree(t.comp)
	node.Right = NewTree(t.comp)
	return true
}

func (t *Tree) walk() {
	if t.Left.Value != nil {
		t.Left.walk()
	}
	if t.Right.Value != nil {
		t.Right.walk()
	}
}

func (t *Tree) Walk() {
	t.walk()
}

func main() {
	t := NewTree(func(current, newval interface{}) bool {
		if current.(int) < newval.(int) {
			return false
		}
		return true
	})

	for x := 0; x < 1000000; x++ {
		t.Add(rand.Intn(x + 1))
	}
}
<commit_msg>Removed unused import<commit_after>package main

import (
	"math/rand"
)

type compare func(interface{}, interface{}) bool

type Tree struct {
	Left   *Tree
	Value  interface{}
	Right  *Tree
	comp   compare
	Levels int
}

func NewTree(f compare) *Tree {
	return &Tree{comp: f}
}

func (t *Tree) Add(v interface{}) bool {
	t.Levels++

	var node *Tree = t
	for node.Value != nil {
		if node.comp(node.Value, v) {
			node = node.Left
		} else {
			node = node.Right
		}
	}

	node.Value = v
	node.Left = NewTree(t.comp)
	node.Right = NewTree(t.comp)
	return true
}

func (t *Tree) walk() {
	if t.Left.Value != nil {
		t.Left.walk()
	}
	if t.Right.Value != nil {
		t.Right.walk()
	}
}

func (t *Tree) Walk() {
	t.walk()
}

func main() {
	t := NewTree(func(current, newval interface{}) bool {
		if current.(int) < newval.(int) {
			return false
		}
		return true
	})

	for x := 0; x < 1000000; x++ {
		t.Add(rand.Intn(x + 1))
	}
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			if remainingBytes > 0 {
				if remainingBytes <= n {
					fmt.Println("1")
					newPacket = true
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = n - remainingBytes
					fmt.Println("msg: ", string(msg))
				} else {
					fmt.Println("2")
					newPacket = false
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = remainingBytes - n
				}
			}
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket {
				fmt.Println("5")
				remainingBytes = 0
				newPacket = false
				msg = ""
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", t)
				switch t {
				case query:
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					remainingBytes = remainingBytes - 4
					if remainingBytes > 0 {
						if remainingBytes <= n {
							newPacket = true
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = n - remainingBytes
							fmt.Println(msg)
							goto NewP
						} else {
							newPacket = false
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = remainingBytes - n
						}
						fmt.Printf("Remaining bytes: %d\n", remainingBytes)
					}
				// case rowDescription:
				// case dataRow:
				// case bindComplete:
				// case commandComplete:
				// 	commandTag = CommandTag(r.readCString())
				default:
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<commit_msg>Update<commit_after>package proxy

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"os"

	"github.com/DimShadoWWW/power-pg/common"
)

var (
	connid = uint64(0)
)

// Start function
func Start(localHost, remoteHost *string, powerCallback common.Callback) {
	fmt.Printf("Proxying from %v to %v\n", localHost, remoteHost)

	localAddr, remoteAddr := getResolvedAddresses(localHost, remoteHost)
	listener := getListener(localAddr)

	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			fmt.Printf("Failed to accept connection '%s'\n", err)
			continue
		}
		connid++

		p := &proxy{
			lconn:  *conn,
			laddr:  localAddr,
			raddr:  remoteAddr,
			erred:  false,
			errsig: make(chan bool),
			prefix: fmt.Sprintf("Connection #%03d ", connid),
		}
		go p.start(powerCallback)
	}
}

func getResolvedAddresses(localHost, remoteHost *string) (*net.TCPAddr, *net.TCPAddr) {
	laddr, err := net.ResolveTCPAddr("tcp", *localHost)
	check(err)
	raddr, err := net.ResolveTCPAddr("tcp", *remoteHost)
	check(err)
	return laddr, raddr
}

func getListener(addr *net.TCPAddr) *net.TCPListener {
	listener, err := net.ListenTCP("tcp", addr)
	check(err)
	return listener
}

type proxy struct {
	sentBytes     uint64
	receivedBytes uint64
	laddr, raddr  *net.TCPAddr
	lconn, rconn  net.TCPConn
	erred         bool
	errsig        chan bool
	prefix        string
}

func (p *proxy) err(s string, err error) {
	if p.erred {
		return
	}
	if err != io.EOF {
		warn(p.prefix+s, err)
	}
	p.errsig <- true
	p.erred = true
}

func (p *proxy) start(powerCallback common.Callback) {
	// defer p.lconn.conn.Close()
	//connect to remote
	rconn, err := net.DialTCP("tcp", nil, p.raddr)
	if err != nil {
		p.err("Remote connection failed: %s", err)
		return
	}
	p.rconn = *rconn
	// p.rconn.alive = true
	// defer p.rconn.conn.Close()
	//bidirectional copy
	go p.pipe(p.lconn, p.rconn, powerCallback)
	go p.pipe(p.rconn, p.lconn, nil)
	//wait for close...
	<-p.errsig
}

func (p *proxy) pipe(src, dst net.TCPConn, powerCallback common.Callback) {
	//data direction
	islocal := src == p.lconn
	//directional copy (64k buffer)
	buff := make(readBuf, 0xffff)
	newPacket := true
	var msg string
	remainingBytes := 0
	if islocal {
		for {
			if remainingBytes == 0 {
				newPacket = true
			}
			var r readBuf
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			fmt.Printf("Readed bytes: %d\n", n)
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}

			r = buff[:n]
			// fmt.Printf("%#v", string(buff[:n]))
			fmt.Printf("Remaining bytes: %d\n", remainingBytes)
			if remainingBytes > 0 {
				if remainingBytes <= n {
					fmt.Println("1")
					newPacket = true
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = n - remainingBytes
					fmt.Println("msg: ", string(msg))
				} else {
					fmt.Println("2")
					newPacket = false
					msg = msg + string(r.next(remainingBytes))
					remainingBytes = remainingBytes - n
				}
			}
			fmt.Println("3")
		NewP:
			fmt.Println("4")
			if newPacket {
				fmt.Println("5")
				remainingBytes = 0
				newPacket = false
				msg = ""
				t := r.byte()
				n = n - 1
				fmt.Println("t: ", t)
				switch t {
				case query:
					// c.rxReadyForQuery(r)
					remainingBytes = r.int32()
					remainingBytes = remainingBytes - 4
					if remainingBytes > 0 {
						if remainingBytes <= n {
							newPacket = true
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = n - remainingBytes
							fmt.Println(msg)
							goto NewP
						} else {
							newPacket = false
							msg = msg + string(r.next(remainingBytes))
							remainingBytes = remainingBytes - n
						}
						fmt.Printf("Remaining bytes: %d\n", remainingBytes)
					}
				// case rowDescription:
				// case dataRow:
				// case bindComplete:
				// case commandComplete:
				// 	commandTag = CommandTag(r.readCString())
				default:
					// if e := c.processContextFreeMsg(t, r); e != nil && softErr == nil {
					// 	softErr = e
					// }
				}
			}
			// r = append(r, buff[:]...)

			// fmt.Println("a")
			// c := src
			// c.reader = bufio.NewReader(src.conn)
			// c.mr.reader = c.reader
			//
			// var t byte
			// var r *msgReader
			// fmt.Println("b")
			// t, r, err := c.rxMsg()
			// fmt.Println("c")
			// if err != nil {
			// 	fmt.Println(err)
			// 	return
			// }
			// fmt.Println("d")
			//
			// fmt.Printf("t: %#v\n", t)

			// n, err := src.Read(buff)
			// if err != nil {
			// 	p.err("Read failed '%s'\n", err)
			// 	return
			// }
			// b := buff[:n]
			// //show output
			//
			//
			// b = getModifiedBuffer(b, powerCallback)
			// n, err = dst.Write(b)
			// //
			// //write out result
			// n, err = dst.Write(b)
			// if err != nil {
			// 	p.err("Write failed '%s'\n", err)
			// 	return
			// }
		}
	} else {
		for {
			n, err := src.Read(buff)
			if err != nil {
				p.err("Read failed '%s'\n", err)
				return
			}
			b := buff[:n]
			//write out result
			n, err = dst.Write(b)
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
			if err != nil {
				p.err("Write failed '%s'\n", err)
				return
			}
		}
	}
}

func getModifiedBuffer(buffer []byte, powerCallback common.Callback) []byte {
	if powerCallback == nil || len(buffer) < 1 || string(buffer[0]) != "Q" || string(buffer[5:11]) != "power:" {
		return buffer
	}
	query := powerCallback(string(buffer[5:]))
	return makeMessage(query)
}

func makeMessage(query string) []byte {
	queryArray := make([]byte, 0, 6+len(query))
	queryArray = append(queryArray, 'Q', 0, 0, 0, 0)
	queryArray = append(queryArray, query...)
	queryArray = append(queryArray, 0)
	binary.BigEndian.PutUint32(queryArray[1:], uint32(len(queryArray)-1))
	return queryArray

}

func check(err error) {
	if err != nil {
		warn(err.Error())
		os.Exit(1)
	}
}

func warn(f string, args ...interface{}) {
	fmt.Printf(f+"\n", args...)
}
<|endoftext|>


<commit_before>package ai

import (
	"fmt"
	"io"
	"text/tabwriter"

	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/tak"
)

const (
	endgameCutoff = 7
)

type FlatScores struct {
	Hard, Soft int
}

type Weights struct {
	TopFlat     int
	EndgameFlat int
	Standing    int
	Capstone    int

	FlatCaptives     FlatScores
	StandingCaptives FlatScores
	CapstoneCaptives FlatScores

	Liberties      int
	GroupLiberties int

	Groups [8]int
}

var defaultWeights = Weights{
	TopFlat:     400,
	EndgameFlat: 800,
	Standing:    200,
	Capstone:    300,

	FlatCaptives: FlatScores{
		Hard: 125,
		Soft: -75,
	},
	StandingCaptives: FlatScores{
		Hard: 125,
		Soft: -50,
	},
	CapstoneCaptives: FlatScores{
		Hard: 150,
		Soft: -50,
	},

	Liberties: 20,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
	},
}

var defaultWeights6 = Weights{
	TopFlat:     400,
	EndgameFlat: 800,
	Standing:    200,
	Capstone:    300,

	FlatCaptives: FlatScores{
		Hard: 125,
		Soft: -200,
	},
	StandingCaptives: FlatScores{
		Hard: 125,
		Soft: -150,
	},
	CapstoneCaptives: FlatScores{
		Hard: 150,
		Soft: -75,
	},

	Liberties: 20,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
		500, // 5
	},
}

var DefaultWeights = []Weights{
	defaultWeights,  // 0
	defaultWeights,  // 1
	defaultWeights,  // 2
	defaultWeights,  // 3
	defaultWeights,  // 4
	defaultWeights,  // 5
	defaultWeights6, // 6
	defaultWeights,  // 7
	defaultWeights,  // 8
}

func MakeEvaluator(size int, w *Weights) EvaluationFunc {
	if w == nil {
		w = &DefaultWeights[size]
	}
	return func(c *bitboard.Constants, p *tak.Position) int64 {
		return evaluate(c, w, p)
	}
}

const moveScale = 100

func evaluateTerminal(p *tak.Position, winner tak.Color) int64 {
	var pieces int64
	if winner == tak.White {
		pieces = int64(p.WhiteStones())
	} else {
		pieces = int64(p.BlackStones())
	}
	switch winner {
	case tak.NoColor:
		return 0
	case p.ToMove():
		return MaxEval - moveScale*int64(p.MoveNumber()) + pieces
	default:
		return MinEval + moveScale*int64(p.MoveNumber()) - pieces
	}
}

func EvaluateWinner(_ *bitboard.Constants, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		return evaluateTerminal(p, winner)
	}
	return 0
}

func evaluate(c *bitboard.Constants, w *Weights, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		return evaluateTerminal(p, winner)
	}

	var ws, bs int64

	analysis := p.Analysis()

	left := p.WhiteStones()
	if p.BlackStones() < left {
		left = p.BlackStones()
	}
	if left > endgameCutoff {
		left = endgameCutoff
	}
	flat := w.TopFlat + ((endgameCutoff-left)*w.EndgameFlat)/endgameCutoff
	if p.ToMove() == tak.White {
		ws += int64(flat/2) + 50
	} else {
		bs += int64(flat/2) + 50
	}

	ws += int64(bitboard.Popcount(p.White&^(p.Caps|p.Standing)) * flat)
	bs += int64(bitboard.Popcount(p.Black&^(p.Caps|p.Standing)) * flat)
	ws += int64(bitboard.Popcount(p.White&p.Standing) * w.Standing)
	bs += int64(bitboard.Popcount(p.Black&p.Standing) * w.Standing)
	ws += int64(bitboard.Popcount(p.White&p.Caps) * w.Capstone)
	bs += int64(bitboard.Popcount(p.Black&p.Caps) * w.Capstone)

	for i, h := range p.Height {
		if h <= 1 {
			continue
		}
		bit := uint64(1 << uint(i))
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		var hf, sf int
		var ptr *int64
		if p.White&bit != 0 {
			sf = bitboard.Popcount(s)
			hf = int(h) - sf - 1
			ptr = &ws
		} else {
			hf = bitboard.Popcount(s)
			sf = int(h) - hf - 1
			ptr = &bs
		}

		switch {
		case p.Standing&(1<<uint(i)) != 0:
			*ptr += (int64(hf*w.StandingCaptives.Hard) +
				int64(sf*w.StandingCaptives.Soft))
		case p.Caps&(1<<uint(i)) != 0:
			*ptr += (int64(hf*w.CapstoneCaptives.Hard) +
				int64(sf*w.CapstoneCaptives.Soft))
		default:
			*ptr += (int64(hf*w.FlatCaptives.Hard) +
				int64(sf*w.FlatCaptives.Soft))
		}
	}

	ws += int64(scoreGroups(c, analysis.WhiteGroups, w, p.Black|p.Standing))
	bs += int64(scoreGroups(c, analysis.BlackGroups, w, p.White|p.Standing))

	if w.Liberties != 0 {
		wr := p.White &^ p.Standing
		br := p.Black &^ p.Standing
		wl := bitboard.Popcount(bitboard.Grow(c, ^p.Black, wr) &^ p.White)
		bl := bitboard.Popcount(bitboard.Grow(c, ^p.White, br) &^ p.Black)
		ws += int64(w.Liberties * wl)
		bs += int64(w.Liberties * bl)
	}

	if p.ToMove() == tak.White {
		return ws - bs
	}
	return bs - ws
}

func scoreGroups(c *bitboard.Constants, gs []uint64, ws *Weights, other uint64) int {
	sc := 0
	var allg uint64
	for _, g := range gs {
		w, h := bitboard.Dimensions(c, g)

		sc += ws.Groups[w]
		sc += ws.Groups[h]
		allg |= g
	}
	if ws.GroupLiberties != 0 {
		libs := bitboard.Popcount(bitboard.Grow(c, ^other, allg) &^ allg)
		sc += libs * ws.GroupLiberties
	}

	return sc
}

func ExplainScore(m *MinimaxAI, out io.Writer, p *tak.Position) {
	tw := tabwriter.NewWriter(out, 4, 8, 1, '\t', 0)
	fmt.Fprintf(tw, "\twhite\tblack\n")
	var scores [2]struct {
		flats    int
		standing int
		caps     int

		stones   int
		captured int
	}

	scores[0].flats = bitboard.Popcount(p.White &^ (p.Caps | p.Standing))
	scores[1].flats = bitboard.Popcount(p.Black &^ (p.Caps | p.Standing))
	scores[0].standing = bitboard.Popcount(p.White & p.Standing)
	scores[1].standing = bitboard.Popcount(p.Black & p.Standing)
	scores[0].caps = bitboard.Popcount(p.White & p.Caps)
	scores[1].caps = bitboard.Popcount(p.Black & p.Caps)

	for i, h := range p.Height {
		if h <= 1 {
			continue
		}
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		bf := bitboard.Popcount(s)
		wf := int(h) - bf - 1
		scores[0].stones += wf
		scores[1].stones += bf

		captured := int(h - 1)
		if captured > p.Size()-1 {
			captured = p.Size() - 1
		}
		if p.White&(1<<uint(i)) != 0 {
			scores[0].captured += captured
		} else {
			scores[1].captured += captured
		}
	}

	fmt.Fprintf(tw, "flats\t%d\t%d\n", scores[0].flats, scores[1].flats)
	fmt.Fprintf(tw, "standing\t%d\t%d\n", scores[0].standing, scores[1].standing)
	fmt.Fprintf(tw, "caps\t%d\t%d\n", scores[0].caps, scores[1].caps)
	fmt.Fprintf(tw, "captured\t%d\t%d\n", scores[0].captured, scores[1].captured)
	fmt.Fprintf(tw, "stones\t%d\t%d\n", scores[0].stones, scores[1].stones)

	analysis := p.Analysis()

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ p.White)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ p.Black)

	fmt.Fprintf(tw, "liberties\t%d\t%d\n", wl, bl)

	var allg uint64
	for i, g := range analysis.WhiteGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t%dx%x\n", i, w, h)
		allg |= g
	}
	wgl := bitboard.Popcount(bitboard.Grow(&m.c, m.c.Mask&^(p.Black|p.Standing), allg) &^ allg)
	allg = 0
	for i, g := range analysis.BlackGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t\t%dx%x\n", i, w, h)
		allg |= g
	}
	bgl := bitboard.Popcount(bitboard.Grow(&m.c, m.c.Mask&^(p.White|p.Standing), allg) &^ allg)
	fmt.Fprintf(tw, "gl\t%d\t%d\n", wgl, bgl)
	tw.Flush()
}
<commit_msg>evaluate: store one score variable<commit_after>package ai

import (
	"fmt"
	"io"
	"text/tabwriter"

	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/tak"
)

const (
	endgameCutoff = 7
)

type FlatScores struct {
	Hard, Soft int
}

type Weights struct {
	TopFlat     int
	EndgameFlat int
	Standing    int
	Capstone    int

	FlatCaptives     FlatScores
	StandingCaptives FlatScores
	CapstoneCaptives FlatScores

	Liberties      int
	GroupLiberties int

	Groups [8]int
}

var defaultWeights = Weights{
	TopFlat:     400,
	EndgameFlat: 800,
	Standing:    200,
	Capstone:    300,

	FlatCaptives: FlatScores{
		Hard: 125,
		Soft: -75,
	},
	StandingCaptives: FlatScores{
		Hard: 125,
		Soft: -50,
	},
	CapstoneCaptives: FlatScores{
		Hard: 150,
		Soft: -50,
	},

	Liberties: 20,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
	},
}

var defaultWeights6 = Weights{
	TopFlat:     400,
	EndgameFlat: 800,
	Standing:    200,
	Capstone:    300,

	FlatCaptives: FlatScores{
		Hard: 125,
		Soft: -200,
	},
	StandingCaptives: FlatScores{
		Hard: 125,
		Soft: -150,
	},
	CapstoneCaptives: FlatScores{
		Hard: 150,
		Soft: -75,
	},

	Liberties: 20,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
		500, // 5
	},
}

var DefaultWeights = []Weights{
	defaultWeights,  // 0
	defaultWeights,  // 1
	defaultWeights,  // 2
	defaultWeights,  // 3
	defaultWeights,  // 4
	defaultWeights,  // 5
	defaultWeights6, // 6
	defaultWeights,  // 7
	defaultWeights,  // 8
}

func MakeEvaluator(size int, w *Weights) EvaluationFunc {
	if w == nil {
		w = &DefaultWeights[size]
	}
	return func(c *bitboard.Constants, p *tak.Position) int64 {
		return evaluate(c, w, p)
	}
}

const moveScale = 100

func evaluateTerminal(p *tak.Position, winner tak.Color) int64 {
	var pieces int64
	if winner == tak.White {
		pieces = int64(p.WhiteStones())
	} else {
		pieces = int64(p.BlackStones())
	}
	switch winner {
	case tak.NoColor:
		return 0
	case p.ToMove():
		return MaxEval - moveScale*int64(p.MoveNumber()) + pieces
	default:
		return MinEval + moveScale*int64(p.MoveNumber()) - pieces
	}
}

func EvaluateWinner(_ *bitboard.Constants, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		return evaluateTerminal(p, winner)
	}
	return 0
}

func evaluate(c *bitboard.Constants, w *Weights, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		return evaluateTerminal(p, winner)
	}

	var score int64

	analysis := p.Analysis()

	left := p.WhiteStones()
	if p.BlackStones() < left {
		left = p.BlackStones()
	}
	if left > endgameCutoff {
		left = endgameCutoff
	}
	flat := w.TopFlat + ((endgameCutoff-left)*w.EndgameFlat)/endgameCutoff
	if p.ToMove() == tak.White {
		score += int64(flat/2) + 50
	} else {
		score -= int64(flat/2) + 50
	}

	score += int64(bitboard.Popcount(p.White&^(p.Caps|p.Standing)) * flat)
	score -= int64(bitboard.Popcount(p.Black&^(p.Caps|p.Standing)) * flat)
	score += int64(bitboard.Popcount(p.White&p.Standing) * w.Standing)
	score -= int64(bitboard.Popcount(p.Black&p.Standing) * w.Standing)
	score += int64(bitboard.Popcount(p.White&p.Caps) * w.Capstone)
	score -= int64(bitboard.Popcount(p.Black&p.Caps) * w.Capstone)

	for i, h := range p.Height {
		if h <= 1 {
			continue
		}
		bit := uint64(1 << uint(i))
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		var hf, sf int
		var sign int64
		if p.White&bit != 0 {
			sf = bitboard.Popcount(s)
			hf = int(h) - sf - 1
			sign = 1
		} else {
			hf = bitboard.Popcount(s)
			sf = int(h) - hf - 1
			sign = -1
		}

		switch {
		case p.Standing&(1<<uint(i)) != 0:
			score += sign * (int64(hf*w.StandingCaptives.Hard) +
				int64(sf*w.StandingCaptives.Soft))
		case p.Caps&(1<<uint(i)) != 0:
			score += sign * (int64(hf*w.CapstoneCaptives.Hard) +
				int64(sf*w.CapstoneCaptives.Soft))
		default:
			score += sign * (int64(hf*w.FlatCaptives.Hard) +
				int64(sf*w.FlatCaptives.Soft))
		}
	}

	score += int64(scoreGroups(c, analysis.WhiteGroups, w, p.Black|p.Standing))
	score -= int64(scoreGroups(c, analysis.BlackGroups, w, p.White|p.Standing))

	if w.Liberties != 0 {
		wr := p.White &^ p.Standing
		br := p.Black &^ p.Standing
		wl := bitboard.Popcount(bitboard.Grow(c, ^p.Black, wr) &^ p.White)
		bl := bitboard.Popcount(bitboard.Grow(c, ^p.White, br) &^ p.Black)
		score += int64(w.Liberties * wl)
		score -= int64(w.Liberties * bl)
	}

	if p.ToMove() == tak.White {
		return score
	}
	return -score
}

func scoreGroups(c *bitboard.Constants, gs []uint64, ws *Weights, other uint64) int {
	sc := 0
	var allg uint64
	for _, g := range gs {
		w, h := bitboard.Dimensions(c, g)

		sc += ws.Groups[w]
		sc += ws.Groups[h]
		allg |= g
	}
	if ws.GroupLiberties != 0 {
		libs := bitboard.Popcount(bitboard.Grow(c, ^other, allg) &^ allg)
		sc += libs * ws.GroupLiberties
	}

	return sc
}

func ExplainScore(m *MinimaxAI, out io.Writer, p *tak.Position) {
	tw := tabwriter.NewWriter(out, 4, 8, 1, '\t', 0)
	fmt.Fprintf(tw, "\twhite\tblack\n")
	var scores [2]struct {
		flats    int
		standing int
		caps     int

		stones   int
		captured int
	}

	scores[0].flats = bitboard.Popcount(p.White &^ (p.Caps | p.Standing))
	scores[1].flats = bitboard.Popcount(p.Black &^ (p.Caps | p.Standing))
	scores[0].standing = bitboard.Popcount(p.White & p.Standing)
	scores[1].standing = bitboard.Popcount(p.Black & p.Standing)
	scores[0].caps = bitboard.Popcount(p.White & p.Caps)
	scores[1].caps = bitboard.Popcount(p.Black & p.Caps)

	for i, h := range p.Height {
		if h <= 1 {
			continue
		}
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		bf := bitboard.Popcount(s)
		wf := int(h) - bf - 1
		scores[0].stones += wf
		scores[1].stones += bf

		captured := int(h - 1)
		if captured > p.Size()-1 {
			captured = p.Size() - 1
		}
		if p.White&(1<<uint(i)) != 0 {
			scores[0].captured += captured
		} else {
			scores[1].captured += captured
		}
	}

	fmt.Fprintf(tw, "flats\t%d\t%d\n", scores[0].flats, scores[1].flats)
	fmt.Fprintf(tw, "standing\t%d\t%d\n", scores[0].standing, scores[1].standing)
	fmt.Fprintf(tw, "caps\t%d\t%d\n", scores[0].caps, scores[1].caps)
	fmt.Fprintf(tw, "captured\t%d\t%d\n", scores[0].captured, scores[1].captured)
	fmt.Fprintf(tw, "stones\t%d\t%d\n", scores[0].stones, scores[1].stones)

	analysis := p.Analysis()

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ p.White)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ p.Black)

	fmt.Fprintf(tw, "liberties\t%d\t%d\n", wl, bl)

	var allg uint64
	for i, g := range analysis.WhiteGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t%dx%x\n", i, w, h)
		allg |= g
	}
	wgl := bitboard.Popcount(bitboard.Grow(&m.c, m.c.Mask&^(p.Black|p.Standing), allg) &^ allg)
	allg = 0
	for i, g := range analysis.BlackGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t\t%dx%x\n", i, w, h)
		allg |= g
	}
	bgl := bitboard.Popcount(bitboard.Grow(&m.c, m.c.Mask&^(p.White|p.Standing), allg) &^ allg)
	fmt.Fprintf(tw, "gl\t%d\t%d\n", wgl, bgl)
	tw.Flush()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"sync"
	"time"
)

type CoalesceWriteCloser struct {
	wc     io.WriteCloser
	buf    *bufio.Writer
	cancel chan bool
	write  chan writeReq
	err    error
}

type writeRes struct {
	n   int
	err error
}

type writeReq struct {
	buf []byte
	res chan writeRes
}

func NewCoalesceWriteCloser(wc io.WriteCloser) *CoalesceWriteCloser {
	c := &CoalesceWriteCloser{
		wc:     wc,
		buf:    bufio.NewWriterSize(wc, 9216),
		cancel: make(chan bool),
		write:  make(chan writeReq),
	}

	go func() {
		ticker := time.NewTicker(500 * time.Millisecond)
		for {
			select {
			case <-c.cancel:
				ticker.Stop()
				return
			case <-ticker.C:
				c.err = c.buf.Flush()
				if c.err != nil {
					ticker.Stop()
					return
				}
			case wreq := <-c.write:
				n, err := c.buf.Write(wreq.buf)
				wreq.res <- writeRes{n, err}
			}
		}
	}()

	return c
}

func (c *CoalesceWriteCloser) Write(p []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}

	wreq := writeReq{p, make(chan writeRes)}
	c.write <- wreq
	wres := <-wreq.res
	return wres.n, wres.err
}

func (c *CoalesceWriteCloser) Close() error {
	c.cancel <- true
	if c.err != nil {
		c.wc.Close()
		return c.err
	}

	err := c.buf.Flush()
	if err != nil {
		c.wc.Close()
		return err
	}

	return c.wc.Close()
}

type TimeoutWriter struct {
	Timeout <-chan time.Time
	timer   *time.Timer
	w       io.WriteCloser
	d       time.Duration
}

func NewTimeoutWriter(w io.WriteCloser, d time.Duration) *TimeoutWriter {
	timer := time.NewTimer(d)

	return &TimeoutWriter{
		Timeout: timer.C,
		timer:   timer,
		w:       w,
		d:       d,
	}
}

func (tw *TimeoutWriter) Write(p []byte) (int, error) {
	tw.timer.Reset(tw.d)

	return tw.w.Write(p)
}

func (tw *TimeoutWriter) Close() error {
	tw.timer.Stop()

	return tw.w.Close()
}

type addReq struct {
	added int
	done  chan bool
}

type LimitWriter struct {
	w            io.WriteCloser
	limit        int64
	current      int64
	currentMutex sync.Mutex
	LimitReached chan bool
	add          chan addReq
}

func NewLimitWriter(w io.WriteCloser, limit int64) *LimitWriter {
	return &LimitWriter{
		w:            w,
		limit:        limit,
		LimitReached: make(chan bool, 1),
		add:          make(chan addReq),
	}
}

func (lw *LimitWriter) Write(p []byte) (int, error) {
	lw.currentMutex.Lock()
	defer lw.currentMutex.Unlock()

	lw.current += int64(len(p))
	if lw.current > lw.limit {
		lw.LimitReached <- true
	}

	return lw.w.Write(p)
}

func (lw *LimitWriter) Close() error {
	return lw.w.Close()
}
<commit_msg>refactor(ioutils): remove the coordinating goroutine in CoalesceWriteCloser<commit_after>package main

import (
	"bufio"
	"io"
	"sync"
	"time"
)

type CoalesceWriteCloser struct {
	wc     io.WriteCloser
	buf    *bufio.Writer
	cancel chan bool
	mutex  sync.Mutex
	err    error
}

type writeRes struct {
	n   int
	err error
}

type writeReq struct {
	buf []byte
	res chan writeRes
}

func NewCoalesceWriteCloser(wc io.WriteCloser) *CoalesceWriteCloser {
	c := &CoalesceWriteCloser{
		wc:     wc,
		buf:    bufio.NewWriterSize(wc, 9216),
		cancel: make(chan bool),
	}

	go func() {
		for {
			select {
			case <-c.cancel:
				return
			default:
			}

			c.err = c.flush()
			if c.err != nil {
				return
			}
			time.Sleep(500 * time.Millisecond)
		}
	}()

	return c
}

func (c *CoalesceWriteCloser) Write(p []byte) (int, error) {
	if c.err != nil {
		return 0, c.err
	}

	c.mutex.Lock()
	defer c.mutex.Unlock()

	return c.buf.Write(p)
}

func (c *CoalesceWriteCloser) Close() error {
	c.cancel <- true
	if c.err != nil {
		c.wc.Close()
		return c.err
	}

	err := c.flush()
	if err != nil {
		c.wc.Close()
		return err
	}

	return c.wc.Close()
}

func (c *CoalesceWriteCloser) flush() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	return c.buf.Flush()
}

type TimeoutWriter struct {
	Timeout <-chan time.Time
	timer   *time.Timer
	w       io.WriteCloser
	d       time.Duration
}

func NewTimeoutWriter(w io.WriteCloser, d time.Duration) *TimeoutWriter {
	timer := time.NewTimer(d)

	return &TimeoutWriter{
		Timeout: timer.C,
		timer:   timer,
		w:       w,
		d:       d,
	}
}

func (tw *TimeoutWriter) Write(p []byte) (int, error) {
	tw.timer.Reset(tw.d)

	return tw.w.Write(p)
}

func (tw *TimeoutWriter) Close() error {
	tw.timer.Stop()

	return tw.w.Close()
}

type addReq struct {
	added int
	done  chan bool
}

type LimitWriter struct {
	w            io.WriteCloser
	limit        int64
	current      int64
	currentMutex sync.Mutex
	LimitReached chan bool
	add          chan addReq
}

func NewLimitWriter(w io.WriteCloser, limit int64) *LimitWriter {
	return &LimitWriter{
		w:            w,
		limit:        limit,
		LimitReached: make(chan bool, 1),
		add:          make(chan addReq),
	}
}

func (lw *LimitWriter) Write(p []byte) (int, error) {
	lw.currentMutex.Lock()
	defer lw.currentMutex.Unlock()

	lw.current += int64(len(p))
	if lw.current > lw.limit {
		lw.LimitReached <- true
	}

	return lw.w.Write(p)
}

func (lw *LimitWriter) Close() error {
	return lw.w.Close()
}
<|endoftext|>