<commit_before>package expr

import (
	"fmt"
	"math"

	"github.com/grafana/metrictank/api/models"
	schema "gopkg.in/raintank/schema.v1"
)

type FuncKeepLastValue struct {
	in    GraphiteFunc
	limit int64
}

func NewKeepLastValue() GraphiteFunc {
	return &FuncKeepLastValue{limit: math.MaxInt64}
}

func (s *FuncKeepLastValue) Signature() ([]Arg, []Arg) {
	var stub string
	return []Arg{
			ArgSeriesList{val: &s.in},
			ArgIn{key: "limit",
				opt: true,
				args: []Arg{
					ArgInt{val: &s.limit},
					ArgString{val: &stub}, // Allow user to specify 'INF' as value. if so, will fall back to maxInt
				},
			},
		},
		[]Arg{ArgSeriesList{}}
}

func (s *FuncKeepLastValue) Context(context Context) Context {
	return context
}

func (s *FuncKeepLastValue) Exec(cache map[Req][]models.Series) ([]models.Series, error) {
	series, err := s.in.Exec(cache)
	if err != nil {
		return nil, err
	}
	limit := int(s.limit)
	outSeries := make([]models.Series, len(series))
	for i, serie := range series {
		serie.Target = fmt.Sprintf("keepLastValue(%s)", serie.Target)
		serie.QueryPatt = serie.Target

		out := pointSlicePool.Get().([]schema.Point)

		var consecutiveNaNs int
		lastVal := math.NaN()

		for i, p := range serie.Datapoints {
			out = append(out, p)
			if math.IsNaN(p.Val) {
				consecutiveNaNs++
				continue
			}
			if 0 < consecutiveNaNs && consecutiveNaNs <= limit && !math.IsNaN(lastVal) {
				for j := i - consecutiveNaNs; j < i; j++ {
					out[j].Val = lastVal
				}
			}
			consecutiveNaNs = 0
			lastVal = p.Val
		}

		if 0 < consecutiveNaNs && consecutiveNaNs <= limit && !math.IsNaN(lastVal) {
			for i := len(out) - consecutiveNaNs; i < len(out); i++ {
				out[i].Val = lastVal
			}
		}

		serie.Datapoints = out
		outSeries[i] = serie
	}
	cache[Req{}] = append(cache[Req{}], outSeries...)
	return outSeries, nil
}
<commit_msg>make function stable<commit_after>package expr

import (
	"fmt"
	"math"

	"github.com/grafana/metrictank/api/models"
	schema "gopkg.in/raintank/schema.v1"
)

type FuncKeepLastValue struct {
	in    GraphiteFunc
	limit int64
}

func NewKeepLastValue() GraphiteFunc {
	return &FuncKeepLastValue{limit: math.MaxInt64}
}

func (s *FuncKeepLastValue) Signature() ([]Arg, []Arg) {
	var stub string
	return []Arg{
			ArgSeriesList{val: &s.in},
			ArgIn{key: "limit",
				opt: true,
				args: []Arg{
					ArgInt{val: &s.limit},
					// Allow user to specify 'INF' as value. if so, will fall back to maxInt
					// Ignores any other strings (just like Graphite)
					ArgString{val: &stub},
				},
			},
		},
		[]Arg{ArgSeriesList{}}
}

func (s *FuncKeepLastValue) Context(context Context) Context {
	return context
}

func (s *FuncKeepLastValue) Exec(cache map[Req][]models.Series) ([]models.Series, error) {
	series, err := s.in.Exec(cache)
	if err != nil {
		return nil, err
	}
	limit := int(s.limit)
	outSeries := make([]models.Series, len(series))
	for i, serie := range series {
		serie.Target = fmt.Sprintf("keepLastValue(%s)", serie.Target)
		serie.QueryPatt = serie.Target

		out := pointSlicePool.Get().([]schema.Point)

		var consecutiveNaNs int
		lastVal := math.NaN()

		for i, p := range serie.Datapoints {
			out = append(out, p)
			if math.IsNaN(p.Val) {
				consecutiveNaNs++
				continue
			}
			if 0 < consecutiveNaNs && consecutiveNaNs <= limit && !math.IsNaN(lastVal) {
				for j := i - consecutiveNaNs; j < i; j++ {
					out[j].Val = lastVal
				}
			}
			consecutiveNaNs = 0
			lastVal = p.Val
		}

		if 0 < consecutiveNaNs && consecutiveNaNs <= limit && !math.IsNaN(lastVal) {
			for i := len(out) - consecutiveNaNs; i < len(out); i++ {
				out[i].Val = lastVal
			}
		}

		serie.Datapoints = out
		outSeries[i] = serie
	}
	cache[Req{}] = append(cache[Req{}], outSeries...)
	return outSeries, nil
}
<|endoftext|>


<commit_before>package binserde

import (
	"bytes"

	"github.com/quux00/ogonori/obinary/rw"
	"github.com/quux00/ogonori/oerror"
	"github.com/quux00/ogonori/oschema"
)

//
// There is apparently a second "binary serialization" system
// in OrientDB that has inconsistent type constants with the
// other one.
//
// Until I understand it better, for now I'm calling it, the
// "typeserializer" though that is misleading since the binserde.go
// Serializer also reads/writes (de/serializes) types.
//

// from Java client code base where all these extend
// com.orientechnologies.common.serialization.types.OBinarySerializer
const (
	BooleanSerializer                       = 1
	ByteSerializer                          = 2
	CharSerializer                          = 3
	DateSerializer                          = 4
	DateTimeSerializer                      = 5
	DoubleSerializer                        = 6
	FloatSerializer                         = 7
	IntegerSerializer                       = 8
	LinkSerializer                          = 9
	LongSerializer                          = 10
	NullSerializer                          = 11
	ShortSerializer                         = 12
	StringSerializer                        = 13
	CompositeKeySerializer                  = 14
	SimpleKeySerializer                     = 15
	StreamSerializerRID                     = 16
	BinaryTypeSerializer                    = 17
	DecimalSerializer                       = 18
	StreamSerializerListRID                 = 19
	StreamSerializerOldRIDContainer         = 20
	StreamSerializerSBTreeIndexRIDContainer = 21
	PhysicalPositionSerializer              = 50
)

type OBinaryTypeSerializer interface {
	Deserialize(buf *bytes.Buffer) (interface{}, error)
	Serialize(val interface{}) ([]byte, error)
}

var TypeSerializers [21]OBinaryTypeSerializer

type OLinkSerializer struct{}

func (ols OLinkSerializer) Deserialize(buf *bytes.Buffer) (interface{}, error) {
	clusterID, err := rw.ReadShort(buf)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}

	clusterPos, err := rw.ReadLong(buf)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}

	rid := oschema.ORID{ClusterID: clusterID, ClusterPos: clusterPos}
	return &oschema.OLink{RID: rid}, nil
}

//
// Serialize serilializes a *oschema.OLink into the binary format
// required by the OrientDB server.  If the `val` passed in is
// not a *oschema.OLink, the method will panic.
//
func (ols OLinkSerializer) Serialize(val interface{}) ([]byte, error) {
	lnk := val.(*oschema.OLink)
	// TODO: do the big-endian translation of shorts and longs to byte arrays
	//       here manually - rather than creating another bytes.Buffer and passing
	//       it to rw.WriteShort, which passes it to binary.Write

	bs := make([]byte, 2+8) // sz of short + long
	buf := bytes.NewBuffer(bs)
	err := rw.WriteShort(buf, lnk.RID.ClusterID)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}
	err = rw.WriteLong(buf, lnk.RID.ClusterPos)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}
	return buf.Bytes(), nil
}

func init() {
	TypeSerializers[LinkSerializer] = OLinkSerializer{}
}
<commit_msg>Clean up TODO item - serialize without buffer<commit_after>package binserde

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"reflect"

	"github.com/quux00/ogonori/obinary/rw"
	"github.com/quux00/ogonori/oerror"
	"github.com/quux00/ogonori/oschema"
)

//
// There is apparently a second "binary serialization" system
// in OrientDB that has inconsistent type constants with the
// other one.
//
// Until I understand it better, for now I'm calling it, the
// "typeserializer" though that is misleading since the binserde.go
// Serializer also reads/writes (de/serializes) types.
//

// from Java client code base where all these extend
// com.orientechnologies.common.serialization.types.OBinarySerializer
const (
	BooleanSerializer                       = 1
	ByteSerializer                          = 2
	CharSerializer                          = 3
	DateSerializer                          = 4
	DateTimeSerializer                      = 5
	DoubleSerializer                        = 6
	FloatSerializer                         = 7
	IntegerSerializer                       = 8
	LinkSerializer                          = 9
	LongSerializer                          = 10
	NullSerializer                          = 11
	ShortSerializer                         = 12
	StringSerializer                        = 13
	CompositeKeySerializer                  = 14
	SimpleKeySerializer                     = 15
	StreamSerializerRID                     = 16
	BinaryTypeSerializer                    = 17
	DecimalSerializer                       = 18
	StreamSerializerListRID                 = 19
	StreamSerializerOldRIDContainer         = 20
	StreamSerializerSBTreeIndexRIDContainer = 21
	PhysicalPositionSerializer              = 50
)

type OBinaryTypeSerializer interface {
	Deserialize(buf *bytes.Buffer) (interface{}, error)
	Serialize(val interface{}) ([]byte, error)
}

var TypeSerializers [21]OBinaryTypeSerializer

type OLinkSerializer struct{}

func (ols OLinkSerializer) Deserialize(buf *bytes.Buffer) (interface{}, error) {
	clusterID, err := rw.ReadShort(buf)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}

	clusterPos, err := rw.ReadLong(buf)
	if err != nil {
		return nil, oerror.NewTrace(err)
	}

	rid := oschema.ORID{ClusterID: clusterID, ClusterPos: clusterPos}
	return &oschema.OLink{RID: rid}, nil
}

//
// Serialize serilializes a *oschema.OLink into the binary format
// required by the OrientDB server.  If the `val` passed in is
// not a *oschema.OLink, the method will panic.
//
func (ols OLinkSerializer) Serialize(val interface{}) ([]byte, error) {
	lnk, ok := val.(*oschema.OLink)
	if !ok {
		return nil, fmt.Errorf("Invalid LINK should be oschema.OLink, got %s", reflect.TypeOf(val))
	}

	bs := make([]byte, 2+8) // sz of short + long
	binary.BigEndian.PutUint16(bs[0:2], uint16(lnk.RID.ClusterID))
	binary.BigEndian.PutUint64(bs[2:10], uint64(lnk.RID.ClusterPos))
	return bs, nil
}

func init() {
	TypeSerializers[LinkSerializer] = OLinkSerializer{}
}
<|endoftext|>


<commit_before>package handlers

import (
	"context"
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"

	"strconv"

	"github.com/alexellis/faas/gateway/requests"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/client"
)

// DefaultMaxReplicas is the amount of replicas a service will auto-scale up to.
const DefaultMaxReplicas = 20

// MakeAlertHandler handles alerts from Prometheus Alertmanager
func MakeAlertHandler(c *client.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Println("Alert received.")
		body, readErr := ioutil.ReadAll(r.Body)

		log.Println(string(body))

		if readErr != nil {
			log.Println(readErr)
			return
		}

		var req requests.PrometheusAlert
		err := json.Unmarshal(body, &req)
		if err != nil {
			log.Println(err)
			return
		}

		if len(req.Alerts) > 0 {
			if err := scaleService(req, c); err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusInternalServerError)
			} else {
				w.WriteHeader(http.StatusOK)
			}
		}
	}
}

func scaleService(req requests.PrometheusAlert, c *client.Client) error {
	var err error
	serviceName := req.Alerts[0].Labels.FunctionName

	if len(serviceName) > 0 {

		service, _, inspectErr := c.ServiceInspectWithRaw(context.Background(), serviceName)
		if inspectErr == nil {

			currentReplicas := *service.Spec.Mode.Replicated.Replicas
			status := req.Status

			replicaLabel := service.Spec.TaskTemplate.ContainerSpec.Labels["com.faas.max_replicas"]
			maxReplicas := DefaultMaxReplicas
			if len(replicaLabel) > 0 {
				maxReplicas, err = strconv.Atoi(replicaLabel)
				if err != nil {
					log.Printf("Bad replica count: %s, should be uint.\n", replicaLabel)
				}
			}
			newReplicas := CalculateReplicas(status, currentReplicas, uint64(maxReplicas))

			if newReplicas == currentReplicas {
				return nil
			}

			log.Printf("Scaling %s to %d replicas.\n", serviceName, newReplicas)
			service.Spec.Mode.Replicated.Replicas = &newReplicas
			updateOpts := types.ServiceUpdateOptions{}
			updateOpts.RegistryAuthFrom = types.RegistryAuthFromSpec

			response, updateErr := c.ServiceUpdate(context.Background(), service.ID, service.Version, service.Spec, updateOpts)
			if updateErr != nil {
				err = updateErr
			}
			log.Println(response)

		} else {
			err = inspectErr
		}
	}
	return err
}

// CalculateReplicas decides what replica count to set depending on current/desired amount
func CalculateReplicas(status string, currentReplicas uint64, maxReplicas uint64) uint64 {
	newReplicas := currentReplicas
	const step = 5

	if status == "firing" {
		if currentReplicas == 1 {
			// First jump is from 1 to "step" i.e. 1->5
			newReplicas = step
		} else {
			if currentReplicas+step > maxReplicas {
				newReplicas = maxReplicas
			} else {
				newReplicas = currentReplicas + step
			}
		}
	} else { // Resolved event.
		newReplicas = 1
	}
	return newReplicas
}
<commit_msg>Fix for Docker upstream<commit_after>package handlers

import (
	"context"
	"encoding/json"
	"io/ioutil"
	"log"
	"net/http"

	"strconv"

	"github.com/alexellis/faas/gateway/requests"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/client"
)

// DefaultMaxReplicas is the amount of replicas a service will auto-scale up to.
const DefaultMaxReplicas = 20

// MakeAlertHandler handles alerts from Prometheus Alertmanager
func MakeAlertHandler(c *client.Client) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Println("Alert received.")
		body, readErr := ioutil.ReadAll(r.Body)

		log.Println(string(body))

		if readErr != nil {
			log.Println(readErr)
			return
		}

		var req requests.PrometheusAlert
		err := json.Unmarshal(body, &req)
		if err != nil {
			log.Println(err)
			return
		}

		if len(req.Alerts) > 0 {
			if err := scaleService(req, c); err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusInternalServerError)
			} else {
				w.WriteHeader(http.StatusOK)
			}
		}
	}
}

func scaleService(req requests.PrometheusAlert, c *client.Client) error {
	var err error
	serviceName := req.Alerts[0].Labels.FunctionName

	if len(serviceName) > 0 {
		opts := types.ServiceInspectOptions{
			InsertDefaults: true,
		}

		service, _, inspectErr := c.ServiceInspectWithRaw(context.Background(), serviceName, opts)
		if inspectErr == nil {

			currentReplicas := *service.Spec.Mode.Replicated.Replicas
			status := req.Status

			replicaLabel := service.Spec.TaskTemplate.ContainerSpec.Labels["com.faas.max_replicas"]
			maxReplicas := DefaultMaxReplicas
			if len(replicaLabel) > 0 {
				maxReplicas, err = strconv.Atoi(replicaLabel)
				if err != nil {
					log.Printf("Bad replica count: %s, should be uint.\n", replicaLabel)
				}
			}
			newReplicas := CalculateReplicas(status, currentReplicas, uint64(maxReplicas))

			if newReplicas == currentReplicas {
				return nil
			}

			log.Printf("Scaling %s to %d replicas.\n", serviceName, newReplicas)
			service.Spec.Mode.Replicated.Replicas = &newReplicas
			updateOpts := types.ServiceUpdateOptions{}
			updateOpts.RegistryAuthFrom = types.RegistryAuthFromSpec

			response, updateErr := c.ServiceUpdate(context.Background(), service.ID, service.Version, service.Spec, updateOpts)
			if updateErr != nil {
				err = updateErr
			}
			log.Println(response)

		} else {
			err = inspectErr
		}
	}
	return err
}

// CalculateReplicas decides what replica count to set depending on current/desired amount
func CalculateReplicas(status string, currentReplicas uint64, maxReplicas uint64) uint64 {
	newReplicas := currentReplicas
	const step = 5

	if status == "firing" {
		if currentReplicas == 1 {
			// First jump is from 1 to "step" i.e. 1->5
			newReplicas = step
		} else {
			if currentReplicas+step > maxReplicas {
				newReplicas = maxReplicas
			} else {
				newReplicas = currentReplicas + step
			}
		}
	} else { // Resolved event.
		newReplicas = 1
	}
	return newReplicas
}
<|endoftext|>


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

package gcsproxy_test

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease/mock_lease"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

func TestMutableContent(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func bufferIs(buf []byte) Matcher {
	return NewMatcher(
		func(candidate interface{}) error {
			p := candidate.([]byte)

			// Compare.
			if &buf[0] != &p[0] {
				return fmt.Errorf(
					"Differing first bytes: %p vs. %p",
					&buf[0],
					&p[0])
			}

			if len(buf) != len(p) {
				return fmt.Errorf(
					"Differing lengths: %d vs. %d",
					len(buf),
					len(p))
			}

			return nil
		},
		fmt.Sprintf("Buffer matches"))
}

////////////////////////////////////////////////////////////////////////
// Invariant-checking mutable content
////////////////////////////////////////////////////////////////////////

// A wrapper around MutableContent that calls CheckInvariants whenever
// invariants should hold. For catching logic errors early in the test.
type checkingMutableContent struct {
	ctx     context.Context
	wrapped *gcsproxy.MutableContent
}

func (mc *checkingMutableContent) Stat() (gcsproxy.StatResult, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.Stat(mc.ctx)
}

func (mc *checkingMutableContent) ReadAt(b []byte, o int64) (int, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.ReadAt(mc.ctx, b, o)
}

func (mc *checkingMutableContent) WriteAt(b []byte, o int64) (int, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.WriteAt(mc.ctx, b, o)
}

func (mc *checkingMutableContent) Truncate(n int64) error {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.Truncate(mc.ctx, n)
}

func (mc *checkingMutableContent) Destroy() {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	mc.wrapped.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const initialContentSize = 11

type mutableContentTest struct {
	ctx context.Context

	initialContent mock_lease.MockReadProxy
	rwl            mock_lease.MockReadWriteLease
	clock          timeutil.SimulatedClock

	mc checkingMutableContent
}

var _ SetUpInterface = &mutableContentTest{}

func (t *mutableContentTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Set up the mock initial contents, including a default size.
	t.initialContent = mock_lease.NewMockReadProxy(
		ti.MockController,
		"initialContent")

	ExpectCall(t.initialContent, "Size")().
		WillRepeatedly(Return(initialContentSize))

	// Set up a mock read/write lease.
	t.rwl = mock_lease.NewMockReadWriteLease(
		ti.MockController,
		"rwl")

	// Ignore uninteresting calls.
	ExpectCall(t.initialContent, "CheckInvariants")().
		WillRepeatedly(Return())

	// Set up the clock.
	t.clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))

	// And the mutable content.
	t.mc.ctx = ti.Ctx
	t.mc.wrapped = gcsproxy.NewMutableContent(
		t.initialContent,
		&t.clock)
}

////////////////////////////////////////////////////////////////////////
// Clean state
////////////////////////////////////////////////////////////////////////

type CleanTest struct {
	mutableContentTest
}

func init() { RegisterTestSuite(&CleanTest{}) }

func (t *CleanTest) ReadAt_CallsProxy() {
	buf := make([]byte, 1)

	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(t.ctx, bufferIs(buf), 17).
		WillOnce(Return(0, errors.New("")))

	// Call
	t.mc.ReadAt(buf, 17)
}

func (t *CleanTest) ReadAt_ProxyFails() {
	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(Any(), Any(), Any()).
		WillOnce(Return(17, errors.New("taco")))

	// Call
	n, err := t.mc.ReadAt(make([]byte, 1), 0)

	ExpectEq(17, n)
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) ReadAt_ProxySuceeds() {
	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(Any(), Any(), Any()).
		WillOnce(Return(17, nil))

	// Call
	n, err := t.mc.ReadAt(make([]byte, 1), 0)

	ExpectEq(17, n)
	ExpectEq(nil, err)
}

func (t *CleanTest) Stat() {
	sr, err := t.mc.Stat()

	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.Size)
	ExpectEq(initialContentSize, sr.DirtyThreshold)
	ExpectEq(nil, sr.Mtime)
}

func (t *CleanTest) WriteAt_UpgradeFails() {
	// Upgrade
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, err := t.mc.WriteAt(make([]byte, 1), 0)

	ExpectThat(err, Error(HasSubstr("Upgrade")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) WriteAt_UpgradeSucceeds() {
	// Upgrade -- succeed.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	// The read/write lease should be called.
	ExpectCall(t.rwl, "WriteAt")(Any(), 17).
		WillOnce(Return(0, errors.New("")))

	// Call.
	t.mc.WriteAt(make([]byte, 1), 17)

	// A further call should go right through to the read/write lease again.
	ExpectCall(t.rwl, "WriteAt")(Any(), 19).
		WillOnce(Return(0, errors.New("")))

	t.mc.WriteAt(make([]byte, 1), 19)
}

func (t *CleanTest) Truncate_UpgradeFails() {
	// Upgrade
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	err := t.mc.Truncate(0)

	ExpectThat(err, Error(HasSubstr("Upgrade")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) Truncate_UpgradeSucceeds() {
	// Upgrade -- succeed.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	// The read/write lease should be called.
	ExpectCall(t.rwl, "Truncate")(17).
		WillOnce(Return(errors.New("")))

	// Call.
	t.mc.Truncate(17)

	// A further call should go right through to the read/write lease again.
	ExpectCall(t.rwl, "Truncate")(19).
		WillOnce(Return(errors.New("")))

	t.mc.Truncate(19)
}

////////////////////////////////////////////////////////////////////////
// Dirty state
////////////////////////////////////////////////////////////////////////

type DirtyTest struct {
	mutableContentTest

	setUpTime time.Time
}

func init() { RegisterTestSuite(&DirtyTest{}) }

func (t *DirtyTest) SetUp(ti *TestInfo) {
	t.mutableContentTest.SetUp(ti)
	t.setUpTime = t.clock.Now()

	// Simulate a successful upgrade.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(nil))

	err := t.mc.Truncate(initialContentSize)
	AssertEq(nil, err)

	// Change the time.
	t.clock.AdvanceTime(time.Second)
}

func (t *DirtyTest) ReadAt_CallsLease() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) ReadAt_LeaseFails() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) ReadAt_LeaseSuceeds() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) Stat_LeaseFails() {
	// Lease
	ExpectCall(t.rwl, "Size")().
		WillOnce(Return(0, errors.New("taco")))

	// Call
	_, err := t.mc.Stat()
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *DirtyTest) Stat_LeaseSucceeds() {
	// Lease
	ExpectCall(t.rwl, "Size")().
		WillOnce(Return(17, nil))

	// Call
	sr, err := t.mc.Stat()
	AssertEq(nil, err)

	// Check the initial state.
	ExpectEq(17, sr.Size)
	ExpectEq(initialContentSize, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.setUpTime)))
}

func (t *DirtyTest) WriteAt_CallsLease() {
	buf := make([]byte, 4)
	const offset = 17

	// Lease
	ExpectCall(t.rwl, "WriteAt")(bufferIs(buf), offset).
		WillOnce(Return(0, errors.New("")))

	// Call
	t.mc.WriteAt(buf, offset)
}

func (t *DirtyTest) WriteAt_LeaseFails() {
	const offset = initialContentSize - 2

	// Lease
	ExpectCall(t.rwl, "WriteAt")(Any(), Any()).
		WillOnce(Return(13, errors.New("taco")))

	// Call
	n, err := t.mc.WriteAt([]byte{}, offset)

	ExpectEq(13, n)
	ExpectThat(err, Error(HasSubstr("taco")))

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(initialContentSize, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(offset, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) WriteAt_LeaseSucceeds() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) WriteAt_DirtyThreshold() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) Truncate_CallsLease() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(17).
		WillOnce(Return(errors.New("")))

	// Call
	t.mc.Truncate(17)
}

func (t *DirtyTest) Truncate_LeaseFails() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(errors.New("taco")))

	// Call
	err := t.mc.Truncate(1)
	ExpectThat(err, Error(HasSubstr("taco")))

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(0, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(1, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) Truncate_LeaseSucceeds() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(nil))

	// Call
	err := t.mc.Truncate(1)
	ExpectEq(nil, err)

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(0, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(1, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) Truncate_DirtyThreshold() {
	var sr gcsproxy.StatResult
	var err error

	// Simulate successful truncations and size requests.
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillRepeatedly(Return(nil))

	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(100, nil))

	// Truncating to the same size should not affect the dirty threshold.
	err = t.mc.Truncate(initialContentSize)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.DirtyThreshold)

	// Nor should truncating upward.
	err = t.mc.Truncate(initialContentSize + 100)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.DirtyThreshold)

	// But truncating downward should.
	err = t.mc.Truncate(initialContentSize - 1)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize-1, sr.DirtyThreshold)
}
<commit_msg>DirtyTest.WriteAt_LeaseSucceeds<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcsproxy_test

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/googlecloudplatform/gcsfuse/gcsproxy"
	"github.com/googlecloudplatform/gcsfuse/lease/mock_lease"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

func TestMutableContent(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func bufferIs(buf []byte) Matcher {
	return NewMatcher(
		func(candidate interface{}) error {
			p := candidate.([]byte)

			// Compare.
			if &buf[0] != &p[0] {
				return fmt.Errorf(
					"Differing first bytes: %p vs. %p",
					&buf[0],
					&p[0])
			}

			if len(buf) != len(p) {
				return fmt.Errorf(
					"Differing lengths: %d vs. %d",
					len(buf),
					len(p))
			}

			return nil
		},
		fmt.Sprintf("Buffer matches"))
}

////////////////////////////////////////////////////////////////////////
// Invariant-checking mutable content
////////////////////////////////////////////////////////////////////////

// A wrapper around MutableContent that calls CheckInvariants whenever
// invariants should hold. For catching logic errors early in the test.
type checkingMutableContent struct {
	ctx     context.Context
	wrapped *gcsproxy.MutableContent
}

func (mc *checkingMutableContent) Stat() (gcsproxy.StatResult, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.Stat(mc.ctx)
}

func (mc *checkingMutableContent) ReadAt(b []byte, o int64) (int, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.ReadAt(mc.ctx, b, o)
}

func (mc *checkingMutableContent) WriteAt(b []byte, o int64) (int, error) {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.WriteAt(mc.ctx, b, o)
}

func (mc *checkingMutableContent) Truncate(n int64) error {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	return mc.wrapped.Truncate(mc.ctx, n)
}

func (mc *checkingMutableContent) Destroy() {
	mc.wrapped.CheckInvariants()
	defer mc.wrapped.CheckInvariants()
	mc.wrapped.Destroy()
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const initialContentSize = 11

type mutableContentTest struct {
	ctx context.Context

	initialContent mock_lease.MockReadProxy
	rwl            mock_lease.MockReadWriteLease
	clock          timeutil.SimulatedClock

	mc checkingMutableContent
}

var _ SetUpInterface = &mutableContentTest{}

func (t *mutableContentTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx

	// Set up the mock initial contents, including a default size.
	t.initialContent = mock_lease.NewMockReadProxy(
		ti.MockController,
		"initialContent")

	ExpectCall(t.initialContent, "Size")().
		WillRepeatedly(Return(initialContentSize))

	// Set up a mock read/write lease.
	t.rwl = mock_lease.NewMockReadWriteLease(
		ti.MockController,
		"rwl")

	// Ignore uninteresting calls.
	ExpectCall(t.initialContent, "CheckInvariants")().
		WillRepeatedly(Return())

	// Set up the clock.
	t.clock.SetTime(time.Date(2012, 8, 15, 22, 56, 0, 0, time.Local))

	// And the mutable content.
	t.mc.ctx = ti.Ctx
	t.mc.wrapped = gcsproxy.NewMutableContent(
		t.initialContent,
		&t.clock)
}

////////////////////////////////////////////////////////////////////////
// Clean state
////////////////////////////////////////////////////////////////////////

type CleanTest struct {
	mutableContentTest
}

func init() { RegisterTestSuite(&CleanTest{}) }

func (t *CleanTest) ReadAt_CallsProxy() {
	buf := make([]byte, 1)

	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(t.ctx, bufferIs(buf), 17).
		WillOnce(Return(0, errors.New("")))

	// Call
	t.mc.ReadAt(buf, 17)
}

func (t *CleanTest) ReadAt_ProxyFails() {
	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(Any(), Any(), Any()).
		WillOnce(Return(17, errors.New("taco")))

	// Call
	n, err := t.mc.ReadAt(make([]byte, 1), 0)

	ExpectEq(17, n)
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) ReadAt_ProxySuceeds() {
	// Proxy
	ExpectCall(t.initialContent, "ReadAt")(Any(), Any(), Any()).
		WillOnce(Return(17, nil))

	// Call
	n, err := t.mc.ReadAt(make([]byte, 1), 0)

	ExpectEq(17, n)
	ExpectEq(nil, err)
}

func (t *CleanTest) Stat() {
	sr, err := t.mc.Stat()

	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.Size)
	ExpectEq(initialContentSize, sr.DirtyThreshold)
	ExpectEq(nil, sr.Mtime)
}

func (t *CleanTest) WriteAt_UpgradeFails() {
	// Upgrade
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, err := t.mc.WriteAt(make([]byte, 1), 0)

	ExpectThat(err, Error(HasSubstr("Upgrade")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) WriteAt_UpgradeSucceeds() {
	// Upgrade -- succeed.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	// The read/write lease should be called.
	ExpectCall(t.rwl, "WriteAt")(Any(), 17).
		WillOnce(Return(0, errors.New("")))

	// Call.
	t.mc.WriteAt(make([]byte, 1), 17)

	// A further call should go right through to the read/write lease again.
	ExpectCall(t.rwl, "WriteAt")(Any(), 19).
		WillOnce(Return(0, errors.New("")))

	t.mc.WriteAt(make([]byte, 1), 19)
}

func (t *CleanTest) Truncate_UpgradeFails() {
	// Upgrade
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	err := t.mc.Truncate(0)

	ExpectThat(err, Error(HasSubstr("Upgrade")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *CleanTest) Truncate_UpgradeSucceeds() {
	// Upgrade -- succeed.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	// The read/write lease should be called.
	ExpectCall(t.rwl, "Truncate")(17).
		WillOnce(Return(errors.New("")))

	// Call.
	t.mc.Truncate(17)

	// A further call should go right through to the read/write lease again.
	ExpectCall(t.rwl, "Truncate")(19).
		WillOnce(Return(errors.New("")))

	t.mc.Truncate(19)
}

////////////////////////////////////////////////////////////////////////
// Dirty state
////////////////////////////////////////////////////////////////////////

type DirtyTest struct {
	mutableContentTest

	setUpTime time.Time
}

func init() { RegisterTestSuite(&DirtyTest{}) }

func (t *DirtyTest) SetUp(ti *TestInfo) {
	t.mutableContentTest.SetUp(ti)
	t.setUpTime = t.clock.Now()

	// Simulate a successful upgrade.
	ExpectCall(t.initialContent, "Upgrade")(Any()).
		WillOnce(Return(t.rwl, nil))

	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(nil))

	err := t.mc.Truncate(initialContentSize)
	AssertEq(nil, err)

	// Change the time.
	t.clock.AdvanceTime(time.Second)
}

func (t *DirtyTest) ReadAt_CallsLease() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) ReadAt_LeaseFails() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) ReadAt_LeaseSuceeds() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) Stat_LeaseFails() {
	// Lease
	ExpectCall(t.rwl, "Size")().
		WillOnce(Return(0, errors.New("taco")))

	// Call
	_, err := t.mc.Stat()
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *DirtyTest) Stat_LeaseSucceeds() {
	// Lease
	ExpectCall(t.rwl, "Size")().
		WillOnce(Return(17, nil))

	// Call
	sr, err := t.mc.Stat()
	AssertEq(nil, err)

	// Check the initial state.
	ExpectEq(17, sr.Size)
	ExpectEq(initialContentSize, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.setUpTime)))
}

func (t *DirtyTest) WriteAt_CallsLease() {
	buf := make([]byte, 4)
	const offset = 17

	// Lease
	ExpectCall(t.rwl, "WriteAt")(bufferIs(buf), offset).
		WillOnce(Return(0, errors.New("")))

	// Call
	t.mc.WriteAt(buf, offset)
}

func (t *DirtyTest) WriteAt_LeaseFails() {
	const offset = initialContentSize - 2

	// Lease
	ExpectCall(t.rwl, "WriteAt")(Any(), Any()).
		WillOnce(Return(13, errors.New("taco")))

	// Call
	n, err := t.mc.WriteAt([]byte{}, offset)

	ExpectEq(13, n)
	ExpectThat(err, Error(HasSubstr("taco")))

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(initialContentSize, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(offset, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) WriteAt_LeaseSucceeds() {
	const offset = initialContentSize - 2

	// Lease
	ExpectCall(t.rwl, "WriteAt")(Any(), Any()).
		WillOnce(Return(13, nil))

	// Call
	n, err := t.mc.WriteAt([]byte{}, offset)

	ExpectEq(13, n)
	ExpectEq(nil, err)

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(initialContentSize, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(offset, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) WriteAt_DirtyThreshold() {
	AssertTrue(false, "TODO")
}

func (t *DirtyTest) Truncate_CallsLease() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(17).
		WillOnce(Return(errors.New("")))

	// Call
	t.mc.Truncate(17)
}

func (t *DirtyTest) Truncate_LeaseFails() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(errors.New("taco")))

	// Call
	err := t.mc.Truncate(1)
	ExpectThat(err, Error(HasSubstr("taco")))

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(0, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(1, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) Truncate_LeaseSucceeds() {
	// Lease
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillOnce(Return(nil))

	// Call
	err := t.mc.Truncate(1)
	ExpectEq(nil, err)

	// The dirty threshold and mtime should have been updated.
	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(0, nil))

	sr, err := t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(1, sr.DirtyThreshold)
	ExpectThat(sr.Mtime, Pointee(timeutil.TimeEq(t.clock.Now())))
}

func (t *DirtyTest) Truncate_DirtyThreshold() {
	var sr gcsproxy.StatResult
	var err error

	// Simulate successful truncations and size requests.
	ExpectCall(t.rwl, "Truncate")(Any()).
		WillRepeatedly(Return(nil))

	ExpectCall(t.rwl, "Size")().
		WillRepeatedly(Return(100, nil))

	// Truncating to the same size should not affect the dirty threshold.
	err = t.mc.Truncate(initialContentSize)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.DirtyThreshold)

	// Nor should truncating upward.
	err = t.mc.Truncate(initialContentSize + 100)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize, sr.DirtyThreshold)

	// But truncating downward should.
	err = t.mc.Truncate(initialContentSize - 1)
	AssertEq(nil, err)

	sr, err = t.mc.Stat()
	AssertEq(nil, err)
	ExpectEq(initialContentSize-1, sr.DirtyThreshold)
}
<|endoftext|>


<commit_before>package protocol

import (
	log "github.com/sirupsen/logrus"

	"github.com/bytom/errors"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
	"github.com/bytom/protocol/state"
	"github.com/bytom/protocol/validation"
)

var (
	// ErrBadBlock is returned when a block is invalid.
	ErrBadBlock = errors.New("invalid block")

	// ErrBadStateRoot is returned when the computed assets merkle root
	// disagrees with the one declared in a block header.
	ErrBadStateRoot = errors.New("invalid state merkle root")
)

// BlockExist check is a block in chain or orphan
func (c *Chain) BlockExist(hash *bc.Hash) bool {
	return c.orphanManage.BlockExist(hash) || c.store.BlockExist(hash)
}

// GetBlockByHash return a block by given hash
func (c *Chain) GetBlockByHash(hash *bc.Hash) (*types.Block, error) {
	return c.store.GetBlock(hash)
}

// GetBlockByHeight return a block by given height
func (c *Chain) GetBlockByHeight(height uint64) (*types.Block, error) {
	c.state.cond.L.Lock()
	hash, ok := c.state.mainChain[height]
	c.state.cond.L.Unlock()
	if !ok {
		return nil, errors.New("can't find block in given hight")
	}
	return c.GetBlockByHash(hash)
}

// ConnectBlock append block to end of chain
func (c *Chain) ConnectBlock(block *types.Block) error {
	c.state.cond.L.Lock()
	defer c.state.cond.L.Unlock()
	return c.connectBlock(block)
}

func (c *Chain) connectBlock(block *types.Block) (err error) {
	bcBlock := types.MapBlock(block)
	utxoView := state.NewUtxoViewpoint()
	bcBlock.TransactionStatus, err = c.store.GetTransactionStatus(&bcBlock.ID)
	if err != nil {
		return err
	}

	if err := c.store.GetTransactionsUtxo(utxoView, bcBlock.Transactions); err != nil {
		return err
	}
	if err := utxoView.ApplyBlock(bcBlock, bcBlock.TransactionStatus); err != nil {
		return err
	}

	blockHash := block.Hash()
	if err := c.setState(block, utxoView, map[uint64]*bc.Hash{block.Height: &blockHash}); err != nil {
		return err
	}

	for _, tx := range block.Transactions {
		c.txPool.RemoveTransaction(&tx.Tx.ID)
	}
	return nil
}

func (c *Chain) getReorganizeBlocks(block *types.Block) ([]*types.Block, []*types.Block) {
	attachBlocks := []*types.Block{}
	detachBlocks := []*types.Block{}
	ancestor := block

	for !c.inMainchain(ancestor) {
		attachBlocks = append([]*types.Block{ancestor}, attachBlocks...)
		ancestor, _ = c.GetBlockByHash(&ancestor.PreviousBlockHash)
	}

	for d := c.state.block; d.Hash() != ancestor.Hash(); d, _ = c.GetBlockByHash(&d.PreviousBlockHash) {
		detachBlocks = append(detachBlocks, d)
	}

	return attachBlocks, detachBlocks
}

func (c *Chain) reorganizeChain(block *types.Block) error {
	attachBlocks, detachBlocks := c.getReorganizeBlocks(block)
	utxoView := state.NewUtxoViewpoint()
	chainChanges := map[uint64]*bc.Hash{}

	for _, d := range detachBlocks {
		detachBlock := types.MapBlock(d)
		if err := c.store.GetTransactionsUtxo(utxoView, detachBlock.Transactions); err != nil {
			return err
		}
		txStatus, err := c.GetTransactionStatus(&detachBlock.ID)
		if err != nil {
			return err
		}
		if err := utxoView.DetachBlock(detachBlock, txStatus); err != nil {
			return err
		}
	}

	for _, a := range attachBlocks {
		attachBlock := types.MapBlock(a)
		if err := c.store.GetTransactionsUtxo(utxoView, attachBlock.Transactions); err != nil {
			return err
		}
		txStatus, err := c.GetTransactionStatus(&attachBlock.ID)
		if err != nil {
			return err
		}

		if err := utxoView.ApplyBlock(attachBlock, txStatus); err != nil {
			return err
		}
		chainChanges[a.Height] = &attachBlock.ID
	}

	return c.setState(block, utxoView, chainChanges)
}

// SaveBlock will validate and save block into storage
func (c *Chain) SaveBlock(block *types.Block) error {
	preBlock, err := c.GetBlockByHash(&block.PreviousBlockHash)
	if err != nil {
		return err
	}

	blockEnts := types.MapBlock(block)
	prevEnts := types.MapBlock(preBlock)

	seed, err := c.GetSeed(block.Height, &block.PreviousBlockHash)
	if err != nil {
		return err
	}

	if err := validation.ValidateBlock(blockEnts, prevEnts, seed, c.store); err != nil {
		return errors.Sub(ErrBadBlock, err)
	}

	if err := c.store.SaveBlock(block, blockEnts.TransactionStatus, seed); err != nil {
		return err
	}

	blockHash := block.Hash()
	log.WithFields(log.Fields{"height": block.Height, "hash": blockHash.String()}).Info("Block saved on disk")
	return nil
}

func (c *Chain) findBestChainTail(block *types.Block) (bestBlock *types.Block) {
	bestBlock = block
	blockHash := block.Hash()
	preorphans, ok := c.orphanManage.preOrphans[blockHash]
	if !ok {
		return
	}

	for _, preorphan := range preorphans {
		orphanBlock, ok := c.orphanManage.Get(preorphan)
		if !ok {
			continue
		}

		if err := c.SaveBlock(orphanBlock); err != nil {
			log.WithFields(log.Fields{
				"height": block.Height,
				"hash":   blockHash.String(),
			}).Errorf("findBestChainTail fail on save block %v", err)
			continue
		}

		if subResult := c.findBestChainTail(orphanBlock); subResult.Height > bestBlock.Height {
			bestBlock = subResult
		}
	}

	c.orphanManage.Delete(&blockHash)
	return
}

// ProcessBlock is the entry for handle block insert
func (c *Chain) ProcessBlock(block *types.Block) (bool, error) {
	blockHash := block.Hash()
	if c.BlockExist(&blockHash) {
		log.WithField("hash", blockHash.String()).Info("Skip process due to block already been handled")
		return false, nil
	}
	if !c.store.BlockExist(&block.PreviousBlockHash) {
		c.orphanManage.Add(block)
		return true, nil
	}
	if err := c.SaveBlock(block); err != nil {
		return false, err
	}

	bestBlock := c.findBestChainTail(block)
	c.state.cond.L.Lock()
	defer c.state.cond.L.Unlock()
	if c.state.block.Hash() == bestBlock.PreviousBlockHash {
		return false, c.connectBlock(bestBlock)
	}

	if bestBlock.Height > c.state.block.Height && bestBlock.Bits >= c.state.block.Bits {
		return false, c.reorganizeChain(bestBlock)
	}

	return false, nil
}
<commit_msg>remove error check of GetBlockByHash in SaveBlock (#485)<commit_after>package protocol

import (
	log "github.com/sirupsen/logrus"

	"github.com/bytom/errors"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
	"github.com/bytom/protocol/state"
	"github.com/bytom/protocol/validation"
)

var (
	// ErrBadBlock is returned when a block is invalid.
	ErrBadBlock = errors.New("invalid block")

	// ErrBadStateRoot is returned when the computed assets merkle root
	// disagrees with the one declared in a block header.
	ErrBadStateRoot = errors.New("invalid state merkle root")
)

// BlockExist check is a block in chain or orphan
func (c *Chain) BlockExist(hash *bc.Hash) bool {
	return c.orphanManage.BlockExist(hash) || c.store.BlockExist(hash)
}

// GetBlockByHash return a block by given hash
func (c *Chain) GetBlockByHash(hash *bc.Hash) (*types.Block, error) {
	return c.store.GetBlock(hash)
}

// GetBlockByHeight return a block by given height
func (c *Chain) GetBlockByHeight(height uint64) (*types.Block, error) {
	c.state.cond.L.Lock()
	hash, ok := c.state.mainChain[height]
	c.state.cond.L.Unlock()
	if !ok {
		return nil, errors.New("can't find block in given hight")
	}
	return c.GetBlockByHash(hash)
}

// ConnectBlock append block to end of chain
func (c *Chain) ConnectBlock(block *types.Block) error {
	c.state.cond.L.Lock()
	defer c.state.cond.L.Unlock()
	return c.connectBlock(block)
}

func (c *Chain) connectBlock(block *types.Block) (err error) {
	bcBlock := types.MapBlock(block)
	utxoView := state.NewUtxoViewpoint()
	bcBlock.TransactionStatus, err = c.store.GetTransactionStatus(&bcBlock.ID)
	if err != nil {
		return err
	}

	if err := c.store.GetTransactionsUtxo(utxoView, bcBlock.Transactions); err != nil {
		return err
	}
	if err := utxoView.ApplyBlock(bcBlock, bcBlock.TransactionStatus); err != nil {
		return err
	}

	blockHash := block.Hash()
	if err := c.setState(block, utxoView, map[uint64]*bc.Hash{block.Height: &blockHash}); err != nil {
		return err
	}

	for _, tx := range block.Transactions {
		c.txPool.RemoveTransaction(&tx.Tx.ID)
	}
	return nil
}

func (c *Chain) getReorganizeBlocks(block *types.Block) ([]*types.Block, []*types.Block) {
	attachBlocks := []*types.Block{}
	detachBlocks := []*types.Block{}
	ancestor := block

	for !c.inMainchain(ancestor) {
		attachBlocks = append([]*types.Block{ancestor}, attachBlocks...)
		ancestor, _ = c.GetBlockByHash(&ancestor.PreviousBlockHash)
	}

	for d := c.state.block; d.Hash() != ancestor.Hash(); d, _ = c.GetBlockByHash(&d.PreviousBlockHash) {
		detachBlocks = append(detachBlocks, d)
	}

	return attachBlocks, detachBlocks
}

func (c *Chain) reorganizeChain(block *types.Block) error {
	attachBlocks, detachBlocks := c.getReorganizeBlocks(block)
	utxoView := state.NewUtxoViewpoint()
	chainChanges := map[uint64]*bc.Hash{}

	for _, d := range detachBlocks {
		detachBlock := types.MapBlock(d)
		if err := c.store.GetTransactionsUtxo(utxoView, detachBlock.Transactions); err != nil {
			return err
		}
		txStatus, err := c.GetTransactionStatus(&detachBlock.ID)
		if err != nil {
			return err
		}
		if err := utxoView.DetachBlock(detachBlock, txStatus); err != nil {
			return err
		}
	}

	for _, a := range attachBlocks {
		attachBlock := types.MapBlock(a)
		if err := c.store.GetTransactionsUtxo(utxoView, attachBlock.Transactions); err != nil {
			return err
		}
		txStatus, err := c.GetTransactionStatus(&attachBlock.ID)
		if err != nil {
			return err
		}

		if err := utxoView.ApplyBlock(attachBlock, txStatus); err != nil {
			return err
		}
		chainChanges[a.Height] = &attachBlock.ID
	}

	return c.setState(block, utxoView, chainChanges)
}

// SaveBlock will validate and save block into storage
func (c *Chain) SaveBlock(block *types.Block) error {
	preBlock, _ := c.GetBlockByHash(&block.PreviousBlockHash)

	blockEnts := types.MapBlock(block)
	prevEnts := types.MapBlock(preBlock)

	seed, err := c.GetSeed(block.Height, &block.PreviousBlockHash)
	if err != nil {
		return err
	}

	if err := validation.ValidateBlock(blockEnts, prevEnts, seed, c.store); err != nil {
		return errors.Sub(ErrBadBlock, err)
	}

	if err := c.store.SaveBlock(block, blockEnts.TransactionStatus, seed); err != nil {
		return err
	}

	blockHash := block.Hash()
	log.WithFields(log.Fields{"height": block.Height, "hash": blockHash.String()}).Info("Block saved on disk")
	return nil
}

func (c *Chain) findBestChainTail(block *types.Block) (bestBlock *types.Block) {
	bestBlock = block
	blockHash := block.Hash()
	preorphans, ok := c.orphanManage.preOrphans[blockHash]
	if !ok {
		return
	}

	for _, preorphan := range preorphans {
		orphanBlock, ok := c.orphanManage.Get(preorphan)
		if !ok {
			continue
		}

		if err := c.SaveBlock(orphanBlock); err != nil {
			log.WithFields(log.Fields{
				"height": block.Height,
				"hash":   blockHash.String(),
			}).Errorf("findBestChainTail fail on save block %v", err)
			continue
		}

		if subResult := c.findBestChainTail(orphanBlock); subResult.Height > bestBlock.Height {
			bestBlock = subResult
		}
	}

	c.orphanManage.Delete(&blockHash)
	return
}

// ProcessBlock is the entry for handle block insert
func (c *Chain) ProcessBlock(block *types.Block) (bool, error) {
	blockHash := block.Hash()
	if c.BlockExist(&blockHash) {
		log.WithField("hash", blockHash.String()).Info("Skip process due to block already been handled")
		return false, nil
	}
	if !c.store.BlockExist(&block.PreviousBlockHash) {
		c.orphanManage.Add(block)
		return true, nil
	}
	if err := c.SaveBlock(block); err != nil {
		return false, err
	}

	bestBlock := c.findBestChainTail(block)
	c.state.cond.L.Lock()
	defer c.state.cond.L.Unlock()
	if c.state.block.Hash() == bestBlock.PreviousBlockHash {
		return false, c.connectBlock(bestBlock)
	}

	if bestBlock.Height > c.state.block.Height && bestBlock.Bits >= c.state.block.Bits {
		return false, c.reorganizeChain(bestBlock)
	}

	return false, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production
//
// This is a test template for the AccountDelete engine.

package engine

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
)

func TestAccountDelete(t *testing.T) {
	tc := SetupEngineTest(t, "acct")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "acct")

	ctx := &Context{}
	eng := NewAccountDelete(tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	_, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(tc.G, fu.Username))
	if err == nil {
		t.Fatal("no error loading deleted user")
	}
	if _, ok := err.(libkb.DeletedError); !ok {
		t.Errorf("loading deleted user error type: %T, expected libkb.DeletedError", err)
	}
}

func TestAccountDeleteAfterRestart(t *testing.T) {
	tc := SetupEngineTest(t, "acct")
	defer tc.Cleanup()

	fu := SignupFakeUserStoreSecret(tc, "acct")

	simulateServiceRestart(t, tc, fu)

	ctx := &Context{}
	eng := NewAccountDelete(tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatalf("AccountDelete after restart was broken but it looks like you've fixed it. Please make this test expect nil error here and uncomment the rest.")
	}
	require.Equal(t, "LoginSession is nil", err.Error())

	// _, err = libkb.LoadUser(libkb.NewLoadUserByNameArg(tc.G, fu.Username))
	// if err == nil {
	// 	t.Fatal("no error loading deleted user")
	// }
	// if _, ok := err.(libkb.DeletedError); !ok {
	// 	t.Errorf("loading deleted user error type: %T, expected libkb.DeletedError", err)
	// }
}

func TestAccountDeleteIdentify(t *testing.T) {
	tc := SetupEngineTest(t, "acct")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "acct")
	u, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(tc.G, fu.Username))
	if err != nil {
		t.Fatal(err)
	}

	ctx := &Context{}
	eng := NewAccountDelete(tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	i := newIdentify2WithUIDTester(tc.G)
	tc.G.Services = i
	arg := &keybase1.Identify2Arg{
		Uid:              u.GetUID(),
		IdentifyBehavior: keybase1.TLFIdentifyBehavior_CLI,
	}
	ieng := NewIdentify2WithUID(tc.G, arg)
	ictx := &Context{IdentifyUI: i}

	err = RunEngine(ieng, ictx)
	if err == nil {
		t.Fatal("identify2 ran successfully on deleted user")
	}
	if _, ok := err.(libkb.DeletedError); !ok {
		t.Errorf("identify2 error: %T, expected libkb.DeletedError", err)
	}
}
<commit_msg>Remove TestAccountDeleteAfterRestart which was doing nothing<commit_after>// Copyright 2016 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production
//
// This is a test template for the AccountDelete engine.

package engine

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
)

func TestAccountDelete(t *testing.T) {
	tc := SetupEngineTest(t, "acct")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "acct")

	ctx := &Context{}
	eng := NewAccountDelete(tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	_, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(tc.G, fu.Username))
	if err == nil {
		t.Fatal("no error loading deleted user")
	}
	if _, ok := err.(libkb.DeletedError); !ok {
		t.Errorf("loading deleted user error type: %T, expected libkb.DeletedError", err)
	}
}

func TestAccountDeleteIdentify(t *testing.T) {
	tc := SetupEngineTest(t, "acct")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "acct")
	u, err := libkb.LoadUser(libkb.NewLoadUserByNameArg(tc.G, fu.Username))
	if err != nil {
		t.Fatal(err)
	}

	ctx := &Context{}
	eng := NewAccountDelete(tc.G)
	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	i := newIdentify2WithUIDTester(tc.G)
	tc.G.Services = i
	arg := &keybase1.Identify2Arg{
		Uid:              u.GetUID(),
		IdentifyBehavior: keybase1.TLFIdentifyBehavior_CLI,
	}
	ieng := NewIdentify2WithUID(tc.G, arg)
	ictx := &Context{IdentifyUI: i}

	err = RunEngine(ieng, ictx)
	if err == nil {
		t.Fatal("identify2 ran successfully on deleted user")
	}
	if _, ok := err.(libkb.DeletedError); !ok {
		t.Errorf("identify2 error: %T, expected libkb.DeletedError", err)
	}
}
<|endoftext|>


<commit_before>package control

import (
	"context"
	"fmt"
	"io"
	"log"
	"sync"

	"github.com/moby/vpnkit/go/pkg/libproxy"
	"github.com/moby/vpnkit/go/pkg/vpnkit"
	"github.com/moby/vpnkit/go/pkg/vpnkit/forward"
	"github.com/moby/vpnkit/go/pkg/vpnkit/transport"
	"github.com/pkg/errors"
)

type Control struct {
	mux       libproxy.Multiplexer
	muxM      *sync.Mutex
	muxC      *sync.Cond
	forwards  map[string]forward.Forward
	forwardsM *sync.Mutex
}

func Make() *Control {
	var muxM sync.Mutex
	muxC := sync.NewCond(&muxM)
	var outsidesM sync.Mutex
	outsides := make(map[string]forward.Forward)
	return &Control{
		muxM:      &muxM,
		muxC:      muxC,
		forwards:  outsides,
		forwardsM: &outsidesM,
	}
}

func (c *Control) SetMux(m libproxy.Multiplexer) {
	c.muxM.Lock()
	defer c.muxM.Unlock()
	log.Println("ready to forward incoming data connections")
	c.mux = m
	c.muxC.Broadcast()
}

func (c *Control) Mux() libproxy.Multiplexer {
	c.muxM.Lock()
	defer c.muxM.Unlock()
	for {
		if c.mux != nil {
			return c.mux
		}
		c.muxC.Wait()
	}
}

func portKey(port *vpnkit.Port) string {
	return port.String()
}

func (c *Control) Expose(_ context.Context, port *vpnkit.Port) error {
	if port == nil {
		return errors.New("cannot expose a nil Port")
	}
	key := portKey(port)
	c.forwardsM.Lock()
	defer c.forwardsM.Unlock()
	if _, ok := c.forwards[key]; ok {
		return errors.New("port already exposed: " + port.String())
	}
	forward, err := forward.Make(c, *port)
	if err != nil {
		// This error (e.g. EADDRINUSE) is special and we want to show it to the user
		return &vpnkit.ExposeError{
			Message: err.Error(),
		}
	}
	c.forwards[key] = forward
	log.Printf("adding port-forward %s\n", port.String())
	go forward.Run()
	return nil
}

func (c *Control) Unexpose(_ context.Context, port *vpnkit.Port) error {
	if port == nil {
		return errors.New("cannot unexpose a nil Port")
	}
	key := portKey(port)
	c.forwardsM.Lock()
	defer c.forwardsM.Unlock()
	forward, ok := c.forwards[key]
	if !ok {
		// make it idempotent
		return nil
	}
	log.Printf("stopping port-forward %s\n", port.String())
	forward.Stop()
	delete(c.forwards, key)
	return nil
}

func (c *Control) ListExposed(_ context.Context) ([]vpnkit.Port, error) {
	var results []vpnkit.Port
	for _, forward := range c.forwards {
		results = append(results, forward.Port())
	}
	return results, nil
}

var _ vpnkit.Implementation = &Control{}
var _ vpnkit.Control = &Control{}

// Listen for incoming data connections
func (c *Control) Listen(path string, quit chan struct{}) {
	t := transport.NewVsockTransport()
	l, err := t.Listen(path)
	if err != nil {
		log.Fatalf("unable to create a data server on AF_VSOCK port %s: %s", path, err)
	}
	for {
		// listen for one connection at a time
		log.Printf("listening on AF_VSOCK port %s for data connection", path)
		conn, err := l.Accept()
		if err != nil {
			log.Printf("unable to accept connection on AF_VSOCK port %s: %s", path, err)
			continue
		}
		log.Printf("accepted data connection on AF_VSOCK port: %s", path)
		c.handleDataConn(conn, quit)
	}
}

// Connect a data connection
func (c *Control) Connect(path string, quit chan struct{}) error {
	for {
		t := transport.NewVsockTransport()
		log.Printf("dialing AF_VSOCK port %s for data connection", path)
		conn, err := t.Dial(context.Background(), path)
		if err != nil {
			return fmt.Errorf("unable to connect data on AF_VSOCK port %s: %s", path, err)
		}
		log.Printf("connected data connection on AF_VSOCK port: %s", path)
		c.handleDataConn(conn, quit)
	}
}

// handle data-plane forwarding
func (c *Control) handleDataConn(rw io.ReadWriteCloser, quit chan struct{}) {
	defer rw.Close()

	mux := libproxy.NewMultiplexer("local", rw)
	mux.Run()
	c.SetMux(mux)
	for {
		conn, destination, err := mux.Accept()
		if err != nil {
			log.Printf("Error accepting subconnection: %v", err)
			return
		}
		go libproxy.Forward(conn, *destination, quit)
	}
}
<commit_msg>vpnkit-forwarder: keep trying to (re-connect)<commit_after>package control

import (
	"context"
	"io"
	"log"
	"sync"
	"time"

	"github.com/moby/vpnkit/go/pkg/libproxy"
	"github.com/moby/vpnkit/go/pkg/vpnkit"
	"github.com/moby/vpnkit/go/pkg/vpnkit/forward"
	"github.com/moby/vpnkit/go/pkg/vpnkit/transport"
	"github.com/pkg/errors"
)

type Control struct {
	mux       libproxy.Multiplexer
	muxM      *sync.Mutex
	muxC      *sync.Cond
	forwards  map[string]forward.Forward
	forwardsM *sync.Mutex
}

func Make() *Control {
	var muxM sync.Mutex
	muxC := sync.NewCond(&muxM)
	var outsidesM sync.Mutex
	outsides := make(map[string]forward.Forward)
	return &Control{
		muxM:      &muxM,
		muxC:      muxC,
		forwards:  outsides,
		forwardsM: &outsidesM,
	}
}

func (c *Control) SetMux(m libproxy.Multiplexer) {
	c.muxM.Lock()
	defer c.muxM.Unlock()
	log.Println("ready to forward incoming data connections")
	c.mux = m
	c.muxC.Broadcast()
}

func (c *Control) Mux() libproxy.Multiplexer {
	c.muxM.Lock()
	defer c.muxM.Unlock()
	for {
		if c.mux != nil {
			return c.mux
		}
		c.muxC.Wait()
	}
}

func portKey(port *vpnkit.Port) string {
	return port.String()
}

func (c *Control) Expose(_ context.Context, port *vpnkit.Port) error {
	if port == nil {
		return errors.New("cannot expose a nil Port")
	}
	key := portKey(port)
	c.forwardsM.Lock()
	defer c.forwardsM.Unlock()
	if _, ok := c.forwards[key]; ok {
		return errors.New("port already exposed: " + port.String())
	}
	forward, err := forward.Make(c, *port)
	if err != nil {
		// This error (e.g. EADDRINUSE) is special and we want to show it to the user
		return &vpnkit.ExposeError{
			Message: err.Error(),
		}
	}
	c.forwards[key] = forward
	log.Printf("adding port-forward %s\n", port.String())
	go forward.Run()
	return nil
}

func (c *Control) Unexpose(_ context.Context, port *vpnkit.Port) error {
	if port == nil {
		return errors.New("cannot unexpose a nil Port")
	}
	key := portKey(port)
	c.forwardsM.Lock()
	defer c.forwardsM.Unlock()
	forward, ok := c.forwards[key]
	if !ok {
		// make it idempotent
		return nil
	}
	log.Printf("stopping port-forward %s\n", port.String())
	forward.Stop()
	delete(c.forwards, key)
	return nil
}

func (c *Control) ListExposed(_ context.Context) ([]vpnkit.Port, error) {
	var results []vpnkit.Port
	for _, forward := range c.forwards {
		results = append(results, forward.Port())
	}
	return results, nil
}

var _ vpnkit.Implementation = &Control{}
var _ vpnkit.Control = &Control{}

// Listen for incoming data connections
func (c *Control) Listen(path string, quit chan struct{}) {
	t := transport.NewVsockTransport()
	l, err := t.Listen(path)
	if err != nil {
		log.Fatalf("unable to create a data server on AF_VSOCK port %s: %s", path, err)
	}
	for {
		// listen for one connection at a time
		log.Printf("listening on AF_VSOCK port %s for data connection", path)
		conn, err := l.Accept()
		if err != nil {
			log.Printf("unable to accept connection on AF_VSOCK port %s: %s", path, err)
			continue
		}
		log.Printf("accepted data connection on AF_VSOCK port: %s", path)
		c.handleDataConn(conn, quit)
	}
}

// Connect a data connection
func (c *Control) Connect(path string, quit chan struct{}) error {
	for {
		t := transport.NewVsockTransport()
		log.Printf("dialing AF_VSOCK port %s for data connection", path)
		conn, err := t.Dial(context.Background(), path)
		if err != nil {
			// This can happen if the server is restarting
			log.Printf("unable to connect data on AF_VSOCK port %s: %s. Is the server restarting? Will retry in 1s.", path, err)
			time.Sleep(time.Second)
			continue
		}
		log.Printf("connected data connection on AF_VSOCK port: %s", path)
		c.handleDataConn(conn, quit)
	}
}

// handle data-plane forwarding
func (c *Control) handleDataConn(rw io.ReadWriteCloser, quit chan struct{}) {
	defer rw.Close()

	mux := libproxy.NewMultiplexer("local", rw)
	mux.Run()
	c.SetMux(mux)
	for {
		conn, destination, err := mux.Accept()
		if err != nil {
			log.Printf("Error accepting subconnection: %v", err)
			return
		}
		go libproxy.Forward(conn, *destination, quit)
	}
}
<|endoftext|>


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

package worker

import (
	"fmt"
	"strconv"
	"strings"
	"testing"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	rpcproto "github.com/youtube/vitess/go/rpcwrap/proto"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/logutil"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
	"github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/wrangler/testlib"
	"github.com/youtube/vitess/go/vt/zktopo"
)

// This is a local SqlQuery RCP implementation to support the tests
type SqlQuery struct {
	t *testing.T
}

func (sq *SqlQuery) GetSessionId(sessionParams *proto.SessionParams, sessionInfo *proto.SessionInfo) error {
	return nil
}

func (sq *SqlQuery) StreamExecute(ctx *rpcproto.Context, query *proto.Query, sendReply func(reply interface{}) error) error {
	// Custom parsing of the query we expect
	min := 100
	max := 200
	var err error
	parts := strings.Split(query.Sql, " ")
	for _, part := range parts {
		if strings.HasPrefix(part, "id>=") {
			min, err = strconv.Atoi(part[4:])
			if err != nil {
				return err
			}
		} else if strings.HasPrefix(part, "id<") {
			max, err = strconv.Atoi(part[3:])
		}
	}
	sq.t.Logf("SqlQuery: got query: %v with min %v max %v", *query, min, max)

	// Send the headers
	if err := sendReply(&mproto.QueryResult{
		Fields: []mproto.Field{
			mproto.Field{
				Name: "id",
				Type: mproto.VT_LONGLONG,
			},
			mproto.Field{
				Name: "msg",
				Type: mproto.VT_VARCHAR,
			},
			mproto.Field{
				Name: "keyspace_id",
				Type: mproto.VT_LONGLONG,
			},
		},
	}); err != nil {
		return err
	}

	// Send the values
	ksids := []uint64{0x2000000000000000, 0x6000000000000000}
	for i := min; i < max; i++ {
		if err := sendReply(&mproto.QueryResult{
			Rows: [][]sqltypes.Value{
				[]sqltypes.Value{
					sqltypes.MakeString([]byte(fmt.Sprintf("%v", i))),
					sqltypes.MakeString([]byte(fmt.Sprintf("Text for %v", i))),
					sqltypes.MakeString([]byte(fmt.Sprintf("%v", ksids[i%2]))),
				},
			},
		}); err != nil {
			return err
		}
	}
	// SELECT id, msg, keyspace_id FROM table1 WHERE id>=180 AND id<190 ORDER BY id
	return nil
}

type ExpectedExecuteFetch struct {
	Query       string
	MaxRows     int
	WantFields  bool
	QueryResult *mproto.QueryResult
	Error       error
}

// FakePoolConnection implements dbconnpool.PoolConnection
type FakePoolConnection struct {
	t      *testing.T
	Closed bool

	ExpectedExecuteFetch      []ExpectedExecuteFetch
	ExpectedExecuteFetchIndex int
}

func NewFakePoolConnectionQueryBinlogOff(t *testing.T, query string) *FakePoolConnection {
	return &FakePoolConnection{
		t: t,
		ExpectedExecuteFetch: []ExpectedExecuteFetch{
			ExpectedExecuteFetch{
				Query:       "SET sql_log_bin = OFF",
				QueryResult: &mproto.QueryResult{},
			},
			ExpectedExecuteFetch{
				Query:       query,
				QueryResult: &mproto.QueryResult{},
			},
			ExpectedExecuteFetch{
				Query:       "SET sql_log_bin = ON",
				QueryResult: &mproto.QueryResult{},
			},
		},
	}
}

func (fpc *FakePoolConnection) ExecuteFetch(query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	if fpc.ExpectedExecuteFetchIndex >= len(fpc.ExpectedExecuteFetch) {
		fpc.t.Errorf("got unexpected out of bound fetch: %v >= %v", fpc.ExpectedExecuteFetchIndex, len(fpc.ExpectedExecuteFetch))
		return nil, fmt.Errorf("unexpected out of bound fetch")
	}
	expected := fpc.ExpectedExecuteFetch[fpc.ExpectedExecuteFetchIndex].Query
	if strings.HasSuffix(expected, "*") {
		if !strings.HasPrefix(query, expected[0:len(expected)-1]) {
			fpc.t.Errorf("got unexpected query start: %v != %v", query, expected)
			return nil, fmt.Errorf("unexpected query")
		}
	} else {
		if query != expected {
			fpc.t.Errorf("got unexpected query: %v != %v", query, expected)
			return nil, fmt.Errorf("unexpected query")
		}
	}
	fpc.t.Logf("ExecuteFetch: %v", query)
	defer func() {
		fpc.ExpectedExecuteFetchIndex++
	}()
	return fpc.ExpectedExecuteFetch[fpc.ExpectedExecuteFetchIndex].QueryResult, nil
}

func (fpc *FakePoolConnection) ExecuteStreamFetch(query string, callback func(*mproto.QueryResult) error, streamBufferSize int) error {
	return nil
}

func (fpc *FakePoolConnection) Id() int64 {
	return 1
}

func (fpc *FakePoolConnection) Close() {
	fpc.Closed = true
}

func (fpc *FakePoolConnection) IsClosed() bool {
	return fpc.Closed
}

func (fpc *FakePoolConnection) Recycle() {
}

// on the source rdonly guy, should only have one query to find min & max
func SourceRdonlyFactory(t *testing.T) func() (dbconnpool.PoolConnection, error) {
	return func() (dbconnpool.PoolConnection, error) {
		return &FakePoolConnection{
			t: t,
			ExpectedExecuteFetch: []ExpectedExecuteFetch{
				ExpectedExecuteFetch{
					Query: "SELECT MIN(id), MAX(id) FROM vt_ks.table1",
					QueryResult: &mproto.QueryResult{
						Fields: []mproto.Field{
							mproto.Field{
								Name: "min",
								Type: mproto.VT_LONGLONG,
							},
							mproto.Field{
								Name: "max",
								Type: mproto.VT_LONGLONG,
							},
						},
						Rows: [][]sqltypes.Value{
							[]sqltypes.Value{
								sqltypes.MakeString([]byte("100")),
								sqltypes.MakeString([]byte("200")),
							},
						},
					},
				},
			},
		}, nil
	}
}

// on the destinations
func DestinationsFactory(t *testing.T, rowCount int) func() (dbconnpool.PoolConnection, error) {
	queryIndex := 0
	return func() (dbconnpool.PoolConnection, error) {
		defer func() {
			queryIndex++
		}()
		switch {
		case queryIndex == 0:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE DATABASE `vt_ks` /*!40100 DEFAULT CHARACTER SET utf8 */"), nil
		case queryIndex == 1:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE TABLE `vt_ks`.`resharding1` (\n"+
				"  `id` bigint(20) NOT NULL,\n"+
				"  `msg` varchar(64) DEFAULT NULL,\n"+
				"  `keyspace_id` bigint(20) unsigned NOT NULL,\n"+
				"  PRIMARY KEY (`id`),\n"+
				"  KEY `by_msg` (`msg`)\n"+
				") ENGINE=InnoDB DEFAULT CHARSET=utf8"), nil
		case queryIndex >= 2 && queryIndex < rowCount+2:
			return NewFakePoolConnectionQueryBinlogOff(t, "INSERT INTO `vt_ks`.table1(id, msg, keyspace_id) VALUES (*"), nil
		case queryIndex == rowCount+2:
			return NewFakePoolConnectionQueryBinlogOff(t, "ALTER TABLE `vt_ks`.`table1` MODIFY   `id` bigint(20) NOT NULL AUTO_INCREMENT"), nil
		case queryIndex == rowCount+3:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE DATABASE IF NOT EXISTS _vt"), nil
		case queryIndex == rowCount+4:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE TABLE IF NOT EXISTS _vt.blp_checkpoint (\n"+
				"  source_shard_uid INT(10) UNSIGNED NOT NULL,\n"+
				"  pos VARCHAR(250) DEFAULT NULL,\n"+
				"  time_updated BIGINT UNSIGNED NOT NULL,\n"+
				"  transaction_timestamp BIGINT UNSIGNED NOT NULL,\n"+
				"  flags VARCHAR(250) DEFAULT NULL,\n"+
				"  PRIMARY KEY (source_shard_uid)) ENGINE=InnoDB"), nil
		case queryIndex == rowCount+5:
			return NewFakePoolConnectionQueryBinlogOff(t, "INSERT INTO _vt.blp_checkpoint (source_shard_uid, pos, time_updated, transaction_timestamp, flags) VALUES (0, 'MariaDB/12-34-5678', *"), nil
		}

		return nil, fmt.Errorf("Unexpected connection")
	}
}

func TestSplitClone(t *testing.T) {
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Minute, time.Second)

	sourceMaster := testlib.NewFakeTablet(t, wr, "cell1", 0,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-80"))
	sourceRdonly := testlib.NewFakeTablet(t, wr, "cell1", 1,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "-80"),
		testlib.TabletParent(sourceMaster.Tablet.Alias))

	leftMaster := testlib.NewFakeTablet(t, wr, "cell1", 10,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-40"))
	leftRdonly := testlib.NewFakeTablet(t, wr, "cell1", 11,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "-40"),
		testlib.TabletParent(leftMaster.Tablet.Alias))

	rightMaster := testlib.NewFakeTablet(t, wr, "cell1", 20,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "40-80"))
	rightRdonly := testlib.NewFakeTablet(t, wr, "cell1", 21,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "40-80"),
		testlib.TabletParent(rightMaster.Tablet.Alias))

	for _, ft := range []*testlib.FakeTablet{sourceMaster, sourceRdonly, leftMaster, leftRdonly, rightMaster, rightRdonly} {
		ft.StartActionLoop(t, wr)
		defer ft.StopActionLoop(t)
	}

	// add the topo and schema data we'll need
	if err := topo.CreateShard(ts, "ks", "80-"); err != nil {
		t.Fatalf("CreateShard(\"-80\") failed: %v", err)
	}
	if err := wr.SetKeyspaceShardingInfo("ks", "keyspace_id", key.KIT_UINT64, 4, false); err != nil {
		t.Fatalf("SetKeyspaceShardingInfo failed: %v", err)
	}
	if err := wr.RebuildKeyspaceGraph("ks", nil); err != nil {
		t.Fatalf("RebuildKeyspaceGraph failed: %v", err)
	}
	sourceRdonly.FakeMysqlDaemon.Schema = &myproto.SchemaDefinition{
		DatabaseSchema: "CREATE DATABASE `{{.DatabaseName}}` /*!40100 DEFAULT CHARACTER SET utf8 */",
		TableDefinitions: []*myproto.TableDefinition{
			&myproto.TableDefinition{
				Name:              "table1",
				Schema:            "CREATE TABLE `resharding1` (\n  `id` bigint(20) NOT NULL AUTO_INCREMENT,\n  `msg` varchar(64) DEFAULT NULL,\n  `keyspace_id` bigint(20) unsigned NOT NULL,\n  PRIMARY KEY (`id`),\n  KEY `by_msg` (`msg`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8",
				Columns:           []string{"id", "msg", "keyspace_id"},
				PrimaryKeyColumns: []string{"id"},
				Type:              myproto.TABLE_BASE_TABLE,
				DataLength:        2048,
				RowCount:          100,
			},
		},
		Version: "unused",
	}
	sourceRdonly.FakeMysqlDaemon.DbaConnectionFactory = SourceRdonlyFactory(t)
	sourceRdonly.FakeMysqlDaemon.CurrentSlaveStatus = &myproto.ReplicationStatus{
		Position: myproto.ReplicationPosition{
			GTIDSet: myproto.MariadbGTID{Domain: 12, Server: 34, Sequence: 5678},
		},
	}
	sourceRdonly.RpcServer.Register(&SqlQuery{t: t})
	leftMaster.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	leftRdonly.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	rightMaster.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	rightRdonly.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)

	wrk := NewSplitCloneWorker(wr, "cell1", "ks", "-80", nil, "populateBlpCheckpoint", 10, 1, 10).(*SplitCloneWorker)
	wrk.Run()
	status := wrk.StatusAsText()
	t.Logf("Got status: %v", status)
	if wrk.err != nil || wrk.state != stateSCDone {
		t.Errorf("Worker run failed")
	}
}
<commit_msg>Removing flakiness in this test. The queryIndex variable was being accessed concurrently.<commit_after>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"fmt"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	mproto "github.com/youtube/vitess/go/mysql/proto"
	rpcproto "github.com/youtube/vitess/go/rpcwrap/proto"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/logutil"
	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
	"github.com/youtube/vitess/go/vt/tabletserver/proto"
	"github.com/youtube/vitess/go/vt/topo"
	"github.com/youtube/vitess/go/vt/wrangler"
	"github.com/youtube/vitess/go/vt/wrangler/testlib"
	"github.com/youtube/vitess/go/vt/zktopo"
)

// This is a local SqlQuery RCP implementation to support the tests
type SqlQuery struct {
	t *testing.T
}

func (sq *SqlQuery) GetSessionId(sessionParams *proto.SessionParams, sessionInfo *proto.SessionInfo) error {
	return nil
}

func (sq *SqlQuery) StreamExecute(ctx *rpcproto.Context, query *proto.Query, sendReply func(reply interface{}) error) error {
	// Custom parsing of the query we expect
	min := 100
	max := 200
	var err error
	parts := strings.Split(query.Sql, " ")
	for _, part := range parts {
		if strings.HasPrefix(part, "id>=") {
			min, err = strconv.Atoi(part[4:])
			if err != nil {
				return err
			}
		} else if strings.HasPrefix(part, "id<") {
			max, err = strconv.Atoi(part[3:])
		}
	}
	sq.t.Logf("SqlQuery: got query: %v with min %v max %v", *query, min, max)

	// Send the headers
	if err := sendReply(&mproto.QueryResult{
		Fields: []mproto.Field{
			mproto.Field{
				Name: "id",
				Type: mproto.VT_LONGLONG,
			},
			mproto.Field{
				Name: "msg",
				Type: mproto.VT_VARCHAR,
			},
			mproto.Field{
				Name: "keyspace_id",
				Type: mproto.VT_LONGLONG,
			},
		},
	}); err != nil {
		return err
	}

	// Send the values
	ksids := []uint64{0x2000000000000000, 0x6000000000000000}
	for i := min; i < max; i++ {
		if err := sendReply(&mproto.QueryResult{
			Rows: [][]sqltypes.Value{
				[]sqltypes.Value{
					sqltypes.MakeString([]byte(fmt.Sprintf("%v", i))),
					sqltypes.MakeString([]byte(fmt.Sprintf("Text for %v", i))),
					sqltypes.MakeString([]byte(fmt.Sprintf("%v", ksids[i%2]))),
				},
			},
		}); err != nil {
			return err
		}
	}
	// SELECT id, msg, keyspace_id FROM table1 WHERE id>=180 AND id<190 ORDER BY id
	return nil
}

type ExpectedExecuteFetch struct {
	Query       string
	MaxRows     int
	WantFields  bool
	QueryResult *mproto.QueryResult
	Error       error
}

// FakePoolConnection implements dbconnpool.PoolConnection
type FakePoolConnection struct {
	t      *testing.T
	Closed bool

	ExpectedExecuteFetch      []ExpectedExecuteFetch
	ExpectedExecuteFetchIndex int
}

func NewFakePoolConnectionQueryBinlogOff(t *testing.T, query string) *FakePoolConnection {
	return &FakePoolConnection{
		t: t,
		ExpectedExecuteFetch: []ExpectedExecuteFetch{
			ExpectedExecuteFetch{
				Query:       "SET sql_log_bin = OFF",
				QueryResult: &mproto.QueryResult{},
			},
			ExpectedExecuteFetch{
				Query:       query,
				QueryResult: &mproto.QueryResult{},
			},
			ExpectedExecuteFetch{
				Query:       "SET sql_log_bin = ON",
				QueryResult: &mproto.QueryResult{},
			},
		},
	}
}

func (fpc *FakePoolConnection) ExecuteFetch(query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
	if fpc.ExpectedExecuteFetchIndex >= len(fpc.ExpectedExecuteFetch) {
		fpc.t.Errorf("got unexpected out of bound fetch: %v >= %v", fpc.ExpectedExecuteFetchIndex, len(fpc.ExpectedExecuteFetch))
		return nil, fmt.Errorf("unexpected out of bound fetch")
	}
	expected := fpc.ExpectedExecuteFetch[fpc.ExpectedExecuteFetchIndex].Query
	if strings.HasSuffix(expected, "*") {
		if !strings.HasPrefix(query, expected[0:len(expected)-1]) {
			fpc.t.Errorf("got unexpected query start: %v != %v", query, expected)
			return nil, fmt.Errorf("unexpected query")
		}
	} else {
		if query != expected {
			fpc.t.Errorf("got unexpected query: %v != %v", query, expected)
			return nil, fmt.Errorf("unexpected query")
		}
	}
	fpc.t.Logf("ExecuteFetch: %v", query)
	defer func() {
		fpc.ExpectedExecuteFetchIndex++
	}()
	return fpc.ExpectedExecuteFetch[fpc.ExpectedExecuteFetchIndex].QueryResult, nil
}

func (fpc *FakePoolConnection) ExecuteStreamFetch(query string, callback func(*mproto.QueryResult) error, streamBufferSize int) error {
	return nil
}

func (fpc *FakePoolConnection) Id() int64 {
	return 1
}

func (fpc *FakePoolConnection) Close() {
	fpc.Closed = true
}

func (fpc *FakePoolConnection) IsClosed() bool {
	return fpc.Closed
}

func (fpc *FakePoolConnection) Recycle() {
}

// on the source rdonly guy, should only have one query to find min & max
func SourceRdonlyFactory(t *testing.T) func() (dbconnpool.PoolConnection, error) {
	return func() (dbconnpool.PoolConnection, error) {
		return &FakePoolConnection{
			t: t,
			ExpectedExecuteFetch: []ExpectedExecuteFetch{
				ExpectedExecuteFetch{
					Query: "SELECT MIN(id), MAX(id) FROM vt_ks.table1",
					QueryResult: &mproto.QueryResult{
						Fields: []mproto.Field{
							mproto.Field{
								Name: "min",
								Type: mproto.VT_LONGLONG,
							},
							mproto.Field{
								Name: "max",
								Type: mproto.VT_LONGLONG,
							},
						},
						Rows: [][]sqltypes.Value{
							[]sqltypes.Value{
								sqltypes.MakeString([]byte("100")),
								sqltypes.MakeString([]byte("200")),
							},
						},
					},
				},
			},
		}, nil
	}
}

// on the destinations
func DestinationsFactory(t *testing.T, rowCount int64) func() (dbconnpool.PoolConnection, error) {
	var queryIndex int64 = -1
	return func() (dbconnpool.PoolConnection, error) {
		qi := atomic.AddInt64(&queryIndex, 1)
		switch {
		case qi == 0:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE DATABASE `vt_ks` /*!40100 DEFAULT CHARACTER SET utf8 */"), nil
		case qi == 1:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE TABLE `vt_ks`.`resharding1` (\n"+
				"  `id` bigint(20) NOT NULL,\n"+
				"  `msg` varchar(64) DEFAULT NULL,\n"+
				"  `keyspace_id` bigint(20) unsigned NOT NULL,\n"+
				"  PRIMARY KEY (`id`),\n"+
				"  KEY `by_msg` (`msg`)\n"+
				") ENGINE=InnoDB DEFAULT CHARSET=utf8"), nil
		case qi >= 2 && qi < rowCount+2:
			return NewFakePoolConnectionQueryBinlogOff(t, "INSERT INTO `vt_ks`.table1(id, msg, keyspace_id) VALUES (*"), nil
		case qi == rowCount+2:
			return NewFakePoolConnectionQueryBinlogOff(t, "ALTER TABLE `vt_ks`.`table1` MODIFY   `id` bigint(20) NOT NULL AUTO_INCREMENT"), nil
		case qi == rowCount+3:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE DATABASE IF NOT EXISTS _vt"), nil
		case qi == rowCount+4:
			return NewFakePoolConnectionQueryBinlogOff(t, "CREATE TABLE IF NOT EXISTS _vt.blp_checkpoint (\n"+
				"  source_shard_uid INT(10) UNSIGNED NOT NULL,\n"+
				"  pos VARCHAR(250) DEFAULT NULL,\n"+
				"  time_updated BIGINT UNSIGNED NOT NULL,\n"+
				"  transaction_timestamp BIGINT UNSIGNED NOT NULL,\n"+
				"  flags VARCHAR(250) DEFAULT NULL,\n"+
				"  PRIMARY KEY (source_shard_uid)) ENGINE=InnoDB"), nil
		case qi == rowCount+5:
			return NewFakePoolConnectionQueryBinlogOff(t, "INSERT INTO _vt.blp_checkpoint (source_shard_uid, pos, time_updated, transaction_timestamp, flags) VALUES (0, 'MariaDB/12-34-5678', *"), nil
		}

		return nil, fmt.Errorf("Unexpected connection")
	}
}

func TestSplitClone(t *testing.T) {
	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Minute, time.Second)

	sourceMaster := testlib.NewFakeTablet(t, wr, "cell1", 0,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-80"))
	sourceRdonly := testlib.NewFakeTablet(t, wr, "cell1", 1,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "-80"),
		testlib.TabletParent(sourceMaster.Tablet.Alias))

	leftMaster := testlib.NewFakeTablet(t, wr, "cell1", 10,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-40"))
	leftRdonly := testlib.NewFakeTablet(t, wr, "cell1", 11,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "-40"),
		testlib.TabletParent(leftMaster.Tablet.Alias))

	rightMaster := testlib.NewFakeTablet(t, wr, "cell1", 20,
		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "40-80"))
	rightRdonly := testlib.NewFakeTablet(t, wr, "cell1", 21,
		topo.TYPE_RDONLY, testlib.TabletKeyspaceShard(t, "ks", "40-80"),
		testlib.TabletParent(rightMaster.Tablet.Alias))

	for _, ft := range []*testlib.FakeTablet{sourceMaster, sourceRdonly, leftMaster, leftRdonly, rightMaster, rightRdonly} {
		ft.StartActionLoop(t, wr)
		defer ft.StopActionLoop(t)
	}

	// add the topo and schema data we'll need
	if err := topo.CreateShard(ts, "ks", "80-"); err != nil {
		t.Fatalf("CreateShard(\"-80\") failed: %v", err)
	}
	if err := wr.SetKeyspaceShardingInfo("ks", "keyspace_id", key.KIT_UINT64, 4, false); err != nil {
		t.Fatalf("SetKeyspaceShardingInfo failed: %v", err)
	}
	if err := wr.RebuildKeyspaceGraph("ks", nil); err != nil {
		t.Fatalf("RebuildKeyspaceGraph failed: %v", err)
	}
	sourceRdonly.FakeMysqlDaemon.Schema = &myproto.SchemaDefinition{
		DatabaseSchema: "CREATE DATABASE `{{.DatabaseName}}` /*!40100 DEFAULT CHARACTER SET utf8 */",
		TableDefinitions: []*myproto.TableDefinition{
			&myproto.TableDefinition{
				Name:              "table1",
				Schema:            "CREATE TABLE `resharding1` (\n  `id` bigint(20) NOT NULL AUTO_INCREMENT,\n  `msg` varchar(64) DEFAULT NULL,\n  `keyspace_id` bigint(20) unsigned NOT NULL,\n  PRIMARY KEY (`id`),\n  KEY `by_msg` (`msg`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8",
				Columns:           []string{"id", "msg", "keyspace_id"},
				PrimaryKeyColumns: []string{"id"},
				Type:              myproto.TABLE_BASE_TABLE,
				DataLength:        2048,
				RowCount:          100,
			},
		},
		Version: "unused",
	}
	sourceRdonly.FakeMysqlDaemon.DbaConnectionFactory = SourceRdonlyFactory(t)
	sourceRdonly.FakeMysqlDaemon.CurrentSlaveStatus = &myproto.ReplicationStatus{
		Position: myproto.ReplicationPosition{
			GTIDSet: myproto.MariadbGTID{Domain: 12, Server: 34, Sequence: 5678},
		},
	}
	sourceRdonly.RpcServer.Register(&SqlQuery{t: t})
	leftMaster.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	leftRdonly.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	rightMaster.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)
	rightRdonly.FakeMysqlDaemon.DbaConnectionFactory = DestinationsFactory(t, 50)

	wrk := NewSplitCloneWorker(wr, "cell1", "ks", "-80", nil, "populateBlpCheckpoint", 10, 1, 10).(*SplitCloneWorker)
	wrk.Run()
	status := wrk.StatusAsText()
	t.Logf("Got status: %v", status)
	if wrk.err != nil || wrk.state != stateSCDone {
		t.Errorf("Worker run failed")
	}
}
<|endoftext|>


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

package gcsx

import (
	"mime"
	"path"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
)

// NewContentTypeBucket creates a wrapper bucket that guesses MIME types for
// newly created or composed objects when an explicit type is not already set.
func NewContentTypeBucket(b gcs.Bucket) gcs.Bucket {
	return contentTypeBucket{b}
}

type contentTypeBucket struct {
	gcs.Bucket
}

func (b contentTypeBucket) CreateObject(
	ctx context.Context,
	req *gcs.CreateObjectRequest) (o *gcs.Object, err error) {
	// Guess a content type if necessary.
	if req.ContentType == "" {
		req.ContentType = mime.TypeByExtension(path.Ext(req.Name))
	}

	// Pass on the request.
	o, err = b.Bucket.CreateObject(ctx, req)
	return
}
<commit_msg>contentTypeBucket.ComposeObjects<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcsx

import (
	"mime"
	"path"

	"github.com/jacobsa/gcloud/gcs"
	"golang.org/x/net/context"
)

// NewContentTypeBucket creates a wrapper bucket that guesses MIME types for
// newly created or composed objects when an explicit type is not already set.
func NewContentTypeBucket(b gcs.Bucket) gcs.Bucket {
	return contentTypeBucket{b}
}

type contentTypeBucket struct {
	gcs.Bucket
}

func (b contentTypeBucket) CreateObject(
	ctx context.Context,
	req *gcs.CreateObjectRequest) (o *gcs.Object, err error) {
	// Guess a content type if necessary.
	if req.ContentType == "" {
		req.ContentType = mime.TypeByExtension(path.Ext(req.Name))
	}

	// Pass on the request.
	o, err = b.Bucket.CreateObject(ctx, req)
	return
}

func (b contentTypeBucket) ComposeObjects(
	ctx context.Context,
	req *gcs.ComposeObjectsRequest) (o *gcs.Object, err error) {
	// Guess a content type if necessary.
	if req.ContentType == "" {
		req.ContentType = mime.TypeByExtension(path.Ext(req.DstName))
	}

	// Pass on the request.
	o, err = b.Bucket.ComposeObjects(ctx, req)
	return
}
<|endoftext|>