<commit_before>package appdash

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
)

const (
	dbName              string = "appdash" // InfluxDB db name.
	spanMeasurementName string = "spans"   // InfluxDB container name for trace spans.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

type InfluxDBStore struct {
	con    *influxDBClient.Client // InfluxDB client connection.
	server *influxDBServer.Server
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	// Current strategy is to remove existing span and save new one
	// instead of updating the existing one.
	// TODO: explore a more efficient alternative strategy.
	if err := in.removeSpanIfExists(id); err != nil {
		return err
	}
	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are uselater on queries.
	tags := make(map[string]string, 3)
	tags["trace_id"] = id.Trace.String()
	tags["span_id"] = id.Span.String()
	tags["parent_id"] = id.Parent.String()
	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}
	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{
		influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   //indexed metadata
			Fields:      fields, //non-indexed metadata
			Time:        time.Now(),
			Precision:   "s",
		},
	}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, err := in.con.Write(bps)
	if err != nil {
		return err
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	t := &Trace{}
	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := influxDBClient.Query{
		Command:  fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id),
		Database: dbName,
	}
	response, err := in.con.Query(q)
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}
	// Expecting one result, since a single query is executed:
	// "SELECT * FROM spans ...".
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of influxdb query response result")
	}
	// Slice series contains all the spans.
	if len(response.Results[0].Series) == 0 {
		return nil, errors.New("trace not found")
	}
	var isRootSpan bool
	// Iterate over series(spans) to create & set trace fields.
	for _, s := range response.Results[0].Series {
		traceID, err := ParseID(s.Tags["trace_id"])
		if err != nil {
			return nil, err
		}
		spanID, err := ParseID(s.Tags["span_id"])
		if err != nil {
			return nil, err
		}
		parentID, err := ParseID(s.Tags["parent_id"])
		if err != nil {
			return nil, err
		}
		if parentID == 0 && isRootSpan {
			// Must be a single root span.
			return nil, errors.New("unexpected multiple root spans")
		}
		if parentID == 0 && !isRootSpan {
			isRootSpan = true
		}
		span := Span{
			ID: SpanID{
				Trace:  ID(traceID),
				Span:   ID(spanID),
				Parent: ID(parentID),
			},
		}
		// s.Values[n] is a slice of span's annotation values
		// len(s.Values) might be greater than one - meaning there are
		// some to drop, see: InfluxDBStore.Collect(...).
		// if so last one is use.
		var fields []interface{}
		if len(s.Values) == 1 {
			fields = s.Values[0]
		}
		if len(s.Values) > 1 {
			fields = s.Values[len(s.Values)-1]
		}
		annotations := make(Annotations, len(fields))
		// Iterates over span's annotation values.
		for i, field := range fields {
			// It is safe to do column[0] (eg. 'Server.Request.Method')
			// matches fields[0] (eg. 'GET')
			key := s.Columns[i]
			var value []byte
			switch field.(type) {
			case string:
				value = []byte(field.(string))
			case nil:
			default:
				return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
			}
			a := Annotation{
				Key:   key,
				Value: value,
			}
			annotations = append(annotations, a)
		}
		span.Annotations = annotations
		if isRootSpan {
			t.Span = span
		} else { // children
			t.Sub = append(t.Sub, &Trace{Span: span})
		}
	}
	return t, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	//TODO: implementation
	return nil, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	v := url.Values{}
	v.Set("q", fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS ", dbName))
	url, err := url.Parse(fmt.Sprintf("%s/%s?%s", in.con.Addr(), "query", v.Encode()))
	if err != nil {
		return err
	}
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		return err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	if c := resp.StatusCode; c < 200 || c > 299 {
		defer resp.Body.Close()
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		return fmt.Errorf("failed to create appdash database, response body: %s", string(b))
	}
	return nil
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	return nil
}

func (in *InfluxDBStore) removeSpanIfExists(id SpanID) error {
	cmd := fmt.Sprintf(`
		DROP SERIES FROM spans WHERE trace_id = '%s' AND span_id = '%s' AND parent_id = '%s'
	`, id.Trace.String(), id.Span.String(), id.Parent.String())
	q := influxDBClient.Query{
		Command:  cmd,
		Database: dbName,
	}
	_, err := in.con.Query(q)
	if err != nil {
		return err
	}
	return nil
}

func NewInfluxDBStore(c *influxDBServer.Config, bi *influxDBServer.BuildInfo) (*InfluxDBStore, error) {
	s, err := influxDBServer.NewServer(c, bi)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<commit_msg>adds Traces implementation & cleanups InfluxDBStore<commit_after>package appdash

import (
	"errors"
	"fmt"
	"net/url"
	"reflect"
	"time"

	influxDBClient "github.com/influxdb/influxdb/client"
	influxDBServer "github.com/influxdb/influxdb/cmd/influxd/run"
	influxDBModels "github.com/influxdb/influxdb/models"
)

const (
	dbName               string = "appdash" // InfluxDB db name.
	spanMeasurementName  string = "spans"   // InfluxDB container name for trace spans.
	defaultTracesPerPage int    = 10        // Default number of traces per page.
)

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*InfluxDBStore)(nil)

type InfluxDBStore struct {
	con           *influxDBClient.Client // InfluxDB client connection.
	server        *influxDBServer.Server // InfluxDB API server.
	tracesPerPage int                    // Number of traces per page.
}

func (in *InfluxDBStore) Collect(id SpanID, anns ...Annotation) error {
	// Current strategy is to remove existing span and save new one
	// instead of updating the existing one.
	// TODO: explore a more efficient alternative strategy.
	if err := in.removeSpanIfExists(id); err != nil {
		return err
	}
	// trace_id, span_id & parent_id are set as tags
	// because InfluxDB tags are indexed & those values
	// are uselater on queries.
	tags := make(map[string]string, 3)
	tags["trace_id"] = id.Trace.String()
	tags["span_id"] = id.Span.String()
	tags["parent_id"] = id.Parent.String()
	// Saving annotations as InfluxDB measurement spans fields
	// which are not indexed.
	fields := make(map[string]interface{}, len(anns))
	for _, ann := range anns {
		fields[ann.Key] = string(ann.Value)
	}
	// InfluxDB point represents a single span.
	pts := []influxDBClient.Point{
		influxDBClient.Point{
			Measurement: spanMeasurementName,
			Tags:        tags,   // indexed metadata.
			Fields:      fields, // non-indexed metadata.
			Time:        time.Now(),
			Precision:   "s",
		},
	}
	bps := influxDBClient.BatchPoints{
		Points:          pts,
		Database:        dbName,
		RetentionPolicy: "default",
	}
	_, err := in.con.Write(bps)
	if err != nil {
		return err
	}
	return nil
}

func (in *InfluxDBStore) Trace(id ID) (*Trace, error) {
	trace := &Trace{}
	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans WHERE trace_id='%s' GROUP BY *", id)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return nil, errors.New("trace not found")
	}
	var isRootSpan bool
	// Iterate over series(spans) to create trace children's & set trace fields.
	for _, s := range result.Series {
		span, err := newSpan(s.Tags)
		if err != nil {
			return nil, err
		}
		if span.ID.Parent == 0 && isRootSpan {
			// Must be a single root span.
			return nil, errors.New("unexpected multiple root spans")
		}
		if span.ID.Parent == 0 && !isRootSpan {
			isRootSpan = true
		}
		annotations, err := annotations(&s)
		if err != nil {
			return trace, nil
		}
		span.Annotations = *annotations
		if isRootSpan { // root span.
			trace.Span = *span
		} else { // children span.
			trace.Sub = append(trace.Sub, &Trace{Span: *span})
		}
	}
	return trace, nil
}

func (in *InfluxDBStore) Traces() ([]*Trace, error) {
	traces := make([]*Trace, 0)
	// GROUP BY * -> meaning group by all tags(trace_id, span_id & parent_id)
	// grouping by all tags includes those and it's values on the query response.
	q := fmt.Sprintf("SELECT * FROM spans GROUP BY * LIMIT %d", in.tracesPerPage)
	result, err := in.executeOneQuery(q)
	if err != nil {
		return nil, err
	}
	// result.Series -> A slice containing all the spans.
	if len(result.Series) == 0 {
		return traces, nil
	}
	// Cache to keep track of traces to be returned.
	tracesCache := make(map[ID]*Trace, 0)
	// Iterate over series(spans) to create traces.
	for _, s := range result.Series {
		var isRootSpan bool
		span, err := newSpan(s.Tags)
		if err != nil {
			return nil, err
		}
		annotations, err := annotations(&s)
		if err != nil {
			return nil, err
		}
		if span.ID.Parent == 0 {
			isRootSpan = true
		}
		span.Annotations = *annotations
		if isRootSpan { // root span.
			trace, present := tracesCache[span.ID.Trace]
			if !present {
				tracesCache[span.ID.Trace] = &Trace{Span: *span}
			} else { // trace already added just update the span.
				trace.Span = *span
			}
		} else { // children span.
			trace, present := tracesCache[span.ID.Trace]
			if !present { // root trace not added yet.
				tracesCache[span.ID.Trace] = &Trace{Sub: []*Trace{&Trace{Span: *span}}}
			} else { // root trace already added so append a sub trace.
				trace.Sub = append(trace.Sub, &Trace{Span: *span})
			}
		}
	}
	for _, t := range tracesCache {
		traces = append(traces, t)
	}
	return traces, nil
}

func (in *InfluxDBStore) Close() {
	in.server.Close()
}

func (in *InfluxDBStore) createDBIfNotExists() error {
	// If no errors query execution was successfully - either DB was created or already exists.
	response, err := in.con.Query(influxDBClient.Query{
		Command: fmt.Sprintf("%s %s", "CREATE DATABASE IF NOT EXISTS", dbName),
	})
	if err != nil {
		return err
	}
	if response.Error() != nil {
		return response.Error()
	}
	return nil
}

func (in *InfluxDBStore) executeOneQuery(command string) (*influxDBClient.Result, error) {
	response, err := in.con.Query(influxDBClient.Query{
		Command:  command,
		Database: dbName,
	})
	if err != nil {
		return nil, err
	}
	if response.Error() != nil {
		return nil, response.Error()
	}
	// Expecting one result, since a single query is executed.
	if len(response.Results) != 1 {
		return nil, errors.New("unexpected number of results for an influxdb single query")
	}
	return &response.Results[0], nil
}

func (in *InfluxDBStore) init(server *influxDBServer.Server) error {
	in.server = server
	url, err := url.Parse(fmt.Sprintf("http://%s:%d", influxDBClient.DefaultHost, influxDBClient.DefaultPort))
	if err != nil {
		return err
	}
	con, err := influxDBClient.NewClient(influxDBClient.Config{URL: *url})
	if err != nil {
		return err
	}
	in.con = con
	if err := in.createDBIfNotExists(); err != nil {
		return err
	}
	in.tracesPerPage = defaultTracesPerPage
	return nil
}

func (in *InfluxDBStore) removeSpanIfExists(id SpanID) error {
	cmd := fmt.Sprintf(`
		DROP SERIES FROM spans WHERE trace_id = '%s' AND span_id = '%s' AND parent_id = '%s'
	`, id.Trace.String(), id.Span.String(), id.Parent.String())
	_, err := in.executeOneQuery(cmd)
	if err != nil {
		return err
	}
	return nil
}

func annotations(s *influxDBModels.Row) (*Annotations, error) {
	// Actually an influxDBModels.Row represents a single InfluxDB serie.
	// s.Values[n] is a slice containing span's annotation values.
	var fields []interface{}
	if len(s.Values) == 1 {
		fields = s.Values[0]
	}
	// len(s.Values) might be greater than one - meaning there are
	// some spans to drop, see: InfluxDBStore.Collect(...).
	// If so last one is picked.
	if len(s.Values) > 1 {
		fields = s.Values[len(s.Values)-1]
	}
	annotations := make(Annotations, len(fields))
	// Iterates over fields which represent span's annotation values.
	for i, field := range fields {
		// It is safe to do column[0] (eg. 'Server.Request.Method')
		// matches fields[0] (eg. 'GET')
		key := s.Columns[i]
		var value []byte
		switch field.(type) {
		case string:
			value = []byte(field.(string))
		case nil:
		default:
			return nil, fmt.Errorf("unexpected field type: %v", reflect.TypeOf(field))
		}
		a := Annotation{
			Key:   key,
			Value: value,
		}
		annotations = append(annotations, a)
	}
	return &annotations, nil
}

func newSpan(tags map[string]string) (*Span, error) {
	s := &Span{}
	traceID, err := ParseID(tags["trace_id"])
	if err != nil {
		return nil, err
	}
	spanID, err := ParseID(tags["span_id"])
	if err != nil {
		return nil, err
	}
	parentID, err := ParseID(tags["parent_id"])
	if err != nil {
		return nil, err
	}
	s.ID = SpanID{
		Trace:  ID(traceID),
		Span:   ID(spanID),
		Parent: ID(parentID),
	}
	return s, nil
}

func NewInfluxDBStore(c *influxDBServer.Config, bi *influxDBServer.BuildInfo) (*InfluxDBStore, error) {
	//TODO: add Authentication.
	s, err := influxDBServer.NewServer(c, bi)
	if err != nil {
		return nil, err
	}
	if err := s.Open(); err != nil {
		return nil, err
	}
	var in InfluxDBStore
	if err := in.init(s); err != nil {
		return nil, err
	}
	return &in, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 eventbox batches incoming events for a single Datastore entity
// for processing.
package eventbox

import (
	"context"
	"fmt"
	"time"

	"github.com/google/uuid"
	"golang.org/x/sync/errgroup"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/trace"
	"go.chromium.org/luci/gae/service/datastore"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/eventbox/dsset"
)

// Emit emits a new event with provided value and auto-generated unique ID.
func Emit(ctx context.Context, value []byte, to *datastore.Key) error {
	d := dsset.Set{Parent: to} // TombstonesDelay doesn't matter for Add.
	id := uuid.New().String()
	if err := d.Add(ctx, []dsset.Item{{ID: id, Value: value}}); err != nil {
		return errors.Annotate(err, "failed to send event").Err()
	}
	return nil
}

// TombstonesDelay is exposed to mitigate frequent errors in CV e2e tests when
// tasks are run in parallel with fake clock.
var TombstonesDelay = 5 * time.Minute

// List returns unprocessed events. For use in tests only.
func List(ctx context.Context, recipient *datastore.Key) (Events, error) {
	d := dsset.Set{
		Parent:          recipient,
		TombstonesDelay: TombstonesDelay,
	}
	const effectivelyUnlimited = 1000000
	switch l, err := d.List(ctx, effectivelyUnlimited); {
	case err != nil:
		return nil, err
	case len(l.Items) == effectivelyUnlimited:
		panic(fmt.Errorf("fetched possibly not all events (limit: %d)", effectivelyUnlimited))
	default:
		return toEvents(l.Items), nil
	}
}

// ErrContention indicates Datastore contention, usually on the mailbox
// recipient entity itself.
var ErrContention = errors.New("Datastore Contention")

// IsErrContention checks if error, possibly wrapped, is ErrContention.
func IsErrContention(err error) bool {
	ret := false
	errors.WalkLeaves(err, func(e error) bool {
		ret = (e == ErrContention)
		return !ret // return false means stop traversing.
	})
	return ret
}

// ProcessBatch reliably processes outstanding events, while transactionally modifying state
// and performing arbitrary side effects.
//
// Returns:
//  - a slice of non-nil post process functions which SHOULD be executed
//    immediately after calling this function. Those are generally extra work
//    that needs to be done as the result of state modification.
//  - error while processing events. Returns wrapped ErrContention
//    if entity's EVersion has changed or there is contention on Datastore
//    entities involved in a transaction.
func ProcessBatch(ctx context.Context, recipient *datastore.Key, p Processor, maxEvents int) ([]PostProcessFn, error) {
	ctx, span := trace.StartSpan(ctx, "go.chromium.org/luci/cv/internal/eventbox/ProcessBatch")
	var err error
	span.Attribute("recipient", recipient.String())
	defer func() { span.End(err) }()
	postProcessFn, err := processBatch(ctx, recipient, p, maxEvents)
	return postProcessFn, err
}

func processBatch(ctx context.Context, recipient *datastore.Key, p Processor, maxEvents int) ([]PostProcessFn, error) {
	var state State
	var expectedEV EVersion
	eg, ectx := errgroup.WithContext(ctx)
	eg.Go(func() (err error) {
		state, expectedEV, err = p.LoadState(ectx)
		return
	})
	d := dsset.Set{
		Parent:          recipient,
		TombstonesDelay: TombstonesDelay,
	}
	var listing *dsset.Listing
	eg.Go(func() (err error) {
		if listing, err = d.List(ectx, maxEvents); err == nil {
			err = dsset.CleanupGarbage(ectx, listing.Garbage)
		}
		return
	})
	if err := eg.Wait(); err != nil {
		return nil, err
	}

	// Compute resulting state before transaction.
	transitions, garbage, err := p.PrepareMutation(ctx, toEvents(listing.Items), state)
	if gErr := deleteSemanticGarbage(ctx, &d, garbage); gErr != nil {
		return nil, gErr
	}
	if err != nil {
		return nil, err
	}
	transitions = withoutNoops(transitions, state)
	if len(transitions) == 0 {
		return nil, nil // nothing to do.
	}

	var innerErr error
	var postProcessFns []PostProcessFn
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) (err error) {
		defer func() { innerErr = err }()
		postProcessFns = nil //  reset, since this func can be retried

		switch latestEV, err := p.FetchEVersion(ctx); {
		case err != nil:
			return err
		case latestEV != expectedEV:
			return errors.Annotate(ErrContention, "EVersion read %d, but expected %d", latestEV, expectedEV).Tag(transient.Tag).Err()
		}
		popOp, err := d.BeginPop(ctx, listing)
		if err != nil {
			return err
		}
		var newState State
		eventsConsumed := 0
		for _, t := range transitions {
			if err := t.apply(ctx, popOp); err != nil {
				return err
			}
			newState = t.TransitionTo
			if t.PostProcessFn != nil {
				postProcessFns = append(postProcessFns, t.PostProcessFn)
			}
			eventsConsumed += len(t.Events)
		}

		logging.Debugf(ctx, "%d transitions, %d events", len(transitions), eventsConsumed)
		if newState != state {
			if err := p.SaveState(ctx, newState, expectedEV+1); err != nil {
				return err
			}
		} else {
			logging.Debugf(ctx, "state wasn't modified")
		}
		return dsset.FinishPop(ctx, popOp)
	}, nil)
	switch {
	case innerErr != nil:
		return nil, innerErr
	case common.IsDatastoreContention(err):
		return nil, errors.Annotate(ErrContention, "failed to commit mutation: %s", err).Tag(transient.Tag).Err()
	case err != nil:
		return nil, errors.Annotate(err, "failed to commit mutation").Tag(transient.Tag).Err()
	default:
		return postProcessFns, nil
	}
}

// Processor defines safe way to process events in a batch.
type Processor interface {
	// LoadState is called to load the state before a transaction.
	LoadState(context.Context) (State, EVersion, error)
	// PrepareMutation is called before a transaction to compute transitions based
	// on a batch of events.
	//
	// The events in a batch are an arbitrary subset of all outstanding events.
	// Because loading of events isn't synchronized with event senders,
	// a recipient of events may see them in different order than the origination
	// order, even if events were produced by a single sender.
	//
	// All actions that must be done atomically with updating state must be
	// encapsulated inside Transition.SideEffectFn callback.
	//
	// Garbage events will be deleted non-transactionally before executing
	// transactional transitions. These events may still be processed by a
	// concurrent invocation of a Processor. The garbage events slice may re-use
	// the given events slice. The garbage will be deleted even if PrepareMutation returns
	// non-nil error.
	//
	// For correctness, two concurrent invocation of a Processor must choose the
	// same events to be deleted as garbage. Consider scenario of 2 events A and B
	// deemed semantically the same and 2 concurrent Processor invocations:
	//   P1: let me delete A and hope to transactionally process B.
	//   P2:  ............ B and ............................... A.
	// Then, it's a real possibility that A and B are both deleted AND no neither
	// P1 nor P2 commits a transaction, thus forever forgetting about A and B.
	PrepareMutation(context.Context, Events, State) (transitions []Transition, garbage Events, err error)
	// FetchEVersion is called at the beginning of a transaction.
	//
	// The returned EVersion is compared against the one associated with a state
	// loaded via GetState. If different, the transaction is aborted and new state
	// isn't saved.
	FetchEVersion(ctx context.Context) (EVersion, error)
	// SaveState is called in a transaction to save the state if it has changed.
	//
	// The passed eversion is incremented value of eversion of what GetState
	// returned before.
	SaveState(context.Context, State, EVersion) error
}

// Event is an incoming event.
type Event dsset.Item

// Events are incoming events.
type Events []Event

// toEvents is an annoying redundant malloc to avoid exposing dsset.Item :(
func toEvents(items []dsset.Item) Events {
	es := make(Events, len(items))
	for i, item := range items {
		es[i] = Event(item)
	}
	return es
}

func deleteSemanticGarbage(ctx context.Context, d *dsset.Set, events Events) error {
	l := len(events)
	if l == 0 {
		return nil
	}
	logging.Debugf(ctx, "eventbox deleting %d semantic garbage events before transaction", l)
	i := -1
	err := d.Delete(ctx, func() string {
		i++
		if i < l {
			return events[i].ID
		}
		return ""
	})
	if err != nil {
		return errors.Annotate(err, "failed to delete %d semantic garbage events before transaction", l).Err()
	}
	return nil
}

// State is an arbitrary object.
//
// Use a pointer to an actual state.
type State interface{}

// EVersion is recipient entity version.
type EVersion int

// PostProcessFn should be executed after event processing completes.
type PostProcessFn func(context.Context) error

// SideEffectFn performs side effects with a Datastore transaction context.
// See Transition.SideEffectFn doc.
type SideEffectFn func(context.Context) error

// Chain combines several SideEffectFn.
//
// NOTE: modifies incoming ... slice.
func Chain(fs ...SideEffectFn) SideEffectFn {
	nonNil := fs[:0]
	for _, f := range fs {
		if f != nil {
			nonNil = append(nonNil, f)
		}
	}
	if len(nonNil) == 0 {
		return nil
	}
	return func(ctx context.Context) error {
		for _, f := range nonNil {
			if err := f(ctx); err != nil {
				return err
			}
		}
		return nil
	}
}

// Transition is a state transition.
type Transition struct {
	// SideEffectFn is called in a transaction to atomically with the state change
	// execute any side effects of a state transition.
	//
	// Typical use is notifying other CV components via TQ tasks.
	// Can be nil, meaning there no side effects to execute.
	//
	// TODO(tandrii): introduce error tag to indicate that failure was clean and
	// should be treated as if Transition wasn't started, s.t. progress of all
	// transitions before can be saved.
	SideEffectFn SideEffectFn
	// Events to consume with this transition.
	Events Events
	// TransitionTo is a state to transition to.
	//
	// It's allowed to transition to the exact same state.
	TransitionTo State
	// PostProcessFn is the function to be called by the eventbox user after
	// event processing completes.
	//
	// Note that it will be called outside of the transaction of all state
	// transitions, so the operation inside this function is not expected
	// to be atomic with this state transition.
	PostProcessFn PostProcessFn
}

func (t *Transition) apply(ctx context.Context, p *dsset.PopOp) error {
	if t.SideEffectFn != nil {
		if err := t.SideEffectFn(ctx); err != nil {
			return err
		}
	}
	for _, e := range t.Events {
		_ = p.Pop(e.ID) // Silently ignore if event has already been consumed.
	}
	return nil
}

// isNoop returns true if the Transition can be skipped entirely.
func (t *Transition) isNoop(oldState State) bool {
	return t.SideEffectFn == nil && len(t.Events) == 0 && t.TransitionTo == oldState && t.PostProcessFn == nil
}

// withoutNoops returns only actionable transitions in the original order.
//
// Modifies incoming slice.
func withoutNoops(all []Transition, s State) []Transition {
	ret := all[:0]
	for _, t := range all {
		if t.isNoop(s) {
			continue
		}
		ret = append(ret, t)
		s = t.TransitionTo
	}
	return ret
}
<commit_msg>[cv] remove logging which is rarely informative.<commit_after>// Copyright 2020 The LUCI 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 eventbox batches incoming events for a single Datastore entity
// for processing.
package eventbox

import (
	"context"
	"fmt"
	"time"

	"github.com/google/uuid"
	"golang.org/x/sync/errgroup"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/common/trace"
	"go.chromium.org/luci/gae/service/datastore"

	"go.chromium.org/luci/cv/internal/common"
	"go.chromium.org/luci/cv/internal/eventbox/dsset"
)

// Emit emits a new event with provided value and auto-generated unique ID.
func Emit(ctx context.Context, value []byte, to *datastore.Key) error {
	d := dsset.Set{Parent: to} // TombstonesDelay doesn't matter for Add.
	id := uuid.New().String()
	if err := d.Add(ctx, []dsset.Item{{ID: id, Value: value}}); err != nil {
		return errors.Annotate(err, "failed to send event").Err()
	}
	return nil
}

// TombstonesDelay is exposed to mitigate frequent errors in CV e2e tests when
// tasks are run in parallel with fake clock.
var TombstonesDelay = 5 * time.Minute

// List returns unprocessed events. For use in tests only.
func List(ctx context.Context, recipient *datastore.Key) (Events, error) {
	d := dsset.Set{
		Parent:          recipient,
		TombstonesDelay: TombstonesDelay,
	}
	const effectivelyUnlimited = 1000000
	switch l, err := d.List(ctx, effectivelyUnlimited); {
	case err != nil:
		return nil, err
	case len(l.Items) == effectivelyUnlimited:
		panic(fmt.Errorf("fetched possibly not all events (limit: %d)", effectivelyUnlimited))
	default:
		return toEvents(l.Items), nil
	}
}

// ErrContention indicates Datastore contention, usually on the mailbox
// recipient entity itself.
var ErrContention = errors.New("Datastore Contention")

// IsErrContention checks if error, possibly wrapped, is ErrContention.
func IsErrContention(err error) bool {
	ret := false
	errors.WalkLeaves(err, func(e error) bool {
		ret = (e == ErrContention)
		return !ret // return false means stop traversing.
	})
	return ret
}

// ProcessBatch reliably processes outstanding events, while transactionally modifying state
// and performing arbitrary side effects.
//
// Returns:
//  - a slice of non-nil post process functions which SHOULD be executed
//    immediately after calling this function. Those are generally extra work
//    that needs to be done as the result of state modification.
//  - error while processing events. Returns wrapped ErrContention
//    if entity's EVersion has changed or there is contention on Datastore
//    entities involved in a transaction.
func ProcessBatch(ctx context.Context, recipient *datastore.Key, p Processor, maxEvents int) ([]PostProcessFn, error) {
	ctx, span := trace.StartSpan(ctx, "go.chromium.org/luci/cv/internal/eventbox/ProcessBatch")
	var err error
	span.Attribute("recipient", recipient.String())
	defer func() { span.End(err) }()
	postProcessFn, err := processBatch(ctx, recipient, p, maxEvents)
	return postProcessFn, err
}

func processBatch(ctx context.Context, recipient *datastore.Key, p Processor, maxEvents int) ([]PostProcessFn, error) {
	var state State
	var expectedEV EVersion
	eg, ectx := errgroup.WithContext(ctx)
	eg.Go(func() (err error) {
		state, expectedEV, err = p.LoadState(ectx)
		return
	})
	d := dsset.Set{
		Parent:          recipient,
		TombstonesDelay: TombstonesDelay,
	}
	var listing *dsset.Listing
	eg.Go(func() (err error) {
		if listing, err = d.List(ectx, maxEvents); err == nil {
			err = dsset.CleanupGarbage(ectx, listing.Garbage)
		}
		return
	})
	if err := eg.Wait(); err != nil {
		return nil, err
	}

	// Compute resulting state before transaction.
	transitions, garbage, err := p.PrepareMutation(ctx, toEvents(listing.Items), state)
	if gErr := deleteSemanticGarbage(ctx, &d, garbage); gErr != nil {
		return nil, gErr
	}
	if err != nil {
		return nil, err
	}
	transitions = withoutNoops(transitions, state)
	if len(transitions) == 0 {
		return nil, nil // nothing to do.
	}

	var innerErr error
	var postProcessFns []PostProcessFn
	err = datastore.RunInTransaction(ctx, func(ctx context.Context) (err error) {
		defer func() { innerErr = err }()
		postProcessFns = nil //  reset, since this func can be retried

		switch latestEV, err := p.FetchEVersion(ctx); {
		case err != nil:
			return err
		case latestEV != expectedEV:
			return errors.Annotate(ErrContention, "EVersion read %d, but expected %d", latestEV, expectedEV).Tag(transient.Tag).Err()
		}
		popOp, err := d.BeginPop(ctx, listing)
		if err != nil {
			return err
		}
		var newState State
		eventsConsumed := 0
		for _, t := range transitions {
			if err := t.apply(ctx, popOp); err != nil {
				return err
			}
			newState = t.TransitionTo
			if t.PostProcessFn != nil {
				postProcessFns = append(postProcessFns, t.PostProcessFn)
			}
			eventsConsumed += len(t.Events)
		}

		if newState != state {
			if err := p.SaveState(ctx, newState, expectedEV+1); err != nil {
				return err
			}
		}
		return dsset.FinishPop(ctx, popOp)
	}, nil)
	switch {
	case innerErr != nil:
		return nil, innerErr
	case common.IsDatastoreContention(err):
		return nil, errors.Annotate(ErrContention, "failed to commit mutation: %s", err).Tag(transient.Tag).Err()
	case err != nil:
		return nil, errors.Annotate(err, "failed to commit mutation").Tag(transient.Tag).Err()
	default:
		return postProcessFns, nil
	}
}

// Processor defines safe way to process events in a batch.
type Processor interface {
	// LoadState is called to load the state before a transaction.
	LoadState(context.Context) (State, EVersion, error)
	// PrepareMutation is called before a transaction to compute transitions based
	// on a batch of events.
	//
	// The events in a batch are an arbitrary subset of all outstanding events.
	// Because loading of events isn't synchronized with event senders,
	// a recipient of events may see them in different order than the origination
	// order, even if events were produced by a single sender.
	//
	// All actions that must be done atomically with updating state must be
	// encapsulated inside Transition.SideEffectFn callback.
	//
	// Garbage events will be deleted non-transactionally before executing
	// transactional transitions. These events may still be processed by a
	// concurrent invocation of a Processor. The garbage events slice may re-use
	// the given events slice. The garbage will be deleted even if PrepareMutation returns
	// non-nil error.
	//
	// For correctness, two concurrent invocation of a Processor must choose the
	// same events to be deleted as garbage. Consider scenario of 2 events A and B
	// deemed semantically the same and 2 concurrent Processor invocations:
	//   P1: let me delete A and hope to transactionally process B.
	//   P2:  ............ B and ............................... A.
	// Then, it's a real possibility that A and B are both deleted AND no neither
	// P1 nor P2 commits a transaction, thus forever forgetting about A and B.
	PrepareMutation(context.Context, Events, State) (transitions []Transition, garbage Events, err error)
	// FetchEVersion is called at the beginning of a transaction.
	//
	// The returned EVersion is compared against the one associated with a state
	// loaded via GetState. If different, the transaction is aborted and new state
	// isn't saved.
	FetchEVersion(ctx context.Context) (EVersion, error)
	// SaveState is called in a transaction to save the state if it has changed.
	//
	// The passed eversion is incremented value of eversion of what GetState
	// returned before.
	SaveState(context.Context, State, EVersion) error
}

// Event is an incoming event.
type Event dsset.Item

// Events are incoming events.
type Events []Event

// toEvents is an annoying redundant malloc to avoid exposing dsset.Item :(
func toEvents(items []dsset.Item) Events {
	es := make(Events, len(items))
	for i, item := range items {
		es[i] = Event(item)
	}
	return es
}

func deleteSemanticGarbage(ctx context.Context, d *dsset.Set, events Events) error {
	l := len(events)
	if l == 0 {
		return nil
	}
	logging.Debugf(ctx, "eventbox deleting %d semantic garbage events before transaction", l)
	i := -1
	err := d.Delete(ctx, func() string {
		i++
		if i < l {
			return events[i].ID
		}
		return ""
	})
	if err != nil {
		return errors.Annotate(err, "failed to delete %d semantic garbage events before transaction", l).Err()
	}
	return nil
}

// State is an arbitrary object.
//
// Use a pointer to an actual state.
type State interface{}

// EVersion is recipient entity version.
type EVersion int

// PostProcessFn should be executed after event processing completes.
type PostProcessFn func(context.Context) error

// SideEffectFn performs side effects with a Datastore transaction context.
// See Transition.SideEffectFn doc.
type SideEffectFn func(context.Context) error

// Chain combines several SideEffectFn.
//
// NOTE: modifies incoming ... slice.
func Chain(fs ...SideEffectFn) SideEffectFn {
	nonNil := fs[:0]
	for _, f := range fs {
		if f != nil {
			nonNil = append(nonNil, f)
		}
	}
	if len(nonNil) == 0 {
		return nil
	}
	return func(ctx context.Context) error {
		for _, f := range nonNil {
			if err := f(ctx); err != nil {
				return err
			}
		}
		return nil
	}
}

// Transition is a state transition.
type Transition struct {
	// SideEffectFn is called in a transaction to atomically with the state change
	// execute any side effects of a state transition.
	//
	// Typical use is notifying other CV components via TQ tasks.
	// Can be nil, meaning there no side effects to execute.
	//
	// TODO(tandrii): introduce error tag to indicate that failure was clean and
	// should be treated as if Transition wasn't started, s.t. progress of all
	// transitions before can be saved.
	SideEffectFn SideEffectFn
	// Events to consume with this transition.
	Events Events
	// TransitionTo is a state to transition to.
	//
	// It's allowed to transition to the exact same state.
	TransitionTo State
	// PostProcessFn is the function to be called by the eventbox user after
	// event processing completes.
	//
	// Note that it will be called outside of the transaction of all state
	// transitions, so the operation inside this function is not expected
	// to be atomic with this state transition.
	PostProcessFn PostProcessFn
}

func (t *Transition) apply(ctx context.Context, p *dsset.PopOp) error {
	if t.SideEffectFn != nil {
		if err := t.SideEffectFn(ctx); err != nil {
			return err
		}
	}
	for _, e := range t.Events {
		_ = p.Pop(e.ID) // Silently ignore if event has already been consumed.
	}
	return nil
}

// isNoop returns true if the Transition can be skipped entirely.
func (t *Transition) isNoop(oldState State) bool {
	return t.SideEffectFn == nil && len(t.Events) == 0 && t.TransitionTo == oldState && t.PostProcessFn == nil
}

// withoutNoops returns only actionable transitions in the original order.
//
// Modifies incoming slice.
func withoutNoops(all []Transition, s State) []Transition {
	ret := all[:0]
	for _, t := range all {
		if t.isNoop(s) {
			continue
		}
		ret = append(ret, t)
		s = t.TransitionTo
	}
	return ret
}
<|endoftext|>


<commit_before>package main

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

	"github.com/couchbaselabs/cbfs/client"
)

var cbfsUrlFlag = flag.String("cbfs", "http://cbfs:8484/", "URL to cbfs base")
var serieslyUrlFlag = flag.String("seriesly", "http://seriesly:3133/",
	"URL to seriesly base")
var pollFreq = flag.Duration("freq", 5*time.Second, "How often to poll cbfs")

var cbfsUrl, serieslyUrl *url.URL
var client *cbfsclient.Client

var nodeLock sync.Mutex
var nodes map[string]cbfsclient.StorageNode

func updateNodes() {
	nodeLock.Lock()
	defer nodeLock.Unlock()

	n, err := client.Nodes()
	lf := log.Printf
	if nodes == nil {
		lf = log.Fatalf
	}
	if err != nil {
		lf("Couldn't update/init nodes: %v", err)
		return
	}

	if len(n) != len(nodes) {
		createDatabases(n)
	}

	nodes = n
}

func createDatabases(m map[string]cbfsclient.StorageNode) {
	for k := range m {
		du := *serieslyUrl
		du.Path = "/" + k

		req, err := http.NewRequest("PUT", du.String(), nil)
		if err != nil {
			log.Fatalf("Error creating DB %q: %v", k, err)
		}

		res, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Printf("Error issuing HTTP request:  %v", err)
			return
		}
		defer res.Body.Close()

		if res.StatusCode != 201 {
			log.Printf("Error creating db: %v", res.Status)
		}
	}
}

func updateNodesLoop() {
	updateNodes()
	for _ = range time.Tick(time.Minute) {
		updateNodes()
	}
}

func httpCopy(dest, src string) error {
	sres, err := http.Get(src)
	if err != nil {
		return err
	}
	defer sres.Body.Close()
	if sres.StatusCode != 200 {
		return fmt.Errorf("HTTP error getting src data from %v: %v", src, sres.Status)
	}

	dres, err := http.Post(dest, sres.Header.Get("Content-Type"), sres.Body)
	if err != nil {
		return err
	}
	defer dres.Body.Close()

	if dres.StatusCode != 201 {
		errmsg, _ := ioutil.ReadAll(io.LimitReader(dres.Body, 512))
		return fmt.Errorf("HTTP Error posting result to %v: %v\n%s",
			dest, dres.StatusCode, errmsg)
	}
	return nil
}

func pollNode(name string, node cbfsclient.StorageNode, t time.Time) {
	du := *serieslyUrl
	du.RawQuery = "ts=" + strconv.FormatInt(t.UnixNano(), 10)
	du.Path = "/" + name

	if err := httpCopy(du.String(), node.URLFor("/.cbfs/debug/")); err != nil {
		log.Printf("Error copying data: %v", err)
	}
}

func poll(t time.Time) {
	for k, v := range nodes {
		go pollNode(k, v, t)
	}
}

func mustParseUrl(ustr string) *url.URL {
	u, e := url.Parse(ustr)
	if e != nil {
		log.Fatalf("Error parsing URL %q: %v", ustr, e)
	}
	return u
}

func parseUrls() {
	cbfsUrl = mustParseUrl(*cbfsUrlFlag)
	serieslyUrl = mustParseUrl(*serieslyUrlFlag)
}

func main() {
	flag.Parse()

	parseUrls()

	var err error
	client, err = cbfsclient.New(cbfsUrl.String())
	if err != nil {
		log.Fatalf("Can't instantiate cbfsclient: %v", err)
	}

	go updateNodesLoop()

	for t := range time.Tick(*pollFreq) {
		poll(t)
	}
}
<commit_msg>Disable keepalives.<commit_after>package main

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

	"github.com/couchbaselabs/cbfs/client"
)

var cbfsUrlFlag = flag.String("cbfs", "http://cbfs:8484/", "URL to cbfs base")
var serieslyUrlFlag = flag.String("seriesly", "http://seriesly:3133/",
	"URL to seriesly base")
var pollFreq = flag.Duration("freq", 5*time.Second, "How often to poll cbfs")

var cbfsUrl, serieslyUrl *url.URL
var client *cbfsclient.Client

var nodeLock sync.Mutex
var nodes map[string]cbfsclient.StorageNode

func updateNodes() {
	nodeLock.Lock()
	defer nodeLock.Unlock()

	n, err := client.Nodes()
	lf := log.Printf
	if nodes == nil {
		lf = log.Fatalf
	}
	if err != nil {
		lf("Couldn't update/init nodes: %v", err)
		return
	}

	if len(n) != len(nodes) {
		createDatabases(n)
	}

	nodes = n
}

func createDatabases(m map[string]cbfsclient.StorageNode) {
	for k := range m {
		du := *serieslyUrl
		du.Path = "/" + k

		req, err := http.NewRequest("PUT", du.String(), nil)
		if err != nil {
			log.Fatalf("Error creating DB %q: %v", k, err)
		}

		res, err := http.DefaultClient.Do(req)
		if err != nil {
			log.Printf("Error issuing HTTP request:  %v", err)
			return
		}
		defer res.Body.Close()

		if res.StatusCode != 201 {
			log.Printf("Error creating db: %v", res.Status)
		}
	}
}

func updateNodesLoop() {
	updateNodes()
	for _ = range time.Tick(time.Minute) {
		updateNodes()
	}
}

func httpCopy(dest, src string) error {
	sres, err := http.Get(src)
	if err != nil {
		return err
	}
	defer sres.Body.Close()
	if sres.StatusCode != 200 {
		return fmt.Errorf("HTTP error getting src data from %v: %v", src, sres.Status)
	}

	dres, err := http.Post(dest, sres.Header.Get("Content-Type"), sres.Body)
	if err != nil {
		return err
	}
	defer dres.Body.Close()

	if dres.StatusCode != 201 {
		errmsg, _ := ioutil.ReadAll(io.LimitReader(dres.Body, 512))
		return fmt.Errorf("HTTP Error posting result to %v: %v\n%s",
			dest, dres.StatusCode, errmsg)
	}
	return nil
}

func pollNode(name string, node cbfsclient.StorageNode, t time.Time) {
	du := *serieslyUrl
	du.RawQuery = "ts=" + strconv.FormatInt(t.UnixNano(), 10)
	du.Path = "/" + name

	if err := httpCopy(du.String(), node.URLFor("/.cbfs/debug/")); err != nil {
		log.Printf("Error copying data: %v", err)
	}
}

func poll(t time.Time) {
	for k, v := range nodes {
		go pollNode(k, v, t)
	}
}

func mustParseUrl(ustr string) *url.URL {
	u, e := url.Parse(ustr)
	if e != nil {
		log.Fatalf("Error parsing URL %q: %v", ustr, e)
	}
	return u
}

func parseUrls() {
	cbfsUrl = mustParseUrl(*cbfsUrlFlag)
	serieslyUrl = mustParseUrl(*serieslyUrlFlag)
}

func initHttp() {
	http.DefaultClient = &http.Client{
		Transport: &http.Transport{
			DisableKeepAlives:     true,
			ResponseHeaderTimeout: time.Millisecond * 100,
		},
	}
}

func main() {
	flag.Parse()

	initHttp()

	parseUrls()

	var err error
	client, err = cbfsclient.New(cbfsUrl.String())
	if err != nil {
		log.Fatalf("Can't instantiate cbfsclient: %v", err)
	}

	go updateNodesLoop()

	for t := range time.Tick(*pollFreq) {
		poll(t)
	}
}
<|endoftext|>


<commit_before>package ssdb

import (
	"bytes"
	"fmt"
	"net"
	"strconv"
)

type Client struct {
	sock     *net.TCPConn
	recv_buf bytes.Buffer
}

func Connect(ip string, port int) (*Client, error) {
	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}
	sock, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return nil, err
	}
	var c Client
	c.sock = sock
	return &c, nil
}

func (c *Client) Do(args ...interface{}) ([]string, error) {
	err := c.send(args)
	if err != nil {
		return nil, err
	}
	resp, err := c.recv()
	return resp, err
}

func (c *Client) Set(key string, val string) (interface{}, error) {
	resp, err := c.Do("set", key, val)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response")
}

// TODO: Will somebody write addition semantic methods?
func (c *Client) Get(key string) (interface{}, error) {
	resp, err := c.Do("get", key)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return resp[1], nil
	}
	if resp[0] == "not_found" {
		return nil, nil
	}
	return nil, fmt.Errorf("bad response")
}

func (c *Client) Del(key string) (interface{}, error) {
	resp, err := c.Do("del", key)
	if err != nil {
		return nil, err
	}

        //response looks like this: [ok 1]
	if len(resp) == 2 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response:resp:%v:", resp)
}

func (c *Client) send(args []interface{}) error {
	var buf bytes.Buffer
	for _, arg := range args {
		var s string
		switch arg := arg.(type) {
		case string:
			s = arg
		case []byte:
			s = string(arg)
		case []string:
			for _, s := range arg {
				buf.WriteString(fmt.Sprintf("%d", len(s)))
				buf.WriteByte('\n')
				buf.WriteString(s)
				buf.WriteByte('\n')
			}
			continue
		case int:
			s = fmt.Sprintf("%d", arg)
		case int64:
			s = fmt.Sprintf("%d", arg)
		case float64:
			s = fmt.Sprintf("%f", arg)
		case bool:
			if arg {
				s = "1"
			} else {
				s = "0"
			}
		case nil:
			s = ""
		default:
			return fmt.Errorf("bad arguments")
		}
		buf.WriteString(fmt.Sprintf("%d", len(s)))
		buf.WriteByte('\n')
		buf.WriteString(s)
		buf.WriteByte('\n')
	}
	buf.WriteByte('\n')
	_, err := c.sock.Write(buf.Bytes())
	return err
}

func (c *Client) recv() ([]string, error) {
	var tmp [8192]byte
	for {
		n, err := c.sock.Read(tmp[0:])
		if err != nil {
			return nil, err
		}
		c.recv_buf.Write(tmp[0:n])
		resp := c.parse()
		if resp == nil || len(resp) > 0 {
			return resp, nil
		}
	}
}

func (c *Client) parse() []string {
	resp := []string{}
	buf := c.recv_buf.Bytes()
	var idx, offset int
	idx = 0
	offset = 0

	for {
		idx = bytes.IndexByte(buf[offset:], '\n')
		if idx == -1 {
			break
		}
		p := buf[offset : offset+idx]
		offset += idx + 1
		//fmt.Printf("> [%s]\n", p);
		if len(p) == 0 || (len(p) == 1 && p[0] == '\r') {
			if len(resp) == 0 {
				continue
			} else {
				c.recv_buf.Next(offset)
				return resp
			}
		}

		size, err := strconv.Atoi(string(p))
		if err != nil || size < 0 {
			return nil
		}
		if offset+size >= c.recv_buf.Len() {
			break
		}

		v := buf[offset : offset+size]
		resp = append(resp, string(v))
		offset += size + 1
	}

	return []string{}
}

// Close The Client Connection
func (c *Client) Close() error {
	return c.sock.Close()
}
<commit_msg>check len(resp)>0 instead of len==2<commit_after>package ssdb

import (
	"bytes"
	"fmt"
	"net"
	"strconv"
)

type Client struct {
	sock     *net.TCPConn
	recv_buf bytes.Buffer
}

func Connect(ip string, port int) (*Client, error) {
	addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", ip, port))
	if err != nil {
		return nil, err
	}
	sock, err := net.DialTCP("tcp", nil, addr)
	if err != nil {
		return nil, err
	}
	var c Client
	c.sock = sock
	return &c, nil
}

func (c *Client) Do(args ...interface{}) ([]string, error) {
	err := c.send(args)
	if err != nil {
		return nil, err
	}
	resp, err := c.recv()
	return resp, err
}

func (c *Client) Set(key string, val string) (interface{}, error) {
	resp, err := c.Do("set", key, val)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response")
}

// TODO: Will somebody write addition semantic methods?
func (c *Client) Get(key string) (interface{}, error) {
	resp, err := c.Do("get", key)
	if err != nil {
		return nil, err
	}
	if len(resp) == 2 && resp[0] == "ok" {
		return resp[1], nil
	}
	if resp[0] == "not_found" {
		return nil, nil
	}
	return nil, fmt.Errorf("bad response")
}

func (c *Client) Del(key string) (interface{}, error) {
	resp, err := c.Do("del", key)
	if err != nil {
		return nil, err
	}

	//response looks like this: [ok 1]
	if len(resp) > 0 && resp[0] == "ok" {
		return true, nil
	}
	return nil, fmt.Errorf("bad response:resp:%v:", resp)
}

func (c *Client) send(args []interface{}) error {
	var buf bytes.Buffer
	for _, arg := range args {
		var s string
		switch arg := arg.(type) {
		case string:
			s = arg
		case []byte:
			s = string(arg)
		case []string:
			for _, s := range arg {
				buf.WriteString(fmt.Sprintf("%d", len(s)))
				buf.WriteByte('\n')
				buf.WriteString(s)
				buf.WriteByte('\n')
			}
			continue
		case int:
			s = fmt.Sprintf("%d", arg)
		case int64:
			s = fmt.Sprintf("%d", arg)
		case float64:
			s = fmt.Sprintf("%f", arg)
		case bool:
			if arg {
				s = "1"
			} else {
				s = "0"
			}
		case nil:
			s = ""
		default:
			return fmt.Errorf("bad arguments")
		}
		buf.WriteString(fmt.Sprintf("%d", len(s)))
		buf.WriteByte('\n')
		buf.WriteString(s)
		buf.WriteByte('\n')
	}
	buf.WriteByte('\n')
	_, err := c.sock.Write(buf.Bytes())
	return err
}

func (c *Client) recv() ([]string, error) {
	var tmp [8192]byte
	for {
		n, err := c.sock.Read(tmp[0:])
		if err != nil {
			return nil, err
		}
		c.recv_buf.Write(tmp[0:n])
		resp := c.parse()
		if resp == nil || len(resp) > 0 {
			return resp, nil
		}
	}
}

func (c *Client) parse() []string {
	resp := []string{}
	buf := c.recv_buf.Bytes()
	var idx, offset int
	idx = 0
	offset = 0

	for {
		idx = bytes.IndexByte(buf[offset:], '\n')
		if idx == -1 {
			break
		}
		p := buf[offset : offset+idx]
		offset += idx + 1
		//fmt.Printf("> [%s]\n", p);
		if len(p) == 0 || (len(p) == 1 && p[0] == '\r') {
			if len(resp) == 0 {
				continue
			} else {
				c.recv_buf.Next(offset)
				return resp
			}
		}

		size, err := strconv.Atoi(string(p))
		if err != nil || size < 0 {
			return nil
		}
		if offset+size >= c.recv_buf.Len() {
			break
		}

		v := buf[offset : offset+size]
		resp = append(resp, string(v))
		offset += size + 1
	}

	return []string{}
}

// Close The Client Connection
func (c *Client) Close() error {
	return c.sock.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"log"
	"os"

	_ "github.com/lib/pq"
)

// initDB connects to the DB and creates the tables if they don't exist
func initDB() *sql.DB {
	// Connect to DB
	dbInfo := os.Getenv("DATABASE_URL")
	if dbInfo == "" {
		dbInfo = "host=localhost port=5432 user=dev password=dev dbname=chat_dev sslmode=disable"
	}

	db, err := sql.Open("postgres", dbInfo)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	err = db.Ping()
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	log.Println("Connected to DB.")

	// Create tables if not exists
	_, err = db.Exec("CREATE EXTENSION IF NOT EXISTS \"pgcrypto\"")
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS messages (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			user_id VARCHAR(255),
			user_name VARCHAR(255),
			user_avatar VARCHAR(255),
			type VARCHAR(255),
			content TEXT,
			date_post TIMESTAMP
		)
	`)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	log.Println("Tables created or already existing.")

	// All good!
	return db
}

// insertMessage inserts a single message into the database
// and returns either the id or an error
func insertMessage(db *sql.DB, msg Message) (string, error) {
	stmt := `
		INSERT INTO messages (
			user_id,
			user_name,
			user_avatar,
			type,
			content,
			date_post
		)
		VALUES ($1, $2, $3, $4, $5, $6)
		RETURNING id
	`

	var id string

	err := db.QueryRow(stmt, msg.UserID, msg.UserName, msg.UserAvatar, msg.Type, msg.Content, msg.Date).Scan(&id)

	if err != nil {
		return "", err
	}
	return id, nil
}

func selectPreviousMessage(db *sql.DB, userID string) (*sql.Rows, error) {
	stmt := `
		(SELECT *
		FROM messages
		WHERE type = 'message'
		ORDER BY date_post DESC
		LIMIT 10)

		UNION

		(SELECT *
		FROM messages
		WHERE type = 'message'
			AND user_id != $1
			AND date_post > (
				SELECT date_post
				FROM messages
				WHERE type = 'logout'
					AND user_id = $1
				ORDER BY date_post DESC
				LIMIT 1
			)
		)
	`

	rows, err := db.Query(stmt, userID)
	if err != nil {
		return &sql.Rows{}, err
	}
	return rows, nil
}

func selectConnectedUsers(db *sql.DB, userID string) (*sql.Rows, error) {
	stmt := `
		SELECT DISTINCT m.user_id, m.user_name, m.user_avatar
		FROM messages m
			INNER JOIN

			(SELECT DISTINCT user_id, MAX(date_post)
			FROM messages
			WHERE type = 'login'
			GROUP BY user_id) users_login
			ON m.user_id = users_login.user_id

			LEFT JOIN

			(SELECT DISTINCT user_id, MAX(date_post)
			FROM messages
			WHERE type = 'logout'
			GROUP BY user_id) users_logout
			ON m.user_id = users_logout.user_id

		WHERE m.user_id != $1
			AND users_login.max > users_logout.max
					OR users_logout.max IS NULL
	`

	rows, err := db.Query(stmt, userID)
	if err != nil {
		return &sql.Rows{}, err
	}
	return rows, nil
}
<commit_msg>Remove duplicates in connected users<commit_after>package main

import (
	"database/sql"
	"log"
	"os"

	_ "github.com/lib/pq"
)

// initDB connects to the DB and creates the tables if they don't exist
func initDB() *sql.DB {
	// Connect to DB
	dbInfo := os.Getenv("DATABASE_URL")
	if dbInfo == "" {
		dbInfo = "host=localhost port=5432 user=dev password=dev dbname=chat_dev sslmode=disable"
	}

	db, err := sql.Open("postgres", dbInfo)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	err = db.Ping()
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	log.Println("Connected to DB.")

	// Create tables if not exists
	_, err = db.Exec("CREATE EXTENSION IF NOT EXISTS \"pgcrypto\"")
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS messages (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			user_id VARCHAR(255),
			user_name VARCHAR(255),
			user_avatar VARCHAR(255),
			type VARCHAR(255),
			content TEXT,
			date_post TIMESTAMP
		)
	`)
	if err != nil {
		log.Fatalf("Error: %v", err)
	}

	log.Println("Tables created or already existing.")

	// All good!
	return db
}

// insertMessage inserts a single message into the database
// and returns either the id or an error
func insertMessage(db *sql.DB, msg Message) (string, error) {
	stmt := `
		INSERT INTO messages (
			user_id,
			user_name,
			user_avatar,
			type,
			content,
			date_post
		)
		VALUES ($1, $2, $3, $4, $5, $6)
		RETURNING id
	`

	var id string

	err := db.QueryRow(stmt, msg.UserID, msg.UserName, msg.UserAvatar, msg.Type, msg.Content, msg.Date).Scan(&id)

	if err != nil {
		return "", err
	}
	return id, nil
}

func selectPreviousMessage(db *sql.DB, userID string) (*sql.Rows, error) {
	stmt := `
		(SELECT *
		FROM messages
		WHERE type = 'message'
		ORDER BY date_post DESC
		LIMIT 10)

		UNION

		(SELECT *
		FROM messages
		WHERE type = 'message'
			AND user_id != $1
			AND date_post > (
				SELECT date_post
				FROM messages
				WHERE type = 'logout'
					AND user_id = $1
				ORDER BY date_post DESC
				LIMIT 1
			)
		)
	`

	rows, err := db.Query(stmt, userID)
	if err != nil {
		return &sql.Rows{}, err
	}
	return rows, nil
}

func selectConnectedUsers(db *sql.DB, userID string) (*sql.Rows, error) {
	stmt := `
		SELECT DISTINCT m.user_id, users_login.user_name, users_login.user_avatar
		FROM messages m
			INNER JOIN

			(SELECT DISTINCT user_id, user_name, user_avatar, MAX(date_post)
			FROM messages
			WHERE type = 'login'
			GROUP BY user_id, user_name, user_avatar) users_login
			ON m.user_id = users_login.user_id

			LEFT JOIN

			(SELECT DISTINCT user_id, user_name, user_avatar, MAX(date_post)
			FROM messages
			WHERE type = 'logout'
			GROUP BY user_id, user_name, user_avatar) users_logout
			ON m.user_id = users_logout.user_id

		WHERE m.user_id != $1
			AND users_login.max > users_logout.max
					OR users_logout.max IS NULL

		ORDER BY users_login.user_name
	`

	rows, err := db.Query(stmt, userID)
	if err != nil {
		return &sql.Rows{}, err
	}
	return rows, nil
}
<|endoftext|>


<commit_before>package sql
<commit_msg>goinstall just doesnt work with cgo<commit_after><|endoftext|>


<commit_before>package state

import (
	"fmt"
	"testing"
)

var simpleFile = []byte(`
{"path": "/tmp/no-exist", "mode": 644}
`)

var simpleFileMeta = Metadata{
	Name: "Simple File",
	Type: "file",
	State: "rendered",
}

func TestStateFactory(t *testing.T) {
	state, err := StateFactory(simpleFileMeta, simpleFile)
	if err != nil {
		fmt.Println(err)
		t.Fail()
	}
	name := state.Meta().Name
	if name != "Simple File" {
		fmt.Printf("Did not load metadata correctly: %s", name)
		t.Fail()
	}
}
<commit_msg>Add source to simpleFile test<commit_after>package state

import (
	"fmt"
	"testing"
)

var simpleFile = []byte(`
{"path": "/tmp/no-exist", "mode": 644, "source": "git:///git@github.com:vektorlab/otter/README.md"}
`)

var simpleFileMeta = Metadata{
	Name: "Simple File",
	Type: "file",
	State: "rendered",
}

func TestStateFactory(t *testing.T) {
	state, err := StateFactory(simpleFileMeta, simpleFile)
	if err != nil {
		fmt.Println(err)
		t.Fail()
	}
	name := state.Meta().Name
	if name != "Simple File" {
		fmt.Printf("Did not load metadata correctly: %s", name)
		t.Fail()
	}
}
<|endoftext|>


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

package fs

import (
	"errors"
	"path"
)

// SkipDir is used as a return value from WalkDirFuncs to indicate that
// the directory named in the call is to be skipped. It is not returned
// as an error by any function.
var SkipDir = errors.New("skip this directory")

// WalkDirFunc is the type of the function called by WalkDir to visit
// each each file or directory.
//
// The path argument contains the argument to Walk as a prefix.
// That is, if Walk is called with root argument "dir" and finds a file
// named "a" in that directory, the walk function will be called with
// argument "dir/a".
//
// The directory and file are joined with Join, which may clean the
// directory name: if Walk is called with the root argument "x/../dir"
// and finds a file named "a" in that directory, the walk function will
// be called with argument "dir/a", not "x/../dir/a".
//
// The d argument is the fs.DirEntry for the named path.
//
// The error result returned by the function controls how WalkDir
// continues. If the function returns the special value SkipDir, WalkDir
// skips the current directory (path if d.IsDir() is true, otherwise
// path's parent directory). Otherwise, if the function returns a non-nil
// error, WalkDir stops entirely and returns that error.
//
// The err argument reports an error related to path, signaling that
// WalkDir will not walk into that directory. The function can decide how
// to handle that error; as described earlier, returning the error will
// cause WalkDir to stop walking the entire tree.
//
// WalkDir calls the function with a non-nil err argument in two cases.
//
// First, if the initial os.Lstat on the root directory fails, WalkDir
// calls the function with path set to root, d set to nil, and err set to
// the error from os.Lstat.
//
// Second, if a directory's ReadDir method fails, WalkDir calls the
// function with path set to the directory's path, d set to an
// fs.DirEntry describing the directory, and err set to the error from
// ReadDir. In this second case, the function is called twice with the
// path of the directory: the first call is before the directory read is
// attempted and has err set to nil, giving the function a chance to
// return SkipDir and avoid the ReadDir entirely. The second call is
// after a failed ReadDir and reports the error from ReadDir.
// (If ReadDir succeeds, there is no second call.)
//
// The differences between WalkDirFunc compared to filepath.WalkFunc are:
//
//   - The second argument has type fs.DirEntry instead of fs.FileInfo.
//   - The function is called before reading a directory, to allow SkipDir
//     to bypass the directory read entirely.
//   - If a directory read fails, the function is called a second time
//     for that directory to report the error.
//
type WalkDirFunc func(path string, entry DirEntry, err error) error

// walkDir recursively descends path, calling walkDirFn.
func walkDir(fsys FS, name string, d DirEntry, walkDirFn WalkDirFunc) error {
	if err := walkDirFn(name, d, nil); err != nil || !d.IsDir() {
		if err == SkipDir && d.IsDir() {
			// Successfully skipped directory.
			err = nil
		}
		return err
	}

	dirs, err := ReadDir(fsys, name)
	if err != nil {
		// Second call, to report ReadDir error.
		err = walkDirFn(name, d, err)
		if err != nil {
			return err
		}
	}

	for _, d1 := range dirs {
		name1 := path.Join(name, d1.Name())
		if err := walkDir(fsys, name1, d1, walkDirFn); err != nil {
			if err == SkipDir {
				break
			}
			return err
		}
	}
	return nil
}

// WalkDir walks the file tree rooted at root, calling fn for each file or
// directory in the tree, including root.
//
// All errors that arise visiting files and directories are filtered by fn:
// see the fs.WalkDirFunc documentation for details.
//
// The files are walked in lexical order, which makes the output deterministic
// but requires WalkDir to read an entire directory into memory before proceeding
// to walk that directory.
//
// WalkDir does not follow symbolic links found in directories,
// but if root itself is a symbolic link, its target will be walked.
func WalkDir(fsys FS, root string, fn WalkDirFunc) error {
	info, err := Stat(fsys, root)
	if err != nil {
		err = fn(root, nil, err)
	} else {
		err = walkDir(fsys, root, &statDirEntry{info}, fn)
	}
	if err == SkipDir {
		return nil
	}
	return err
}

type statDirEntry struct {
	info FileInfo
}

func (d *statDirEntry) Name() string            { return d.info.Name() }
func (d *statDirEntry) IsDir() bool             { return d.info.IsDir() }
func (d *statDirEntry) Type() FileMode          { return d.info.Mode().Type() }
func (d *statDirEntry) Info() (FileInfo, error) { return d.info, nil }
<commit_msg>io/fs: make WalkDirFunc parameter name consistent with doc comment<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fs

import (
	"errors"
	"path"
)

// SkipDir is used as a return value from WalkDirFuncs to indicate that
// the directory named in the call is to be skipped. It is not returned
// as an error by any function.
var SkipDir = errors.New("skip this directory")

// WalkDirFunc is the type of the function called by WalkDir to visit
// each each file or directory.
//
// The path argument contains the argument to Walk as a prefix.
// That is, if Walk is called with root argument "dir" and finds a file
// named "a" in that directory, the walk function will be called with
// argument "dir/a".
//
// The directory and file are joined with Join, which may clean the
// directory name: if Walk is called with the root argument "x/../dir"
// and finds a file named "a" in that directory, the walk function will
// be called with argument "dir/a", not "x/../dir/a".
//
// The d argument is the fs.DirEntry for the named path.
//
// The error result returned by the function controls how WalkDir
// continues. If the function returns the special value SkipDir, WalkDir
// skips the current directory (path if d.IsDir() is true, otherwise
// path's parent directory). Otherwise, if the function returns a non-nil
// error, WalkDir stops entirely and returns that error.
//
// The err argument reports an error related to path, signaling that
// WalkDir will not walk into that directory. The function can decide how
// to handle that error; as described earlier, returning the error will
// cause WalkDir to stop walking the entire tree.
//
// WalkDir calls the function with a non-nil err argument in two cases.
//
// First, if the initial os.Lstat on the root directory fails, WalkDir
// calls the function with path set to root, d set to nil, and err set to
// the error from os.Lstat.
//
// Second, if a directory's ReadDir method fails, WalkDir calls the
// function with path set to the directory's path, d set to an
// fs.DirEntry describing the directory, and err set to the error from
// ReadDir. In this second case, the function is called twice with the
// path of the directory: the first call is before the directory read is
// attempted and has err set to nil, giving the function a chance to
// return SkipDir and avoid the ReadDir entirely. The second call is
// after a failed ReadDir and reports the error from ReadDir.
// (If ReadDir succeeds, there is no second call.)
//
// The differences between WalkDirFunc compared to filepath.WalkFunc are:
//
//   - The second argument has type fs.DirEntry instead of fs.FileInfo.
//   - The function is called before reading a directory, to allow SkipDir
//     to bypass the directory read entirely.
//   - If a directory read fails, the function is called a second time
//     for that directory to report the error.
//
type WalkDirFunc func(path string, d DirEntry, err error) error

// walkDir recursively descends path, calling walkDirFn.
func walkDir(fsys FS, name string, d DirEntry, walkDirFn WalkDirFunc) error {
	if err := walkDirFn(name, d, nil); err != nil || !d.IsDir() {
		if err == SkipDir && d.IsDir() {
			// Successfully skipped directory.
			err = nil
		}
		return err
	}

	dirs, err := ReadDir(fsys, name)
	if err != nil {
		// Second call, to report ReadDir error.
		err = walkDirFn(name, d, err)
		if err != nil {
			return err
		}
	}

	for _, d1 := range dirs {
		name1 := path.Join(name, d1.Name())
		if err := walkDir(fsys, name1, d1, walkDirFn); err != nil {
			if err == SkipDir {
				break
			}
			return err
		}
	}
	return nil
}

// WalkDir walks the file tree rooted at root, calling fn for each file or
// directory in the tree, including root.
//
// All errors that arise visiting files and directories are filtered by fn:
// see the fs.WalkDirFunc documentation for details.
//
// The files are walked in lexical order, which makes the output deterministic
// but requires WalkDir to read an entire directory into memory before proceeding
// to walk that directory.
//
// WalkDir does not follow symbolic links found in directories,
// but if root itself is a symbolic link, its target will be walked.
func WalkDir(fsys FS, root string, fn WalkDirFunc) error {
	info, err := Stat(fsys, root)
	if err != nil {
		err = fn(root, nil, err)
	} else {
		err = walkDir(fsys, root, &statDirEntry{info}, fn)
	}
	if err == SkipDir {
		return nil
	}
	return err
}

type statDirEntry struct {
	info FileInfo
}

func (d *statDirEntry) Name() string            { return d.info.Name() }
func (d *statDirEntry) IsDir() bool             { return d.info.IsDir() }
func (d *statDirEntry) Type() FileMode          { return d.info.Mode().Type() }
func (d *statDirEntry) Info() (FileInfo, error) { return d.info, nil }
<|endoftext|>


<commit_before>package connections

import (
	"errors"
	"sync"

	"github.com/sirupsen/logrus"
)

const firstGroupId = 1

type ConnectionGroupManager struct {
	groups      map[int]*ConnectionGroup
	groupsMutex *sync.RWMutex
	groupLimit  int
	connsLimit  int
	connsCount  int
	logger      logrus.FieldLogger
}

func NewConnectionGroupManager(logger logrus.FieldLogger, groupLimit, connsLimit int) (*ConnectionGroupManager, error) {
	if groupLimit > 0 {
		return &ConnectionGroupManager{
			groups:      map[int]*ConnectionGroup{},
			groupsMutex: &sync.RWMutex{},
			groupLimit:  groupLimit,
			connsLimit:  connsLimit,
			logger:      logger,
		}, nil
	}

	return nil, errors.New("cannot create connection group manager: invalid group limit")
}

func (m *ConnectionGroupManager) unsafeIsFull() bool {
	return len(m.groups) == m.groupLimit
}

func (m *ConnectionGroupManager) IsFull() bool {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	return m.unsafeIsFull()
}

type ErrAddGroup string

func (e ErrAddGroup) Error() string {
	return "cannot add group: " + string(e)
}

var (
	ErrGroupLimitReached = ErrAddGroup("limit group count reached")
	ErrCannotGetID       = ErrAddGroup("cannot get id for group")
	ErrConnsLimitReached = ErrAddGroup("cannot reserve connections for group: connections count reached")
)

func (m *ConnectionGroupManager) Add(group *ConnectionGroup) (int, error) {
	m.groupsMutex.Lock()
	defer m.groupsMutex.Unlock()

	if m.unsafeIsFull() {
		return 0, ErrGroupLimitReached
	}

	if group.GetLimit() > m.connsLimit-m.connsCount {
		if m.connsLimit-m.connsCount < 1 {
			return 0, ErrConnsLimitReached
		}
		group.SetLimit(m.connsLimit - m.connsCount)
	}

	m.connsCount += group.GetLimit()

	for id := firstGroupId; id <= len(m.groups)+firstGroupId; id++ {
		if _, occupied := m.groups[id]; !occupied {
			m.groups[id] = group
			return id, nil
		}
	}

	return 0, ErrCannotGetID
}

type ErrDeleteGroup string

func (e ErrDeleteGroup) Error() string {
	return "cannot delete group: " + string(e)
}

var (
	ErrDeleteNotEmptyGroup = ErrDeleteGroup("group is not empty")
	ErrDeleteNotFoundGroup = ErrDeleteGroup("group not found")
)

func (m *ConnectionGroupManager) Delete(group *ConnectionGroup) error {
	m.groupsMutex.Lock()
	defer m.groupsMutex.Unlock()

	if !group.IsEmpty() {
		return ErrDeleteNotEmptyGroup
	}

	for id := range m.groups {
		if m.groups[id] == group {
			delete(m.groups, id)
			m.connsCount -= group.GetLimit()

			return nil
		}
	}

	return ErrDeleteNotFoundGroup
}

var ErrNotFoundGroup = errors.New("not found group")

func (m *ConnectionGroupManager) Get(id int) (*ConnectionGroup, error) {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()

	if group, ok := m.groups[id]; ok {
		return group, nil
	}

	return nil, ErrNotFoundGroup
}

func (m *ConnectionGroupManager) Groups() map[int]*ConnectionGroup {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	groups := map[int]*ConnectionGroup{}
	for id, group := range m.groups {
		groups[id] = group
	}
	return groups
}

func (m *ConnectionGroupManager) GroupLimit() int {
	return m.groupLimit
}

func (m *ConnectionGroupManager) GroupCount() int {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	return len(m.groups)
}

func (m *ConnectionGroupManager) Capacity() float32 {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()

	var count = 0
	for _, group := range m.groups {
		count += group.GetCount()
	}

	return float32(count) / float32(m.connsLimit)
}
<commit_msg>Create todos marks in connection group manager method Add<commit_after>package connections

import (
	"errors"
	"sync"

	"github.com/sirupsen/logrus"
)

const firstGroupId = 1

type ConnectionGroupManager struct {
	groups      map[int]*ConnectionGroup
	groupsMutex *sync.RWMutex
	groupLimit  int
	connsLimit  int
	connsCount  int
	logger      logrus.FieldLogger
}

func NewConnectionGroupManager(logger logrus.FieldLogger, groupLimit, connsLimit int) (*ConnectionGroupManager, error) {
	if groupLimit > 0 {
		return &ConnectionGroupManager{
			groups:      map[int]*ConnectionGroup{},
			groupsMutex: &sync.RWMutex{},
			groupLimit:  groupLimit,
			connsLimit:  connsLimit,
			logger:      logger,
		}, nil
	}

	return nil, errors.New("cannot create connection group manager: invalid group limit")
}

func (m *ConnectionGroupManager) unsafeIsFull() bool {
	return len(m.groups) == m.groupLimit
}

func (m *ConnectionGroupManager) IsFull() bool {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	return m.unsafeIsFull()
}

type ErrAddGroup string

func (e ErrAddGroup) Error() string {
	return "cannot add group: " + string(e)
}

var (
	ErrGroupLimitReached = ErrAddGroup("limit group count reached")
	ErrCannotGetID       = ErrAddGroup("cannot get id for group")
	ErrConnsLimitReached = ErrAddGroup("cannot reserve connections for group: connections count reached")
)

func (m *ConnectionGroupManager) Add(group *ConnectionGroup) (int, error) {
	// TODO: Fix method to receive group and required conn limit.

	// TODO: Fix method to return (id int, count int, err error), where
	// id is group identifier, count is reserved connection count for the
	// group, and err is error if occurred.

	m.groupsMutex.Lock()
	defer m.groupsMutex.Unlock()

	if m.unsafeIsFull() {
		return 0, ErrGroupLimitReached
	}

	if group.GetLimit() > m.connsLimit-m.connsCount {
		if m.connsLimit-m.connsCount < 1 {
			return 0, ErrConnsLimitReached
		}
		group.SetLimit(m.connsLimit - m.connsCount)
	}

	m.connsCount += group.GetLimit()

	for id := firstGroupId; id <= len(m.groups)+firstGroupId; id++ {
		if _, occupied := m.groups[id]; !occupied {
			m.groups[id] = group
			return id, nil
		}
	}

	return 0, ErrCannotGetID
}

type ErrDeleteGroup string

func (e ErrDeleteGroup) Error() string {
	return "cannot delete group: " + string(e)
}

var (
	ErrDeleteNotEmptyGroup = ErrDeleteGroup("group is not empty")
	ErrDeleteNotFoundGroup = ErrDeleteGroup("group not found")
)

func (m *ConnectionGroupManager) Delete(group *ConnectionGroup) error {
	m.groupsMutex.Lock()
	defer m.groupsMutex.Unlock()

	if !group.IsEmpty() {
		return ErrDeleteNotEmptyGroup
	}

	for id := range m.groups {
		if m.groups[id] == group {
			delete(m.groups, id)
			m.connsCount -= group.GetLimit()

			return nil
		}
	}

	return ErrDeleteNotFoundGroup
}

var ErrNotFoundGroup = errors.New("not found group")

func (m *ConnectionGroupManager) Get(id int) (*ConnectionGroup, error) {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()

	if group, ok := m.groups[id]; ok {
		return group, nil
	}

	return nil, ErrNotFoundGroup
}

func (m *ConnectionGroupManager) Groups() map[int]*ConnectionGroup {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	groups := map[int]*ConnectionGroup{}
	for id, group := range m.groups {
		groups[id] = group
	}
	return groups
}

func (m *ConnectionGroupManager) GroupLimit() int {
	return m.groupLimit
}

func (m *ConnectionGroupManager) GroupCount() int {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()
	return len(m.groups)
}

func (m *ConnectionGroupManager) Capacity() float32 {
	m.groupsMutex.RLock()
	defer m.groupsMutex.RUnlock()

	var count = 0
	for _, group := range m.groups {
		count += group.GetCount()
	}

	return float32(count) / float32(m.connsLimit)
}
<|endoftext|>


<commit_before>//  Copyright 2018 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 mixer

import (
	"context"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"fortio.org/fortio/fhttp"
	"fortio.org/fortio/periodic"

	"istio.io/istio/pkg/test/framework/components/ingress"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/prometheus"
	"istio.io/istio/pkg/test/shell"
	"istio.io/istio/tests/util"
)

const (
	destLabel         = "destination_service"
	responseCodeLabel = "response_code"
	reporterLabel     = "reporter"
)

func GetDestinationLabel() string {
	return destLabel
}

func GetResponseCodeLabel() string {
	return responseCodeLabel
}

func GetReporterCodeLabel() string {
	return reporterLabel
}

func VisitProductPage(ing ingress.Instance, timeout time.Duration, wantStatus int, t *testing.T) error {
	start := time.Now()
	endpointIP := ing.HTTPAddress()
	for {
		response, err := ing.Call(ingress.CallOptions{
			Host:     "",
			Path:     "/productpage",
			CallType: ingress.PlainText,
			Address:  endpointIP})
		if err != nil {
			t.Logf("Unable to connect to product page: %v", err)
		}

		status := response.Code
		if status == wantStatus {
			t.Logf("Got %d response from product page!", wantStatus)
			return nil
		}

		if time.Since(start) > timeout {
			return fmt.Errorf("could not retrieve product page in %v: Last status: %v", timeout, status)
		}

		time.Sleep(3 * time.Second)
	}
}

func ValidateMetric(t *testing.T, prometheus prometheus.Instance, query, metricName string, want float64) {
	got, err := getMetric(t, prometheus, query, metricName)
	if err != nil {
		t.Fatalf("%v", err)
	}
	t.Logf("%s: %f", metricName, got)
	if got < want {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Errorf("bad metric value: got %f, want at least %f", got, want)
	}
}

func getMetric(t *testing.T, prometheus prometheus.Instance, query, metricName string) (float64, error) {
	t.Helper()

	t.Logf("prometheus query: %s", query)
	value, err := prometheus.WaitForQuiesce(query)
	if err != nil {
		return 0, fmt.Errorf("could not get metrics from prometheus: %v", err)
	}

	got, err := prometheus.Sum(value, nil)
	if err != nil {
		t.Logf("value: %s", value.String())
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		return 0, fmt.Errorf("could not find metric value: %v", err)
	}

	return got, nil
}

func Fqdn(service, namespace string) string {
	return fmt.Sprintf("%s.%s.svc.cluster.local", service, namespace)
}

// promDump gets all of the recorded values for a metric by name and generates a report of the values.
// used for debugging of failures to provide a comprehensive view of traffic experienced.
func PromDump(prometheus prometheus.Instance, metric string) string {
	return PromDumpWithAttributes(prometheus, metric, nil)
}

// promDumpWithAttributes is used to get all of the recorded values of a metric for particular attributes.
// Attributes have to be of format %s=\"%s\"
// nolint: unparam
func PromDumpWithAttributes(prometheus prometheus.Instance, metric string, attributes []string) string {
	if value, err := prometheus.WaitForQuiesce(fmt.Sprintf("%s{%s}", metric, strings.Join(attributes, ", "))); err == nil {
		return value.String()
	}

	return ""
}

func SendTraffic(ingress ingress.Instance, t *testing.T, msg, url, extraHeader string, calls int64) *fhttp.HTTPRunnerResults {
	t.Log(msg)
	if url == "" {
		url = fmt.Sprintf("%s/productpage", ingress.HTTPAddress())
	}

	// run at a high enough QPS (here 10) to ensure that enough
	// traffic is generated to trigger 429s from the 1 QPS rate limit rule
	opts := fhttp.HTTPRunnerOptions{
		RunnerOptions: periodic.RunnerOptions{
			QPS:        10,
			Exactly:    calls,     // will make exactly 300 calls, so run for about 30 seconds
			NumThreads: 5,         // get the same number of calls per connection (300/5=60)
			Out:        os.Stderr, // Only needed because of log capture issue
		},
		HTTPOptions: fhttp.HTTPOptions{
			URL: url,
		},
	}
	if extraHeader != "" {
		opts.HTTPOptions.AddAndValidateExtraHeader(extraHeader)
	}
	// productpage should still return 200s when ratings is rate-limited.
	res, err := fhttp.RunHTTPTest(&opts)
	if err != nil {
		t.Fatalf("Generating traffic via fortio failed: %v", err)
	}
	return res
}

func SendTrafficAndWaitForExpectedStatus(ingress ingress.Instance, t *testing.T, msg, url string, calls int64,
	httpStatusCode int) {
	retry := util.Retrier{
		BaseDelay: 15 * time.Second,
		Retries:   3,
	}

	retryFn := func(_ context.Context, i int) error {
		res := SendTraffic(ingress, t, msg, url, "", calls)
		// Verify you get specified http return code.
		if float64(res.RetCodes[httpStatusCode]) == 0 {
			return fmt.Errorf("could not get %v status", httpStatusCode)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Fatalf("Failed with err: %v", err)
	}
}

func GetAndValidateAccessLog(ns namespace.Instance, t *testing.T, labelSelector, container string, validate func(string) error) {
	retry := util.Retrier{
		BaseDelay: 15 * time.Second,
		Retries:   3,
		MaxDelay:  30 * time.Second,
	}

	retryFn := func(_ context.Context, i int) error {
		content, err := shell.Execute(false, "kubectl logs -n %s -l %s -c %s --tail=-1",
			ns.Name(), labelSelector, container)
		if err != nil {
			return fmt.Errorf("unable to get access logs from mixer: %v , content %v", err, content)
		}
		err = validate(content)
		if err != nil {
			return fmt.Errorf("error validating content %v ", err)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Fatalf("Failed with err: %v", err)
	}
}

func FetchRequestCount(t *testing.T, prometheus prometheus.Instance, service, additionalLabels, namespace string,
	totalReqExpected float64) (prior429s float64, prior200s float64) {
	var err error
	t.Log("Establishing metrics baseline for test...")

	retry := util.Retrier{
		BaseDelay: 30 * time.Second,
		Retries:   2,
	}
	metricName := "istio_requests_total"

	retryFn := func(_ context.Context, i int) error {
		t.Helper()
		t.Logf("Trying to find metrics via promql (attempt %d)...", i)
		query := fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace), reporterLabel, "destination", additionalLabels)
		totalReq, err := getMetric(t, prometheus, query, "istio_requests_total")
		t.Logf("Expected Req: %v Got: %v", totalReqExpected, totalReq)
		if totalReqExpected == float64(0) && totalReq == float64(0) {
			return fmt.Errorf("returning 0. totalReqExpected : %v", totalReqExpected)
		}
		if err != nil {
			return err
		}
		if totalReq < totalReqExpected {
			return fmt.Errorf("total Requests: %f less than expected: %f", totalReq, totalReqExpected)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("could not find istio_requests_total  (msg: %v)", err)
		return 0, 0
	}

	query := fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace),
		reporterLabel, "destination", responseCodeLabel, "429", additionalLabels)
	prior429s, err = getMetric(t, prometheus, query, "istio_requests_total")
	if err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("error getting prior 429s, using 0 as value (msg: %v)", err)
		prior429s = 0
	}

	query = fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace),
		reporterLabel, "destination", responseCodeLabel, "200", additionalLabels)
	prior200s, err = getMetric(t, prometheus, query, "istio_requests_total")
	if err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("error getting prior 200s, using 0 as value (msg: %v)", err)
		prior200s = 0
	}
	t.Logf("Baseline established: prior200s = %f, prior429s = %f", prior200s, prior429s)

	return prior429s, prior200s
}

func AllowRuleSync(t *testing.T) {
	t.Log("Sleeping to allow rules to take effect...")
	time.Sleep(15 * time.Second)
}
<commit_msg>Set tail to high number for access log test (#15395)<commit_after>//  Copyright 2018 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 mixer

import (
	"context"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

	"fortio.org/fortio/fhttp"
	"fortio.org/fortio/periodic"

	"istio.io/istio/pkg/test/framework/components/ingress"
	"istio.io/istio/pkg/test/framework/components/namespace"
	"istio.io/istio/pkg/test/framework/components/prometheus"
	"istio.io/istio/pkg/test/shell"
	"istio.io/istio/tests/util"
)

const (
	destLabel         = "destination_service"
	responseCodeLabel = "response_code"
	reporterLabel     = "reporter"
)

func GetDestinationLabel() string {
	return destLabel
}

func GetResponseCodeLabel() string {
	return responseCodeLabel
}

func GetReporterCodeLabel() string {
	return reporterLabel
}

func VisitProductPage(ing ingress.Instance, timeout time.Duration, wantStatus int, t *testing.T) error {
	start := time.Now()
	endpointIP := ing.HTTPAddress()
	for {
		response, err := ing.Call(ingress.CallOptions{
			Host:     "",
			Path:     "/productpage",
			CallType: ingress.PlainText,
			Address:  endpointIP})
		if err != nil {
			t.Logf("Unable to connect to product page: %v", err)
		}

		status := response.Code
		if status == wantStatus {
			t.Logf("Got %d response from product page!", wantStatus)
			return nil
		}

		if time.Since(start) > timeout {
			return fmt.Errorf("could not retrieve product page in %v: Last status: %v", timeout, status)
		}

		time.Sleep(3 * time.Second)
	}
}

func ValidateMetric(t *testing.T, prometheus prometheus.Instance, query, metricName string, want float64) {
	got, err := getMetric(t, prometheus, query, metricName)
	if err != nil {
		t.Fatalf("%v", err)
	}
	t.Logf("%s: %f", metricName, got)
	if got < want {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Errorf("bad metric value: got %f, want at least %f", got, want)
	}
}

func getMetric(t *testing.T, prometheus prometheus.Instance, query, metricName string) (float64, error) {
	t.Helper()

	t.Logf("prometheus query: %s", query)
	value, err := prometheus.WaitForQuiesce(query)
	if err != nil {
		return 0, fmt.Errorf("could not get metrics from prometheus: %v", err)
	}

	got, err := prometheus.Sum(value, nil)
	if err != nil {
		t.Logf("value: %s", value.String())
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		return 0, fmt.Errorf("could not find metric value: %v", err)
	}

	return got, nil
}

func Fqdn(service, namespace string) string {
	return fmt.Sprintf("%s.%s.svc.cluster.local", service, namespace)
}

// promDump gets all of the recorded values for a metric by name and generates a report of the values.
// used for debugging of failures to provide a comprehensive view of traffic experienced.
func PromDump(prometheus prometheus.Instance, metric string) string {
	return PromDumpWithAttributes(prometheus, metric, nil)
}

// promDumpWithAttributes is used to get all of the recorded values of a metric for particular attributes.
// Attributes have to be of format %s=\"%s\"
// nolint: unparam
func PromDumpWithAttributes(prometheus prometheus.Instance, metric string, attributes []string) string {
	if value, err := prometheus.WaitForQuiesce(fmt.Sprintf("%s{%s}", metric, strings.Join(attributes, ", "))); err == nil {
		return value.String()
	}

	return ""
}

func SendTraffic(ingress ingress.Instance, t *testing.T, msg, url, extraHeader string, calls int64) *fhttp.HTTPRunnerResults {
	t.Log(msg)
	if url == "" {
		url = fmt.Sprintf("%s/productpage", ingress.HTTPAddress())
	}

	// run at a high enough QPS (here 10) to ensure that enough
	// traffic is generated to trigger 429s from the 1 QPS rate limit rule
	opts := fhttp.HTTPRunnerOptions{
		RunnerOptions: periodic.RunnerOptions{
			QPS:        10,
			Exactly:    calls,     // will make exactly 300 calls, so run for about 30 seconds
			NumThreads: 5,         // get the same number of calls per connection (300/5=60)
			Out:        os.Stderr, // Only needed because of log capture issue
		},
		HTTPOptions: fhttp.HTTPOptions{
			URL: url,
		},
	}
	if extraHeader != "" {
		opts.HTTPOptions.AddAndValidateExtraHeader(extraHeader)
	}
	// productpage should still return 200s when ratings is rate-limited.
	res, err := fhttp.RunHTTPTest(&opts)
	if err != nil {
		t.Fatalf("Generating traffic via fortio failed: %v", err)
	}
	return res
}

func SendTrafficAndWaitForExpectedStatus(ingress ingress.Instance, t *testing.T, msg, url string, calls int64,
	httpStatusCode int) {
	retry := util.Retrier{
		BaseDelay: 15 * time.Second,
		Retries:   3,
	}

	retryFn := func(_ context.Context, i int) error {
		res := SendTraffic(ingress, t, msg, url, "", calls)
		// Verify you get specified http return code.
		if float64(res.RetCodes[httpStatusCode]) == 0 {
			return fmt.Errorf("could not get %v status", httpStatusCode)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Fatalf("Failed with err: %v", err)
	}
}

func GetAndValidateAccessLog(ns namespace.Instance, t *testing.T, labelSelector, container string, validate func(string) error) {
	retry := util.Retrier{
		BaseDelay: 15 * time.Second,
		Retries:   3,
		MaxDelay:  30 * time.Second,
	}

	retryFn := func(_ context.Context, i int) error {
		// Different kubectl versions seem to return different amounts of logs. To ensure we get them all, set tail to a large number
		content, err := shell.Execute(false, "kubectl logs -n %s -l %s -c %s --tail=10000000",
			ns.Name(), labelSelector, container)
		if err != nil {
			return fmt.Errorf("unable to get access logs from mixer: %v , content %v", err, content)
		}
		err = validate(content)
		if err != nil {
			return fmt.Errorf("error validating content %v ", err)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Fatalf("Failed with err: %v", err)
	}
}

func FetchRequestCount(t *testing.T, prometheus prometheus.Instance, service, additionalLabels, namespace string,
	totalReqExpected float64) (prior429s float64, prior200s float64) {
	var err error
	t.Log("Establishing metrics baseline for test...")

	retry := util.Retrier{
		BaseDelay: 30 * time.Second,
		Retries:   2,
	}
	metricName := "istio_requests_total"

	retryFn := func(_ context.Context, i int) error {
		t.Helper()
		t.Logf("Trying to find metrics via promql (attempt %d)...", i)
		query := fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace), reporterLabel, "destination", additionalLabels)
		totalReq, err := getMetric(t, prometheus, query, "istio_requests_total")
		t.Logf("Expected Req: %v Got: %v", totalReqExpected, totalReq)
		if totalReqExpected == float64(0) && totalReq == float64(0) {
			return fmt.Errorf("returning 0. totalReqExpected : %v", totalReqExpected)
		}
		if err != nil {
			return err
		}
		if totalReq < totalReqExpected {
			return fmt.Errorf("total Requests: %f less than expected: %f", totalReq, totalReqExpected)
		}
		return nil
	}

	if _, err := retry.Retry(context.Background(), retryFn); err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("could not find istio_requests_total  (msg: %v)", err)
		return 0, 0
	}

	query := fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace),
		reporterLabel, "destination", responseCodeLabel, "429", additionalLabels)
	prior429s, err = getMetric(t, prometheus, query, "istio_requests_total")
	if err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("error getting prior 429s, using 0 as value (msg: %v)", err)
		prior429s = 0
	}

	query = fmt.Sprintf("istio_requests_total{%s=\"%s\",%s=\"%s\",%s=\"%s\",%s}", destLabel, Fqdn(service, namespace),
		reporterLabel, "destination", responseCodeLabel, "200", additionalLabels)
	prior200s, err = getMetric(t, prometheus, query, "istio_requests_total")
	if err != nil {
		t.Logf("prometheus values for %s:\n%s", metricName, PromDump(prometheus, metricName))
		t.Logf("error getting prior 200s, using 0 as value (msg: %v)", err)
		prior200s = 0
	}
	t.Logf("Baseline established: prior200s = %f, prior429s = %f", prior200s, prior429s)

	return prior429s, prior200s
}

func AllowRuleSync(t *testing.T) {
	t.Log("Sleeping to allow rules to take effect...")
	time.Sleep(15 * time.Second)
}
<|endoftext|>