<commit_before>// Copyright 2017 The etcd 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 integration

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"testing"
	"time"

	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.etcd.io/etcd/pkg/v3/transport"
	"go.etcd.io/etcd/server/v3/etcdserver"
)

// TestMetricDbSizeBoot checks that the db size metric is set on boot.
func TestMetricDbSizeBoot(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	v, err := clus.Members[0].Metric("etcd_debugging_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}

	if v == "0" {
		t.Fatalf("expected non-zero, got %q", v)
	}
}

func TestMetricDbSizeDefrag(t *testing.T) {
	testMetricDbSizeDefrag(t, "etcd")
}

func TestMetricDbSizeDefragDebugging(t *testing.T) {
	testMetricDbSizeDefrag(t, "etcd_debugging")
}

// testMetricDbSizeDefrag checks that the db size metric is set after defrag.
func testMetricDbSizeDefrag(t *testing.T, name string) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	kvc := toGRPC(clus.Client(0)).KV
	mc := toGRPC(clus.Client(0)).Maintenance

	// expand the db size
	numPuts := 25 // large enough to write more than 1 page
	putreq := &pb.PutRequest{Key: []byte("k"), Value: make([]byte, 4096)}
	for i := 0; i < numPuts; i++ {
		time.Sleep(10 * time.Millisecond) // to execute multiple backend txn
		if _, err := kvc.Put(context.TODO(), putreq); err != nil {
			t.Fatal(err)
		}
	}

	// wait for backend txn sync
	time.Sleep(500 * time.Millisecond)

	expected := numPuts * len(putreq.Value)
	beforeDefrag, err := clus.Members[0].Metric(name + "_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	bv, err := strconv.Atoi(beforeDefrag)
	if err != nil {
		t.Fatal(err)
	}
	if bv < expected {
		t.Fatalf("expected db size greater than %d, got %d", expected, bv)
	}
	beforeDefragInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	biu, err := strconv.Atoi(beforeDefragInUse)
	if err != nil {
		t.Fatal(err)
	}
	if biu < expected {
		t.Fatalf("expected db size in use is greater than %d, got %d", expected, biu)
	}

	// clear out historical keys, in use bytes should free pages
	creq := &pb.CompactionRequest{Revision: int64(numPuts), Physical: true}
	if _, kerr := kvc.Compact(context.TODO(), creq); kerr != nil {
		t.Fatal(kerr)
	}

	validateAfterCompactionInUse := func() error {
		// Put to move PendingPages to FreePages
		if _, err = kvc.Put(context.TODO(), putreq); err != nil {
			t.Fatal(err)
		}
		time.Sleep(500 * time.Millisecond)

		afterCompactionInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
		if err != nil {
			t.Fatal(err)
		}
		aciu, err := strconv.Atoi(afterCompactionInUse)
		if err != nil {
			t.Fatal(err)
		}
		if biu <= aciu {
			return fmt.Errorf("expected less than %d, got %d after compaction", biu, aciu)
		}
		return nil
	}

	// backend rollbacks read transaction asynchronously (PR #10523),
	// which causes the result to be flaky. Retry 3 times.
	maxRetry, retry := 3, 0
	for {
		err := validateAfterCompactionInUse()
		if err == nil {
			break
		}
		retry++
		if retry >= maxRetry {
			t.Fatalf(err.Error())
		}
	}

	// defrag should give freed space back to fs
	mc.Defragment(context.TODO(), &pb.DefragmentRequest{})

	afterDefrag, err := clus.Members[0].Metric(name + "_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	av, err := strconv.Atoi(afterDefrag)
	if err != nil {
		t.Fatal(err)
	}
	if bv <= av {
		t.Fatalf("expected less than %d, got %d after defrag", bv, av)
	}

	afterDefragInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	adiu, err := strconv.Atoi(afterDefragInUse)
	if err != nil {
		t.Fatal(err)
	}
	if adiu > av {
		t.Fatalf("db size in use (%d) is expected less than db size (%d) after defrag", adiu, av)
	}
}

func TestMetricQuotaBackendBytes(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	qs, err := clus.Members[0].Metric("etcd_server_quota_backend_bytes")
	if err != nil {
		t.Fatal(err)
	}
	qv, err := strconv.ParseFloat(qs, 64)
	if err != nil {
		t.Fatal(err)
	}
	if int64(qv) != etcdserver.DefaultQuotaBytes {
		t.Fatalf("expected %d, got %f", etcdserver.DefaultQuotaBytes, qv)
	}
}

func TestMetricsHealth(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	tr, err := transport.NewTransport(transport.TLSInfo{}, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	u := clus.Members[0].ClientURLs[0]
	u.Path = "/health"
	resp, err := tr.RoundTrip(&http.Request{
		Header: make(http.Header),
		Method: http.MethodGet,
		URL:    &u,
	})
	resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}

	hv, err := clus.Members[0].Metric("etcd_server_health_failures")
	if err != nil {
		t.Fatal(err)
	}
	if hv != "0" {
		t.Fatalf("expected '0' from etcd_server_health_failures, got %q", hv)
	}
}
<commit_msg>Fix/remove broken: TestMetricDbSizeDefragDebugging<commit_after>// Copyright 2017 The etcd 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 integration

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"testing"
	"time"

	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.etcd.io/etcd/pkg/v3/transport"
	"go.etcd.io/etcd/server/v3/etcdserver"
)

// TestMetricDbSizeBoot checks that the db size metric is set on boot.
func TestMetricDbSizeBoot(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	v, err := clus.Members[0].Metric("etcd_debugging_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}

	if v == "0" {
		t.Fatalf("expected non-zero, got %q", v)
	}
}

func TestMetricDbSizeDefrag(t *testing.T) {
	testMetricDbSizeDefrag(t, "etcd")
}

// testMetricDbSizeDefrag checks that the db size metric is set after defrag.
func testMetricDbSizeDefrag(t *testing.T, name string) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	kvc := toGRPC(clus.Client(0)).KV
	mc := toGRPC(clus.Client(0)).Maintenance

	// expand the db size
	numPuts := 25 // large enough to write more than 1 page
	putreq := &pb.PutRequest{Key: []byte("k"), Value: make([]byte, 4096)}
	for i := 0; i < numPuts; i++ {
		time.Sleep(10 * time.Millisecond) // to execute multiple backend txn
		if _, err := kvc.Put(context.TODO(), putreq); err != nil {
			t.Fatal(err)
		}
	}

	// wait for backend txn sync
	time.Sleep(500 * time.Millisecond)

	expected := numPuts * len(putreq.Value)
	beforeDefrag, err := clus.Members[0].Metric(name + "_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	bv, err := strconv.Atoi(beforeDefrag)
	if err != nil {
		t.Fatal(err)
	}
	if bv < expected {
		t.Fatalf("expected db size greater than %d, got %d", expected, bv)
	}
	beforeDefragInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	biu, err := strconv.Atoi(beforeDefragInUse)
	if err != nil {
		t.Fatal(err)
	}
	if biu < expected {
		t.Fatalf("expected db size in use is greater than %d, got %d", expected, biu)
	}

	// clear out historical keys, in use bytes should free pages
	creq := &pb.CompactionRequest{Revision: int64(numPuts), Physical: true}
	if _, kerr := kvc.Compact(context.TODO(), creq); kerr != nil {
		t.Fatal(kerr)
	}

	validateAfterCompactionInUse := func() error {
		// Put to move PendingPages to FreePages
		if _, err = kvc.Put(context.TODO(), putreq); err != nil {
			t.Fatal(err)
		}
		time.Sleep(500 * time.Millisecond)

		afterCompactionInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
		if err != nil {
			t.Fatal(err)
		}
		aciu, err := strconv.Atoi(afterCompactionInUse)
		if err != nil {
			t.Fatal(err)
		}
		if biu <= aciu {
			return fmt.Errorf("expected less than %d, got %d after compaction", biu, aciu)
		}
		return nil
	}

	// backend rollbacks read transaction asynchronously (PR #10523),
	// which causes the result to be flaky. Retry 3 times.
	maxRetry, retry := 3, 0
	for {
		err := validateAfterCompactionInUse()
		if err == nil {
			break
		}
		retry++
		if retry >= maxRetry {
			t.Fatalf(err.Error())
		}
	}

	// defrag should give freed space back to fs
	mc.Defragment(context.TODO(), &pb.DefragmentRequest{})

	afterDefrag, err := clus.Members[0].Metric(name + "_mvcc_db_total_size_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	av, err := strconv.Atoi(afterDefrag)
	if err != nil {
		t.Fatal(err)
	}
	if bv <= av {
		t.Fatalf("expected less than %d, got %d after defrag", bv, av)
	}

	afterDefragInUse, err := clus.Members[0].Metric("etcd_mvcc_db_total_size_in_use_in_bytes")
	if err != nil {
		t.Fatal(err)
	}
	adiu, err := strconv.Atoi(afterDefragInUse)
	if err != nil {
		t.Fatal(err)
	}
	if adiu > av {
		t.Fatalf("db size in use (%d) is expected less than db size (%d) after defrag", adiu, av)
	}
}

func TestMetricQuotaBackendBytes(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	qs, err := clus.Members[0].Metric("etcd_server_quota_backend_bytes")
	if err != nil {
		t.Fatal(err)
	}
	qv, err := strconv.ParseFloat(qs, 64)
	if err != nil {
		t.Fatal(err)
	}
	if int64(qv) != etcdserver.DefaultQuotaBytes {
		t.Fatalf("expected %d, got %f", etcdserver.DefaultQuotaBytes, qv)
	}
}

func TestMetricsHealth(t *testing.T) {
	BeforeTest(t)
	clus := NewClusterV3(t, &ClusterConfig{Size: 1})
	defer clus.Terminate(t)

	tr, err := transport.NewTransport(transport.TLSInfo{}, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	u := clus.Members[0].ClientURLs[0]
	u.Path = "/health"
	resp, err := tr.RoundTrip(&http.Request{
		Header: make(http.Header),
		Method: http.MethodGet,
		URL:    &u,
	})
	resp.Body.Close()
	if err != nil {
		t.Fatal(err)
	}

	hv, err := clus.Members[0].Metric("etcd_server_health_failures")
	if err != nil {
		t.Fatal(err)
	}
	if hv != "0" {
		t.Fatalf("expected '0' from etcd_server_health_failures, got %q", hv)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/jessevdk/go-flags"
)

func main() {
	var options struct{} // No common flags

	parser := flags.NewParser(&options, flags.Default)

	parser.AddCommand("search",
		"Search for libraries associated with specific terms",
		"Search for libraries associated with specific terms",
		&SearchCommand{})

	parser.AddCommand("list",
		"List all available libraries",
		"List all available libraries",
		&ListCommand{})

	parser.AddCommand("install",
		"Install named libraries",
		"Install named libraries",
		&InstallCommand{})

	parser.AddCommand("index",
		"Build library index",
		"Build library index",
		&IndexCommand{})

	parser.AddCommand("info",
		"Print information about a specific library",
		"Print information about a specific library",
		&InfoCommand{})

	parser.AddCommand("version",
		"Version information about impact itself",
		"Version information about impact itself",
		&VersionCommand{})

	if _, err := parser.Parse(); err != nil {
		os.Exit(1)
	}
}
<commit_msg>Fixed some build errors<commit_after>package main

import (
	"os"

	"github.com/jessevdk/go-flags"
)

func main() {
	var options struct{} // No common flags

	parser := flags.NewParser(&options, flags.Default)

	parser.AddCommand("search",
		"Search for libraries associated with specific terms",
		"Search for libraries associated with specific terms",
		&SearchCommand{})

	parser.AddCommand("install",
		"Install named libraries",
		"Install named libraries",
		&InstallCommand{})

	parser.AddCommand("index",
		"Build library index",
		"Build library index",
		&IndexCommand{})

	parser.AddCommand("version",
		"Version information about impact itself",
		"Version information about impact itself",
		&VersionCommand{})

	if _, err := parser.Parse(); err != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package goproxy_image

import (
	"bytes"
	. "github.com/elazarl/goproxy"
	"github.com/elazarl/goproxy/regretable"
	"image"
	_ "image/gif"
	"image/jpeg"
	"image/png"
	"io/ioutil"
	"net/http"
)

var RespIsImage = ContentTypeIs("image/gif",
	"image/jpeg",
	"image/pjpeg",
	"application/octet-stream",
	"image/png")

// "image/tiff" tiff support is in external package, and rarely used, so we omitted it

func HandleImage(f func(img image.Image, ctx *ProxyCtx) image.Image) RespHandler {
	return FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
		if !RespIsImage.HandleResp(resp, ctx) {
			return resp
		}
		if resp.StatusCode != 200 {
			// we might get 304 - not modified response without data
			return resp
		}
		contentType := resp.Header.Get("Content-Type")

		regret := regretable.NewRegretableReaderCloser(resp.Body)
		resp.Body = regret
		img, imgType, err := image.Decode(resp.Body)
		if err != nil {
			regret.Regret()
			ctx.Warnf("%s: %s", ctx.Req.Method+" "+ctx.Req.URL.String()+" Image from "+ctx.Req.RequestURI+"content type"+
				contentType+"cannot be decoded returning original image", err)
			return resp
		}
		result := f(img, ctx)
		buf := bytes.NewBuffer([]byte{})
		switch contentType {
		// No gif image encoder in go - convert to png
		case "image/gif", "image/png":
			if err := png.Encode(buf, result); err != nil {
				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
				return resp
			}
			resp.Header.Set("Content-Type", "image/png")
		case "image/jpeg", "image/pjpeg":
			if err := jpeg.Encode(buf, result, nil); err != nil {
				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
				return resp
			}
		case "application/octet-stream":
			switch imgType {
			case "jpeg":
				if err := jpeg.Encode(buf, result, nil); err != nil {
					ctx.Warnf("Cannot encode image as jpeg, returning orig %v %v", ctx.Req.URL.String(), err)
					return resp
				}
			case "png", "gif":
				if err := png.Encode(buf, result); err != nil {
					ctx.Warnf("Cannot encode image as png, returning orig %v %v", ctx.Req.URL.String(), err)
					return resp
				}
			}
		default:
			panic("unhandlable type" + contentType)
		}
		resp.Body = ioutil.NopCloser(buf)
		return resp
	})
}
<commit_msg>fixes #64, bigger buffer size for image extension<commit_after>package goproxy_image

import (
	"bytes"
	"image"
	_ "image/gif"
	"image/jpeg"
	"image/png"
	"io/ioutil"
	"net/http"
	. "github.com/elazarl/goproxy"
	"github.com/elazarl/goproxy/regretable"
)

var RespIsImage = ContentTypeIs("image/gif",
	"image/jpeg",
	"image/pjpeg",
	"application/octet-stream",
	"image/png")

// "image/tiff" tiff support is in external package, and rarely used, so we omitted it

func HandleImage(f func(img image.Image, ctx *ProxyCtx) image.Image) RespHandler {
	return FuncRespHandler(func(resp *http.Response, ctx *ProxyCtx) *http.Response {
		if !RespIsImage.HandleResp(resp, ctx) {
			return resp
		}
		if resp.StatusCode != 200 {
			// we might get 304 - not modified response without data
			return resp
		}
		contentType := resp.Header.Get("Content-Type")

		const kb = 1024
		regret := regretable.NewRegretableReaderCloserSize(resp.Body, 16*kb)
		resp.Body = regret
		img, imgType, err := image.Decode(resp.Body)
		if err != nil {
			regret.Regret()
			ctx.Warnf("%s: %s", ctx.Req.Method+" "+ctx.Req.URL.String()+" Image from "+ctx.Req.RequestURI+"content type"+
				contentType+"cannot be decoded returning original image", err)
			return resp
		}
		result := f(img, ctx)
		buf := bytes.NewBuffer([]byte{})
		switch contentType {
		// No gif image encoder in go - convert to png
		case "image/gif", "image/png":
			if err := png.Encode(buf, result); err != nil {
				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
				return resp
			}
			resp.Header.Set("Content-Type", "image/png")
		case "image/jpeg", "image/pjpeg":
			if err := jpeg.Encode(buf, result, nil); err != nil {
				ctx.Warnf("Cannot encode image, returning orig %v %v", ctx.Req.URL.String(), err)
				return resp
			}
		case "application/octet-stream":
			switch imgType {
			case "jpeg":
				if err := jpeg.Encode(buf, result, nil); err != nil {
					ctx.Warnf("Cannot encode image as jpeg, returning orig %v %v", ctx.Req.URL.String(), err)
					return resp
				}
			case "png", "gif":
				if err := png.Encode(buf, result); err != nil {
					ctx.Warnf("Cannot encode image as png, returning orig %v %v", ctx.Req.URL.String(), err)
					return resp
				}
			}
		default:
			panic("unhandlable type" + contentType)
		}
		resp.Body = ioutil.NopCloser(buf)
		return resp
	})
}
<|endoftext|>


<commit_before>package expect

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"regexp"
	"sync"
	"syscall"
	"time"

	"github.com/kr/pty"
)

// Expect is a program interaction session.
type Expect struct {
	timeout time.Duration
	pty     io.ReadWriteCloser
	Killer  func()
	buffer  []byte

	// channel for receiving read events
	readChan   chan readEvent
	readStatus error

	logger Logger
}

// Match is returned from exp.Expect*() when a match is found.
type Match struct {
	Before string
	Groups []string
}

type readEvent struct {
	buf    []byte
	status error
}

// ErrTimeout is returned from exp.Expect*() when a timeout is reached.
var ErrTimeout = errors.New("Expect Timeout")

// READ_SIZE is the largest amount of data expect attempts to read in a single I/O operation.
// This likely needs some research and tuning.
const READ_SIZE = 4094

// Create an Expect instance from a command.
// Effectively the same as Create(pty.Start(exec.Command(name, args...)))
func Spawn(name string, args ...string) (*Expect, error) {
	cmd := exec.Command(name, args...)
	pty, err := pty.Start(cmd)
	if err != nil {
		return nil, err
	}
	killer := func() {
		cmd.Process.Kill()
	}
	return Create(pty, killer, nil, 0), nil
}

// Create an Expect instance from something that we can do read/writes off of.
//
// Note: Close() must be called to cleanup this process.
func Create(pty io.ReadWriteCloser, killer func(), logger Logger, timeout time.Duration) (exp *Expect) {
	if timeout == 0 {
		timeout = time.Hour * 24 * 356
	}
	rv := Expect{
		timeout:  timeout,
		pty:      pty,
		readChan: make(chan readEvent),
		logger:   logger,
		Killer:   killer,
	}

	// Start up processes
	rv.startReader()

	// Done
	return &rv
}

// Timeout() returns amount of time an Expect() call will wait for the output to appear.
func (exp *Expect) Timeout() time.Duration {
	return exp.timeout
}

// SetTimeout(Duration) sets the amount of time an Expect() call will wait for the output to appear.
func (exp *Expect) SetTimeout(d time.Duration) {
	exp.timeout = d
}

// Return the current buffer.
//
// Note: This is not all data received off the network, but data that has been received for processing.
func (exp *Expect) Buffer() []byte {
	return exp.buffer
}

// Kill & close off process.
//
// Note: This *must* be run to cleanup the process
func (exp *Expect) Close() error {
	exp.Killer()
	err := exp.pty.Close()
	for readEvent := range exp.readChan {
		exp.mergeRead(readEvent)
	}
	return err
}

// Send data to program
func (exp *Expect) Send(s string) error {
	return exp.send([]byte(s), false)
}

// Send data, but mark it as masked to observers. Use this for passwords
func (exp *Expect) SendMasked(s string) error {
	return exp.send([]byte(s), true)
}

// Send several lines data (separated by \n) to the process
func (exp *Expect) SendLn(lines ...string) error {
	for _, l := range lines {
		if err := exp.Send(l + "\n"); err != nil {
			return err
		}
	}
	return nil
}

func (exp *Expect) send(arr []byte, masked bool) error {
	for len(arr) > 0 {
		if n, err := exp.pty.Write(arr); err == nil {
			if masked {
				exp.logger.SendMasked(time.Now(), arr[0:n])
			} else {
				exp.logger.Send(time.Now(), arr[0:n])
			}
			arr = arr[n:]
		} else {
			return err
		}
	}
	return nil
}

// ExpectRegexp searches the I/O read stream for a pattern within .Timeout()
func (exp *Expect) ExpectRegexp(pat *regexp.Regexp) (Match, error) {
	exp.logger.ExpectCall(time.Now(), pat)

	// Read error happened.
	if exp.readStatus != nil {
		exp.logger.ExpectReturn(time.Now(), Match{}, exp.readStatus)
		return Match{}, exp.readStatus
	}

	// Calculate absolute timeout
	giveUpTime := time.Now().Add(exp.timeout)

	// Loop until we match or read some data
	for first := true; first || time.Now().Before(giveUpTime); first = false {
		// Read some data
		if !first {
			if !exp.readData(giveUpTime) {
				return Match{}, io.EOF
			}
		}

		// Check for match
		if m, found := exp.checkForMatch(pat); found {
			exp.logger.ExpectReturn(time.Now(), m, nil)
			return m, nil
		}

		// If no match, check for read error (likely io.EOF)
		if exp.readStatus != nil {
			exp.logger.ExpectReturn(time.Now(), Match{}, exp.readStatus)
			return Match{}, exp.readStatus
		}
	}

	// Time is up
	exp.logger.ExpectReturn(time.Now(), Match{}, ErrTimeout)
	return Match{}, ErrTimeout
}

func (exp *Expect) checkForMatch(pat *regexp.Regexp) (m Match, found bool) {

	matches := pat.FindSubmatchIndex(exp.buffer)
	if matches != nil {
		found = true
		groupCount := len(matches) / 2
		m.Groups = make([]string, groupCount)

		for i := 0; i < groupCount; i++ {
			start := matches[2*i]
			end := matches[2*i+1]
			if start >= 0 && end >= 0 {
				m.Groups[i] = string(exp.buffer[start:end])
			}
		}
		m.Before = string(exp.buffer[0:matches[0]])
		exp.buffer = exp.buffer[matches[1]:]
	}
	return
}

func (exp *Expect) readData(giveUpTime time.Time) bool {
	wait := giveUpTime.Sub(time.Now())
	select {
	case read, ok := <-exp.readChan:
		if ok {
			exp.mergeRead(read)
		} else {
			return false
		}

	case <-time.After(wait):
		// Timeout & return
	}
	return true
}

func (exp *Expect) mergeRead(read readEvent) {
	exp.buffer = append(exp.buffer, read.buf...)
	exp.readStatus = read.status
	exp.fixNewLines()

	if len(read.buf) > 0 {
		exp.logger.Recv(time.Now(), read.buf)
	}

	if read.status == io.EOF {
		exp.logger.RecvEOF(time.Now())
	}
}

var newLineRegexp *regexp.Regexp
var newLineOnce sync.Once

// fixNewLines will change various newlines combinations to \n
func (exp *Expect) fixNewLines() {
	newLineOnce.Do(func() { newLineRegexp = regexp.MustCompile("\r\n") })

	// This code could probably be optimized
	exp.buffer = newLineRegexp.ReplaceAllLiteral(exp.buffer, []byte("\n"))
}

// Expect(s string) is equivalent to exp.ExpectRegexp(regexp.MustCompile(s))
func (exp *Expect) Expect(expr string) (m Match, err error) {
	return exp.ExpectRegexp(regexp.MustCompile(expr))
}

// Wait for EOF
func (exp *Expect) ExpectEOF() error {
	_, err := exp.Expect("$EOF")
	return err
}

func (exp *Expect) startReader() {
	bufferInput := make(chan readEvent)

	// Buffer shim
	go func() {
		queue := make([]readEvent, 0)
		done := false

		// Normal I/O loop
		for !done {
			var sendItem readEvent
			var sendChan chan readEvent = nil

			// Set up send operation if we have data to send
			if len(queue) > 0 {
				sendItem = queue[0]
				sendChan = exp.readChan
			}

			// I/O
			select {
			case sendChan <- sendItem:
				queue = queue[1:]
			case read, ok := <-bufferInput:
				if ok {
					queue = append(queue, read)
				} else {
					done = true
				}
			}
		}

		// Drain buffer
		for _, read := range queue {
			exp.readChan <- read
		}

		// Close output
		close(exp.readChan)
	}()

	// Reader process
	go func() {
		done := false
		for !done {
			buf := make([]byte, READ_SIZE)
			n, err := exp.pty.Read(buf)
			buf = buf[0:n]

			// OSX: Closed FD returns io.EOF
			// Linux: Closed FD returns syscall.EIO, translate to io.EOF
			pathErr, ok := err.(*os.PathError)
			if ok && pathErr.Err == syscall.EIO {
				err = io.EOF
			}

			exp.logger.RecvNet(time.Now(), buf)
			bufferInput <- readEvent{buf, err}

			if err != nil {
				done = true
			}
		}
		close(bufferInput)
	}()

}
<commit_msg>fixNewLines does seem to cause problems...<commit_after>package expect

import (
	"errors"
	"io"
	"os"
	"os/exec"
	"regexp"
	"sync"
	"syscall"
	"time"

	"github.com/kr/pty"
)

// Expect is a program interaction session.
type Expect struct {
	timeout time.Duration
	pty     io.ReadWriteCloser
	Killer  func()
	buffer  []byte

	// channel for receiving read events
	readChan   chan readEvent
	readStatus error

	logger Logger
}

// Match is returned from exp.Expect*() when a match is found.
type Match struct {
	Before string
	Groups []string
}

type readEvent struct {
	buf    []byte
	status error
}

// ErrTimeout is returned from exp.Expect*() when a timeout is reached.
var ErrTimeout = errors.New("Expect Timeout")

// READ_SIZE is the largest amount of data expect attempts to read in a single I/O operation.
// This likely needs some research and tuning.
const READ_SIZE = 4094

// Create an Expect instance from a command.
// Effectively the same as Create(pty.Start(exec.Command(name, args...)))
func Spawn(name string, args ...string) (*Expect, error) {
	cmd := exec.Command(name, args...)
	pty, err := pty.Start(cmd)
	if err != nil {
		return nil, err
	}
	killer := func() {
		cmd.Process.Kill()
	}
	return Create(pty, killer, nil, 0), nil
}

// Create an Expect instance from something that we can do read/writes off of.
//
// Note: Close() must be called to cleanup this process.
func Create(pty io.ReadWriteCloser, killer func(), logger Logger, timeout time.Duration) (exp *Expect) {
	if timeout == 0 {
		timeout = time.Hour * 24 * 356
	}
	rv := Expect{
		timeout:  timeout,
		pty:      pty,
		readChan: make(chan readEvent),
		logger:   logger,
		Killer:   killer,
	}

	// Start up processes
	rv.startReader()

	// Done
	return &rv
}

// Timeout() returns amount of time an Expect() call will wait for the output to appear.
func (exp *Expect) Timeout() time.Duration {
	return exp.timeout
}

// SetTimeout(Duration) sets the amount of time an Expect() call will wait for the output to appear.
func (exp *Expect) SetTimeout(d time.Duration) {
	exp.timeout = d
}

// Return the current buffer.
//
// Note: This is not all data received off the network, but data that has been received for processing.
func (exp *Expect) Buffer() []byte {
	return exp.buffer
}

// Kill & close off process.
//
// Note: This *must* be run to cleanup the process
func (exp *Expect) Close() error {
	exp.Killer()
	err := exp.pty.Close()
	for readEvent := range exp.readChan {
		exp.mergeRead(readEvent)
	}
	return err
}

// Send data to program
func (exp *Expect) Send(s string) error {
	return exp.send([]byte(s), false)
}

// Send data, but mark it as masked to observers. Use this for passwords
func (exp *Expect) SendMasked(s string) error {
	return exp.send([]byte(s), true)
}

// Send several lines data (separated by \n) to the process
func (exp *Expect) SendLn(lines ...string) error {
	for _, l := range lines {
		if err := exp.Send(l + "\n"); err != nil {
			return err
		}
	}
	return nil
}

func (exp *Expect) send(arr []byte, masked bool) error {
	for len(arr) > 0 {
		if n, err := exp.pty.Write(arr); err == nil {
			if masked {
				exp.logger.SendMasked(time.Now(), arr[0:n])
			} else {
				exp.logger.Send(time.Now(), arr[0:n])
			}
			arr = arr[n:]
		} else {
			return err
		}
	}
	return nil
}

// ExpectRegexp searches the I/O read stream for a pattern within .Timeout()
func (exp *Expect) ExpectRegexp(pat *regexp.Regexp) (Match, error) {
	exp.logger.ExpectCall(time.Now(), pat)

	// Read error happened.
	if exp.readStatus != nil {
		exp.logger.ExpectReturn(time.Now(), Match{}, exp.readStatus)
		return Match{}, exp.readStatus
	}

	// Calculate absolute timeout
	giveUpTime := time.Now().Add(exp.timeout)

	// Loop until we match or read some data
	for first := true; first || time.Now().Before(giveUpTime); first = false {
		// Read some data
		if !first {
			if !exp.readData(giveUpTime) {
				return Match{}, io.EOF
			}
		}

		// Check for match
		if m, found := exp.checkForMatch(pat); found {
			exp.logger.ExpectReturn(time.Now(), m, nil)
			return m, nil
		}

		// If no match, check for read error (likely io.EOF)
		if exp.readStatus != nil {
			exp.logger.ExpectReturn(time.Now(), Match{}, exp.readStatus)
			return Match{}, exp.readStatus
		}
	}

	// Time is up
	exp.logger.ExpectReturn(time.Now(), Match{}, ErrTimeout)
	return Match{}, ErrTimeout
}

func (exp *Expect) checkForMatch(pat *regexp.Regexp) (m Match, found bool) {

	matches := pat.FindSubmatchIndex(exp.buffer)
	if matches != nil {
		found = true
		groupCount := len(matches) / 2
		m.Groups = make([]string, groupCount)

		for i := 0; i < groupCount; i++ {
			start := matches[2*i]
			end := matches[2*i+1]
			if start >= 0 && end >= 0 {
				m.Groups[i] = string(exp.buffer[start:end])
			}
		}
		m.Before = string(exp.buffer[0:matches[0]])
		exp.buffer = exp.buffer[matches[1]:]
	}
	return
}

func (exp *Expect) readData(giveUpTime time.Time) bool {
	wait := giveUpTime.Sub(time.Now())
	select {
	case read, ok := <-exp.readChan:
		if ok {
			exp.mergeRead(read)
		} else {
			return false
		}

	case <-time.After(wait):
		// Timeout & return
	}
	return true
}

func (exp *Expect) mergeRead(read readEvent) {
	exp.buffer = append(exp.buffer, read.buf...)
	exp.readStatus = read.status
	//exp.fixNewLines()

	if len(read.buf) > 0 {
		exp.logger.Recv(time.Now(), read.buf)
	}

	if read.status == io.EOF {
		exp.logger.RecvEOF(time.Now())
	}
}

var newLineRegexp *regexp.Regexp
var newLineOnce sync.Once

// fixNewLines will change various newlines combinations to \n
func (exp *Expect) fixNewLines() {
	newLineOnce.Do(func() { newLineRegexp = regexp.MustCompile("\r\n") })

	// This code could probably be optimized
	exp.buffer = newLineRegexp.ReplaceAllLiteral(exp.buffer, []byte("\n"))
}

// Expect(s string) is equivalent to exp.ExpectRegexp(regexp.MustCompile(s))
func (exp *Expect) Expect(expr string) (m Match, err error) {
	return exp.ExpectRegexp(regexp.MustCompile(expr))
}

// Wait for EOF
func (exp *Expect) ExpectEOF() error {
	_, err := exp.Expect("$EOF")
	return err
}

func (exp *Expect) startReader() {
	bufferInput := make(chan readEvent)

	// Buffer shim
	go func() {
		queue := make([]readEvent, 0)
		done := false

		// Normal I/O loop
		for !done {
			var sendItem readEvent
			var sendChan chan readEvent = nil

			// Set up send operation if we have data to send
			if len(queue) > 0 {
				sendItem = queue[0]
				sendChan = exp.readChan
			}

			// I/O
			select {
			case sendChan <- sendItem:
				queue = queue[1:]
			case read, ok := <-bufferInput:
				if ok {
					queue = append(queue, read)
				} else {
					done = true
				}
			}
		}

		// Drain buffer
		for _, read := range queue {
			exp.readChan <- read
		}

		// Close output
		close(exp.readChan)
	}()

	// Reader process
	go func() {
		done := false
		for !done {
			buf := make([]byte, READ_SIZE)
			n, err := exp.pty.Read(buf)
			buf = buf[0:n]

			// OSX: Closed FD returns io.EOF
			// Linux: Closed FD returns syscall.EIO, translate to io.EOF
			pathErr, ok := err.(*os.PathError)
			if ok && pathErr.Err == syscall.EIO {
				err = io.EOF
			}

			exp.logger.RecvNet(time.Now(), buf)
			bufferInput <- readEvent{buf, err}

			if err != nil {
				done = true
			}
		}
		close(bufferInput)
	}()

}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/square/quotaservice/master/LICENSE

package events

import (
	"time"

	"fmt"

	"github.com/square/quotaservice/logging"
)

type EventType int

const (
	EVENT_TOKENS_SERVED EventType = iota
	EVENT_TIMEOUT_SERVING_TOKENS
	EVENT_TOO_MANY_TOKENS_REQUESTED
	EVENT_BUCKET_MISS
	EVENT_BUCKET_CREATED
	EVENT_BUCKET_REMOVED
)

var eventNames = []string{
	EVENT_TOKENS_SERVED:             "EVENT_TOKENS_SERVED",
	EVENT_TIMEOUT_SERVING_TOKENS:    "EVENT_TIMEOUT_SERVING_TOKENS",
	EVENT_TOO_MANY_TOKENS_REQUESTED: "EVENT_TOO_MANY_TOKENS_REQUESTED",
	EVENT_BUCKET_MISS:               "EVENT_BUCKET_MISS",
	EVENT_BUCKET_CREATED:            "EVENT_BUCKET_CREATED",
	EVENT_BUCKET_REMOVED:            "EVENT_BUCKET_REMOVED"}

func (et EventType) String() string {
	name := eventNames[et]
	if name == "" {
		panic(fmt.Sprintf("Don't know event %d", et))
	}

	return name
}

type Event interface {
	EventType() EventType
	Namespace() string
	BucketName() string
	Dynamic() bool
	NumTokens() int64
	WaitTime() time.Duration
}

// EventProducer is a hook into the notification system, to inform listeners that certain events
// take place.
type EventProducer struct {
	c chan Event
}

func (e *EventProducer) Emit(event Event) {
	select {
	case e.c <- event:
	// OK
	default:
		logging.Println("Event buffer full; dropping event.")
	}
}

func (e *EventProducer) notifyListeners(l Listener) {
	for event := range e.c {
		l(event)
	}
}

// Listener is a function that consumes an Event
type Listener func(details Event)

// RegisterListener takes a Listener and a buffer size and
// returns an EventProducer that consumes events and notifies listeners
func RegisterListener(listener Listener, bufsize int) *EventProducer {
	if listener == nil {
		panic("Cannot register a nil listener")
	}

	ep := &EventProducer{make(chan Event, bufsize)}

	go ep.notifyListeners(listener)

	return ep
}

type namedEvent struct {
	eventType             EventType
	namespace, bucketName string
	dynamic               bool
}

func (n *namedEvent) String() string {
	return fmt.Sprintf("namedEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		n.eventType, n.namespace, n.bucketName, n.dynamic, 0, 0)
}

func (n *namedEvent) EventType() EventType {
	return n.eventType
}

func (n *namedEvent) Namespace() string {
	return n.namespace
}

func (n *namedEvent) BucketName() string {
	return n.bucketName
}

func (n *namedEvent) Dynamic() bool {
	return n.dynamic
}

func (n *namedEvent) NumTokens() int64 {
	return 0
}

func (n *namedEvent) WaitTime() time.Duration {
	return 0
}

type tokenEvent struct {
	*namedEvent
	numTokens int64
}

func (t *tokenEvent) String() string {
	return fmt.Sprintf("tokenEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		t.eventType, t.namespace, t.bucketName, t.dynamic, t.numTokens, 0)
}

func (t *tokenEvent) NumTokens() int64 {
	return t.numTokens
}

type tokenWaitEvent struct {
	*tokenEvent
	waitTime time.Duration
}

func (t *tokenWaitEvent) String() string {
	return fmt.Sprintf("tokenWaitEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		t.eventType, t.namespace, t.bucketName, t.dynamic, t.numTokens, t.waitTime)
}

func (t *tokenWaitEvent) WaitTime() time.Duration {
	return t.waitTime
}

// NewTokensServedEvent creates a new event with the type EVENT_TOKENS_SERVED
func NewTokensServedEvent(namespace, bucketName string, dynamic bool, numTokens int64, waitTime time.Duration) Event {
	return &tokenWaitEvent{
		tokenEvent: &tokenEvent{
			namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TOKENS_SERVED),
			numTokens:  numTokens},
		waitTime: waitTime}
}

// NewTimedOutEvent creates a new event with the type EVENT_TIMEOUT_SERVING_TOKENS
func NewTimedOutEvent(namespace, bucketName string, dynamic bool, numTokens int64) Event {
	return &tokenEvent{
		namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TIMEOUT_SERVING_TOKENS),
		numTokens:  numTokens}
}

// NewTooManyTokensRequestedEvent creates a new event with the type EVENT_TOO_MANY_TOKENS_REQUESTED
func NewTooManyTokensRequestedEvent(namespace, bucketName string, dynamic bool, numTokens int64) Event {
	return &tokenEvent{
		namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TOO_MANY_TOKENS_REQUESTED),
		numTokens:  numTokens}
}

// NewBucketMissedEvent creates a new event with the type EVENT_BUCKET_MISS
func NewBucketMissedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_MISS)
}

// NewBucketCreatedEvent creates a new event with the type EVENT_BUCKET_CREATED
func NewBucketCreatedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_CREATED)
}

// NewBucketRemovedEvent creates a new event with the type EVENT_BUCKET_REMOVED
func NewBucketRemovedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_REMOVED)
}

func newNamedEvent(namespace, bucketName string, dynamic bool, eventType EventType) *namedEvent {
	return &namedEvent{
		eventType:  eventType,
		namespace:  namespace,
		bucketName: bucketName,
		dynamic:    dynamic}
}
<commit_msg>Added more context to dropping event log<commit_after>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/square/quotaservice/master/LICENSE

package events

import (
	"time"

	"fmt"

	"github.com/square/quotaservice/logging"
)

type EventType int

const (
	EVENT_TOKENS_SERVED EventType = iota
	EVENT_TIMEOUT_SERVING_TOKENS
	EVENT_TOO_MANY_TOKENS_REQUESTED
	EVENT_BUCKET_MISS
	EVENT_BUCKET_CREATED
	EVENT_BUCKET_REMOVED
)

var eventNames = []string{
	EVENT_TOKENS_SERVED:             "EVENT_TOKENS_SERVED",
	EVENT_TIMEOUT_SERVING_TOKENS:    "EVENT_TIMEOUT_SERVING_TOKENS",
	EVENT_TOO_MANY_TOKENS_REQUESTED: "EVENT_TOO_MANY_TOKENS_REQUESTED",
	EVENT_BUCKET_MISS:               "EVENT_BUCKET_MISS",
	EVENT_BUCKET_CREATED:            "EVENT_BUCKET_CREATED",
	EVENT_BUCKET_REMOVED:            "EVENT_BUCKET_REMOVED"}

func (et EventType) String() string {
	name := eventNames[et]
	if name == "" {
		panic(fmt.Sprintf("Don't know event %d", et))
	}

	return name
}

type Event interface {
	EventType() EventType
	Namespace() string
	BucketName() string
	Dynamic() bool
	NumTokens() int64
	WaitTime() time.Duration
}

// EventProducer is a hook into the notification system, to inform listeners that certain events
// take place.
type EventProducer struct {
	c chan Event
}

func (e *EventProducer) Emit(event Event) {
	select {
	case e.c <- event:
	// OK
	default:
		logging.Printf("Event buffer full; dropping %s event for %s.%s", event.EventType(), event.Namespace(), event.BucketName())
	}
}

func (e *EventProducer) notifyListeners(l Listener) {
	for event := range e.c {
		l(event)
	}
}

// Listener is a function that consumes an Event
type Listener func(details Event)

// RegisterListener takes a Listener and a buffer size and
// returns an EventProducer that consumes events and notifies listeners
func RegisterListener(listener Listener, bufsize int) *EventProducer {
	if listener == nil {
		panic("Cannot register a nil listener")
	}

	ep := &EventProducer{make(chan Event, bufsize)}

	go ep.notifyListeners(listener)

	return ep
}

type namedEvent struct {
	eventType             EventType
	namespace, bucketName string
	dynamic               bool
}

func (n *namedEvent) String() string {
	return fmt.Sprintf("namedEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		n.eventType, n.namespace, n.bucketName, n.dynamic, 0, 0)
}

func (n *namedEvent) EventType() EventType {
	return n.eventType
}

func (n *namedEvent) Namespace() string {
	return n.namespace
}

func (n *namedEvent) BucketName() string {
	return n.bucketName
}

func (n *namedEvent) Dynamic() bool {
	return n.dynamic
}

func (n *namedEvent) NumTokens() int64 {
	return 0
}

func (n *namedEvent) WaitTime() time.Duration {
	return 0
}

type tokenEvent struct {
	*namedEvent
	numTokens int64
}

func (t *tokenEvent) String() string {
	return fmt.Sprintf("tokenEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		t.eventType, t.namespace, t.bucketName, t.dynamic, t.numTokens, 0)
}

func (t *tokenEvent) NumTokens() int64 {
	return t.numTokens
}

type tokenWaitEvent struct {
	*tokenEvent
	waitTime time.Duration
}

func (t *tokenWaitEvent) String() string {
	return fmt.Sprintf("tokenWaitEvent{type: %v, namespace: %v, name: %v, dynamic: %v, numTokens: %v, waitTime: %v}",
		t.eventType, t.namespace, t.bucketName, t.dynamic, t.numTokens, t.waitTime)
}

func (t *tokenWaitEvent) WaitTime() time.Duration {
	return t.waitTime
}

// NewTokensServedEvent creates a new event with the type EVENT_TOKENS_SERVED
func NewTokensServedEvent(namespace, bucketName string, dynamic bool, numTokens int64, waitTime time.Duration) Event {
	return &tokenWaitEvent{
		tokenEvent: &tokenEvent{
			namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TOKENS_SERVED),
			numTokens:  numTokens},
		waitTime: waitTime}
}

// NewTimedOutEvent creates a new event with the type EVENT_TIMEOUT_SERVING_TOKENS
func NewTimedOutEvent(namespace, bucketName string, dynamic bool, numTokens int64) Event {
	return &tokenEvent{
		namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TIMEOUT_SERVING_TOKENS),
		numTokens:  numTokens}
}

// NewTooManyTokensRequestedEvent creates a new event with the type EVENT_TOO_MANY_TOKENS_REQUESTED
func NewTooManyTokensRequestedEvent(namespace, bucketName string, dynamic bool, numTokens int64) Event {
	return &tokenEvent{
		namedEvent: newNamedEvent(namespace, bucketName, dynamic, EVENT_TOO_MANY_TOKENS_REQUESTED),
		numTokens:  numTokens}
}

// NewBucketMissedEvent creates a new event with the type EVENT_BUCKET_MISS
func NewBucketMissedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_MISS)
}

// NewBucketCreatedEvent creates a new event with the type EVENT_BUCKET_CREATED
func NewBucketCreatedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_CREATED)
}

// NewBucketRemovedEvent creates a new event with the type EVENT_BUCKET_REMOVED
func NewBucketRemovedEvent(namespace, bucketName string, dynamic bool) Event {
	return newNamedEvent(namespace, bucketName, dynamic, EVENT_BUCKET_REMOVED)
}

func newNamedEvent(namespace, bucketName string, dynamic bool, eventType EventType) *namedEvent {
	return &namedEvent{
		eventType:  eventType,
		namespace:  namespace,
		bucketName: bucketName,
		dynamic:    dynamic}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Gorka Lerchundi Osa. All rights reserved.
// Use of this source code is governed by the Apache License, Version 2.0
// that can be found in the LICENSE file.
package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/coreos/etcd/client"
	"github.com/glerchundi/setup-etcd-peers-environment/providers"
	"github.com/glerchundi/setup-etcd-peers-environment/providers/aws"
	"github.com/glerchundi/setup-etcd-peers-environment/util"
)

func main() {
	app := cli.NewApp()
	app.Name = "setup-etcd-peers-environment"
	app.Version = "0.1.1"
	app.Usage = "manage etcd cluster peers based on AWS autoscaling groups"
	app.Action = mainAction
	app.Flags = []cli.Flag {
		cli.StringFlag{
			Name: "out, o",
			Value: "/etc/sysconfig/etcd-peers",
			Usage: "etcd peers environment config file destination",
		},
	}
	app.RunAndExitOnError()
}

func mainAction(c *cli.Context) {
	environmentFilePath := c.GlobalString("out")
	if _, err := os.Stat(environmentFilePath); err == nil {
		log.Printf("etcd-peers file %s already created, exiting.", environmentFilePath)
    	return
	}

	tempFilePath := environmentFilePath + ".tmp"
	tempFile, err := os.Create(tempFilePath)
	if err != nil {
		log.Fatal(err)
	}

	defer tempFile.Close()
	if err := writeEnvironment(tempFile); err != nil {
		log.Fatal(err)
	}

	if err := os.Rename(tempFilePath, environmentFilePath); err != nil {
		log.Fatal(err)
	}

	os.Exit(0)
}

func writeEnvironment(w io.Writer) error {
	var buffer bytes.Buffer
	var err error
	var provider providers.Provider = aws.New()

	instanceId, err := provider.GetInstanceId()
	if err != nil {
		return err
	}

	instanceIp, err := provider.GetInstancePrivateAddress()
	if err != nil {
		return err
	}

	clusterMembersByName, err := provider.GetClusterMembers()
	if err != nil {
		return err
	}

	// retrieve current cluster members
	var etcdMembers []client.Member
	var goodEtcdClientURL string = ""
	for _, memberIp := range clusterMembersByName {
		if memberIp == instanceIp {
			continue
		}

		etcdClientURL := util.EtcdClientURLFromIP(memberIp)
		etcdMembers, err = util.EtcdListMembers(etcdClientURL)
		if err == nil {
			goodEtcdClientURL = etcdClientURL
			break
		}
	}

	// etcd parameters
	var initialClusterState string
    var initialCluster string

	// check if instanceId is already member of cluster
	var isMember bool = false
	for _, member := range etcdMembers {
		if member.Name == instanceId {
			isMember = true
			break
		}
	}

	// if i am not already listed as a member of the cluster assume that this is a new cluster
	if etcdMembers != nil && !isMember {
		log.Print("joining an existing cluster, using this client url: ", goodEtcdClientURL)
		initialClusterState = "existing"

		//
		// detect and remove bad peers
		//

		// create a reverse cluster members map
		clusterMembersByIp := make(map[string]string)
		for memberName, memberIp := range clusterMembersByName {
			clusterMembersByIp[memberIp] = memberName
		}

		for _, etcdMember := range etcdMembers {
			peerURL, err := url.Parse(etcdMember.PeerURLs[0])
			if err != nil {
				return err
			}

			peerHost, _, err := net.SplitHostPort(peerURL.Host)
			if err != nil {
				return err
			}

			if _, ok := clusterMembersByIp[peerHost]; !ok {
				log.Print("removing etcd member ", etcdMember.ID)
				err = util.EtcdRemoveMember(goodEtcdClientURL, etcdMember.ID)
				if err != nil {
					return err
				}
			}
		}

		//
		// join an existing cluster
		//

		instancePeerURL := util.EtcdPeerURLFromIP(instanceIp)
		log.Print("adding etcd member ", instancePeerURL)
		member, err := util.EtcdAddMember(goodEtcdClientURL, instancePeerURL)
		if member != nil {
			return err
		}
	} else {
		log.Print("creating new cluster")
		initialClusterState = "new"
	}

	// initial cluster
	kvs := make([]string, 0)
	for memberName, memberIp := range clusterMembersByName {
		kvs = append(kvs, fmt.Sprintf("%s=%s", memberName, util.EtcdPeerURLFromIP(memberIp)))
	}
	initialCluster = strings.Join(kvs, ",")

	// create environment variables
	buffer.WriteString(fmt.Sprintf("ETCD_NAME=%s\n", instanceId))
	buffer.WriteString(fmt.Sprintf("ETCD_INITIAL_CLUSTER_STATE=%s\n", initialClusterState))
	buffer.WriteString(fmt.Sprintf("ETCD_INITIAL_CLUSTER=%s\n", initialCluster))

	if _, err := buffer.WriteTo(w); err != nil {
		return err
	}

	return nil
}<commit_msg>protect against renaming a non-closed file<commit_after>// Copyright (c) 2015 Gorka Lerchundi Osa. All rights reserved.
// Use of this source code is governed by the Apache License, Version 2.0
// that can be found in the LICENSE file.
package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/coreos/etcd/client"
	"github.com/glerchundi/setup-etcd-peers-environment/providers"
	"github.com/glerchundi/setup-etcd-peers-environment/providers/aws"
	"github.com/glerchundi/setup-etcd-peers-environment/util"
)

func main() {
	app := cli.NewApp()
	app.Name = "setup-etcd-peers-environment"
	app.Version = "0.1.1"
	app.Usage = "manage etcd cluster peers based on AWS autoscaling groups"
	app.Action = mainAction
	app.Flags = []cli.Flag {
		cli.StringFlag{
			Name: "out, o",
			Value: "/etc/sysconfig/etcd-peers",
			Usage: "etcd peers environment config file destination",
		},
	}
	app.RunAndExitOnError()
}

func mainAction(c *cli.Context) {
	environmentFilePath := c.GlobalString("out")
	if _, err := os.Stat(environmentFilePath); err == nil {
		log.Printf("etcd-peers file %s already created, exiting.", environmentFilePath)
    	return
	}

	tempFilePath := environmentFilePath + ".tmp"
	tempFile, err := os.Create(tempFilePath)
	if err != nil {
		log.Fatal(err)
	}

	isClosed := false

	defer func() {
		if !isClosed {
			tempFile.Close()
		}
	}()
	if err := writeEnvironment(tempFile); err != nil {
		log.Fatal(err)
	}

	if err := tempFile.Close(); err != nil {
		return err
	}

	isClosed = true

	if err := os.Rename(tempFilePath, environmentFilePath); err != nil {
		log.Fatal(err)
	}

	os.Exit(0)
}

func writeEnvironment(w io.Writer) error {
	var buffer bytes.Buffer
	var err error
	var provider providers.Provider = aws.New()

	instanceId, err := provider.GetInstanceId()
	if err != nil {
		return err
	}

	instanceIp, err := provider.GetInstancePrivateAddress()
	if err != nil {
		return err
	}

	clusterMembersByName, err := provider.GetClusterMembers()
	if err != nil {
		return err
	}

	// retrieve current cluster members
	var etcdMembers []client.Member
	var goodEtcdClientURL string = ""
	for _, memberIp := range clusterMembersByName {
		if memberIp == instanceIp {
			continue
		}

		etcdClientURL := util.EtcdClientURLFromIP(memberIp)
		etcdMembers, err = util.EtcdListMembers(etcdClientURL)
		if err == nil {
			goodEtcdClientURL = etcdClientURL
			break
		}
	}

	// etcd parameters
	var initialClusterState string
    var initialCluster string

	// check if instanceId is already member of cluster
	var isMember bool = false
	for _, member := range etcdMembers {
		if member.Name == instanceId {
			isMember = true
			break
		}
	}

	// if i am not already listed as a member of the cluster assume that this is a new cluster
	if etcdMembers != nil && !isMember {
		log.Print("joining an existing cluster, using this client url: ", goodEtcdClientURL)
		initialClusterState = "existing"

		//
		// detect and remove bad peers
		//

		// create a reverse cluster members map
		clusterMembersByIp := make(map[string]string)
		for memberName, memberIp := range clusterMembersByName {
			clusterMembersByIp[memberIp] = memberName
		}

		for _, etcdMember := range etcdMembers {
			peerURL, err := url.Parse(etcdMember.PeerURLs[0])
			if err != nil {
				return err
			}

			peerHost, _, err := net.SplitHostPort(peerURL.Host)
			if err != nil {
				return err
			}

			if _, ok := clusterMembersByIp[peerHost]; !ok {
				log.Print("removing etcd member ", etcdMember.ID)
				err = util.EtcdRemoveMember(goodEtcdClientURL, etcdMember.ID)
				if err != nil {
					return err
				}
			}
		}

		//
		// join an existing cluster
		//

		instancePeerURL := util.EtcdPeerURLFromIP(instanceIp)
		log.Print("adding etcd member ", instancePeerURL)
		member, err := util.EtcdAddMember(goodEtcdClientURL, instancePeerURL)
		if member != nil {
			return err
		}
	} else {
		log.Print("creating new cluster")
		initialClusterState = "new"
	}

	// initial cluster
	kvs := make([]string, 0)
	for memberName, memberIp := range clusterMembersByName {
		kvs = append(kvs, fmt.Sprintf("%s=%s", memberName, util.EtcdPeerURLFromIP(memberIp)))
	}
	initialCluster = strings.Join(kvs, ",")

	// create environment variables
	buffer.WriteString(fmt.Sprintf("ETCD_NAME=%s\n", instanceId))
	buffer.WriteString(fmt.Sprintf("ETCD_INITIAL_CLUSTER_STATE=%s\n", initialClusterState))
	buffer.WriteString(fmt.Sprintf("ETCD_INITIAL_CLUSTER=%s\n", initialCluster))

	if _, err := buffer.WriteTo(w); err != nil {
		return err
	}

	return nil
}<|endoftext|>


<commit_before>package releaseAnalyzer

import (
	"bytes"
	"fmt"
	"github.com/haskelladdict/mbdr/libmbd"
	"log"
	"math"
	"math/rand"
	"sort"
	"time"
)

const (
	numSyt              = 8  // number of synaptotagmin molecules (with 5 Ca2+ sites each)
	numY                = 16 // number of second sensor (Y) sites
	numActiveSyt        = 2  // how many Ca2+ sites need to be bound for sensors
	numActiveY          = 1  // to become active
	vesicleFusionEnergy = 40
	pulseDuration       = 3e-3 // pulse duration in [s]
)

// type of binding site (syt or second sensor)
const (
	sytSite = iota
	ySite
)

// CaSensor defines a single synaptotagmin and Y sites
type caSensor struct {
	sites    []int // ca sites contributing to sensors
	siteType int   // type of sensor (syt or Y)
}

// ActEvent keeps track of a single activation/deactivation event
type ActEvent struct {
	sensorID  int    // sensor which was activated/deactivated
	vesicleID string // vesicleID were activation event took place
	eventIter int    // iteration when event occured
	activated bool   // activated is set to true and deactivated otherwise
}

// sort infrastructure for sorting ActEvents according to the event time
type byIter []ActEvent

func (e byIter) Len() int {
	return len(e)
}

func (e byIter) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e byIter) Less(i, j int) bool {
	return e[i].eventIter < e[j].eventIter
}

// ReleaseEvent keeps track of vesicle release events
type ReleaseEvent struct {
	sensors   []int  // list of sensors involved in release event
	vesicleID string // id of vesicle which was released
	eventIter uint64 // iteration when event occurred
}

var caSensors []caSensor

// random number generator for Metropolis-Hasting
var rng *rand.Rand

func init() {
	rng = rand.New(rand.NewSource(time.Now().UnixNano()))

	caSensors = make([]caSensor, numSyt+numY)

	// define synaptogamin and Y sites
	caSensors[0] = caSensor{[]int{8, 9, 29, 30, 31}, sytSite}
	caSensors[1] = caSensor{[]int{7, 32, 33, 34, 35}, sytSite}
	caSensors[2] = caSensor{[]int{3, 6, 36, 37, 38}, sytSite}
	caSensors[3] = caSensor{[]int{17, 39, 40, 41, 42}, sytSite}
	caSensors[4] = caSensor{[]int{15, 16, 43, 44, 45}, sytSite}
	caSensors[5] = caSensor{[]int{14, 46, 47, 48, 49}, sytSite}
	caSensors[6] = caSensor{[]int{4, 12, 24, 50, 51}, sytSite}
	caSensors[7] = caSensor{[]int{10, 25, 26, 27, 28}, sytSite}
	caSensors[8] = caSensor{[]int{122}, ySite}
	caSensors[9] = caSensor{[]int{70}, ySite}
	caSensors[10] = caSensor{[]int{126}, ySite}
	caSensors[11] = caSensor{[]int{142}, ySite}
	caSensors[12] = caSensor{[]int{62}, ySite}
	caSensors[13] = caSensor{[]int{118}, ySite}
	caSensors[14] = caSensor{[]int{22}, ySite}
	caSensors[15] = caSensor{[]int{134}, ySite}
	caSensors[16] = caSensor{[]int{110}, ySite}
	caSensors[17] = caSensor{[]int{66}, ySite}
	caSensors[18] = caSensor{[]int{106}, ySite}
	caSensors[19] = caSensor{[]int{130}, ySite}
	caSensors[20] = caSensor{[]int{2}, ySite}
	caSensors[21] = caSensor{[]int{114}, ySite}
	caSensors[22] = caSensor{[]int{42}, ySite}
	caSensors[23] = caSensor{[]int{138}, ySite}
}

// analyze is the main entry point for analyzing the mouse AZ model. It
// determines release events and collects statistics
func analyze(data *libmbd.MCellData, model *ReleaseModel, energyModel bool, seed,
	numPulses, numActiveSites, sytEnergy, yEnergy int) ([]string, error) {

	var releases []*ReleaseEvent
	for _, vesID := range model.VesicleIDs {
		evts, err := extractActivationEvents(data, numPulses, seed, vesID,
			model.SensorTemplateString)
		if err != nil {
			return nil, err
		}
		if evts == nil {
			continue
		}

		rel, err := extractReleaseEvents(evts, data.BlockSize(), energyModel,
			numActiveSites, sytEnergy, yEnergy, vesID)
		if err != nil {
			return nil, err
		}
		if rel != nil {
			releases = append(releases, rel)
		}
	}

	return assembleReleaseMsgs(data, seed, releases), nil
}

// assembleReleaseMsgs creates a slice of strings with summary statistics for all
// released vesicles for a given seed
func assembleReleaseMsgs(data *libmbd.MCellData, seed int, rel []*ReleaseEvent) []string {
	messages := make([]string, 0)
	timeStep := data.StepSize()
	for _, r := range rel {
		buffer := bytes.NewBufferString("")
		channels, err := determineCaChanContrib(data, r)
		if err != nil {
			log.Fatal(err)
		}

		eventTime := float64(r.eventIter) * timeStep
		// figure out if event happened within or between pulses
		pulseID := int(math.Floor(eventTime / isiValue))
		var pulseString string
		if eventTime-float64(pulseID)*isiValue > pulseDuration {
			pulseString = fmt.Sprintf("ISI_%d", pulseID+1)
		} else {
			pulseString = fmt.Sprintf("%d", pulseID+1)
		}

		fmt.Fprintf(buffer, "seed : %d   vesicleID : %s   time : %e   pulseID : %s", seed,
			r.vesicleID, eventTime, pulseString)
		fmt.Fprintf(buffer, "  sensors: [")
		for _, s := range r.sensors {
			fmt.Fprintf(buffer, "%d ", s)
		}
		fmt.Fprintf(buffer, "]")
		fmt.Fprintf(buffer, "  channels: [")
		for n, c := range channels {
			fmt.Fprintf(buffer, "%s : %d  ", n, int(c))
		}
		fmt.Fprintf(buffer, "]")
		messages = append(messages, buffer.String())
	}
	return messages
}

// extractActivationEvents returns a slice with actvation and deactivation events
// for the given vesicle and active zone
func extractActivationEvents(data *libmbd.MCellData, numPulses, seed int,
	vesicleID, template string) ([]ActEvent, error) {

	var events []ActEvent
	// analyze the activation/deactivation status of each ca sensor.
	// NOTE: for now we merge the binding data for individual pulses into one
	for id := 0; id < len(caSensors); id++ {
		sensor := caSensors[id]
		sensorString := "sensor"
		actThresh := numActiveSyt
		if sensor.siteType == ySite {
			sensorString = "sensor_Y"
			actThresh = numActiveY
		}

		sensorData := make([]int, data.BlockSize())
		for _, s := range sensor.sites {
			for p := 0; p < numPulses; p++ {
				dataName := fmt.Sprintf(template, vesicleID, sensorString, s, p+1, seed)
				bd, err := data.BlockDataByName(dataName)
				if err != nil {
					return nil, err
				}

				if len(bd.Col) != 1 {
					return nil, fmt.Errorf("data set %s had more than one data column",
						dataName)
				}
				for i := 0; i < len(sensorData); i++ {
					sensorData[i] += int(bd.Col[0][i])
				}
			}
		}

		// check for activation events
		active := false
		for i, b := range sensorData {
			if !active && b >= actThresh {
				active = true
				events = append(events, ActEvent{id, vesicleID, i, active})
			} else if active && b < actThresh {
				active = false
				events = append(events, ActEvent{id, vesicleID, i, active})
			}
		}
	}
	return events, nil
}

// extractReleaseEvents determines if the given vesicle was released given
// a list of sensor activation events. If no release took place returns nil.
func extractReleaseEvents(evts []ActEvent, maxIter uint64, energyModel bool,
	numActiveSites, sytEnergy, yEnergy int, vesicleID string) (*ReleaseEvent, error) {

	sort.Sort(byIter(evts))
	activeEvts := make(map[int]struct{})
	for i, e := range evts {
		_, present := activeEvts[e.sensorID]
		if e.activated {
			if present {
				return nil, fmt.Errorf("trying to add active event that already exists")
			}
			activeEvts[e.sensorID] = struct{}{}
		} else {
			if !present {
				return nil, fmt.Errorf("trying to remove a nonexisting active event")
			}
			delete(activeEvts, e.sensorID)
		}

		// special case: If the next event happens simultaneously we apply it right away
		if i+1 < len(evts) && evts[i+1].eventIter == e.eventIter {
			continue
		}

		var rel *ReleaseEvent
		var relError error
		if energyModel {
			// use the energy model to determine release
			energy := getEnergy(activeEvts, sytEnergy, yEnergy)
			nextIter := getNextIter(i, maxIter, evts)
			rel, relError = checkForEnergyRelease(energy, vesicleID, e, activeEvts, nextIter)
		} else {
			// use the deterministic model to determine release
			rel, relError = checkForDeterministicRelease(vesicleID, numActiveSites, e, activeEvts)
		}
		if relError != nil {
			return nil, relError
		}
		if rel != nil {
			return rel, nil
		}
	}
	return nil, nil
}

// getEnergy computes the total energy corresponding to the current number
// of active synaptotagmin and Y sites. Also returns the number of active syts
func getEnergy(events map[int]struct{}, sytEnergy, yEnergy int) int {
	var energy int
	for s, _ := range events {
		if caSensors[s].siteType == sytSite {
			energy += sytEnergy
		} else {
			energy += yEnergy
		}
	}
	return energy
}

// getNextIter determines the iteration of the next event that will happen in
// the event queue
func getNextIter(iter int, maxIter uint64, evts []ActEvent) uint64 {
	nextIter := maxIter
	if iter < len(evts)-1 {
		nextIter = uint64(evts[iter+1].eventIter)
	}
	return nextIter
}

// checkForDeterministicRelease tests if vesicles are released according
// to a deterministic critertion, i.e. as soon as numActiveSites syt or
// Y sites are active
func checkForDeterministicRelease(vesID string, numActiveSites int, evt ActEvent,
	activeEvts map[int]struct{}) (*ReleaseEvent, error) {
	if len(activeEvts) == numActiveSites {
		var sensors []int
		for a, _ := range activeEvts {
			sensors = append(sensors, a)
		}
		return &ReleaseEvent{sensors: sensors, vesicleID: vesID,
			eventIter: uint64(evt.eventIter)}, nil
	}
	return nil, nil
}

// checkForEnergyRelease tests if an energy release according to specified
// syt and y site energies takes place. Check for releases given the current
// energy until next event or the end of simulation. To do this we basically
// test for each iteration between now and the next event if a release takes
// place using the Metrolpolis-Hasting algorithm
func checkForEnergyRelease(energy int, vesID string, evt ActEvent,
	activeEvts map[int]struct{}, nextIter uint64) (*ReleaseEvent, error) {

	numIters := nextIter - uint64(evt.eventIter)
	if nextIter < uint64(evt.eventIter) {
		return nil, fmt.Errorf("encountered out of order release event")
	}
	if iter, ok := checkForRelease(energy, numIters); ok {
		var sensors []int
		for a, _ := range activeEvts {
			sensors = append(sensors, a)
		}
		return &ReleaseEvent{sensors: sensors, vesicleID: vesID,
			eventIter: uint64(evt.eventIter) + iter}, nil
	}
	return nil, nil
}

// checkForReleases uses a Metropolis-Hasting scheme to test numIter times
// if vesicle release happens given the provided bound sensor energy
func checkForRelease(energy int, numIters uint64) (uint64, bool) {

	if energy >= vesicleFusionEnergy {
		return 0, true
	}

	prob := math.Exp(float64(energy - vesicleFusionEnergy))
	if prob >= 1 {
		log.Fatal("probability out of bounds")
	}
	for i := uint64(0); i < numIters; i++ {
		if rng.Float64() < prob {
			return i, true
		}
	}
	return 0, false
}
<commit_msg>Added a sanity check.<commit_after>package releaseAnalyzer

import (
	"bytes"
	"fmt"
	"github.com/haskelladdict/mbdr/libmbd"
	"log"
	"math"
	"math/rand"
	"sort"
	"time"
)

const (
	numSyt              = 8  // number of synaptotagmin molecules (with 5 Ca2+ sites each)
	numY                = 16 // number of second sensor (Y) sites
	numActiveSyt        = 2  // how many Ca2+ sites need to be bound for sensors
	numActiveY          = 1  // to become active
	vesicleFusionEnergy = 40
	pulseDuration       = 3e-3 // pulse duration in [s]
)

// type of binding site (syt or second sensor)
const (
	sytSite = iota
	ySite
)

// CaSensor defines a single synaptotagmin and Y sites
type caSensor struct {
	sites    []int // ca sites contributing to sensors
	siteType int   // type of sensor (syt or Y)
}

// ActEvent keeps track of a single activation/deactivation event
type ActEvent struct {
	sensorID  int    // sensor which was activated/deactivated
	vesicleID string // vesicleID were activation event took place
	eventIter int    // iteration when event occured
	activated bool   // activated is set to true and deactivated otherwise
}

// sort infrastructure for sorting ActEvents according to the event time
type byIter []ActEvent

func (e byIter) Len() int {
	return len(e)
}

func (e byIter) Swap(i, j int) {
	e[i], e[j] = e[j], e[i]
}

func (e byIter) Less(i, j int) bool {
	return e[i].eventIter < e[j].eventIter
}

// ReleaseEvent keeps track of vesicle release events
type ReleaseEvent struct {
	sensors   []int  // list of sensors involved in release event
	vesicleID string // id of vesicle which was released
	eventIter uint64 // iteration when event occurred
}

var caSensors []caSensor

// random number generator for Metropolis-Hasting
var rng *rand.Rand

func init() {
	rng = rand.New(rand.NewSource(time.Now().UnixNano()))

	caSensors = make([]caSensor, numSyt+numY)

	// define synaptogamin and Y sites
	caSensors[0] = caSensor{[]int{8, 9, 29, 30, 31}, sytSite}
	caSensors[1] = caSensor{[]int{7, 32, 33, 34, 35}, sytSite}
	caSensors[2] = caSensor{[]int{3, 6, 36, 37, 38}, sytSite}
	caSensors[3] = caSensor{[]int{17, 39, 40, 41, 42}, sytSite}
	caSensors[4] = caSensor{[]int{15, 16, 43, 44, 45}, sytSite}
	caSensors[5] = caSensor{[]int{14, 46, 47, 48, 49}, sytSite}
	caSensors[6] = caSensor{[]int{4, 12, 24, 50, 51}, sytSite}
	caSensors[7] = caSensor{[]int{10, 25, 26, 27, 28}, sytSite}
	caSensors[8] = caSensor{[]int{122}, ySite}
	caSensors[9] = caSensor{[]int{70}, ySite}
	caSensors[10] = caSensor{[]int{126}, ySite}
	caSensors[11] = caSensor{[]int{142}, ySite}
	caSensors[12] = caSensor{[]int{62}, ySite}
	caSensors[13] = caSensor{[]int{118}, ySite}
	caSensors[14] = caSensor{[]int{22}, ySite}
	caSensors[15] = caSensor{[]int{134}, ySite}
	caSensors[16] = caSensor{[]int{110}, ySite}
	caSensors[17] = caSensor{[]int{66}, ySite}
	caSensors[18] = caSensor{[]int{106}, ySite}
	caSensors[19] = caSensor{[]int{130}, ySite}
	caSensors[20] = caSensor{[]int{2}, ySite}
	caSensors[21] = caSensor{[]int{114}, ySite}
	caSensors[22] = caSensor{[]int{42}, ySite}
	caSensors[23] = caSensor{[]int{138}, ySite}
}

// analyze is the main entry point for analyzing the mouse AZ model. It
// determines release events and collects statistics
func analyze(data *libmbd.MCellData, model *ReleaseModel, energyModel bool, seed,
	numPulses, numActiveSites, sytEnergy, yEnergy int) ([]string, error) {

	var releases []*ReleaseEvent
	for _, vesID := range model.VesicleIDs {
		evts, err := extractActivationEvents(data, numPulses, seed, vesID,
			model.SensorTemplateString)
		if err != nil {
			return nil, err
		}
		if evts == nil {
			continue
		}

		rel, err := extractReleaseEvents(evts, data.BlockSize(), energyModel,
			numActiveSites, sytEnergy, yEnergy, vesID)
		if err != nil {
			return nil, err
		}
		if rel != nil {
			releases = append(releases, rel)
		}
	}

	return assembleReleaseMsgs(data, seed, releases), nil
}

// assembleReleaseMsgs creates a slice of strings with summary statistics for all
// released vesicles for a given seed
func assembleReleaseMsgs(data *libmbd.MCellData, seed int, rel []*ReleaseEvent) []string {
	messages := make([]string, 0)
	timeStep := data.StepSize()
	for _, r := range rel {
		buffer := bytes.NewBufferString("")
		channels, err := determineCaChanContrib(data, r)
		if err != nil {
			log.Fatal(err)
		}
		if err := checkCaNumbers(channels, r); err != nil {
			fmt.Printf("In seed %d, vesicle %s, time %f\n", seed, r.vesicleID,
				float64(r.eventIter)*data.StepSize())
			log.Fatal(err)
		}

		eventTime := float64(r.eventIter) * timeStep
		// figure out if event happened within or between pulses
		pulseID := int(math.Floor(eventTime / isiValue))
		var pulseString string
		if eventTime-float64(pulseID)*isiValue > pulseDuration {
			pulseString = fmt.Sprintf("ISI_%d", pulseID+1)
		} else {
			pulseString = fmt.Sprintf("%d", pulseID+1)
		}

		fmt.Fprintf(buffer, "seed : %d   vesicleID : %s   time : %e   pulseID : %s", seed,
			r.vesicleID, eventTime, pulseString)
		fmt.Fprintf(buffer, "  sensors: [")
		for _, s := range r.sensors {
			fmt.Fprintf(buffer, "%d ", s)
		}
		fmt.Fprintf(buffer, "]")
		fmt.Fprintf(buffer, "  channels: [")
		for n, c := range channels {
			fmt.Fprintf(buffer, "%s : %d  ", n, int(c))
		}
		fmt.Fprintf(buffer, "]")
		messages = append(messages, buffer.String())
	}
	return messages
}

// extractActivationEvents returns a slice with actvation and deactivation events
// for the given vesicle and active zone
func extractActivationEvents(data *libmbd.MCellData, numPulses, seed int,
	vesicleID, template string) ([]ActEvent, error) {

	var events []ActEvent
	// analyze the activation/deactivation status of each ca sensor.
	// NOTE: for now we merge the binding data for individual pulses into one
	for id := 0; id < len(caSensors); id++ {
		sensor := caSensors[id]
		sensorString := "sensor"
		actThresh := numActiveSyt
		if sensor.siteType == ySite {
			sensorString = "sensor_Y"
			actThresh = numActiveY
		}

		sensorData := make([]int, data.BlockSize())
		for _, s := range sensor.sites {
			for p := 0; p < numPulses; p++ {
				dataName := fmt.Sprintf(template, vesicleID, sensorString, s, p+1, seed)
				bd, err := data.BlockDataByName(dataName)
				if err != nil {
					return nil, err
				}

				if len(bd.Col) != 1 {
					return nil, fmt.Errorf("data set %s had more than one data column",
						dataName)
				}
				for i := 0; i < len(sensorData); i++ {
					sensorData[i] += int(bd.Col[0][i])
				}
			}
		}

		// check for activation events
		active := false
		for i, b := range sensorData {
			if !active && b >= actThresh {
				active = true
				events = append(events, ActEvent{id, vesicleID, i, active})
			} else if active && b < actThresh {
				active = false
				events = append(events, ActEvent{id, vesicleID, i, active})
			}
		}
	}
	return events, nil
}

// extractReleaseEvents determines if the given vesicle was released given
// a list of sensor activation events. If no release took place returns nil.
func extractReleaseEvents(evts []ActEvent, maxIter uint64, energyModel bool,
	numActiveSites, sytEnergy, yEnergy int, vesicleID string) (*ReleaseEvent, error) {

	sort.Sort(byIter(evts))
	activeEvts := make(map[int]struct{})
	for i, e := range evts {
		_, present := activeEvts[e.sensorID]
		if e.activated {
			if present {
				return nil, fmt.Errorf("trying to add active event that already exists")
			}
			activeEvts[e.sensorID] = struct{}{}
		} else {
			if !present {
				return nil, fmt.Errorf("trying to remove a nonexisting active event")
			}
			delete(activeEvts, e.sensorID)
		}

		// special case: If the next event happens simultaneously we apply it right away
		if i+1 < len(evts) && evts[i+1].eventIter == e.eventIter {
			continue
		}

		var rel *ReleaseEvent
		var relError error
		if energyModel {
			// use the energy model to determine release
			energy := getEnergy(activeEvts, sytEnergy, yEnergy)
			nextIter := getNextIter(i, maxIter, evts)
			rel, relError = checkForEnergyRelease(energy, vesicleID, e, activeEvts, nextIter)
		} else {
			// use the deterministic model to determine release
			rel, relError = checkForDeterministicRelease(vesicleID, numActiveSites, e, activeEvts)
		}
		if relError != nil {
			return nil, relError
		}
		if rel != nil {
			return rel, nil
		}
	}
	return nil, nil
}

// getEnergy computes the total energy corresponding to the current number
// of active synaptotagmin and Y sites. Also returns the number of active syts
func getEnergy(events map[int]struct{}, sytEnergy, yEnergy int) int {
	var energy int
	for s, _ := range events {
		if caSensors[s].siteType == sytSite {
			energy += sytEnergy
		} else {
			energy += yEnergy
		}
	}
	return energy
}

// getNextIter determines the iteration of the next event that will happen in
// the event queue
func getNextIter(iter int, maxIter uint64, evts []ActEvent) uint64 {
	nextIter := maxIter
	if iter < len(evts)-1 {
		nextIter = uint64(evts[iter+1].eventIter)
	}
	return nextIter
}

// checkForDeterministicRelease tests if vesicles are released according
// to a deterministic critertion, i.e. as soon as numActiveSites syt or
// Y sites are active
func checkForDeterministicRelease(vesID string, numActiveSites int, evt ActEvent,
	activeEvts map[int]struct{}) (*ReleaseEvent, error) {
	if len(activeEvts) == numActiveSites {
		var sensors []int
		for a, _ := range activeEvts {
			sensors = append(sensors, a)
		}
		return &ReleaseEvent{sensors: sensors, vesicleID: vesID,
			eventIter: uint64(evt.eventIter)}, nil
	}
	return nil, nil
}

// checkForEnergyRelease tests if an energy release according to specified
// syt and y site energies takes place. Check for releases given the current
// energy until next event or the end of simulation. To do this we basically
// test for each iteration between now and the next event if a release takes
// place using the Metrolpolis-Hasting algorithm
func checkForEnergyRelease(energy int, vesID string, evt ActEvent,
	activeEvts map[int]struct{}, nextIter uint64) (*ReleaseEvent, error) {

	numIters := nextIter - uint64(evt.eventIter)
	if nextIter < uint64(evt.eventIter) {
		return nil, fmt.Errorf("encountered out of order release event")
	}
	if iter, ok := checkForRelease(energy, numIters); ok {
		var sensors []int
		for a, _ := range activeEvts {
			sensors = append(sensors, a)
		}
		return &ReleaseEvent{sensors: sensors, vesicleID: vesID,
			eventIter: uint64(evt.eventIter) + iter}, nil
	}
	return nil, nil
}

// checkForReleases uses a Metropolis-Hasting scheme to test numIter times
// if vesicle release happens given the provided bound sensor energy
func checkForRelease(energy int, numIters uint64) (uint64, bool) {

	if energy >= vesicleFusionEnergy {
		return 0, true
	}

	prob := math.Exp(float64(energy - vesicleFusionEnergy))
	if prob >= 1 {
		log.Fatal("probability out of bounds")
	}
	for i := uint64(0); i < numIters; i++ {
		if rng.Float64() < prob {
			return i, true
		}
	}
	return 0, false
}

// checkCaNumbers does a sanity check to ensure that the number of bound
// calcium ions is equal or larger than what is expected based on the activated
// syt and Y sites
func checkCaNumbers(channels map[string]float64, r *ReleaseEvent) error {
	var expected int
	for _, s := range r.sensors {
		if caSensors[s].siteType == sytSite {
			expected += 2
		} else if caSensors[s].siteType == ySite {
			expected += 1
		} else {
			return fmt.Errorf("in checkCaNumbers: Encountered incorrect binding site "+
				"type %d", caSensors[s].siteType)
		}
	}

	var actual int
	for _, c := range channels {
		actual += int(c)
	}

	if actual < expected {
		return fmt.Errorf("The number of bound Ca ions (%d) is smaller than expected "+
			"based on the activation status (%d)", actual, expected)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package grpctmclient

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"runtime"
	"runtime/pprof"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/nettest"
	"google.golang.org/grpc"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/vttablet/grpctmserver"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager"
	"vitess.io/vitess/go/vt/vttablet/tmrpctest"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

func grpcTestServer(t testing.TB, tm tabletmanager.RPCTM) (*net.TCPAddr, func()) {
	t.Helper()

	lis, err := nettest.NewLocalListener("tcp")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	s := grpc.NewServer()
	grpctmserver.RegisterForTest(s, tm)
	go s.Serve(lis)

	var shutdownOnce sync.Once

	return lis.Addr().(*net.TCPAddr), func() {
		shutdownOnce.Do(func() {
			s.Stop()
			lis.Close()
		})
	}
}

func BenchmarkCachedConnClient(b *testing.B) {
	tmserv := tmrpctest.NewFakeRPCTM(b)
	tablets := make([]*topodatapb.Tablet, 4)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(b, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	b.ResetTimer()
	client := NewCachedConnClient(5)
	defer client.Close()

	for i := 0; i < b.N; i++ {
		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100)
		x := rand.Intn(len(tablets))
		err := client.Ping(ctx, tablets[x])
		if err != nil {
			b.Errorf("error pinging tablet %v: %w", tablets[x].Hostname, err)
		}

		cancel()
	}
}

func TestCachedConnClient(t *testing.T) {
	t.Parallel()

	if os.Getenv("VT_PPROF_TEST") != "" {
		file, err := os.Create(fmt.Sprintf("%s.profile.out", t.Name()))
		require.NoError(t, err)
		defer file.Close()
		if err := pprof.StartCPUProfile(file); err != nil {
			t.Errorf("failed to start cpu profile: %v", err)
			return
		}
		defer pprof.StopCPUProfile()
	}

	testCtx, testCancel := context.WithCancel(context.Background())
	wg := sync.WaitGroup{}
	procs := 0

	wg.Add(1)
	go func() {
		defer wg.Done()
		procs = runtime.NumGoroutine()

		for {
			select {
			case <-testCtx.Done():
				return
			case <-time.After(time.Millisecond * 100):
				newProcs := runtime.NumGoroutine()
				if newProcs > procs {
					procs = newProcs
				}
			}
		}
	}()

	numTablets := 100
	numGoroutines := 8

	tmserv := tmrpctest.NewFakeRPCTM(t)
	tablets := make([]*topodatapb.Tablet, numTablets)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(t, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	poolSize := int(float64(numTablets) * 0.5)
	client := NewCachedConnClient(poolSize)
	defer client.Close()

	dialAttempts := sync2.NewAtomicInt64(0)
	dialErrors := sync2.NewAtomicInt64(0)

	longestDials := make(chan time.Duration, numGoroutines)

	for i := 0; i < numGoroutines; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			attempts := 0
			jitter := time.Second * 0
			longestDial := time.Duration(0)

			for {
				select {
				case <-testCtx.Done():
					dialAttempts.Add(int64(attempts))
					longestDials <- longestDial
					return
				case <-time.After(jitter):
					jitter = time.Millisecond * (time.Duration(rand.Intn(11) + 50))
					attempts++

					tablet := tablets[rand.Intn(len(tablets))]
					start := time.Now()
					_, closer, err := client.dialer.dial(context.Background(), tablet)
					if err != nil {
						dialErrors.Add(1)
						continue
					}

					dialDuration := time.Since(start)
					if dialDuration > longestDial {
						longestDial = dialDuration
					}

					closer.Close()
				}
			}
		}()
	}

	time.Sleep(time.Minute)
	testCancel()
	wg.Wait()
	close(longestDials)

	longestDial := time.Duration(0)
	for dialDuration := range longestDials {
		if dialDuration > longestDial {
			longestDial = dialDuration
		}
	}

	attempts, errors := dialAttempts.Get(), dialErrors.Get()
	assert.Less(t, float64(errors)/float64(attempts), 0.001, fmt.Sprintf("fewer than 0.1%% of dial attempts should fail (attempts = %d, errors = %d, max running procs = %d)", attempts, errors, procs))
	assert.Less(t, errors, int64(1), "at least one dial attempt failed (attempts = %d, errors = %d)", attempts, errors)
	assert.Less(t, longestDial.Milliseconds(), int64(50))
}

func TestCachedConnClient_evictions(t *testing.T) {
	tmserv := tmrpctest.NewFakeRPCTM(t)
	tablets := make([]*topodatapb.Tablet, 5)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(t, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	testCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	connHoldContext, connHoldCancel := context.WithCancel(testCtx)

	client := NewCachedConnClient(len(tablets) - 1)
	for i := 0; i < len(tablets)-1; i++ {
		_, closer, err := client.dialer.dial(context.Background(), tablets[i])
		t.Logf("holding connection open to %d", tablets[i].Alias.Uid)
		require.NoError(t, err)

		ctx := testCtx
		if i == 0 {
			ctx = connHoldContext
		}
		go func(ctx context.Context, closer io.Closer) {
			// Hold on to one connection until the test is done.
			// In the case of tablets[0], hold on to the connection until we
			// signal to close it.
			<-ctx.Done()
			closer.Close()
		}(ctx, closer)
	}

	dialCtx, dialCancel := context.WithTimeout(testCtx, time.Millisecond*50)
	defer dialCancel()

	err := client.Ping(dialCtx, tablets[0]) // this should take the rlock_fast path
	assert.NoError(t, err, "could not redial on inuse cached connection")

	err = client.Ping(dialCtx, tablets[4]) // this will enter the poll loop until context timeout
	assert.Error(t, err, "should have timed out waiting for an eviction, while all conns were held")

	// free up a connection
	connHoldCancel()

	dialCtx, dialCancel = context.WithTimeout(testCtx, time.Millisecond*100)
	defer dialCancel()

	err = client.Ping(dialCtx, tablets[4]) // this will enter the poll loop and evict a connection
	assert.NoError(t, err, "should have evicted a conn and succeeded to dial")
}
<commit_msg>Add steadystate benchmarks<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package grpctmclient

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"runtime"
	"runtime/pprof"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/nettest"
	"google.golang.org/grpc"

	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/vttablet/grpctmserver"
	"vitess.io/vitess/go/vt/vttablet/tabletmanager"
	"vitess.io/vitess/go/vt/vttablet/tmrpctest"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

func grpcTestServer(t testing.TB, tm tabletmanager.RPCTM) (*net.TCPAddr, func()) {
	t.Helper()

	lis, err := nettest.NewLocalListener("tcp")
	if err != nil {
		t.Fatalf("Cannot listen: %v", err)
	}

	s := grpc.NewServer()
	grpctmserver.RegisterForTest(s, tm)
	go s.Serve(lis)

	var shutdownOnce sync.Once

	return lis.Addr().(*net.TCPAddr), func() {
		shutdownOnce.Do(func() {
			s.Stop()
			lis.Close()
		})
	}
}

func BenchmarkCachedConnClient(b *testing.B) {
	tmserv := tmrpctest.NewFakeRPCTM(b)
	tablets := make([]*topodatapb.Tablet, 4)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(b, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	b.ResetTimer()
	client := NewCachedConnClient(5)
	defer client.Close()

	for i := 0; i < b.N; i++ {
		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100)
		x := rand.Intn(len(tablets))
		err := client.Ping(ctx, tablets[x])
		if err != nil {
			b.Errorf("error pinging tablet %v: %w", tablets[x].Hostname, err)
		}

		cancel()
	}
}

func BenchmarkCachedConnClientSteadyState(b *testing.B) {
	tmserv := tmrpctest.NewFakeRPCTM(b)
	tablets := make([]*topodatapb.Tablet, 1000)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(b, tmserv)
		// b.Cleanup(shutdown)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	client := NewCachedConnClient(100)
	// b.Cleanup(client.Close)
	defer client.Close()

	// fill the pool
	for i := 0; i < 100; i++ {
		err := client.Ping(context.Background(), tablets[i])
		require.NoError(b, err)
	}

	procs := runtime.GOMAXPROCS(0) / 4
	if procs == 0 {
		procs = 2
	}

	pingsPerProc := len(tablets) / procs
	if pingsPerProc == 0 {
		pingsPerProc = 2
	}

	b.ResetTimer()

	// Begin the benchmark
	for i := 0; i < b.N; i++ {
		ctx, cancel := context.WithCancel(context.Background())

		var wg sync.WaitGroup
		for j := 0; j < procs; j++ {
			wg.Add(1)
			go func() {
				defer wg.Done()

				for k := 0; k < pingsPerProc; k++ {
					func() {
						ctx, cancel := context.WithTimeout(ctx, time.Second*5)
						defer cancel()

						x := rand.Intn(len(tablets))
						err := client.Ping(ctx, tablets[x])
						assert.NoError(b, err)
					}()
				}
			}()
		}

		wg.Wait()
		cancel()
	}
}

func BenchmarkCachedConnClientSteadyStateRedials(b *testing.B) {
	tmserv := tmrpctest.NewFakeRPCTM(b)
	tablets := make([]*topodatapb.Tablet, 1000)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(b, tmserv)
		// b.Cleanup(shutdown)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	client := NewCachedConnClient(1000)
	// b.Cleanup(client.Close)
	defer client.Close()

	// fill the pool
	for i := 0; i < 1000; i++ {
		err := client.Ping(context.Background(), tablets[i])
		require.NoError(b, err)
	}

	procs := runtime.GOMAXPROCS(0) / 4
	if procs == 0 {
		procs = 2
	}

	pingsPerProc := len(tablets) / procs
	if pingsPerProc == 0 {
		pingsPerProc = 2
	}

	b.ResetTimer()

	// Begin the benchmark
	for i := 0; i < b.N; i++ {
		ctx, cancel := context.WithCancel(context.Background())

		var wg sync.WaitGroup
		for j := 0; j < procs; j++ {
			wg.Add(1)
			go func() {
				defer wg.Done()

				for k := 0; k < pingsPerProc; k++ {
					func() {
						ctx, cancel := context.WithTimeout(ctx, time.Second*5)
						defer cancel()

						x := rand.Intn(len(tablets))
						err := client.Ping(ctx, tablets[x])
						assert.NoError(b, err)
					}()
				}
			}()
		}

		wg.Wait()
		cancel()
	}
}

func BenchmarkCachedConnClientSteadyStateEvictions(b *testing.B) {
	tmserv := tmrpctest.NewFakeRPCTM(b)
	tablets := make([]*topodatapb.Tablet, 1000)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(b, tmserv)
		b.Cleanup(shutdown)

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	client := NewCachedConnClient(100)
	b.Cleanup(client.Close)

	// fill the pool
	for i := 0; i < 100; i++ {
		err := client.Ping(context.Background(), tablets[i])
		require.NoError(b, err)
	}

	assert.Equal(b, len(client.dialer.(*cachedConnDialer).conns), 100)

	procs := runtime.GOMAXPROCS(0) / 4
	if procs == 0 {
		procs = 2
	}

	start := 100
	b.ResetTimer()

	// Begin the benchmark
	for i := 0; i < b.N; i++ {
		ctx, cancel := context.WithCancel(context.Background())
		ch := make(chan int, 100) // 100 dials per iteration

		var wg sync.WaitGroup
		for j := 0; j < procs; j++ {
			wg.Add(1)
			go func() {
				defer wg.Done()

				for idx := range ch {
					func() {
						ctx, cancel := context.WithTimeout(ctx, time.Second*5)
						defer cancel()

						err := client.Ping(ctx, tablets[idx])
						assert.NoError(b, err)
					}()
				}
			}()
		}

		for j := 0; j < cap(ch); j++ {
			start = (start + j) % 1000 // go in increasing order, wrapping around
			ch <- start
		}

		close(ch)
		wg.Wait()
		cancel()
	}
}

func TestCachedConnClient(t *testing.T) {
	t.Parallel()

	if os.Getenv("VT_PPROF_TEST") != "" {
		file, err := os.Create(fmt.Sprintf("%s.profile.out", t.Name()))
		require.NoError(t, err)
		defer file.Close()
		if err := pprof.StartCPUProfile(file); err != nil {
			t.Errorf("failed to start cpu profile: %v", err)
			return
		}
		defer pprof.StopCPUProfile()
	}

	testCtx, testCancel := context.WithCancel(context.Background())
	wg := sync.WaitGroup{}
	procs := 0

	wg.Add(1)
	go func() {
		defer wg.Done()
		procs = runtime.NumGoroutine()

		for {
			select {
			case <-testCtx.Done():
				return
			case <-time.After(time.Millisecond * 100):
				newProcs := runtime.NumGoroutine()
				if newProcs > procs {
					procs = newProcs
				}
			}
		}
	}()

	numTablets := 100
	numGoroutines := 8

	tmserv := tmrpctest.NewFakeRPCTM(t)
	tablets := make([]*topodatapb.Tablet, numTablets)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(t, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	poolSize := int(float64(numTablets) * 0.5)
	client := NewCachedConnClient(poolSize)
	defer client.Close()

	dialAttempts := sync2.NewAtomicInt64(0)
	dialErrors := sync2.NewAtomicInt64(0)

	longestDials := make(chan time.Duration, numGoroutines)

	for i := 0; i < numGoroutines; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			attempts := 0
			jitter := time.Second * 0
			longestDial := time.Duration(0)

			for {
				select {
				case <-testCtx.Done():
					dialAttempts.Add(int64(attempts))
					longestDials <- longestDial
					return
				case <-time.After(jitter):
					jitter = time.Millisecond * (time.Duration(rand.Intn(11) + 50))
					attempts++

					tablet := tablets[rand.Intn(len(tablets))]
					start := time.Now()
					_, closer, err := client.dialer.dial(context.Background(), tablet)
					if err != nil {
						dialErrors.Add(1)
						continue
					}

					dialDuration := time.Since(start)
					if dialDuration > longestDial {
						longestDial = dialDuration
					}

					closer.Close()
				}
			}
		}()
	}

	time.Sleep(time.Minute)
	testCancel()
	wg.Wait()
	close(longestDials)

	longestDial := time.Duration(0)
	for dialDuration := range longestDials {
		if dialDuration > longestDial {
			longestDial = dialDuration
		}
	}

	attempts, errors := dialAttempts.Get(), dialErrors.Get()
	assert.Less(t, float64(errors)/float64(attempts), 0.001, fmt.Sprintf("fewer than 0.1%% of dial attempts should fail (attempts = %d, errors = %d, max running procs = %d)", attempts, errors, procs))
	assert.Less(t, errors, int64(1), "at least one dial attempt failed (attempts = %d, errors = %d)", attempts, errors)
	assert.Less(t, longestDial.Milliseconds(), int64(50))
}

func TestCachedConnClient_evictions(t *testing.T) {
	tmserv := tmrpctest.NewFakeRPCTM(t)
	tablets := make([]*topodatapb.Tablet, 5)
	for i := 0; i < len(tablets); i++ {
		addr, shutdown := grpcTestServer(t, tmserv)
		defer shutdown()

		tablets[i] = &topodatapb.Tablet{
			Alias: &topodatapb.TabletAlias{
				Cell: "test",
				Uid:  uint32(addr.Port),
			},
			Hostname: addr.IP.String(),
			PortMap: map[string]int32{
				"grpc": int32(addr.Port),
			},
		}
	}

	testCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	connHoldContext, connHoldCancel := context.WithCancel(testCtx)

	client := NewCachedConnClient(len(tablets) - 1)
	for i := 0; i < len(tablets)-1; i++ {
		_, closer, err := client.dialer.dial(context.Background(), tablets[i])
		t.Logf("holding connection open to %d", tablets[i].Alias.Uid)
		require.NoError(t, err)

		ctx := testCtx
		if i == 0 {
			ctx = connHoldContext
		}
		go func(ctx context.Context, closer io.Closer) {
			// Hold on to one connection until the test is done.
			// In the case of tablets[0], hold on to the connection until we
			// signal to close it.
			<-ctx.Done()
			closer.Close()
		}(ctx, closer)
	}

	dialCtx, dialCancel := context.WithTimeout(testCtx, time.Millisecond*50)
	defer dialCancel()

	err := client.Ping(dialCtx, tablets[0]) // this should take the rlock_fast path
	assert.NoError(t, err, "could not redial on inuse cached connection")

	err = client.Ping(dialCtx, tablets[4]) // this will enter the poll loop until context timeout
	assert.Error(t, err, "should have timed out waiting for an eviction, while all conns were held")

	// free up a connection
	connHoldCancel()

	dialCtx, dialCancel = context.WithTimeout(testCtx, time.Millisecond*100)
	defer dialCancel()

	err = client.Ping(dialCtx, tablets[4]) // this will enter the poll loop and evict a connection
	assert.NoError(t, err, "should have evicted a conn and succeeded to dial")
}
<|endoftext|>


<commit_before>package controllers

import (
	"html/template"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/qor/action_bar"
	"github.com/qor/qor-example/app/models"
	"github.com/qor/qor-example/config"
	"github.com/qor/qor-example/config/admin"
	"github.com/qor/seo"
)

func ProductShow(ctx *gin.Context) {
	var (
		product        models.Product
		colorVariation models.ColorVariation
		seoSetting     models.SEOSetting
		codes          = strings.Split(ctx.Param("code"), "_")
		productCode    = codes[0]
		colorCode      string
	)

	if len(codes) > 1 {
		colorCode = codes[1]
	}

	if DB(ctx).Where(&models.Product{Code: productCode}).First(&product).RecordNotFound() {
		http.Redirect(ctx.Writer, ctx.Request, "/", http.StatusFound)
	}

	DB(ctx).Preload("Product").Preload("Color").Preload("SizeVariations.Size").Where(&models.ColorVariation{ProductID: product.ID, ColorCode: colorCode}).First(&colorVariation)
	DB(ctx).First(&seoSetting)

	editButton := admin.ActionBar.RenderEditButton(ctx.Writer, ctx.Request, product)
	config.View.Funcs(funcsMap(ctx)).Execute(
		"product_show",
		gin.H{
			"ActionBarTag":   admin.ActionBar.Render(ctx.Writer, ctx.Request, action_bar.Config{InlineActions: []template.HTML{editButton}}),
			"Product":        product,
			"ColorVariation": colorVariation,
			"SeoTag":         seoSetting.ProductPage.Render(seoSetting, product),
			"MicroProduct": seo.MicroProduct{
				Name:        product.Name,
				Description: product.Description,
				BrandName:   product.Category.Name,
				SKU:         product.Code,
				Price:       float64(product.Price),
				Image:       colorVariation.MainImageURL(),
			}.Render(),
			"CurrentUser":   CurrentUser(ctx),
			"CurrentLocale": CurrentLocale(ctx),
		},
		ctx.Request,
		ctx.Writer,
	)
}

func funcsMap(ctx *gin.Context) template.FuncMap {
	funcMaps := map[string]interface{}{
		"related_products": func(cv models.ColorVariation) []models.Product {
			var products []models.Product
			DB(ctx).Preload("ColorVariations").Limit(4).Find(&products, "id <> ?", cv.ProductID)
			return products
		},
		"other_also_bought": func(cv models.ColorVariation) []models.Product {
			var products []models.Product
			DB(ctx).Preload("ColorVariations").Order("id ASC").Limit(8).Find(&products, "id <> ?", cv.ProductID)
			return products
		},
	}
	for key, value := range I18nFuncMap(ctx) {
		funcMaps[key] = value
	}
	for key, value := range admin.ActionBar.FuncMap(ctx.Writer, ctx.Request) {
		funcMaps[key] = value
	}
	return funcMaps
}
<commit_msg>seo: update ActionBar api<commit_after>package controllers

import (
	"html/template"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/qor/action_bar"
	"github.com/qor/qor-example/app/models"
	"github.com/qor/qor-example/config"
	"github.com/qor/qor-example/config/admin"
	"github.com/qor/seo"
)

func ProductShow(ctx *gin.Context) {
	var (
		product        models.Product
		colorVariation models.ColorVariation
		seoSetting     models.SEOSetting
		codes          = strings.Split(ctx.Param("code"), "_")
		productCode    = codes[0]
		colorCode      string
	)

	if len(codes) > 1 {
		colorCode = codes[1]
	}

	if DB(ctx).Where(&models.Product{Code: productCode}).First(&product).RecordNotFound() {
		http.Redirect(ctx.Writer, ctx.Request, "/", http.StatusFound)
	}

	DB(ctx).Preload("Product").Preload("Color").Preload("SizeVariations.Size").Where(&models.ColorVariation{ProductID: product.ID, ColorCode: colorCode}).First(&colorVariation)
	DB(ctx).First(&seoSetting)

	editButton := admin.ActionBar.RenderEditButtonWithResource(ctx.Writer, ctx.Request, product)
	config.View.Funcs(funcsMap(ctx)).Execute(
		"product_show",
		gin.H{
			"ActionBarTag":   admin.ActionBar.Render(ctx.Writer, ctx.Request, action_bar.Config{InlineActions: []template.HTML{editButton}}),
			"Product":        product,
			"ColorVariation": colorVariation,
			"SeoTag":         seoSetting.ProductPage.Render(seoSetting, product),
			"MicroProduct": seo.MicroProduct{
				Name:        product.Name,
				Description: product.Description,
				BrandName:   product.Category.Name,
				SKU:         product.Code,
				Price:       float64(product.Price),
				Image:       colorVariation.MainImageURL(),
			}.Render(),
			"CurrentUser":   CurrentUser(ctx),
			"CurrentLocale": CurrentLocale(ctx),
		},
		ctx.Request,
		ctx.Writer,
	)
}

func funcsMap(ctx *gin.Context) template.FuncMap {
	funcMaps := map[string]interface{}{
		"related_products": func(cv models.ColorVariation) []models.Product {
			var products []models.Product
			DB(ctx).Preload("ColorVariations").Limit(4).Find(&products, "id <> ?", cv.ProductID)
			return products
		},
		"other_also_bought": func(cv models.ColorVariation) []models.Product {
			var products []models.Product
			DB(ctx).Preload("ColorVariations").Order("id ASC").Limit(8).Find(&products, "id <> ?", cv.ProductID)
			return products
		},
	}
	for key, value := range I18nFuncMap(ctx) {
		funcMaps[key] = value
	}
	for key, value := range admin.ActionBar.FuncMap(ctx.Writer, ctx.Request) {
		funcMaps[key] = value
	}
	return funcMaps
}
<|endoftext|>