<commit_before>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package frontgate

import (
	"context"
	"crypto/tls"
	"fmt"
	"os"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb/metadata/frontgate"
	"openpitrix.io/openpitrix/pkg/pb/metadata/types"
	"openpitrix.io/openpitrix/pkg/service/metadata/pilot/pilotutil"
)

type Server struct {
	cfg            *ConfigManager
	tlsPilotConfig *tls.Config
	etcd           *EtcdClientManager

	ch   *pilotutil.FrameChannel
	conn *grpc.ClientConn
	err  error
}

func Serve(cfg *ConfigManager, tlsPilotConfig *tls.Config) {
	p := &Server{
		cfg:            cfg,
		tlsPilotConfig: tlsPilotConfig,
		etcd:           NewEtcdClientManager(),
	}

	go ServeReverseRpcServerForPilot(cfg.Get(), tlsPilotConfig, p)
	go func() {
		err := pbfrontgate.ListenAndServeFrontgateService("tcp",
			fmt.Sprintf(":%d", cfg.Get().ListenPort),
			p,
		)
		if err != nil {
			logger.Critical(nil, "%+v", err)
			os.Exit(1)
		}
	}()

	<-make(chan bool)
}

func ServeReverseRpcServerForPilot(
	cfg *pbtypes.FrontgateConfig, tlsConfig *tls.Config,
	service pbfrontgate.FrontgateService,
) {
	logger.Info(nil, "ReverseRpcServerForPilot begin")
	defer logger.Info(nil, "ReverseRpcServerForPilot end")

	var lastErrCode = codes.OK

	for {
		ch, conn, err := pilotutil.DialFrontgateChannelTLS(
			context.Background(), fmt.Sprintf("%s:%d", cfg.PilotHost, cfg.PilotPort),
			tlsConfig,
		)
		if err != nil {
			gerr, ok := status.FromError(err)
			if !ok {
				logger.Error(nil, "err shoule be grpc error type")
				time.Sleep(time.Second)
				continue
			}

			if gerr.Code() != lastErrCode {
				logger.Error(nil, "did not connect: %v", gerr.Err())
			}

			lastErrCode = gerr.Code()
			continue
		} else {
			if lastErrCode == codes.Unavailable {
				logger.Info(nil, "pilot connect ok")
			}

			lastErrCode = codes.OK
		}

		pbfrontgate.ServeFrontgateService(ch, service)
		conn.Close()
	}
}
<commit_msg>frontgate: add withBlock and withKeepalive in DialFrontgateChannelTLS<commit_after>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package frontgate

import (
	"context"
	"crypto/tls"
	"fmt"
	"os"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"

	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb/metadata/frontgate"
	"openpitrix.io/openpitrix/pkg/pb/metadata/types"
	"openpitrix.io/openpitrix/pkg/service/metadata/pilot/pilotutil"
)

type Server struct {
	cfg            *ConfigManager
	tlsPilotConfig *tls.Config
	etcd           *EtcdClientManager

	ch   *pilotutil.FrameChannel
	conn *grpc.ClientConn
	err  error
}

func Serve(cfg *ConfigManager, tlsPilotConfig *tls.Config) {
	p := &Server{
		cfg:            cfg,
		tlsPilotConfig: tlsPilotConfig,
		etcd:           NewEtcdClientManager(),
	}

	go ServeReverseRpcServerForPilot(cfg.Get(), tlsPilotConfig, p)
	go func() {
		err := pbfrontgate.ListenAndServeFrontgateService("tcp",
			fmt.Sprintf(":%d", cfg.Get().ListenPort),
			p,
		)
		if err != nil {
			logger.Critical(nil, "%+v", err)
			os.Exit(1)
		}
	}()

	<-make(chan bool)
}

func ServeReverseRpcServerForPilot(
	cfg *pbtypes.FrontgateConfig, tlsConfig *tls.Config,
	service pbfrontgate.FrontgateService,
) {
	logger.Info(nil, "ReverseRpcServerForPilot begin")
	defer logger.Info(nil, "ReverseRpcServerForPilot end")

	var lastErrCode = codes.OK

	for {
		ch, conn, err := pilotutil.DialFrontgateChannelTLS(
			context.Background(), fmt.Sprintf("%s:%d", cfg.PilotHost, cfg.PilotPort),
			tlsConfig,
			grpc.WithBlock(),
			grpc.WithKeepaliveParams(keepalive.ClientParameters{
				Time:                30 * time.Second,
				Timeout:             10 * time.Second,
				PermitWithoutStream: true,
			}),
			grpc.WithTimeout(20*time.Second),
		)
		if err != nil {
			gerr, ok := status.FromError(err)
			if !ok {
				logger.Error(nil, "err shoule be grpc error type")
				time.Sleep(time.Second)
				continue
			}

			if gerr.Code() != lastErrCode {
				logger.Error(nil, "did not connect: %v", gerr.Err())
			}

			lastErrCode = gerr.Code()
			continue
		} else {
			if lastErrCode == codes.Unavailable {
				logger.Info(nil, "pilot connect ok")
			}

			lastErrCode = codes.OK
		}

		pbfrontgate.ServeFrontgateService(ch, service)
		conn.Close()
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The Cockroach 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.
//
// Author: Matt Tracy (matt@cockroachlabs.com)

package storage_test

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/base"
	"github.com/cockroachdb/cockroach/pkg/internal/client"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/server"
	"github.com/cockroachdb/cockroach/pkg/storage"
	"github.com/cockroachdb/cockroach/pkg/storage/engine"
	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
	"github.com/cockroachdb/cockroach/pkg/ts"
	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
	"github.com/cockroachdb/cockroach/pkg/util"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
	"github.com/kr/pretty"
	"github.com/pkg/errors"
)

type modelTimeSeriesDataStore struct {
	syncutil.Mutex
	t                  testing.TB
	containsCalled     int
	pruneCalled        int
	pruneSeenStartKeys map[string]struct{}
	pruneSeenEndKeys   map[string]struct{}
}

func (m *modelTimeSeriesDataStore) ContainsTimeSeries(start, end roachpb.RKey) bool {
	if !start.Less(end) {
		m.t.Fatalf("ContainsTimeSeries passed start key %v which is not less than end key %v", start, end)
	}
	m.Lock()
	defer m.Unlock()
	m.containsCalled++
	return true
}

func (m *modelTimeSeriesDataStore) PruneTimeSeries(
	ctx context.Context,
	snapshot engine.Reader,
	start, end roachpb.RKey,
	db *client.DB,
	now hlc.Timestamp,
) error {
	if snapshot == nil {
		m.t.Fatal("PruneTimeSeries was passed a nil snapshot")
	}
	if db == nil {
		m.t.Fatal("PruneTimeSeries was passed a nil client.DB")
	}
	if !start.Less(end) {
		m.t.Fatalf("PruneTimeSeries passed start key %v which is not less than end key %v", start, end)
	}

	m.Lock()
	defer m.Unlock()
	m.pruneCalled++
	m.pruneSeenStartKeys[start.String()] = struct{}{}
	m.pruneSeenEndKeys[end.String()] = struct{}{}
	return nil
}

// TestTimeSeriesMaintenanceQueue verifies shouldQueue and process method
// pass the correct data to the store's TimeSeriesData
func TestTimeSeriesMaintenanceQueue(t *testing.T) {
	defer leaktest.AfterTest(t)()

	model := &modelTimeSeriesDataStore{
		t:                  t,
		pruneSeenStartKeys: make(map[string]struct{}),
		pruneSeenEndKeys:   make(map[string]struct{}),
	}

	manual := hlc.NewManualClock(1)
	cfg := storage.TestStoreConfig(hlc.NewClock(manual.UnixNano, time.Nanosecond))
	cfg.TimeSeriesDataStore = model
	cfg.TestingKnobs.DisableScanner = true
	cfg.TestingKnobs.DisableSplitQueue = true

	store, stopper := createTestStoreWithConfig(t, cfg)
	defer stopper.Stop()

	// Generate several splits.
	splitKeys := []roachpb.Key{roachpb.Key("c"), roachpb.Key("b"), roachpb.Key("a")}
	for _, k := range splitKeys {
		repl := store.LookupReplica(roachpb.RKey(k), nil)
		args := adminSplitArgs(k, k)
		if _, pErr := client.SendWrappedWith(context.Background(), store, roachpb.Header{
			RangeID: repl.RangeID,
		}, &args); pErr != nil {
			t.Fatal(pErr)
		}
	}

	// Generate a list of start/end keys the model should have been passed by
	// the queue. This consists of all split keys, with KeyMin as an additional
	// start and KeyMax as an additional end.
	expectedStartKeys := make(map[string]struct{})
	expectedEndKeys := make(map[string]struct{})
	expectedStartKeys[roachpb.KeyMin.String()] = struct{}{}
	expectedEndKeys[roachpb.KeyMax.String()] = struct{}{}
	for _, expected := range splitKeys {
		expectedStartKeys[expected.String()] = struct{}{}
		expectedEndKeys[expected.String()] = struct{}{}
	}

	// Wait for splits to complete and system config to be available.
	util.SucceedsSoon(t, func() error {
		if a, e := store.ReplicaCount(), len(expectedEndKeys); a != e {
			return fmt.Errorf("expected %d replicas in store; found %d", a, e)
		}
		if _, ok := store.Gossip().GetSystemConfig(); !ok {
			return fmt.Errorf("system config not yet available")
		}
		return nil
	})

	// Force replica scan to run, which will populate the model.
	now := store.Clock().Now()
	store.ForceTimeSeriesMaintenanceQueueProcess()

	// Wait for processing to complete.
	util.SucceedsSoon(t, func() error {
		model.Lock()
		defer model.Unlock()
		if a, e := model.containsCalled, len(expectedStartKeys); a != e {
			return fmt.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
		}
		if a, e := model.pruneCalled, len(expectedStartKeys); a != e {
			return fmt.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
		}
		return nil
	})

	model.Lock()
	if a, e := model.pruneSeenStartKeys, expectedStartKeys; !reflect.DeepEqual(a, e) {
		t.Errorf("start keys seen by PruneTimeSeries did not match expectation: %s", pretty.Diff(a, e))
	}
	if a, e := model.pruneSeenEndKeys, expectedEndKeys; !reflect.DeepEqual(a, e) {
		t.Errorf("end keys seen by PruneTimeSeries did not match expectation: %s", pretty.Diff(a, e))
	}
	model.Unlock()

	util.SucceedsSoon(t, func() error {
		keys := []roachpb.RKey{roachpb.RKeyMin}
		for _, k := range splitKeys {
			keys = append(keys, roachpb.RKey(k))
		}
		for _, key := range keys {
			repl := store.LookupReplica(key, nil)
			ts, err := repl.GetQueueLastProcessed(context.TODO(), "timeSeriesMaintenance")
			if err != nil {
				return err
			}
			if ts.Less(now) {
				return errors.Errorf("expected last processed %s > %s", ts, now)
			}
		}
		return nil
	})

	// Force replica scan to run. But because we haven't moved the
	// clock forward, no pruning will take place on second invocation.
	store.ForceTimeSeriesMaintenanceQueueProcess()
	if a, e := model.containsCalled, len(expectedStartKeys); a != e {
		t.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
	}
	if a, e := model.pruneCalled, len(expectedStartKeys); a != e {
		t.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
	}

	// Move clock forward and force to scan again.
	manual.Increment(storage.TimeSeriesMaintenanceInterval.Nanoseconds())
	store.ForceTimeSeriesMaintenanceQueueProcess()
	util.SucceedsSoon(t, func() error {
		if a, e := model.containsCalled, len(expectedStartKeys)*2; a != e {
			return errors.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
		}
		if a, e := model.pruneCalled, len(expectedStartKeys)*2; a != e {
			return errors.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
		}
		return nil
	})
}

// TestTimeSeriesMaintenanceQueueServer verifies that the time series
// maintenance queue runs correctly on a test server.
func TestTimeSeriesMaintenanceQueueServer(t *testing.T) {
	defer leaktest.AfterTest(t)()

	s, _, db := serverutils.StartServer(t, base.TestServerArgs{})
	defer s.Stopper().Stop()
	tsrv := s.(*server.TestServer)
	tsdb := tsrv.TsDB()

	// Populate time series data into the server. One time series, with one
	// datapoint at the current time and two datapoints older than the pruning
	// threshold. Datapoint timestamps are set to the midpoint of sample duration
	// periods; this simplifies verification.
	seriesName := "test.metric"
	sourceName := "source1"
	now := tsrv.Clock().PhysicalNow()
	nearPast := now - (ts.Resolution10s.PruneThreshold() * 2)
	farPast := now - (ts.Resolution10s.PruneThreshold() * 4)
	sampleDuration := ts.Resolution10s.SampleDuration()
	datapoints := []tspb.TimeSeriesDatapoint{
		{
			TimestampNanos: farPast - farPast%sampleDuration + sampleDuration/2,
			Value:          100.0,
		},
		{
			TimestampNanos: nearPast - (nearPast)%sampleDuration + sampleDuration/2,
			Value:          200.0,
		},
		{
			TimestampNanos: now - now%sampleDuration + sampleDuration/2,
			Value:          300.0,
		},
	}
	if err := tsdb.StoreData(context.TODO(), ts.Resolution10s, []tspb.TimeSeriesData{
		{
			Name:       seriesName,
			Source:     sourceName,
			Datapoints: datapoints,
		},
	}); err != nil {
		t.Fatal(err)
	}

	// Generate a split key at a timestamp halfway between near past and far past.
	splitKey := ts.MakeDataKey(
		seriesName, sourceName, ts.Resolution10s, farPast+(nearPast-farPast)/2,
	)

	// Force a range split in between near past and far past. This guarantees
	// that the pruning operation will issue a DeleteRange which spans ranges.
	if err := db.AdminSplit(context.TODO(), splitKey); err != nil {
		t.Fatal(err)
	}

	// getDatapoints queries all datapoints in the series from the beginning
	// of time to a point in the near future.
	getDatapoints := func() ([]tspb.TimeSeriesDatapoint, error) {
		dps, _, err := tsdb.Query(
			context.TODO(),
			tspb.Query{Name: seriesName},
			ts.Resolution10s,
			ts.Resolution10s.SampleDuration(),
			0,
			now+ts.Resolution10s.SlabDuration(),
		)
		return dps, err
	}

	// Verify the datapoints are all present.
	actualDatapoints, err := getDatapoints()
	if err != nil {
		t.Fatal(err)
	}
	if a, e := actualDatapoints, datapoints; !reflect.DeepEqual(a, e) {
		t.Fatalf("got datapoints %v, expected %v, diff: %s", a, e, pretty.Diff(a, e))
	}

	// Force pruning.
	storeID := roachpb.StoreID(1)
	store, err := tsrv.Stores().GetStore(roachpb.StoreID(1))
	if err != nil {
		t.Fatalf("error retrieving store %d: %s", storeID, err)
	}
	store.ForceTimeSeriesMaintenanceQueueProcess()

	// Verify the older datapoint has been pruned.
	util.SucceedsSoon(t, func() error {
		actualDatapoints, err = getDatapoints()
		if err != nil {
			return err
		}
		if a, e := actualDatapoints, datapoints[2:]; !reflect.DeepEqual(a, e) {
			return fmt.Errorf("got datapoints %v, expected %v, diff: %s", a, e, pretty.Diff(a, e))
		}
		return nil
	})
}
<commit_msg>storage: fix race in time series maintenance test<commit_after>// Copyright 2016 The Cockroach 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.
//
// Author: Matt Tracy (matt@cockroachlabs.com)

package storage_test

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/base"
	"github.com/cockroachdb/cockroach/pkg/internal/client"
	"github.com/cockroachdb/cockroach/pkg/roachpb"
	"github.com/cockroachdb/cockroach/pkg/server"
	"github.com/cockroachdb/cockroach/pkg/storage"
	"github.com/cockroachdb/cockroach/pkg/storage/engine"
	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
	"github.com/cockroachdb/cockroach/pkg/ts"
	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
	"github.com/cockroachdb/cockroach/pkg/util"
	"github.com/cockroachdb/cockroach/pkg/util/hlc"
	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
	"github.com/kr/pretty"
	"github.com/pkg/errors"
)

type modelTimeSeriesDataStore struct {
	syncutil.Mutex
	t                  testing.TB
	containsCalled     int
	pruneCalled        int
	pruneSeenStartKeys map[string]struct{}
	pruneSeenEndKeys   map[string]struct{}
}

func (m *modelTimeSeriesDataStore) ContainsTimeSeries(start, end roachpb.RKey) bool {
	if !start.Less(end) {
		m.t.Fatalf("ContainsTimeSeries passed start key %v which is not less than end key %v", start, end)
	}
	m.Lock()
	defer m.Unlock()
	m.containsCalled++
	return true
}

func (m *modelTimeSeriesDataStore) PruneTimeSeries(
	ctx context.Context,
	snapshot engine.Reader,
	start, end roachpb.RKey,
	db *client.DB,
	now hlc.Timestamp,
) error {
	if snapshot == nil {
		m.t.Fatal("PruneTimeSeries was passed a nil snapshot")
	}
	if db == nil {
		m.t.Fatal("PruneTimeSeries was passed a nil client.DB")
	}
	if !start.Less(end) {
		m.t.Fatalf("PruneTimeSeries passed start key %v which is not less than end key %v", start, end)
	}

	m.Lock()
	defer m.Unlock()
	m.pruneCalled++
	m.pruneSeenStartKeys[start.String()] = struct{}{}
	m.pruneSeenEndKeys[end.String()] = struct{}{}
	return nil
}

// TestTimeSeriesMaintenanceQueue verifies shouldQueue and process method
// pass the correct data to the store's TimeSeriesData
func TestTimeSeriesMaintenanceQueue(t *testing.T) {
	defer leaktest.AfterTest(t)()

	model := &modelTimeSeriesDataStore{
		t:                  t,
		pruneSeenStartKeys: make(map[string]struct{}),
		pruneSeenEndKeys:   make(map[string]struct{}),
	}

	manual := hlc.NewManualClock(1)
	cfg := storage.TestStoreConfig(hlc.NewClock(manual.UnixNano, time.Nanosecond))
	cfg.TimeSeriesDataStore = model
	cfg.TestingKnobs.DisableScanner = true
	cfg.TestingKnobs.DisableSplitQueue = true

	store, stopper := createTestStoreWithConfig(t, cfg)
	defer stopper.Stop()

	// Generate several splits.
	splitKeys := []roachpb.Key{roachpb.Key("c"), roachpb.Key("b"), roachpb.Key("a")}
	for _, k := range splitKeys {
		repl := store.LookupReplica(roachpb.RKey(k), nil)
		args := adminSplitArgs(k, k)
		if _, pErr := client.SendWrappedWith(context.Background(), store, roachpb.Header{
			RangeID: repl.RangeID,
		}, &args); pErr != nil {
			t.Fatal(pErr)
		}
	}

	// Generate a list of start/end keys the model should have been passed by
	// the queue. This consists of all split keys, with KeyMin as an additional
	// start and KeyMax as an additional end.
	expectedStartKeys := make(map[string]struct{})
	expectedEndKeys := make(map[string]struct{})
	expectedStartKeys[roachpb.KeyMin.String()] = struct{}{}
	expectedEndKeys[roachpb.KeyMax.String()] = struct{}{}
	for _, expected := range splitKeys {
		expectedStartKeys[expected.String()] = struct{}{}
		expectedEndKeys[expected.String()] = struct{}{}
	}

	// Wait for splits to complete and system config to be available.
	util.SucceedsSoon(t, func() error {
		if a, e := store.ReplicaCount(), len(expectedEndKeys); a != e {
			return fmt.Errorf("expected %d replicas in store; found %d", a, e)
		}
		if _, ok := store.Gossip().GetSystemConfig(); !ok {
			return fmt.Errorf("system config not yet available")
		}
		return nil
	})

	// Force replica scan to run, which will populate the model.
	now := store.Clock().Now()
	store.ForceTimeSeriesMaintenanceQueueProcess()

	// Wait for processing to complete.
	util.SucceedsSoon(t, func() error {
		model.Lock()
		defer model.Unlock()
		if a, e := model.containsCalled, len(expectedStartKeys); a != e {
			return fmt.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
		}
		if a, e := model.pruneCalled, len(expectedStartKeys); a != e {
			return fmt.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
		}
		return nil
	})

	model.Lock()
	if a, e := model.pruneSeenStartKeys, expectedStartKeys; !reflect.DeepEqual(a, e) {
		t.Errorf("start keys seen by PruneTimeSeries did not match expectation: %s", pretty.Diff(a, e))
	}
	if a, e := model.pruneSeenEndKeys, expectedEndKeys; !reflect.DeepEqual(a, e) {
		t.Errorf("end keys seen by PruneTimeSeries did not match expectation: %s", pretty.Diff(a, e))
	}
	model.Unlock()

	util.SucceedsSoon(t, func() error {
		keys := []roachpb.RKey{roachpb.RKeyMin}
		for _, k := range splitKeys {
			keys = append(keys, roachpb.RKey(k))
		}
		for _, key := range keys {
			repl := store.LookupReplica(key, nil)
			ts, err := repl.GetQueueLastProcessed(context.TODO(), "timeSeriesMaintenance")
			if err != nil {
				return err
			}
			if ts.Less(now) {
				return errors.Errorf("expected last processed %s > %s", ts, now)
			}
		}
		return nil
	})

	// Force replica scan to run. But because we haven't moved the
	// clock forward, no pruning will take place on second invocation.
	store.ForceTimeSeriesMaintenanceQueueProcess()
	model.Lock()
	if a, e := model.containsCalled, len(expectedStartKeys); a != e {
		t.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
	}
	if a, e := model.pruneCalled, len(expectedStartKeys); a != e {
		t.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
	}
	model.Unlock()

	// Move clock forward and force to scan again.
	manual.Increment(storage.TimeSeriesMaintenanceInterval.Nanoseconds())
	store.ForceTimeSeriesMaintenanceQueueProcess()
	util.SucceedsSoon(t, func() error {
		model.Lock()
		defer model.Unlock()
		if a, e := model.containsCalled, len(expectedStartKeys)*2; a != e {
			return errors.Errorf("ContainsTimeSeries called %d times; expected %d", a, e)
		}
		if a, e := model.pruneCalled, len(expectedStartKeys)*2; a != e {
			return errors.Errorf("PruneTimeSeries called %d times; expected %d", a, e)
		}
		return nil
	})
}

// TestTimeSeriesMaintenanceQueueServer verifies that the time series
// maintenance queue runs correctly on a test server.
func TestTimeSeriesMaintenanceQueueServer(t *testing.T) {
	defer leaktest.AfterTest(t)()

	s, _, db := serverutils.StartServer(t, base.TestServerArgs{})
	defer s.Stopper().Stop()
	tsrv := s.(*server.TestServer)
	tsdb := tsrv.TsDB()

	// Populate time series data into the server. One time series, with one
	// datapoint at the current time and two datapoints older than the pruning
	// threshold. Datapoint timestamps are set to the midpoint of sample duration
	// periods; this simplifies verification.
	seriesName := "test.metric"
	sourceName := "source1"
	now := tsrv.Clock().PhysicalNow()
	nearPast := now - (ts.Resolution10s.PruneThreshold() * 2)
	farPast := now - (ts.Resolution10s.PruneThreshold() * 4)
	sampleDuration := ts.Resolution10s.SampleDuration()
	datapoints := []tspb.TimeSeriesDatapoint{
		{
			TimestampNanos: farPast - farPast%sampleDuration + sampleDuration/2,
			Value:          100.0,
		},
		{
			TimestampNanos: nearPast - (nearPast)%sampleDuration + sampleDuration/2,
			Value:          200.0,
		},
		{
			TimestampNanos: now - now%sampleDuration + sampleDuration/2,
			Value:          300.0,
		},
	}
	if err := tsdb.StoreData(context.TODO(), ts.Resolution10s, []tspb.TimeSeriesData{
		{
			Name:       seriesName,
			Source:     sourceName,
			Datapoints: datapoints,
		},
	}); err != nil {
		t.Fatal(err)
	}

	// Generate a split key at a timestamp halfway between near past and far past.
	splitKey := ts.MakeDataKey(
		seriesName, sourceName, ts.Resolution10s, farPast+(nearPast-farPast)/2,
	)

	// Force a range split in between near past and far past. This guarantees
	// that the pruning operation will issue a DeleteRange which spans ranges.
	if err := db.AdminSplit(context.TODO(), splitKey); err != nil {
		t.Fatal(err)
	}

	// getDatapoints queries all datapoints in the series from the beginning
	// of time to a point in the near future.
	getDatapoints := func() ([]tspb.TimeSeriesDatapoint, error) {
		dps, _, err := tsdb.Query(
			context.TODO(),
			tspb.Query{Name: seriesName},
			ts.Resolution10s,
			ts.Resolution10s.SampleDuration(),
			0,
			now+ts.Resolution10s.SlabDuration(),
		)
		return dps, err
	}

	// Verify the datapoints are all present.
	actualDatapoints, err := getDatapoints()
	if err != nil {
		t.Fatal(err)
	}
	if a, e := actualDatapoints, datapoints; !reflect.DeepEqual(a, e) {
		t.Fatalf("got datapoints %v, expected %v, diff: %s", a, e, pretty.Diff(a, e))
	}

	// Force pruning.
	storeID := roachpb.StoreID(1)
	store, err := tsrv.Stores().GetStore(roachpb.StoreID(1))
	if err != nil {
		t.Fatalf("error retrieving store %d: %s", storeID, err)
	}
	store.ForceTimeSeriesMaintenanceQueueProcess()

	// Verify the older datapoint has been pruned.
	util.SucceedsSoon(t, func() error {
		actualDatapoints, err = getDatapoints()
		if err != nil {
			return err
		}
		if a, e := actualDatapoints, datapoints[2:]; !reflect.DeepEqual(a, e) {
			return fmt.Errorf("got datapoints %v, expected %v, diff: %s", a, e, pretty.Diff(a, e))
		}
		return nil
	})
}
<|endoftext|>


<commit_before>package disk_test

import (
	"errors"
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	fakeboshaction "github.com/cloudfoundry/bosh-agent/agent/action/fakes"
	. "github.com/cloudfoundry/bosh-agent/platform/disk"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	fakesys "github.com/cloudfoundry/bosh-utils/system/fakes"
)

const devSdaSfdiskEmptyDump = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        0, size=    0, Id= 0
/dev/sda2 : start=        0, size=    0, Id= 0
/dev/sda3 : start=        0, size=    0, Id= 0
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devSdaSfdiskNotableDumpStderr = `
sfdisk: ERROR: sector 0 does not have an msdos signature
 /dev/sda: unrecognized partition table type
No partitions found`

const devSdaSfdiskDump = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        1, size= xxxx, Id=82
/dev/sda2 : start=     xxxx, size= xxxx, Id=83
/dev/sda3 : start=     xxxx, size= xxxx, Id=83
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devSdaSfdiskDumpOnePartition = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        1, size= xxxx, Id=83
/dev/sda2 : start=     xxxx, size= xxxx, Id=83
/dev/sda3 : start=        0, size=    0, Id= 0
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devMapperSfdiskDumpOnePartition = `# partition table of /dev/mapper/xxxxxx
unit: sectors

/dev/mapper/xxxxxx1 : start=        1, size= xxxx , Id=83
/dev/mapper/xxxxxx2 : start=        0, size=        0, Id= 0
/dev/mapper/xxxxxx3 : start=        0, size=        0, Id= 0
/dev/mapper/xxxxxx4 : start=        0, size=        0, Id= 0
`

const expectedDmSetupLs = `
xxxxxx-part1	(252:1)
xxxxxx	(252:0)
`

var _ = Describe("sfdiskPartitioner", func() {
	var (
		runner      *fakesys.FakeCmdRunner
		partitioner Partitioner
		fakeclock   *fakeboshaction.FakeClock
	)

	BeforeEach(func() {
		runner = fakesys.NewFakeCmdRunner()
		logger := boshlog.NewLogger(boshlog.LevelNone)
		fakeclock = &fakeboshaction.FakeClock{}

		partitioner = NewSfdiskPartitioner(logger, runner, fakeclock)
	})

	It("sfdisk partition", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskEmptyDump})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: "1048576"})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	Context("when we get an error occurs", func() {
		Context("during get partitions", func() {
			It("raises error", func() {
				runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Error: errors.New("Some weird error")})

				partitions := []Partition{
					{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
					{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
					{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
				}

				err := partitioner.Partition("/dev/sda", partitions)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("Some weird error"))
			})
		})

		Context("when getting device size", func() {
			It("raises error", func() {
				runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
				runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Error: errors.New("Another weird error")})

				partitions := []Partition{
					{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
				}

				err := partitioner.Partition("/dev/sda", partitions)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("Another weird error"))
			})
		})
	})

	It("sfdisk partition with no partition table", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stderr: devSdaSfdiskNotableDumpStderr})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	It("sfdisk partition for multipath", func() {
		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/mapper/xxxxxx", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/mapper/xxxxxx"}))
		Expect(22).To(Equal(len(runner.RunCommands)))
		Expect(runner.RunCommands[1]).To(Equal([]string{"/etc/init.d/open-iscsi", "restart"}))
	})

	It("sfdisk get device size in mb", func() {
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 40000*1024)})

		size, err := partitioner.GetDeviceSizeInBytes("/dev/sda")
		Expect(err).ToNot(HaveOccurred())

		Expect(size).To(Equal(uint64(40000 * 1024 * 1024)))
	})

	It("sfdisk partition when partitions already match", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDump})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 525*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1020*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda3", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 500*1024)})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(0))
	})

	It("sfdisk partition when partitions already match for mutlitpath", func() {
		runner.AddCmdResult("sfdisk -d /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: devMapperSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024+7000)})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx-part1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
		}

		partitioner.Partition("/dev/mapper/xxxxxx", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(0))
	})

	It("sfdisk partition with last partition not matching size", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 512*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(1))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	It("sfdisk partition with last partition filling disk", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(0).To(Equal(len(runner.RunCommandsWithInput)))
	})

	It("sfdisk command is retried 20 times", func() {
		for i := 0; i < 19; i++ {
			testError := fmt.Errorf("test error")
			runner.AddCmdResult(",,L\n sfdisk -uM /dev/sda", fakesys.FakeCmdResult{ExitStatus: 1, Error: testError})
		}
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: "1048576"})

		runner.AddCmdResult(",,L\n sfdisk -uM /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})

		partitions := []Partition{
			{Type: PartitionTypeLinux},
		}

		err := partitioner.Partition("/dev/sda", partitions)
		Expect(err).To(BeNil())
		Expect(fakeclock.SleepCallCount()).To(Equal(19))
		Expect(len(runner.RunCommandsWithInput)).To(Equal(20))
	})

	It("dmsetup command is retried 20 times", func() {
		runner.AddCmdResult("sfdisk -d /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		for i := 0; i < 19; i++ {
			testError := fmt.Errorf("test error")
			runner.AddCmdResult("dmsetup ls", fakesys.FakeCmdResult{ExitStatus: 1, Error: testError})
		}
		runner.AddCmdResult("dmsetup ls", fakesys.FakeCmdResult{Stdout: expectedDmSetupLs})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: "1048576"})

		partitions := []Partition{
			{Type: PartitionTypeLinux},
		}

		err := partitioner.Partition("/dev/mapper/xxxxxx", partitions)
		Expect(err).To(BeNil())
		Expect(fakeclock.SleepCallCount()).To(Equal(19))
		Expect(len(runner.RunCommands)).To(Equal(25))
	})
})
<commit_msg>Correct a typo for multipath<commit_after>package disk_test

import (
	"errors"
	"fmt"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	fakeboshaction "github.com/cloudfoundry/bosh-agent/agent/action/fakes"
	. "github.com/cloudfoundry/bosh-agent/platform/disk"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	fakesys "github.com/cloudfoundry/bosh-utils/system/fakes"
)

const devSdaSfdiskEmptyDump = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        0, size=    0, Id= 0
/dev/sda2 : start=        0, size=    0, Id= 0
/dev/sda3 : start=        0, size=    0, Id= 0
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devSdaSfdiskNotableDumpStderr = `
sfdisk: ERROR: sector 0 does not have an msdos signature
 /dev/sda: unrecognized partition table type
No partitions found`

const devSdaSfdiskDump = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        1, size= xxxx, Id=82
/dev/sda2 : start=     xxxx, size= xxxx, Id=83
/dev/sda3 : start=     xxxx, size= xxxx, Id=83
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devSdaSfdiskDumpOnePartition = `# partition table of /dev/sda
unit: sectors

/dev/sda1 : start=        1, size= xxxx, Id=83
/dev/sda2 : start=     xxxx, size= xxxx, Id=83
/dev/sda3 : start=        0, size=    0, Id= 0
/dev/sda4 : start=        0, size=    0, Id= 0
`

const devMapperSfdiskDumpOnePartition = `# partition table of /dev/mapper/xxxxxx
unit: sectors

/dev/mapper/xxxxxx1 : start=        1, size= xxxx , Id=83
/dev/mapper/xxxxxx2 : start=        0, size=        0, Id= 0
/dev/mapper/xxxxxx3 : start=        0, size=        0, Id= 0
/dev/mapper/xxxxxx4 : start=        0, size=        0, Id= 0
`

const expectedDmSetupLs = `
xxxxxx-part1	(252:1)
xxxxxx	(252:0)
`

var _ = Describe("sfdiskPartitioner", func() {
	var (
		runner      *fakesys.FakeCmdRunner
		partitioner Partitioner
		fakeclock   *fakeboshaction.FakeClock
	)

	BeforeEach(func() {
		runner = fakesys.NewFakeCmdRunner()
		logger := boshlog.NewLogger(boshlog.LevelNone)
		fakeclock = &fakeboshaction.FakeClock{}

		partitioner = NewSfdiskPartitioner(logger, runner, fakeclock)
	})

	It("sfdisk partition", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskEmptyDump})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: "1048576"})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	Context("when we get an error occurs", func() {
		Context("during get partitions", func() {
			It("raises error", func() {
				runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Error: errors.New("Some weird error")})

				partitions := []Partition{
					{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
					{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
					{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
				}

				err := partitioner.Partition("/dev/sda", partitions)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("Some weird error"))
			})
		})

		Context("when getting device size", func() {
			It("raises error", func() {
				runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
				runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Error: errors.New("Another weird error")})

				partitions := []Partition{
					{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
				}

				err := partitioner.Partition("/dev/sda", partitions)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("Another weird error"))
			})
		})
	})

	It("sfdisk partition with no partition table", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stderr: devSdaSfdiskNotableDumpStderr})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	It("sfdisk partition for multipath", func() {
		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/mapper/xxxxxx", partitions)

		Expect(1).To(Equal(len(runner.RunCommandsWithInput)))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",512,S\n,1024,L\n,,L\n", "sfdisk", "-uM", "/dev/mapper/xxxxxx"}))
		Expect(22).To(Equal(len(runner.RunCommands)))
		Expect(runner.RunCommands[1]).To(Equal([]string{"/etc/init.d/open-iscsi", "restart"}))
	})

	It("sfdisk get device size in mb", func() {
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 40000*1024)})

		size, err := partitioner.GetDeviceSizeInBytes("/dev/sda")
		Expect(err).ToNot(HaveOccurred())

		Expect(size).To(Equal(uint64(40000 * 1024 * 1024)))
	})

	It("sfdisk partition when partitions already match", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDump})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 525*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1020*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda3", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 500*1024)})

		partitions := []Partition{
			{Type: PartitionTypeSwap, SizeInBytes: 512 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux, SizeInBytes: 512 * 1024 * 1024},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(0))
	})

	It("sfdisk partition when partitions already match for multipath", func() {
		runner.AddCmdResult("sfdisk -d /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: devMapperSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024+7000)})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx-part1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
		}

		partitioner.Partition("/dev/mapper/xxxxxx", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(0))
	})

	It("sfdisk partition with last partition not matching size", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 512*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(len(runner.RunCommandsWithInput)).To(Equal(1))
		Expect(runner.RunCommandsWithInput[0]).To(Equal([]string{",1024,L\n,,L\n", "sfdisk", "-uM", "/dev/sda"}))
	})

	It("sfdisk partition with last partition filling disk", func() {
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 2048*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda1", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})
		runner.AddCmdResult("sfdisk -s /dev/sda2", fakesys.FakeCmdResult{Stdout: fmt.Sprintf("%d\n", 1024*1024)})

		partitions := []Partition{
			{Type: PartitionTypeLinux, SizeInBytes: 1024 * 1024 * 1024},
			{Type: PartitionTypeLinux},
		}

		partitioner.Partition("/dev/sda", partitions)

		Expect(0).To(Equal(len(runner.RunCommandsWithInput)))
	})

	It("sfdisk command is retried 20 times", func() {
		for i := 0; i < 19; i++ {
			testError := fmt.Errorf("test error")
			runner.AddCmdResult(",,L\n sfdisk -uM /dev/sda", fakesys.FakeCmdResult{ExitStatus: 1, Error: testError})
		}
		runner.AddCmdResult("sfdisk -d /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		runner.AddCmdResult("sfdisk -s /dev/sda", fakesys.FakeCmdResult{Stdout: "1048576"})

		runner.AddCmdResult(",,L\n sfdisk -uM /dev/sda", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})

		partitions := []Partition{
			{Type: PartitionTypeLinux},
		}

		err := partitioner.Partition("/dev/sda", partitions)
		Expect(err).To(BeNil())
		Expect(fakeclock.SleepCallCount()).To(Equal(19))
		Expect(len(runner.RunCommandsWithInput)).To(Equal(20))
	})

	It("dmsetup command is retried 20 times", func() {
		runner.AddCmdResult("sfdisk -d /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: devSdaSfdiskDumpOnePartition})
		for i := 0; i < 19; i++ {
			testError := fmt.Errorf("test error")
			runner.AddCmdResult("dmsetup ls", fakesys.FakeCmdResult{ExitStatus: 1, Error: testError})
		}
		runner.AddCmdResult("dmsetup ls", fakesys.FakeCmdResult{Stdout: expectedDmSetupLs})
		runner.AddCmdResult("sfdisk -s /dev/mapper/xxxxxx", fakesys.FakeCmdResult{Stdout: "1048576"})

		partitions := []Partition{
			{Type: PartitionTypeLinux},
		}

		err := partitioner.Partition("/dev/mapper/xxxxxx", partitions)
		Expect(err).To(BeNil())
		Expect(fakeclock.SleepCallCount()).To(Equal(19))
		Expect(len(runner.RunCommands)).To(Equal(25))
	})
})
<|endoftext|>


<commit_before>/* https://leetcode.com/problems/remove-linked-list-elements/#/description
Remove all elements from a linked list of integers that have value val.

Example
Given: 1 --> 2 --> 6 --> 3 --> 4 --> 5 --> 6, val = 6
Return: 1 --> 2 --> 3 --> 4 --> 5
*/

package leetcode

func removeElements(head *ListNode, val int) *ListNode {
	var r, cur, temp *ListNode
	for ; head != nil; head = head.Next {
		if head.Val != val {
			temp = &ListNode{Val: head.Val}
			if r == nil {
				r = temp
				cur = r
			} else {
				cur.Next = temp
				cur = temp
			}
		}
	}
	return r
}
<commit_msg>enhance code<commit_after>/* https://leetcode.com/problems/remove-linked-list-elements/#/description
Remove all elements from a linked list of integers that have value val.

Example
Given: 1 --> 2 --> 6 --> 3 --> 4 --> 5 --> 6, val = 6
Return: 1 --> 2 --> 3 --> 4 --> 5
*/

package leetcode

func removeElements(head *ListNode, val int) *ListNode {
	var r, cur, temp *ListNode
	for ; head != nil; head = head.Next {
		if head.Val != val {
			temp = &ListNode{Val: head.Val}
			if r == nil {
				r = temp
			} else {
				cur.Next = temp
			}
			cur = temp
		}
	}
	return r
}
<|endoftext|>


<commit_before>package cli

import "fmt"
import "strings"

type paramable struct {
	*writer
	params       paramsList
	paramValues  map[string]Value
	paramsParsed bool
}

// Arg returns the i'th argument.  Arg(0) is the first remaining argument
// after flags have been processed.
func (cmd *paramable) Param(key string) Value {
	value, ok := cmd.paramValues[key]
	if ok {
		return value
	}
	var emptyString stringValue
	emptyString = ""
	return &emptyString
}

// Args returns the non-flag arguments.
func (cmd *paramable) Params() map[string]Value {
	return cmd.paramValues
}

// UsageString returns the params usage as a string
func (cmd *paramable) UsageString() string {
	var formattednames []string
	for i := 0; i < len(cmd.params); i++ {
		param := cmd.params[i]
		formattednames = append(formattednames, fmt.Sprintf("<%s>", param.Name))
	}
	return strings.Join(formattednames, " ")
}

// Set Param names from strings
func (cmd *paramable) DefineParams(names ...string) {
	var params []*Param
	for i := 0; i < len(names); i++ {
		name := names[i]
		param := &Param{Name: name}
		params = append(params, param)
	}
	cmd.params = params
}

func (cmd *paramable) parse(args []string) []string {
	var seenParams paramsList

	if len(cmd.params) == 0 {
		return args
	}
	i := 0
	for i < len(args) && i < len(cmd.params) {
		param := cmd.params[i]
		seenParams = append(seenParams, param)
		str := ""
		if cmd.paramValues == nil {
			cmd.paramValues = make(map[string]Value)
		}
		cmd.paramValues[param.Name] = newStringValue(args[i], &str)
		i++
	}
	missingParams := cmd.params.Compare(seenParams)
	if len(missingParams) > 0 {
		var msg string
		if len(missingParams) == 1 {
			msg = "missing param"
		} else {
			msg = "missing params"
		}
		cmd.errf("%s: %s", msg, strings.Join(missingParams.Names(), ", "))
	}

	return args[i:]
}
<commit_msg>raise an error if an invalid param is requested<commit_after>package cli

import "fmt"
import "strings"

type paramable struct {
	*writer
	params       paramsList
	paramValues  map[string]Value
	paramsParsed bool
}

// Param returns named param
func (cmd *paramable) Param(key string) Value {
	value, ok := cmd.paramValues[key]
	if !ok {
		cmd.errf("invalid param: %s", key)
	}
	return value
}

// Args returns the non-flag arguments.
func (cmd *paramable) Params() map[string]Value {
	return cmd.paramValues
}

// UsageString returns the params usage as a string
func (cmd *paramable) UsageString() string {
	var formattednames []string
	for i := 0; i < len(cmd.params); i++ {
		param := cmd.params[i]
		formattednames = append(formattednames, fmt.Sprintf("<%s>", param.Name))
	}
	return strings.Join(formattednames, " ")
}

// Set Param names from strings
func (cmd *paramable) DefineParams(names ...string) {
	var params []*Param
	for i := 0; i < len(names); i++ {
		name := names[i]
		param := &Param{Name: name}
		params = append(params, param)
	}
	cmd.params = params
}

func (cmd *paramable) parse(args []string) []string {
	var seenParams paramsList

	if len(cmd.params) == 0 {
		return args
	}
	i := 0
	for i < len(args) && i < len(cmd.params) {
		param := cmd.params[i]
		seenParams = append(seenParams, param)
		str := ""
		if cmd.paramValues == nil {
			cmd.paramValues = make(map[string]Value)
		}
		cmd.paramValues[param.Name] = newStringValue(args[i], &str)
		i++
	}
	missingParams := cmd.params.Compare(seenParams)
	if len(missingParams) > 0 {
		var msg string
		if len(missingParams) == 1 {
			msg = "missing param"
		} else {
			msg = "missing params"
		}
		cmd.errf("%s: %s", msg, strings.Join(missingParams.Names(), ", "))
	}

	return args[i:]
}
<|endoftext|>


<commit_before>package mongotrainer

import (
	// "code.google.com/p/goprotobuf/proto"
	"fmt"
	dt "github.com/ajtulloch/decisiontrees"
	pb "github.com/ajtulloch/decisiontrees/protobufs"
	"github.com/golang/glog"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"time"
)

// MongoTrainer polls a MongoDB collection for changes
// and spins up training jobs based on these changes
type MongoTrainer struct {
	Collection *mgo.Collection
}

type trainingTask struct {
	objectID bson.ObjectId
	row      *pb.TrainingRow
}

type idRow struct {
	ID bson.ObjectId `bson:"_id,omitempty"`
}

func (m *MongoTrainer) pollTasks(c chan *trainingTask) {
	getUnclaimedTask := func() {
		id := idRow{}
		err := m.Collection.Find(bson.M{
			"trainingStatus": pb.TrainingStatus_UNCLAIMED.Enum(),
		}).Select(bson.M{
			"_id": 1,
		}).One(&id)
		if err != nil {
			glog.Error(err)
			return
		}

		t := &pb.TrainingRow{}
		err = m.Collection.FindId(id.ID).One(t)
		if err != nil {
			glog.Error(err)
			return
		}
		c <- &trainingTask{
			objectID: id.ID,
			row:      t,
		}
	}

	getUnclaimedTask()
	for {
		select {
		case <-time.After(*MongoPollTime):
			getUnclaimedTask()
		}
	}
}

// Loop starts the polling thread, and selects on the channel of
// potential tasks
func (m *MongoTrainer) Loop() {
	taskChannel := make(chan *trainingTask)
	go func() { m.pollTasks(taskChannel) }()
	for {
		select {
		case task := <-taskChannel:
			err := m.runTask(task)
			glog.Info("Starting task %v", task)
			if err != nil {
				glog.Errorf("Got error %v running task %v", err, task)
				continue
			}
			glog.Info("Successfully trained task %v", task)
		}
	}
}

func (m *MongoTrainer) runTraining(task *trainingTask) error {
	dataSource, err := NewDataSource(task.row.GetDataSourceConfig(), m.Collection.Database.Session)
	if err != nil {
		return err
	}
	trainingData, err := dataSource.GetTrainingData()
	if err != nil {
		return err
	}

	generator, err := dt.NewForestGenerator(task.row.GetForestConfig())
	if err != nil {
		return err
	}
	task.row.Forest = generator.ConstructForest(trainingData.GetTrain())
	task.row.TrainingResults = dt.LearningCurve(task.row.Forest, trainingData.GetTest())
	return nil
}

func (m *MongoTrainer) claimTask(task *trainingTask) error {
	return m.cas(task.objectID, pb.TrainingStatus_UNCLAIMED, pb.TrainingStatus_PROCESSING)
}

func (m *MongoTrainer) runTask(task *trainingTask) error {
	err := m.claimTask(task)
	if err != nil {
		return err
	}

	err = m.runTraining(task)
	if err != nil {
		return err
	}

	err = m.finalizeTask(task)
	if err != nil {
		return err
	}
	return nil
}

// cas atomically compares-and-swaps the given objectId between the given training statuses
func (m *MongoTrainer) cas(objectID bson.ObjectId, from, to pb.TrainingStatus) error {
	newRow := &pb.TrainingRow{}
	changeInfo, err := m.Collection.Find(bson.M{
		"_id":            objectID,
		"trainingStatus": from.Enum(),
	}).Apply(mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"trainingStatus": to.Enum(),
			},
		},
		ReturnNew: true,
	}, newRow)

	if err != nil {
		return err
	}

	if changeInfo.Updated != 1 {
		return fmt.Errorf("failed CAS'ing task %v from state %v to state %v", objectID, from, to)
	}
	glog.Infof("Updated objectId %v from state %v to state %v", objectID, from, to)
	return nil
}

func (m *MongoTrainer) finalizeTask(task *trainingTask) error {
	err := m.cas(task.objectID, pb.TrainingStatus_PROCESSING, pb.TrainingStatus_FINISHED)
	if err != nil {
		return err
	}
	return m.Collection.UpdateId(task.objectID, bson.M{
		"$set": bson.M{
			"forest":          task.row.Forest,
			"trainingResults": task.row.TrainingResults,
		},
	})
}
<commit_msg>[Decision Trees] Minor MongoTrainer polling refactoring<commit_after>package mongotrainer

import (
	// "code.google.com/p/goprotobuf/proto"
	"fmt"
	dt "github.com/ajtulloch/decisiontrees"
	pb "github.com/ajtulloch/decisiontrees/protobufs"
	"github.com/golang/glog"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"time"
)

// MongoTrainer polls a MongoDB collection for changes
// and spins up training jobs based on these changes
type MongoTrainer struct {
	Collection *mgo.Collection
}

type trainingTask struct {
	objectID bson.ObjectId
	row      *pb.TrainingRow
}

type idRow struct {
	ID bson.ObjectId `bson:"_id,omitempty"`
}

func (m *MongoTrainer) pollTasks(c chan *trainingTask) {
	sendUnclaimedTask := func() {
		id := idRow{}
		err := m.Collection.Find(bson.M{
			"trainingStatus": pb.TrainingStatus_UNCLAIMED.Enum(),
		}).Select(bson.M{
			"_id": 1,
		}).One(&id)
		if err != nil {
			glog.Error(err)
			return
		}

		t := &pb.TrainingRow{}
		err = m.Collection.FindId(id.ID).One(t)
		if err != nil {
			glog.Error(err)
			return
		}
		c <- &trainingTask{
			objectID: id.ID,
			row:      t,
		}
	}

	for {
		sendUnclaimedTask()
		time.Sleep(*MongoPollTime)
	}
}

// Loop starts the polling thread, and selects on the channel of
// potential tasks
func (m *MongoTrainer) Loop() {
	taskChannel := make(chan *trainingTask)
	go func() { m.pollTasks(taskChannel) }()
	for {
		select {
		case task := <-taskChannel:
			err := m.runTask(task)
			glog.Info("Starting task %v", task)
			if err != nil {
				glog.Errorf("Got error %v running task %v", err, task)
				continue
			}
			glog.Info("Successfully trained task %v", task)
		}
	}
}

func (m *MongoTrainer) runTraining(task *trainingTask) error {
	dataSource, err := NewDataSource(task.row.GetDataSourceConfig(), m.Collection.Database.Session)
	if err != nil {
		return err
	}
	trainingData, err := dataSource.GetTrainingData()
	if err != nil {
		return err
	}

	generator, err := dt.NewForestGenerator(task.row.GetForestConfig())
	if err != nil {
		return err
	}
	task.row.Forest = generator.ConstructForest(trainingData.GetTrain())
	task.row.TrainingResults = dt.LearningCurve(task.row.Forest, trainingData.GetTest())
	return nil
}

func (m *MongoTrainer) claimTask(task *trainingTask) error {
	return m.cas(task.objectID, pb.TrainingStatus_UNCLAIMED, pb.TrainingStatus_PROCESSING)
}

func (m *MongoTrainer) runTask(task *trainingTask) error {
	err := m.claimTask(task)
	if err != nil {
		return err
	}

	err = m.runTraining(task)
	if err != nil {
		return err
	}

	err = m.finalizeTask(task)
	if err != nil {
		return err
	}
	return nil
}

// cas atomically compares-and-swaps the given objectId between the given training statuses
func (m *MongoTrainer) cas(objectID bson.ObjectId, from, to pb.TrainingStatus) error {
	newRow := &pb.TrainingRow{}
	changeInfo, err := m.Collection.Find(bson.M{
		"_id":            objectID,
		"trainingStatus": from.Enum(),
	}).Apply(mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"trainingStatus": to.Enum(),
			},
		},
		ReturnNew: true,
	}, newRow)

	if err != nil {
		return err
	}

	if changeInfo.Updated != 1 {
		return fmt.Errorf("failed CAS'ing task %v from state %v to state %v", objectID, from, to)
	}
	glog.Infof("Updated objectId %v from state %v to state %v", objectID, from, to)
	return nil
}

func (m *MongoTrainer) finalizeTask(task *trainingTask) error {
	err := m.cas(task.objectID, pb.TrainingStatus_PROCESSING, pb.TrainingStatus_FINISHED)
	if err != nil {
		return err
	}
	return m.Collection.UpdateId(task.objectID, bson.M{
		"$set": bson.M{
			"forest":          task.row.Forest,
			"trainingResults": task.row.TrainingResults,
		},
	})
}
<|endoftext|>


<commit_before>package client

import (
	"log"
	"strconv"
	"testing"
)

// NOTE, before running the test execute the server.sh script in the run folder.
// The scripts runs two cache servers on ports 8080 and 8081
// The 8081 server is run with authentication psw=123

const (
	connectionString string = "http://localhost:8080"
	connectionStringAuth string = "http://localhost:8081"
	psw string = "123"
)

func TestClient_SetGetDel(t *testing.T) {

	const key = "key"
	const value = "value"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)
	test_SetGetDel(client, key, value, t)
}

func TestClient_SetGetDel_WithAuth(t *testing.T) {

	const key = "key"
	const value = "value"

	conns := Connections{
		{connectionStringAuth, psw},
	}

	client := NewClient(conns)
	test_SetGetDel(client, key, value, t)
}

func TestClient_SetGetDel_Sharded(t *testing.T) {
	conns := Connections{
		{connectionString, ""},
		{connectionStringAuth, psw},
	}

	client := NewClient(conns)

	for i := 0; i < 10; i++ {
		test_SetGetDel(client, strconv.Itoa(i), "value", t)
	}
}

func test_SetGetDel(client *Client, key string, value string, t *testing.T) {

	err := client.Set(key, value, 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	returnedValue, err := client.Get(key)

	if err != nil {
		t.Error("Failed to get the key", err)
	}

	if returnedValue != value {
		t.Error("Set value", value, "is not equal to returned value", returnedValue)
	}

	err = client.Del(key)

	if err != nil {
		t.Error("Failed to delete the key", err)
	}
}

func TestClient_Keys(t *testing.T) {

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	keys, err := client.Keys()

	const key1 = "key1"
	const key2 = "key2"

	if err != nil {
		t.Error("Failed to get keys", err)
	}

	if len(keys) != 0 {
		t.Errorf("There should be no keys, but there are %d keys", len(keys))
	}

	// Insert key1
	err = client.Set(key1, "value", 5)

	if err != nil {
		t.Error("Failed to set the key1", err)
	}

	// Insert key2
	err = client.Set(key2, "value", 5)

	if err != nil {
		t.Error("Failed to set the key2", err)
	}

	keys, err = client.Keys()

	if err != nil {
		t.Error("Failed to get keys", err)
	}

	if len(keys) != 2 {
		log.Println("Keys", keys)
		t.Errorf("There should be only one keys, but there are %d keys", len(keys))
	}

	if !contains(keys, key1) || !contains(keys, key2) {
		t.Errorf("Keys contains unexpected key. Keys=%s", keys)
	}

	// Tear down
	client.Del(key1)
	client.Del(key2)

}

func TestClient_Update(t *testing.T) {

	const key = "key"
	const updatedValue = "updated"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.Update(key, "value")

	if err == nil {
		t.Error("Expected: key not found")
	}

	// Insert key
	err = client.Set(key, "value", 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	err = client.Update(key, updatedValue)

	if (err != nil) {
		t.Error("Failed to update. Err = ", err)
	}

	value, err := client.Get(key)

	if value != updatedValue {
		t.Errorf("Update value '%s' does not equal to returned value '%s'", updatedValue, value)
	}

	// Tear down
	client.Del(key)

}

func TestClient_UpdateWithTtl(t *testing.T) {

	const key = "key"
	const updatedValue = "updated"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.UpdateWithTtl(key, "value", 5)

	if err == nil {
		t.Error("Expected: key not found")
	}

	// Insert key
	err = client.Set(key, "value", 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	// Update
	err = client.UpdateWithTtl(key, updatedValue, 25)

	// Assertions
	if (err != nil) {
		t.Error("Failed to update. Err = ", err)
	}

	value, err := client.Get(key)

	if value != updatedValue {
		t.Errorf("Update value '%s' does not equal to returned value '%s'", updatedValue, value)
	}

	// Tear down
	client.Del(key)

}

func TestClient_HSet_HGET(t *testing.T) {

	const key = "key"
	const hashKey = "hashKey"
	const value = "value"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.HSet(key, hashKey, value)

	if err != nil {
		t.Errorf("Failed to hset '%s' with hash key '%s' and value '%s'. Err = %s", key, hashKey, value, err)
	}

	returnedValue, err := client.HGet(key, hashKey)

	if err != nil {
		t.Error("Failed to get the key", err)
	}

	if returnedValue != value {
		t.Error("Value", value, "is not equal to returned value", returnedValue)
	}

	// Tear down
	client.Del(key)
}

func TestClient_LRange_LPUSH_LPOP(t *testing.T) {
	const listKey = "rangelistKey"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	// LPUSH 10 items
	for i := 0; i < 10; i++ {
		err := client.LPush(listKey, strconv.Itoa(i))
		if err != nil {
			t.Errorf("Failed to lpush. ListKey = '%s'. Error = %s", listKey, err)
		}
	}

	values, err := client.LRange(listKey, 2, 4)

	if err != nil {
		t.Fatalf("LRange failed. Err = %s", err)
	}

	log.Println("LRANGE", values)

	for _, value := range values {
		if value != "2" && value != "3" && value != "4" {
			t.Errorf("Values contain unexpected value '%s", value)
		}
	}

	// Tear down
	for i := 0; i < 10; i++ {
		_, err := client.LPop(listKey)
		if err != nil {
			t.Errorf("Failed to lpop. ListKey = '%s'. Error = %s", listKey, err)
		}
	}

	// Tear down
	client.Del(listKey)
}

func TestClient_RPush_RPop(t *testing.T) {
	//TODO://
}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
<commit_msg>TestClient_Keys_Sharded<commit_after>package client

import (
	"log"
	"strconv"
	"testing"
)

// NOTE, before running the test execute the server.sh script in the run folder.
// The scripts runs two cache servers on ports 8080 and 8081
// The 8081 server is run with authentication psw=123

const (
	connectionString string = "http://localhost:8080"
	connectionStringAuth string = "http://localhost:8081"
	psw string = "123"
)

func TestClient_SetGetDel(t *testing.T) {

	const key = "key"
	const value = "value"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)
	test_SetGetDel(client, key, value, t)
}

func TestClient_SetGetDel_WithAuth(t *testing.T) {

	const key = "key"
	const value = "value"

	conns := Connections{
		{connectionStringAuth, psw},
	}

	client := NewClient(conns)
	test_SetGetDel(client, key, value, t)
}

func TestClient_Keys_Sharded(t *testing.T) {

	const n  = 10

	// Two shards
	conns := Connections{
		{connectionString, ""},
		{connectionStringAuth, psw},
	}

	client := NewClient(conns)

	for i := 0; i < n; i++ {
		err := client.Set(strconv.Itoa(i), "value", 5)

		if err != nil {
			t.Error("Failed to set the key", err)
		}
	}

	keys, err := client.Keys()

	if err != nil {
		t.Error("Failed to get keys", err)
	}

	if len(keys) != n {
		log.Println("Keys", keys)
		t.Errorf("There should be %d keys, but there are %d keys", n, len(keys))
	}

	//Tear down
	for i := 0; i < n; i++ {
		client.Del(strconv.Itoa(i))
	}

}

func test_SetGetDel(client *Client, key string, value string, t *testing.T) {

	err := client.Set(key, value, 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	returnedValue, err := client.Get(key)

	if err != nil {
		t.Error("Failed to get the key", err)
	}

	if returnedValue != value {
		t.Error("Set value", value, "is not equal to returned value", returnedValue)
	}

	err = client.Del(key)

	if err != nil {
		t.Error("Failed to delete the key", err)
	}
}

func TestClient_Keys(t *testing.T) {

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	keys, err := client.Keys()

	const key1 = "key1"
	const key2 = "key2"

	if err != nil {
		t.Error("Failed to get keys", err)
	}

	if len(keys) != 0 {
		t.Errorf("There should be no keys, but there are %d keys", len(keys))
	}

	// Insert key1
	err = client.Set(key1, "value", 5)

	if err != nil {
		t.Error("Failed to set the key1", err)
	}

	// Insert key2
	err = client.Set(key2, "value", 5)

	if err != nil {
		t.Error("Failed to set the key2", err)
	}

	keys, err = client.Keys()

	if err != nil {
		t.Error("Failed to get keys", err)
	}

	if len(keys) != 2 {
		log.Println("Keys", keys)
		t.Errorf("There should be '%d' keys, but there are %d keys", 2, len(keys))
	}

	if !contains(keys, key1) || !contains(keys, key2) {
		t.Errorf("Keys contains unexpected key. Keys=%s", keys)
	}

	// Tear down
	client.Del(key1)
	client.Del(key2)

}

func TestClient_Update(t *testing.T) {

	const key = "key"
	const updatedValue = "updated"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.Update(key, "value")

	if err == nil {
		t.Error("Expected: key not found")
	}

	// Insert key
	err = client.Set(key, "value", 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	err = client.Update(key, updatedValue)

	if (err != nil) {
		t.Error("Failed to update. Err = ", err)
	}

	value, err := client.Get(key)

	if value != updatedValue {
		t.Errorf("Update value '%s' does not equal to returned value '%s'", updatedValue, value)
	}

	// Tear down
	client.Del(key)

}

func TestClient_UpdateWithTtl(t *testing.T) {

	const key = "key"
	const updatedValue = "updated"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.UpdateWithTtl(key, "value", 5)

	if err == nil {
		t.Error("Expected: key not found")
	}

	// Insert key
	err = client.Set(key, "value", 5)

	if err != nil {
		t.Error("Failed to set the key", err)
	}

	// Update
	err = client.UpdateWithTtl(key, updatedValue, 25)

	// Assertions
	if (err != nil) {
		t.Error("Failed to update. Err = ", err)
	}

	value, err := client.Get(key)

	if value != updatedValue {
		t.Errorf("Update value '%s' does not equal to returned value '%s'", updatedValue, value)
	}

	// Tear down
	client.Del(key)

}

func TestClient_HSet_HGET(t *testing.T) {

	const key = "key"
	const hashKey = "hashKey"
	const value = "value"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	err := client.HSet(key, hashKey, value)

	if err != nil {
		t.Errorf("Failed to hset '%s' with hash key '%s' and value '%s'. Err = %s", key, hashKey, value, err)
	}

	returnedValue, err := client.HGet(key, hashKey)

	if err != nil {
		t.Error("Failed to get the key", err)
	}

	if returnedValue != value {
		t.Error("Value", value, "is not equal to returned value", returnedValue)
	}

	// Tear down
	client.Del(key)
}

func TestClient_LRange_LPUSH_LPOP(t *testing.T) {
	const listKey = "rangelistKey"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	// LPUSH 10 items
	for i := 0; i < 10; i++ {
		err := client.LPush(listKey, strconv.Itoa(i))
		if err != nil {
			t.Errorf("Failed to lpush. ListKey = '%s'. Error = %s", listKey, err)
		}
	}

	values, err := client.LRange(listKey, 2, 4)

	if err != nil {
		t.Fatalf("LRange failed. Err = %s", err)
	}

	log.Println("LRANGE", values)

	for _, value := range values {
		if value != "2" && value != "3" && value != "4" {
			t.Errorf("Values contain unexpected value '%s", value)
		}
	}

	// Tear down
	for i := 0; i < 10; i++ {
		_, err := client.LPop(listKey)
		if err != nil {
			t.Errorf("Failed to lpop. ListKey = '%s'. Error = %s", listKey, err)
		}
	}

	// Tear down
	client.Del(listKey)
}

func TestClient_RPush_RPop(t *testing.T) {

	const key = "rpushpopkey"
	const value = "value"

	conns := Connections{
		{connectionString, ""},
	}

	client := NewClient(conns)

	// RPush
	err := client.RPush(key, value)
	if err != nil {
		t.Errorf("Failed to RPush. ListKey = '%s'. Error = %s", key, err)
	}

	// RPop
	returnedValue, err := client.RPop(key)
	if err != nil {
		t.Errorf("Failed to RPop. ListKey = '%s'. Error = %s", key, err)
	}

	if returnedValue != value {
		t.Error("Value", value, "is not equal to returned value", returnedValue)
	}

	// Tear down
	client.Del(key)

}

func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// +build !linux nonsystemd

package journald

import (
	"context"

	"github.com/inconshreveable/log15"
)

var Supported bool = false

type JournaldReader interface {
	Start()
	Stop()
	Entries() chan map[string]string
}

type reader struct {
	entries chan map[string]string
}

func NewReader(ctx context.Context, logger log15.Logger) (JournaldReader, error) {
	r := &reader{}
	r.entries = make(chan map[string]string)
	return r, nil
}

func (r *reader) Start() {}
func (r *reader) Stop()  {}

func (r *reader) Entries() chan map[string]string {
	return r.entries
}
<commit_msg>fix compile on macosx<commit_after>// +build !linux nonsystemd

package journald

import (
	"github.com/inconshreveable/log15"
)

var Supported bool = false

type JournaldReader interface {
	Start(coding string)
	Stop()
	Shutdown()
	Entries() chan map[string]string
}

type reader struct {
	entries chan map[string]string
}

func NewReader(logger log15.Logger) (JournaldReader, error) {
	r := &reader{}
	r.entries = make(chan map[string]string)
	return r, nil
}

func (r *reader) Start(coding string) {}
func (r *reader) Stop()               {}
func (r *reader) Shutdown()           {}

func (r *reader) Entries() chan map[string]string {
	return r.entries
}
<|endoftext|>


<commit_before>package lexer

// LexemeStream provides a constant stream of lexemes, allowing
// for an n-Lookahead
type LexemeStream interface {
	EOF() bool

	Next() (Lexeme, error)
	Get(n int) ([]Lexeme, []error)

	Peek() (Lexeme, error)
	Lookahead(n int) ([]Lexeme, error)
}

// RuneStreamPosition represents some column and row in the
// rune stream
type RuneStreamPosition struct {
	Row    int
	Column int
}

// RuneStream provides a constant stream of bytes, potentially from a file
type RuneStream interface {
	EOF() bool

	Next() rune
	Get(n int) []rune

	Peek() rune
	Lookahead(n int) []rune

	Position() RuneStreamPosition
}

//
// ByteStream for string
//

type StringRuneStream struct {
	Source string

	at        int
	lookahead []rune
	x         int
	y         int
}

// Position returns the current column the row is pointed at; e.g.,
// the rune returned at Peek() would be at this location.
func (s *StringRuneStream) Position() RuneStreamPosition {
	return RuneStreamPosition{Row: s.y + 1, Column: s.x + 1}
}

func (s *StringRuneStream) updatePosition(r rune) {
	if r == '\n' {
		s.y++
		s.x = 0
	} else {
		s.x++
	}
}

func (s *StringRuneStream) EOF() bool {
	return len(s.lookahead) == 0 && s.at >= len(s.Source)
}

func (s *StringRuneStream) getOne() rune {
	if s.at >= len(s.Source) {
		return 0
	}

	next := []rune(s.Source)[s.at]

	s.at++
	return next
}

func (s *StringRuneStream) getN(n int) []rune {
	bytes := make([]rune, n)

	for i := 0; i < n; i++ {
		bytes[i] = s.getOne()
	}

	return bytes
}

func (s *StringRuneStream) Next() rune {
	if len(s.lookahead) > 0 {
		next := s.lookahead[0]
		s.lookahead = s.lookahead[1:]
		s.updatePosition(next)
		return next
	}

	next := s.getOne()
	s.updatePosition(next)
	return next
}

func (s *StringRuneStream) Get(n int) []rune {
	bytes := make([]rune, n)

	for i := 0; i < n; i++ {
		bytes[i] = s.Next()
	}

	return bytes
}

func (s *StringRuneStream) Peek() rune {
	return s.Lookahead(1)[0]
}

func (s *StringRuneStream) Lookahead(n int) []rune {
	rem := n - len(s.lookahead)

	if rem > 0 {
		s.lookahead = append(s.lookahead, s.getN(rem)...)
	}

	return s.lookahead[0:n]
}
<commit_msg>Fix LexemeStream interface<commit_after>package lexer

// LexemeStream provides a constant stream of lexemes, allowing
// for an n-Lookahead
type LexemeStream interface {
	EOF() bool

	Next() (Lexeme, error)
	Get(n int) ([]Lexeme, []error)

	Peek() (Lexeme, error)
	Lookahead(n int) ([]Lexeme, []error)
}

// RuneStreamPosition represents some column and row in the
// rune stream
type RuneStreamPosition struct {
	Row    int
	Column int
}

// RuneStream provides a constant stream of bytes, potentially from a file
type RuneStream interface {
	EOF() bool

	Next() rune
	Get(n int) []rune

	Peek() rune
	Lookahead(n int) []rune

	Position() RuneStreamPosition
}

//
// ByteStream for string
//

type StringRuneStream struct {
	Source string

	at        int
	lookahead []rune
	x         int
	y         int
}

// Position returns the current column the row is pointed at; e.g.,
// the rune returned at Peek() would be at this location.
func (s *StringRuneStream) Position() RuneStreamPosition {
	return RuneStreamPosition{Row: s.y + 1, Column: s.x + 1}
}

func (s *StringRuneStream) updatePosition(r rune) {
	if r == '\n' {
		s.y++
		s.x = 0
	} else {
		s.x++
	}
}

func (s *StringRuneStream) EOF() bool {
	return len(s.lookahead) == 0 && s.at >= len(s.Source)
}

func (s *StringRuneStream) getOne() rune {
	if s.at >= len(s.Source) {
		return 0
	}

	next := []rune(s.Source)[s.at]

	s.at++
	return next
}

func (s *StringRuneStream) getN(n int) []rune {
	bytes := make([]rune, n)

	for i := 0; i < n; i++ {
		bytes[i] = s.getOne()
	}

	return bytes
}

func (s *StringRuneStream) Next() rune {
	if len(s.lookahead) > 0 {
		next := s.lookahead[0]
		s.lookahead = s.lookahead[1:]
		s.updatePosition(next)
		return next
	}

	next := s.getOne()
	s.updatePosition(next)
	return next
}

func (s *StringRuneStream) Get(n int) []rune {
	bytes := make([]rune, n)

	for i := 0; i < n; i++ {
		bytes[i] = s.Next()
	}

	return bytes
}

func (s *StringRuneStream) Peek() rune {
	return s.Lookahead(1)[0]
}

func (s *StringRuneStream) Lookahead(n int) []rune {
	rem := n - len(s.lookahead)

	if rem > 0 {
		s.lookahead = append(s.lookahead, s.getN(rem)...)
	}

	return s.lookahead[0:n]
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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 main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	// Injection related imports.
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/injection"
	revisioninformer "knative.dev/serving/pkg/client/injection/informers/serving/v1alpha1/revision"

	"github.com/kelseyhightower/envconfig"
	perrors "github.com/pkg/errors"
	"go.opencensus.io/stats/view"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/util/wait"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection/sharedmain"
	pkglogging "knative.dev/pkg/logging"
	"knative.dev/pkg/logging/logkey"
	"knative.dev/pkg/metrics"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/profiling"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/system"
	"knative.dev/pkg/tracing"
	tracingconfig "knative.dev/pkg/tracing/config"
	"knative.dev/pkg/version"
	"knative.dev/pkg/websocket"
	"knative.dev/serving/pkg/activator"
	activatorconfig "knative.dev/serving/pkg/activator/config"
	activatorhandler "knative.dev/serving/pkg/activator/handler"
	activatornet "knative.dev/serving/pkg/activator/net"
	"knative.dev/serving/pkg/apis/networking"
	"knative.dev/serving/pkg/autoscaler"
	"knative.dev/serving/pkg/goversion"
	pkghttp "knative.dev/serving/pkg/http"
	"knative.dev/serving/pkg/logging"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/queue"
)

// Fail if using unsupported go version.
var _ = goversion.IsSupported()

const (
	component = "activator"

	// Add a little buffer space between request handling and stat
	// reporting so that latency in the stat pipeline doesn't
	// interfere with request handling.
	statReportingQueueLength = 10

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100

	// The number of requests that are queued on the breaker before the 503s are sent.
	// The value must be adjusted depending on the actual production requirements.
	breakerQueueDepth = 10000

	// The upper bound for concurrent requests sent to the revision.
	// As new endpoints show up, the Breakers concurrency increases up to this value.
	breakerMaxConcurrency = 1000

	// The port on which autoscaler WebSocket server listens.
	autoscalerPort = ":8080"
)

var (
	masterURL = flag.String("master", "", "The address of the Kubernetes API server. "+
		"Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func statReporter(statSink *websocket.ManagedConnection, stopCh <-chan struct{},
	statChan <-chan *autoscaler.StatMessage, logger *zap.SugaredLogger) {
	for {
		select {
		case sm := <-statChan:
			if err := statSink.Send(sm); err != nil {
				logger.Errorw("Error while sending stat", zap.Error(err))
			}
		case <-stopCh:
			// It's a sending connection, so no drainage required.
			statSink.Shutdown()
			return
		}
	}
}

type config struct {
	PodName string `split_words:"true" required:"true"`
}

func main() {
	flag.Parse()

	// Set up a context that we can cancel to tell informers and other subprocesses to stop.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Report stats on Go memory usage every 30 seconds.
	msp := metrics.NewMemStatsAll()
	msp.Start(ctx, 30*time.Second)
	if err := view.Register(msp.DefaultViews()...); err != nil {
		log.Fatalf("Error exporting go memstats view: %v", err)
	}

	cfg, err := sharedmain.GetConfig(*masterURL, *kubeconfig)
	if err != nil {
		log.Fatal("Error building kubeconfig:", err)
	}

	log.Printf("Registering %d clients", len(injection.Default.GetClients()))
	log.Printf("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	log.Printf("Registering %d informers", len(injection.Default.GetInformers()))

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// Set up our logger.
	loggingConfig, err := sharedmain.GetLoggingConfig(ctx)
	if err != nil {
		log.Fatal("Error loading/parsing logging configuration:", err)
	}
	logger, atomicLevel := pkglogging.NewLoggerFromConfig(loggingConfig, component)
	logger = logger.With(zap.String(logkey.ControllerType, component))
	ctx = pkglogging.WithLogger(ctx, logger)
	defer flush(logger)

	kubeClient := kubeclient.Get(ctx)

	// Run informers instead of starting them from the factory to prevent the sync hanging because of empty handler.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", zap.Error(err))
	}

	logger.Info("Starting the knative activator")

	// We sometimes startup faster than we can reach kube-api. Poll on failure to prevent us terminating
	if perr := wait.PollImmediate(time.Second, 60*time.Second, func() (bool, error) {
		if err = version.CheckMinimumVersion(kubeClient.Discovery()); err != nil {
			logger.Errorw("Failed to get k8s version", zap.Error(err))
		}
		return err == nil, nil
	}); perr != nil {
		logger.Fatalw("Timed out attempting to get k8s version", zap.Error(err))
	}

	reporter, err := activator.NewStatsReporter()
	if err != nil {
		logger.Fatalw("Failed to create stats reporter", zap.Error(err))
	}

	statCh := make(chan *autoscaler.StatMessage, statReportingQueueLength)
	defer close(statCh)

	reqCh := make(chan activatorhandler.ReqEvent, requestCountingQueueLength)
	defer close(reqCh)

	params := queue.BreakerParams{QueueDepth: breakerQueueDepth, MaxConcurrency: breakerMaxConcurrency, InitialCapacity: 0}

	// Start throttler.
	throttler := activatornet.NewThrottler(ctx, params)
	go throttler.Run(ctx)

	oct := tracing.NewOpenCensusTracer(tracing.WithExporter(networking.ActivatorServiceName, logger))

	tracerUpdater := configmap.TypeFilter(&tracingconfig.Config{})(func(name string, value interface{}) {
		cfg := value.(*tracingconfig.Config)
		if err := oct.ApplyConfig(cfg); err != nil {
			logger.Errorw("Unable to apply open census tracer config", zap.Error(err))
			return
		}
	})

	// Set up our config store
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace())
	configStore := activatorconfig.NewStore(logger, tracerUpdater)
	configStore.WatchConfigs(configMapWatcher)

	// Open a WebSocket connection to the autoscaler.
	autoscalerEndpoint := fmt.Sprintf("ws://%s.%s.svc.%s%s", "autoscaler", system.Namespace(), pkgnet.GetClusterDomainName(), autoscalerPort)
	logger.Info("Connecting to autoscaler at", autoscalerEndpoint)
	statSink := websocket.NewDurableSendingConnection(autoscalerEndpoint, logger)
	go statReporter(statSink, ctx.Done(), statCh, logger)

	var env config
	if err := envconfig.Process("", &env); err != nil {
		logger.Fatalw("Failed to process env", zap.Error(err))
	}
	podName := env.PodName

	// Create and run our concurrency reporter
	reportTicker := time.NewTicker(time.Second)
	defer reportTicker.Stop()
	cr := activatorhandler.NewConcurrencyReporter(ctx, podName, reqCh,
		reportTicker.C, statCh, reporter)
	go cr.Run(ctx.Done())

	// Create activation handler chain
	// Note: innermost handlers are specified first, ie. the last handler in the chain will be executed first
	var ah http.Handler = activatorhandler.New(
		ctx,
		throttler,
		reporter)
	ah = activatorhandler.NewRequestEventHandler(reqCh, ah)
	ah = tracing.HTTPSpanMiddleware(ah)
	ah = configStore.HTTPMiddleware(ah)
	reqLogHandler, err := pkghttp.NewRequestLogHandler(ah, logging.NewSyncFileWriter(os.Stdout), "",
		requestLogTemplateInputGetter(revisioninformer.Get(ctx).Lister()))
	if err != nil {
		logger.Fatalw("Unable to create request log handler", zap.Error(err))
	}
	ah = reqLogHandler
	ah = &activatorhandler.ProbeHandler{NextHandler: ah}

	// Set up our health check based on the health of stat sink and environmental factors.
	// When drainCh is closed, we should start to drain connections.
	hc, drainCh := newHealthCheck(logger, statSink)
	ah = &activatorhandler.HealthHandler{HealthCheck: hc, NextHandler: ah}

	// NOTE: MetricHandler is being used as the outermost handler for the purpose of measuring the request latency.
	ah = activatorhandler.NewMetricHandler(ctx, reporter, ah)
	ah = network.NewProbeHandler(ah)

	profilingHandler := profiling.NewHandler(logger, false)
	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher.Watch(pkglogging.ConfigMapName(), pkglogging.UpdateLevelFromConfigMap(logger, atomicLevel, component))

	// Watch the observability config map
	configMapWatcher.Watch(metrics.ConfigMapName(),
		metrics.UpdateExporterFromConfigMap(component, logger),
		updateRequestLogFromConfigMap(logger, reqLogHandler),
		profilingHandler.UpdateFromConfigMap)

	if err = configMapWatcher.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start configuration manager", zap.Error(err))
	}

	servers := map[string]*http.Server{
		"http1":   network.NewServer(":"+strconv.Itoa(networking.BackendHTTPPort), ah),
		"h2c":     network.NewServer(":"+strconv.Itoa(networking.BackendHTTP2Port), ah),
		"profile": profiling.NewServer(profilingHandler),
	}

	errCh := make(chan error, len(servers))
	for name, server := range servers {
		go func(name string, s *http.Server) {
			// Don't forward ErrServerClosed as that indicates we're already shutting down.
			if err := s.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				errCh <- perrors.Wrapf(err, "%s server failed", name)
			}
		}(name, server)
	}

	// Wait for the signal to drain.
	select {
	case <-drainCh:
		logger.Info("Received the drain signal.")
	case err := <-errCh:
		logger.Errorw("Failed to run HTTP server", zap.Error(err))
	}

	// The drain has started (we are now failing readiness probes).  Let the effects of this
	// propagate so that new requests are no longer routed our way.
	time.Sleep(30 * time.Second)
	logger.Info("Done waiting, shutting down servers.")

	// Drain outstanding requests, and stop accepting new ones.
	for _, server := range servers {
		server.Shutdown(context.Background())
	}
	logger.Info("Servers shutdown.")
}

func newHealthCheck(logger *zap.SugaredLogger, statSink *websocket.ManagedConnection) (func() error, <-chan struct{}) {
	// When we get SIGTERM (sigCh closes), start failing readiness probes.
	sigCh := signals.SetupSignalHandler()

	// Some duration after our first readiness probe failure (to allow time
	// for the network to reprogram) send the signal to drain connections.
	drainCh := make(chan struct{})
	once := sync.Once{}

	return func() error {
		select {
		case <-sigCh:
			// Signal to start the process of draining.
			once.Do(func() {
				logger.Info("Received SIGTERM")
				close(drainCh)
			})
			return errors.New("received SIGTERM from kubelet")
		default:
			logger.Debug("No signal yet.")
			return statSink.Status()
		}
	}, drainCh
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	os.Stdout.Sync()
	os.Stderr.Sync()
	metrics.FlushExporter()
}
<commit_msg>Unbuffer the channel and send the stat in the go-routine (#5705)<commit_after>/*
Copyright 2018 The Knative 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 main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	// Injection related imports.
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/injection"
	revisioninformer "knative.dev/serving/pkg/client/injection/informers/serving/v1alpha1/revision"

	"github.com/kelseyhightower/envconfig"
	perrors "github.com/pkg/errors"
	"go.opencensus.io/stats/view"
	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/util/wait"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection/sharedmain"
	pkglogging "knative.dev/pkg/logging"
	"knative.dev/pkg/logging/logkey"
	"knative.dev/pkg/metrics"
	pkgnet "knative.dev/pkg/network"
	"knative.dev/pkg/profiling"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/system"
	"knative.dev/pkg/tracing"
	tracingconfig "knative.dev/pkg/tracing/config"
	"knative.dev/pkg/version"
	"knative.dev/pkg/websocket"
	"knative.dev/serving/pkg/activator"
	activatorconfig "knative.dev/serving/pkg/activator/config"
	activatorhandler "knative.dev/serving/pkg/activator/handler"
	activatornet "knative.dev/serving/pkg/activator/net"
	"knative.dev/serving/pkg/apis/networking"
	"knative.dev/serving/pkg/autoscaler"
	"knative.dev/serving/pkg/goversion"
	pkghttp "knative.dev/serving/pkg/http"
	"knative.dev/serving/pkg/logging"
	"knative.dev/serving/pkg/network"
	"knative.dev/serving/pkg/queue"
)

// Fail if using unsupported go version.
var _ = goversion.IsSupported()

const (
	component = "activator"

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100

	// The number of requests that are queued on the breaker before the 503s are sent.
	// The value must be adjusted depending on the actual production requirements.
	breakerQueueDepth = 10000

	// The upper bound for concurrent requests sent to the revision.
	// As new endpoints show up, the Breakers concurrency increases up to this value.
	breakerMaxConcurrency = 1000

	// The port on which autoscaler WebSocket server listens.
	autoscalerPort = ":8080"
)

var (
	masterURL = flag.String("master", "", "The address of the Kubernetes API server. "+
		"Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func statReporter(statSink *websocket.ManagedConnection, stopCh <-chan struct{},
	statChan <-chan *autoscaler.StatMessage, logger *zap.SugaredLogger) {
	for {
		select {
		case sm := <-statChan:
			go func() {
				if err := statSink.Send(sm); err != nil {
					logger.Errorw("Error while sending stat", zap.Error(err))
				}
			}()
		case <-stopCh:
			// It's a sending connection, so no drainage required.
			statSink.Shutdown()
			return
		}
	}
}

type config struct {
	PodName string `split_words:"true" required:"true"`
}

func main() {
	flag.Parse()

	// Set up a context that we can cancel to tell informers and other subprocesses to stop.
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Report stats on Go memory usage every 30 seconds.
	msp := metrics.NewMemStatsAll()
	msp.Start(ctx, 30*time.Second)
	if err := view.Register(msp.DefaultViews()...); err != nil {
		log.Fatalf("Error exporting go memstats view: %v", err)
	}

	cfg, err := sharedmain.GetConfig(*masterURL, *kubeconfig)
	if err != nil {
		log.Fatal("Error building kubeconfig:", err)
	}

	log.Printf("Registering %d clients", len(injection.Default.GetClients()))
	log.Printf("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
	log.Printf("Registering %d informers", len(injection.Default.GetInformers()))

	ctx, informers := injection.Default.SetupInformers(ctx, cfg)

	// Set up our logger.
	loggingConfig, err := sharedmain.GetLoggingConfig(ctx)
	if err != nil {
		log.Fatal("Error loading/parsing logging configuration:", err)
	}
	logger, atomicLevel := pkglogging.NewLoggerFromConfig(loggingConfig, component)
	logger = logger.With(zap.String(logkey.ControllerType, component))
	ctx = pkglogging.WithLogger(ctx, logger)
	defer flush(logger)

	kubeClient := kubeclient.Get(ctx)

	// Run informers instead of starting them from the factory to prevent the sync hanging because of empty handler.
	if err := controller.StartInformers(ctx.Done(), informers...); err != nil {
		logger.Fatalw("Failed to start informers", zap.Error(err))
	}

	logger.Info("Starting the knative activator")

	// We sometimes startup faster than we can reach kube-api. Poll on failure to prevent us terminating
	if perr := wait.PollImmediate(time.Second, 60*time.Second, func() (bool, error) {
		if err = version.CheckMinimumVersion(kubeClient.Discovery()); err != nil {
			logger.Errorw("Failed to get k8s version", zap.Error(err))
		}
		return err == nil, nil
	}); perr != nil {
		logger.Fatalw("Timed out attempting to get k8s version", zap.Error(err))
	}

	reporter, err := activator.NewStatsReporter()
	if err != nil {
		logger.Fatalw("Failed to create stats reporter", zap.Error(err))
	}

	statCh := make(chan *autoscaler.StatMessage)
	defer close(statCh)

	reqCh := make(chan activatorhandler.ReqEvent, requestCountingQueueLength)
	defer close(reqCh)

	params := queue.BreakerParams{QueueDepth: breakerQueueDepth, MaxConcurrency: breakerMaxConcurrency, InitialCapacity: 0}

	// Start throttler.
	throttler := activatornet.NewThrottler(ctx, params)
	go throttler.Run(ctx)

	oct := tracing.NewOpenCensusTracer(tracing.WithExporter(networking.ActivatorServiceName, logger))

	tracerUpdater := configmap.TypeFilter(&tracingconfig.Config{})(func(name string, value interface{}) {
		cfg := value.(*tracingconfig.Config)
		if err := oct.ApplyConfig(cfg); err != nil {
			logger.Errorw("Unable to apply open census tracer config", zap.Error(err))
			return
		}
	})

	// Set up our config store
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace())
	configStore := activatorconfig.NewStore(logger, tracerUpdater)
	configStore.WatchConfigs(configMapWatcher)

	// Open a WebSocket connection to the autoscaler.
	autoscalerEndpoint := fmt.Sprintf("ws://%s.%s.svc.%s%s", "autoscaler", system.Namespace(), pkgnet.GetClusterDomainName(), autoscalerPort)
	logger.Info("Connecting to autoscaler at", autoscalerEndpoint)
	statSink := websocket.NewDurableSendingConnection(autoscalerEndpoint, logger)
	go statReporter(statSink, ctx.Done(), statCh, logger)

	var env config
	if err := envconfig.Process("", &env); err != nil {
		logger.Fatalw("Failed to process env", zap.Error(err))
	}
	podName := env.PodName

	// Create and run our concurrency reporter
	reportTicker := time.NewTicker(time.Second)
	defer reportTicker.Stop()
	cr := activatorhandler.NewConcurrencyReporter(ctx, podName, reqCh,
		reportTicker.C, statCh, reporter)
	go cr.Run(ctx.Done())

	// Create activation handler chain
	// Note: innermost handlers are specified first, ie. the last handler in the chain will be executed first
	var ah http.Handler = activatorhandler.New(
		ctx,
		throttler,
		reporter)
	ah = activatorhandler.NewRequestEventHandler(reqCh, ah)
	ah = tracing.HTTPSpanMiddleware(ah)
	ah = configStore.HTTPMiddleware(ah)
	reqLogHandler, err := pkghttp.NewRequestLogHandler(ah, logging.NewSyncFileWriter(os.Stdout), "",
		requestLogTemplateInputGetter(revisioninformer.Get(ctx).Lister()))
	if err != nil {
		logger.Fatalw("Unable to create request log handler", zap.Error(err))
	}
	ah = reqLogHandler
	ah = &activatorhandler.ProbeHandler{NextHandler: ah}

	// Set up our health check based on the health of stat sink and environmental factors.
	// When drainCh is closed, we should start to drain connections.
	hc, drainCh := newHealthCheck(logger, statSink)
	ah = &activatorhandler.HealthHandler{HealthCheck: hc, NextHandler: ah}

	// NOTE: MetricHandler is being used as the outermost handler for the purpose of measuring the request latency.
	ah = activatorhandler.NewMetricHandler(ctx, reporter, ah)
	ah = network.NewProbeHandler(ah)

	profilingHandler := profiling.NewHandler(logger, false)
	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher.Watch(pkglogging.ConfigMapName(), pkglogging.UpdateLevelFromConfigMap(logger, atomicLevel, component))

	// Watch the observability config map
	configMapWatcher.Watch(metrics.ConfigMapName(),
		metrics.UpdateExporterFromConfigMap(component, logger),
		updateRequestLogFromConfigMap(logger, reqLogHandler),
		profilingHandler.UpdateFromConfigMap)

	if err = configMapWatcher.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start configuration manager", zap.Error(err))
	}

	servers := map[string]*http.Server{
		"http1":   network.NewServer(":"+strconv.Itoa(networking.BackendHTTPPort), ah),
		"h2c":     network.NewServer(":"+strconv.Itoa(networking.BackendHTTP2Port), ah),
		"profile": profiling.NewServer(profilingHandler),
	}

	errCh := make(chan error, len(servers))
	for name, server := range servers {
		go func(name string, s *http.Server) {
			// Don't forward ErrServerClosed as that indicates we're already shutting down.
			if err := s.ListenAndServe(); err != nil && err != http.ErrServerClosed {
				errCh <- perrors.Wrapf(err, "%s server failed", name)
			}
		}(name, server)
	}

	// Wait for the signal to drain.
	select {
	case <-drainCh:
		logger.Info("Received the drain signal.")
	case err := <-errCh:
		logger.Errorw("Failed to run HTTP server", zap.Error(err))
	}

	// The drain has started (we are now failing readiness probes).  Let the effects of this
	// propagate so that new requests are no longer routed our way.
	time.Sleep(30 * time.Second)
	logger.Info("Done waiting, shutting down servers.")

	// Drain outstanding requests, and stop accepting new ones.
	for _, server := range servers {
		server.Shutdown(context.Background())
	}
	logger.Info("Servers shutdown.")
}

func newHealthCheck(logger *zap.SugaredLogger, statSink *websocket.ManagedConnection) (func() error, <-chan struct{}) {
	// When we get SIGTERM (sigCh closes), start failing readiness probes.
	sigCh := signals.SetupSignalHandler()

	// Some duration after our first readiness probe failure (to allow time
	// for the network to reprogram) send the signal to drain connections.
	drainCh := make(chan struct{})
	once := sync.Once{}

	return func() error {
		select {
		case <-sigCh:
			// Signal to start the process of draining.
			once.Do(func() {
				logger.Info("Received SIGTERM")
				close(drainCh)
			})
			return errors.New("received SIGTERM from kubelet")
		default:
			logger.Debug("No signal yet.")
			return statSink.Status()
		}
	}, drainCh
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	os.Stdout.Sync()
	os.Stderr.Sync()
	metrics.FlushExporter()
}
<|endoftext|>