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

package etcdmain

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	"go.etcd.io/etcd/client/pkg/v3/fileutil"
	"go.etcd.io/etcd/client/pkg/v3/logutil"
	"go.etcd.io/etcd/client/pkg/v3/types"
	"go.etcd.io/etcd/pkg/v3/osutil"
	"go.etcd.io/etcd/server/v3/embed"
	"go.etcd.io/etcd/server/v3/etcdserver/api/v2discovery"
	"go.etcd.io/etcd/server/v3/etcdserver/errors"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

type dirType string

var (
	dirMember = dirType("member")
	dirProxy  = dirType("proxy")
	dirEmpty  = dirType("empty")
)

func startEtcdOrProxyV2(args []string) {
	grpc.EnableTracing = false

	cfg := newConfig()
	defaultInitialCluster := cfg.ec.InitialCluster

	err := cfg.parse(args[1:])
	lg := cfg.ec.GetLogger()
	// If we failed to parse the whole configuration, print the error using
	// preferably the resolved logger from the config,
	// but if does not exists, create a new temporary logger.
	if lg == nil {
		var zapError error
		// use this logger
		lg, zapError = logutil.CreateDefaultZapLogger(zap.InfoLevel)
		if zapError != nil {
			fmt.Printf("error creating zap logger %v", zapError)
			os.Exit(1)
		}
	}
	lg.Info("Running: ", zap.Strings("args", args))
	if err != nil {
		lg.Warn("failed to verify flags", zap.Error(err))
		switch err {
		case embed.ErrUnsetAdvertiseClientURLsFlag:
			lg.Warn("advertise client URLs are not set", zap.Error(err))
		}
		os.Exit(1)
	}

	cfg.ec.SetupGlobalLoggers()

	defer func() {
		logger := cfg.ec.GetLogger()
		if logger != nil {
			logger.Sync()
		}
	}()

	defaultHost, dhErr := (&cfg.ec).UpdateDefaultClusterFromName(defaultInitialCluster)
	if defaultHost != "" {
		lg.Info(
			"detected default host for advertise",
			zap.String("host", defaultHost),
		)
	}
	if dhErr != nil {
		lg.Info("failed to detect default host", zap.Error(dhErr))
	}

	if cfg.ec.Dir == "" {
		cfg.ec.Dir = fmt.Sprintf("%v.etcd", cfg.ec.Name)
		lg.Warn(
			"'data-dir' was empty; using default",
			zap.String("data-dir", cfg.ec.Dir),
		)
	}

	var stopped <-chan struct{}
	var errc <-chan error

	which := identifyDataDirOrDie(cfg.ec.GetLogger(), cfg.ec.Dir)
	if which != dirEmpty {
		lg.Info(
			"server has already been initialized",
			zap.String("data-dir", cfg.ec.Dir),
			zap.String("dir-type", string(which)),
		)
		switch which {
		case dirMember:
			stopped, errc, err = startEtcd(&cfg.ec)
		case dirProxy:
			lg.Panic("v2 http proxy has already been deprecated in 3.6", zap.String("dir-type", string(which)))
		default:
			lg.Panic(
				"unknown directory type",
				zap.String("dir-type", string(which)),
			)
		}
	} else {
		lg.Info(
			"Initialize and start etcd server",
			zap.String("data-dir", cfg.ec.Dir),
			zap.String("dir-type", string(which)),
		)
		stopped, errc, err = startEtcd(&cfg.ec)
	}

	if err != nil {
		if derr, ok := err.(*errors.DiscoveryError); ok {
			switch derr.Err {
			case v2discovery.ErrDuplicateID:
				lg.Warn(
					"member has been registered with discovery service",
					zap.String("name", cfg.ec.Name),
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(derr.Err),
				)
				lg.Warn(
					"but could not find valid cluster configuration",
					zap.String("data-dir", cfg.ec.Dir),
				)
				lg.Warn("check data dir if previous bootstrap succeeded")
				lg.Warn("or use a new discovery token if previous bootstrap failed")

			case v2discovery.ErrDuplicateName:
				lg.Warn(
					"member with duplicated name has already been registered",
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(derr.Err),
				)
				lg.Warn("cURL the discovery token URL for details")
				lg.Warn("do not reuse discovery token; generate a new one to bootstrap a cluster")

			default:
				lg.Warn(
					"failed to bootstrap; discovery token was already used",
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(err),
				)
				lg.Warn("do not reuse discovery token; generate a new one to bootstrap a cluster")
			}
			os.Exit(1)
		}

		if strings.Contains(err.Error(), "include") && strings.Contains(err.Error(), "--initial-cluster") {
			lg.Warn("failed to start", zap.Error(err))
			if cfg.ec.InitialCluster == cfg.ec.InitialClusterFromName(cfg.ec.Name) {
				lg.Warn("forgot to set --initial-cluster?")
			}
			if types.URLs(cfg.ec.APUrls).String() == embed.DefaultInitialAdvertisePeerURLs {
				lg.Warn("forgot to set --initial-advertise-peer-urls?")
			}
			if cfg.ec.InitialCluster == cfg.ec.InitialClusterFromName(cfg.ec.Name) && len(cfg.ec.Durl) == 0 && len(cfg.ec.DiscoveryCfg.Endpoints) == 0 {
				lg.Warn("V2 discovery settings (i.e., --discovery) or v3 discovery settings (i.e., --discovery-token, --discovery-endpoints) are not set")
			}
			os.Exit(1)
		}
		lg.Fatal("discovery failed", zap.Error(err))
	}

	osutil.HandleInterrupts(lg)

	// At this point, the initialization of etcd is done.
	// The listeners are listening on the TCP ports and ready
	// for accepting connections. The etcd instance should be
	// joined with the cluster and ready to serve incoming
	// connections.
	notifySystemd(lg)

	select {
	case lerr := <-errc:
		// fatal out on listener errors
		lg.Fatal("listener failed", zap.Error(lerr))
	case <-stopped:
	}

	osutil.Exit(0)
}

// startEtcd runs StartEtcd in addition to hooks needed for standalone etcd.
func startEtcd(cfg *embed.Config) (<-chan struct{}, <-chan error, error) {
	e, err := embed.StartEtcd(cfg)
	if err != nil {
		return nil, nil, err
	}
	osutil.RegisterInterruptHandler(e.Close)
	select {
	case <-e.Server.ReadyNotify(): // wait for e.Server to join the cluster
	case <-e.Server.StopNotify(): // publish aborted from 'ErrStopped'
	case <-time.After(cfg.ExperimentalWaitClusterReadyTimeout):
		e.GetLogger().Warn("startEtcd: timed out waiting for the ready notification")
	}
	return e.Server.StopNotify(), e.Err(), nil
}

// identifyDataDirOrDie returns the type of the data dir.
// Dies if the datadir is invalid.
func identifyDataDirOrDie(lg *zap.Logger, dir string) dirType {
	names, err := fileutil.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return dirEmpty
		}
		lg.Fatal("failed to list data directory", zap.String("dir", dir), zap.Error(err))
	}

	var m, p bool
	for _, name := range names {
		switch dirType(name) {
		case dirMember:
			m = true
		case dirProxy:
			p = true
		default:
			lg.Warn(
				"found invalid file under data directory",
				zap.String("filename", name),
				zap.String("data-dir", dir),
			)
		}
	}

	if m && p {
		lg.Fatal("invalid datadir; both member and proxy directories exist")
	}
	if m {
		return dirMember
	}
	if p {
		return dirProxy
	}
	return dirEmpty
}

func checkSupportArch() {
	lg, err := logutil.CreateDefaultZapLogger(zap.InfoLevel)
	if err != nil {
		panic(err)
	}
	// to add a new platform, check https://github.com/etcd-io/website/blob/main/content/en/docs/next/op-guide/supported-platform.md
	if runtime.GOARCH == "amd64" ||
		runtime.GOARCH == "arm64" ||
		runtime.GOARCH == "ppc64le" ||
		runtime.GOARCH == "s390x" {
		return
	}
	// unsupported arch only configured via environment variable
	// so unset here to not parse through flag
	defer os.Unsetenv("ETCD_UNSUPPORTED_ARCH")
	if env, ok := os.LookupEnv("ETCD_UNSUPPORTED_ARCH"); ok && env == runtime.GOARCH {
		lg.Info("running etcd on unsupported architecture since ETCD_UNSUPPORTED_ARCH is set", zap.String("arch", env))
		return
	}

	lg.Error("running etcd on unsupported architecture since ETCD_UNSUPPORTED_ARCH is set", zap.String("arch", runtime.GOARCH))
	os.Exit(1)
}
<commit_msg>etcdmain: Updates v3.5 supported platform docs<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcdmain

import (
	"fmt"
	"os"
	"runtime"
	"strings"
	"time"

	"go.etcd.io/etcd/client/pkg/v3/fileutil"
	"go.etcd.io/etcd/client/pkg/v3/logutil"
	"go.etcd.io/etcd/client/pkg/v3/types"
	"go.etcd.io/etcd/pkg/v3/osutil"
	"go.etcd.io/etcd/server/v3/embed"
	"go.etcd.io/etcd/server/v3/etcdserver/api/v2discovery"
	"go.etcd.io/etcd/server/v3/etcdserver/errors"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

type dirType string

var (
	dirMember = dirType("member")
	dirProxy  = dirType("proxy")
	dirEmpty  = dirType("empty")
)

func startEtcdOrProxyV2(args []string) {
	grpc.EnableTracing = false

	cfg := newConfig()
	defaultInitialCluster := cfg.ec.InitialCluster

	err := cfg.parse(args[1:])
	lg := cfg.ec.GetLogger()
	// If we failed to parse the whole configuration, print the error using
	// preferably the resolved logger from the config,
	// but if does not exists, create a new temporary logger.
	if lg == nil {
		var zapError error
		// use this logger
		lg, zapError = logutil.CreateDefaultZapLogger(zap.InfoLevel)
		if zapError != nil {
			fmt.Printf("error creating zap logger %v", zapError)
			os.Exit(1)
		}
	}
	lg.Info("Running: ", zap.Strings("args", args))
	if err != nil {
		lg.Warn("failed to verify flags", zap.Error(err))
		switch err {
		case embed.ErrUnsetAdvertiseClientURLsFlag:
			lg.Warn("advertise client URLs are not set", zap.Error(err))
		}
		os.Exit(1)
	}

	cfg.ec.SetupGlobalLoggers()

	defer func() {
		logger := cfg.ec.GetLogger()
		if logger != nil {
			logger.Sync()
		}
	}()

	defaultHost, dhErr := (&cfg.ec).UpdateDefaultClusterFromName(defaultInitialCluster)
	if defaultHost != "" {
		lg.Info(
			"detected default host for advertise",
			zap.String("host", defaultHost),
		)
	}
	if dhErr != nil {
		lg.Info("failed to detect default host", zap.Error(dhErr))
	}

	if cfg.ec.Dir == "" {
		cfg.ec.Dir = fmt.Sprintf("%v.etcd", cfg.ec.Name)
		lg.Warn(
			"'data-dir' was empty; using default",
			zap.String("data-dir", cfg.ec.Dir),
		)
	}

	var stopped <-chan struct{}
	var errc <-chan error

	which := identifyDataDirOrDie(cfg.ec.GetLogger(), cfg.ec.Dir)
	if which != dirEmpty {
		lg.Info(
			"server has already been initialized",
			zap.String("data-dir", cfg.ec.Dir),
			zap.String("dir-type", string(which)),
		)
		switch which {
		case dirMember:
			stopped, errc, err = startEtcd(&cfg.ec)
		case dirProxy:
			lg.Panic("v2 http proxy has already been deprecated in 3.6", zap.String("dir-type", string(which)))
		default:
			lg.Panic(
				"unknown directory type",
				zap.String("dir-type", string(which)),
			)
		}
	} else {
		lg.Info(
			"Initialize and start etcd server",
			zap.String("data-dir", cfg.ec.Dir),
			zap.String("dir-type", string(which)),
		)
		stopped, errc, err = startEtcd(&cfg.ec)
	}

	if err != nil {
		if derr, ok := err.(*errors.DiscoveryError); ok {
			switch derr.Err {
			case v2discovery.ErrDuplicateID:
				lg.Warn(
					"member has been registered with discovery service",
					zap.String("name", cfg.ec.Name),
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(derr.Err),
				)
				lg.Warn(
					"but could not find valid cluster configuration",
					zap.String("data-dir", cfg.ec.Dir),
				)
				lg.Warn("check data dir if previous bootstrap succeeded")
				lg.Warn("or use a new discovery token if previous bootstrap failed")

			case v2discovery.ErrDuplicateName:
				lg.Warn(
					"member with duplicated name has already been registered",
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(derr.Err),
				)
				lg.Warn("cURL the discovery token URL for details")
				lg.Warn("do not reuse discovery token; generate a new one to bootstrap a cluster")

			default:
				lg.Warn(
					"failed to bootstrap; discovery token was already used",
					zap.String("discovery-token", cfg.ec.Durl),
					zap.Error(err),
				)
				lg.Warn("do not reuse discovery token; generate a new one to bootstrap a cluster")
			}
			os.Exit(1)
		}

		if strings.Contains(err.Error(), "include") && strings.Contains(err.Error(), "--initial-cluster") {
			lg.Warn("failed to start", zap.Error(err))
			if cfg.ec.InitialCluster == cfg.ec.InitialClusterFromName(cfg.ec.Name) {
				lg.Warn("forgot to set --initial-cluster?")
			}
			if types.URLs(cfg.ec.APUrls).String() == embed.DefaultInitialAdvertisePeerURLs {
				lg.Warn("forgot to set --initial-advertise-peer-urls?")
			}
			if cfg.ec.InitialCluster == cfg.ec.InitialClusterFromName(cfg.ec.Name) && len(cfg.ec.Durl) == 0 && len(cfg.ec.DiscoveryCfg.Endpoints) == 0 {
				lg.Warn("V2 discovery settings (i.e., --discovery) or v3 discovery settings (i.e., --discovery-token, --discovery-endpoints) are not set")
			}
			os.Exit(1)
		}
		lg.Fatal("discovery failed", zap.Error(err))
	}

	osutil.HandleInterrupts(lg)

	// At this point, the initialization of etcd is done.
	// The listeners are listening on the TCP ports and ready
	// for accepting connections. The etcd instance should be
	// joined with the cluster and ready to serve incoming
	// connections.
	notifySystemd(lg)

	select {
	case lerr := <-errc:
		// fatal out on listener errors
		lg.Fatal("listener failed", zap.Error(lerr))
	case <-stopped:
	}

	osutil.Exit(0)
}

// startEtcd runs StartEtcd in addition to hooks needed for standalone etcd.
func startEtcd(cfg *embed.Config) (<-chan struct{}, <-chan error, error) {
	e, err := embed.StartEtcd(cfg)
	if err != nil {
		return nil, nil, err
	}
	osutil.RegisterInterruptHandler(e.Close)
	select {
	case <-e.Server.ReadyNotify(): // wait for e.Server to join the cluster
	case <-e.Server.StopNotify(): // publish aborted from 'ErrStopped'
	case <-time.After(cfg.ExperimentalWaitClusterReadyTimeout):
		e.GetLogger().Warn("startEtcd: timed out waiting for the ready notification")
	}
	return e.Server.StopNotify(), e.Err(), nil
}

// identifyDataDirOrDie returns the type of the data dir.
// Dies if the datadir is invalid.
func identifyDataDirOrDie(lg *zap.Logger, dir string) dirType {
	names, err := fileutil.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return dirEmpty
		}
		lg.Fatal("failed to list data directory", zap.String("dir", dir), zap.Error(err))
	}

	var m, p bool
	for _, name := range names {
		switch dirType(name) {
		case dirMember:
			m = true
		case dirProxy:
			p = true
		default:
			lg.Warn(
				"found invalid file under data directory",
				zap.String("filename", name),
				zap.String("data-dir", dir),
			)
		}
	}

	if m && p {
		lg.Fatal("invalid datadir; both member and proxy directories exist")
	}
	if m {
		return dirMember
	}
	if p {
		return dirProxy
	}
	return dirEmpty
}

func checkSupportArch() {
	lg, err := logutil.CreateDefaultZapLogger(zap.InfoLevel)
	if err != nil {
		panic(err)
	}
	// to add a new platform, check https://github.com/etcd-io/website/blob/main/content/en/docs/v3.5/op-guide/supported-platform.md
	if runtime.GOARCH == "amd64" ||
		runtime.GOARCH == "arm64" ||
		runtime.GOARCH == "ppc64le" ||
		runtime.GOARCH == "s390x" {
		return
	}
	// unsupported arch only configured via environment variable
	// so unset here to not parse through flag
	defer os.Unsetenv("ETCD_UNSUPPORTED_ARCH")
	if env, ok := os.LookupEnv("ETCD_UNSUPPORTED_ARCH"); ok && env == runtime.GOARCH {
		lg.Info("running etcd on unsupported architecture since ETCD_UNSUPPORTED_ARCH is set", zap.String("arch", env))
		return
	}

	lg.Error("running etcd on unsupported architecture since ETCD_UNSUPPORTED_ARCH is set", zap.String("arch", runtime.GOARCH))
	os.Exit(1)
}
<|endoftext|>


<commit_before>package mmm

import (
	"fmt"
	"reflect"
	"unsafe"
)

// -----------------------------------------------------------------------------

// BytesOf copies the in-memory representation of `v` in `bytes`.
//
// Only numeric, array, struct types and any combination of the above are
// supported.
func BytesOf(v interface{}, bytes []byte) error {
	return bytesOf(v, bytes)
}

func bytesOf(v interface{}, bytes []byte) error {
	t := reflect.TypeOf(v)
	k := t.Kind()
	switch k {
	case reflect.Bool:
		return bytesOfNumericType(v, bytes)
	case reflect.Int:
		return bytesOfNumericType(v, bytes)
	case reflect.Int8:
		return bytesOfNumericType(v, bytes)
	case reflect.Int16:
		return bytesOfNumericType(v, bytes)
	case reflect.Int32:
		return bytesOfNumericType(v, bytes)
	case reflect.Int64:
		return bytesOfNumericType(v, bytes)
	case reflect.Uint:
		return bytesOfNumericType(v, bytes)
	case reflect.Uint8:
		return bytesOfNumericType(v, bytes)
	case reflect.Uint16:
		return bytesOfNumericType(v, bytes)
	case reflect.Uint32:
		return bytesOfNumericType(v, bytes)
	case reflect.Uint64:
		return bytesOfNumericType(v, bytes)
	case reflect.Uintptr:
		return bytesOfNumericType(v, bytes)
	case reflect.Float32:
		return bytesOfNumericType(v, bytes)
	case reflect.Float64:
		return bytesOfNumericType(v, bytes)
	case reflect.Complex64:
		return bytesOfNumericType(v, bytes)
	case reflect.Complex128:
		return bytesOfNumericType(v, bytes)
	case reflect.Array:
		return bytesOfArrayType(v, bytes)
	case reflect.Struct:
		return bytesOfStructType(v, bytes)
	default:
		return Error(fmt.Sprintf("unsuppported type: %#v", k.String()))
	}

	return nil
}

// -----------------------------------------------------------------------------

func extractDataAddress(v interface{}) uintptr {
	// extract data address from the interface
	vbytes := *((*[unsafe.Sizeof(v)]byte)(unsafe.Pointer(&v)))
	vuintptr := *(*uintptr)(unsafe.Pointer(&(vbytes[unsafe.Sizeof(v)-unsafe.Sizeof(uintptr(0))])))

	return vuintptr
}

func bytesOfNumericType(v interface{}, bytes []byte) error {
	// interpret data as a contiguous array of 16 bytes
	vvalue := *((*[16]byte)(unsafe.Pointer(extractDataAddress(v))))
	// copy data byte-representation to `bytes`
	copy(bytes, vvalue[:reflect.ValueOf(v).Type().Size()])

	return nil
}

func bytesOfArrayType(v interface{}, bytes []byte) error {
	vuintptr := extractDataAddress(v)
	size := reflect.ValueOf(v).Type().Size()
	// loop until there are no more data chunks
	for i := uintptr(0); size > 0; i += 128 {
		// interpret data as a contiguous array of 128 bytes
		vvalue := *((*[128]byte)(unsafe.Pointer(vuintptr + i)))
		// copy data byte-representation to `bytes`
		if size > 128 {
			copy(bytes[i:], vvalue[:128])
			size -= 128
		} else {
			copy(bytes[i:], vvalue[:size])
			size = 0
		}
	}

	return nil
}

func bytesOfStructType(v interface{}, bytes []byte) error {
	return bytesOfArrayType(v, bytes)
}
<commit_msg>bytesOf() now uses TypeOf()<commit_after>package mmm

import (
	"reflect"
	"unsafe"
)

// -----------------------------------------------------------------------------

// BytesOf copies the in-memory representation of `v` in `bytes`.
//
// Only numeric, array, struct types and any combination of the above are
// supported.
func BytesOf(v interface{}, bytes []byte) error {
	return bytesOf(v, bytes)
}

func bytesOf(v interface{}, bytes []byte) error {
	t, err := TypeOf(v)
	if err != nil {
		return err
	}

	switch t {
	case TypeNumeric:
		return bytesOfNumericType(v, bytes)
	case TypeArray:
		return bytesOfArrayType(v, bytes)
	case TypeStruct:
		return bytesOfStructType(v, bytes)
	}

	return nil
}

// -----------------------------------------------------------------------------

func extractDataAddress(v interface{}) uintptr {
	// extract data address from the interface
	vbytes := *((*[unsafe.Sizeof(v)]byte)(unsafe.Pointer(&v)))
	vuintptr := *(*uintptr)(unsafe.Pointer(&(vbytes[unsafe.Sizeof(v)-unsafe.Sizeof(uintptr(0))])))

	return vuintptr
}

func bytesOfNumericType(v interface{}, bytes []byte) error {
	// interpret data as a contiguous array of 16 bytes
	vvalue := *((*[16]byte)(unsafe.Pointer(extractDataAddress(v))))
	// copy data byte-representation to `bytes`
	copy(bytes, vvalue[:reflect.ValueOf(v).Type().Size()])

	return nil
}

func bytesOfArrayType(v interface{}, bytes []byte) error {
	vuintptr := extractDataAddress(v)
	size := reflect.ValueOf(v).Type().Size()
	// loop until there are no more data chunks
	for i := uintptr(0); size > 0; i += 128 {
		// interpret data as a contiguous array of 128 bytes
		vvalue := *((*[128]byte)(unsafe.Pointer(vuintptr + i)))
		// copy data byte-representation to `bytes`
		if size > 128 {
			copy(bytes[i:], vvalue[:128])
			size -= 128
		} else {
			copy(bytes[i:], vvalue[:size])
			size = 0
		}
	}

	return nil
}

func bytesOfStructType(v interface{}, bytes []byte) error {
	return bytesOfArrayType(v, bytes)
}
<|endoftext|>


<commit_before>package lm2

import (
	"math/rand"
	"sync"
)

type recordCache struct {
	cache        map[int64]*record
	maxKeyRecord *record
	size         int
	preventPurge bool
	lock         sync.RWMutex
}

func newCache(size int) *recordCache {
	return &recordCache{
		cache:        map[int64]*record{},
		maxKeyRecord: nil,
		size:         size,
	}
}

func (rc *recordCache) findLastLessThan(key string) int64 {
	rc.lock.RLock()
	defer rc.lock.RUnlock()

	if rc.maxKeyRecord != nil {
		if rc.maxKeyRecord.Key < key {
			return rc.maxKeyRecord.Offset
		}
	}
	max := ""
	maxOffset := int64(0)

	for offset, record := range rc.cache {
		if record.Key >= key {
			continue
		}
		if record.Key > max {
			max = record.Key
			maxOffset = offset
		}
	}
	return maxOffset
}

func (rc *recordCache) push(rec *record) {
	rc.lock.RLock()

	if rc.maxKeyRecord == nil || rc.maxKeyRecord.Key < rec.Key {
		rc.lock.RUnlock()

		rc.lock.Lock()
		if rc.maxKeyRecord == nil || rc.maxKeyRecord.Key < rec.Key {
			rc.maxKeyRecord = rec
		}
		rc.lock.Unlock()

		return
	}

	if len(rc.cache) == rc.size && rand.Float32() >= cacheProb {
		rc.lock.RUnlock()
		return
	}

	rc.lock.RUnlock()
	rc.lock.Lock()

	rc.cache[rec.Offset] = rec
	if !rc.preventPurge {
		rc.purge()
	}

	rc.lock.Unlock()
}

func (rc *recordCache) purge() {
	purged := 0
	for len(rc.cache) > rc.size {
		deletedKey := int64(0)
		for k := range rc.cache {
			if k == rc.maxKeyRecord.Offset {
				continue
			}
			deletedKey = k
			break
		}
		delete(rc.cache, deletedKey)
		purged++
	}
}

func (rc *recordCache) flushOffsets(offsets []int64) {
	rc.lock.Lock()
	for _, offset := range offsets {
		delete(rc.cache, offset)
	}
	rc.lock.Unlock()
}
<commit_msg>make cache probability dependent on height (#51)<commit_after>package lm2

import (
	"math/rand"
	"sync"
)

type recordCache struct {
	cache        map[int64]*record
	maxKeyRecord *record
	size         int
	preventPurge bool
	lock         sync.RWMutex
}

func newCache(size int) *recordCache {
	return &recordCache{
		cache:        map[int64]*record{},
		maxKeyRecord: nil,
		size:         size,
	}
}

func (rc *recordCache) findLastLessThan(key string) int64 {
	rc.lock.RLock()
	defer rc.lock.RUnlock()

	if rc.maxKeyRecord != nil {
		if rc.maxKeyRecord.Key < key {
			return rc.maxKeyRecord.Offset
		}
	}
	max := ""
	maxOffset := int64(0)

	for offset, record := range rc.cache {
		if record.Key >= key {
			continue
		}
		if record.Key > max {
			max = record.Key
			maxOffset = offset
		}
	}
	return maxOffset
}

func (rc *recordCache) push(rec *record) {
	rc.lock.RLock()

	if rc.maxKeyRecord == nil || rc.maxKeyRecord.Key < rec.Key {
		rc.lock.RUnlock()

		rc.lock.Lock()
		if rc.maxKeyRecord == nil || rc.maxKeyRecord.Key < rec.Key {
			rc.maxKeyRecord = rec
		}
		rc.lock.Unlock()

		return
	}

	if len(rc.cache) == rc.size && rand.Float32() >= cacheProb/float32(len(rec.Next)) {
		rc.lock.RUnlock()
		return
	}

	rc.lock.RUnlock()
	rc.lock.Lock()

	rc.cache[rec.Offset] = rec
	if !rc.preventPurge {
		rc.purge()
	}

	rc.lock.Unlock()
}

func (rc *recordCache) purge() {
	purged := 0
	for len(rc.cache) > rc.size {
		deletedKey := int64(0)
		for k := range rc.cache {
			if k == rc.maxKeyRecord.Offset {
				continue
			}
			deletedKey = k
			break
		}
		delete(rc.cache, deletedKey)
		purged++
	}
}

func (rc *recordCache) flushOffsets(offsets []int64) {
	rc.lock.Lock()
	for _, offset := range offsets {
		delete(rc.cache, offset)
	}
	rc.lock.Unlock()
}
<|endoftext|>


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

package main

import (
	"io"
	"log"
	"sort"
	"time"
)

type page struct {
	n      offset
	body   []byte
	expire time.Time
	cached bool
}

func newPage(n offset, body []byte) *page {
	return &page{
		n:    n,
		body: body,
	}
}

func (p *page) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(p.body)
	return int64(n), err
}

type group string

type entry struct {
	deadline time.Time
	data     []byte
}

func newEntry(data []byte, d time.Duration) *entry {
	return &entry{
		deadline: time.Now().Add(d),
		data:     data,
	}
}

func (ce *entry) invalid(t time.Time) bool {
	return !ce.deadline.After(t)
}

func (ce *entry) asPage(n offset) *page {
	p := newPage(n, ce.data)
	p.expire = ce.deadline
	return p
}

type entries struct {
	ents map[group]map[offset]*entry
}

func newEntries() *entries {
	return &entries{
		ents: make(map[group]map[offset]*entry),
	}
}

func (e *entries) count() int {
	tot := 0
	for k := range e.ents {
		tot += len(e.ents[k])
	}
	return tot
}

func (e *entries) get(cg group, n offset) (*entry, bool) {
	pages, ok := e.ents[cg]
	if !ok {
		return nil, false
	}
	ce, ok := pages[n]
	return ce, ok
}

func (e *entries) put(cg group, n offset, ce *entry) {
	_, ok := e.ents[cg]
	if !ok {
		e.ents[cg] = make(map[offset]*entry)
	}
	e.ents[cg][n] = ce
}

func (e *entries) has(cg group, n offset, t time.Time) bool {
	_, ok := e.ents[cg]
	if !ok {
		return false
	}
	ent, ok := e.ents[cg][n]
	if !ok {
		return false
	}
	return !ent.invalid(t)
}

// remove assumes the entry exists
func (e *entries) remove(cg group, n offset) {
	delete(e.ents[cg], n)
	if len(e.ents[cg]) == 0 {
		delete(e.ents, cg)
	}
}

func (e *entries) purge(cg group, st *stats) {
	st.mem(-e.sizeof(cg))
	delete(e.ents, cg)
}

func (e *entries) sizeof(cg group) int {
	var s int
	for n := range e.ents[cg] {
		s += len(e.ents[cg][n].data)
	}
	return s
}

func (e *entries) oldestDeadline(cg group) time.Time {
	var t time.Time
	for off := range e.ents[cg] {
		entry := e.ents[cg][off]
		if t.IsZero() || t.After(entry.deadline) {
			t = entry.deadline
		}
	}
	return t
}

func (e *entries) gc(t time.Time, st *stats) {
	for cg, ents := range e.ents {
		for n := range ents {
			if ents[n].invalid(t) {
				st.mem(-len(ents[n].data))
				e.remove(cg, n)
			}
		}
	}
}

type cacheFunc func() error

type cache struct {
	entries *entries
	waits   *waiters
	fetcher *fetcher
	config  *config
	stat    *stats
	events  chan cacheFunc
	debug   func(string, ...interface{})
}

func newCache(f *fetcher, logs *logbuf, cf *config) *cache {
	c := &cache{
		fetcher: f,
		config:  cf,
		events:  make(chan cacheFunc),
		entries: newEntries(),
		waits:   newWaiters(),
		stat:    newStats(),
		debug:   logs.debug,
	}
	go c.gc(cf.gcpause)
	go c.run()
	return c
}

func (c *cache) run() {
	for f := range c.events {
		if err := f(); err != nil {
			log.Print("cache: ", err)
		}
	}
}

func (c *cache) gc(d time.Duration) {
	done := make(chan struct{})
	for {
		time.Sleep(d)
		c.events <- func() error {
			c.debug("running garbage collector cycle, memory is %d", c.stat.Mem)
			c.entries.gc(time.Now(), c.stat)
			c.debug("garbage collection done, memory is %d", c.stat.Mem)
			done <- struct{}{}
			return nil
		}
		<-done
	}
}

type timeGroup struct {
	t  time.Time
	cg group
}

type timeGroups struct {
	entries []timeGroup
}

type byTime []timeGroup

func (a byTime) Len() int           { return len(a) }
func (a byTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byTime) Less(i, j int) bool { return a[i].t.After(a[j].t) }

func makeTimeGroups(e *entries) timeGroups {
	var i int
	tg := make([]timeGroup, len(e.ents))
	for cg := range e.ents {
		tg[i] = timeGroup{
			t:  e.oldestDeadline(cg),
			cg: cg,
		}
		i++
	}
	sort.Sort(byTime(tg))
	return timeGroups{tg}
}

func (tg *timeGroups) purgeOldest(c *cache) {
	entry, ts := tg.entries[len(tg.entries)-1], tg.entries[0:len(tg.entries)-1]
	c.entries.purge(entry.cg, c.stat)
	tg.entries = ts
}

func (tg *timeGroups) empty() bool {
	return len(tg.entries) == 0
}

func (c *cache) oom(target int64) {
	c.debug("OOM called: using %d, limit is %d", c.stat.Mem, target)
	tg := makeTimeGroups(c.entries)
	for {
		tg.purgeOldest(c)
		// Repeat unless we have no more entries or we are using less memory than target
		if tg.empty() || !c.stat.above(target) {
			break
		}
	}
	c.debug("OOM: mem now %d", c.stat.Mem)
}

// put inserts a page into the cache (after it was fetched).
func (c *cache) put(cg group, p *page, err error) {
	c.events <- func() error {
		ce := newEntry(p.body, c.config.lifetime)
		// If an entry exists, it will be overwritten
		if ent, ok := c.entries.get(cg, p.n); ok {
			c.stat.mem(-len(ent.data))
		}
		c.entries.put(cg, p.n, ce)
		c.stat.mem(len(p.body))
		c.debug("added page %s/%d", cg, p.n)
		if c.stat.above(c.config.maxMemory) {
			go func() {
				c.events <- func() error {
					c.oom(c.config.maxMemory)
					return nil
				}
			}()
		}
		// If there were waiters, signal that the wait is over
		c.waits.done(cg, p.n)
		return err
	}
}

// prefetch requests pages (n-m, n) and (n, n+m) if not already fetched
func (c *cache) prefetch(cg group, n, m int, t time.Time) {
	i := n - m
	if i < 0 {
		i = 0
	}
	for ; i < n+m; i++ {
		if i == n {
			// just fetched
			continue
		}
		off := offset(i * c.config.incr)
		if c.entries.has(cg, off, t) || c.waits.has(cg, off) {
			// already fetched or requested
			continue
		}
		c.waits.wait(cg, off)
		res := newResource(c.config.tmpl, cg, off)
		c.fetcher.request(newJob(res, c))
	}
}

func (c *cache) request(cg group, n int, t time.Time) chan struct{} {
	off := offset(n * c.config.incr)
	wait := c.waits.wait(cg, off)
	res := newResource(c.config.tmpl, cg, off)
	c.fetcher.request(newJob(res, c))
	c.prefetch(cg, n, c.config.npref, t)
	return wait
}

func (c *cache) stats() *stats {
	var st *stats
	wait := make(chan struct{})
	c.events <- func() error {
		c.stat.Entries = c.entries.count()
		c.stat.Waiters = c.waits.count()
		st = c.stat.clone()
		wait <- struct{}{}
		return nil
	}
	<-wait
	return st
}

func (c *cache) get(cg group, n int) *page {
	var (
		page *page
		wait chan struct{}
	)
	cached := true
	requested := make(chan struct{})
	off := offset(n * c.config.incr)
	c.debug("%s/%d: requesting from cache", cg, off)
	for {
		wait = nil
		c.events <- func() error {
			defer func() { requested <- struct{}{} }()
			var now time.Time
			ce, ok := c.entries.get(cg, off)
			if ok {
				now = time.Now()
			}
			if !ok || ce.invalid(now) {
				c.debug("%s/%d: not cached, requested", cg, off)
				wait = c.request(cg, n, now)
				return nil
			}
			c.debug("%s/%d: found", cg, off)
			c.prefetch(cg, n, c.config.npref, now)
			c.stat.hit(cached)
			page = ce.asPage(off)
			return nil
		}
		<-requested
		// content was already in cache, return it
		if wait == nil {
			page.cached = cached
			return page
		}
		// We needed to request the object, it was not cached
		cached = false
		// content is being fetched, wait and try to get again
		<-wait
	}
}
<commit_msg>cache: wait for already requested items<commit_after>// Copyright 2017 Giulio Iotti. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"io"
	"log"
	"sort"
	"time"
)

type page struct {
	n      offset
	body   []byte
	expire time.Time
	cached bool
}

func newPage(n offset, body []byte) *page {
	return &page{
		n:    n,
		body: body,
	}
}

func (p *page) WriteTo(w io.Writer) (int64, error) {
	n, err := w.Write(p.body)
	return int64(n), err
}

type group string

type entry struct {
	deadline time.Time
	data     []byte
}

func newEntry(data []byte, d time.Duration) *entry {
	return &entry{
		deadline: time.Now().Add(d),
		data:     data,
	}
}

func (ce *entry) invalid(t time.Time) bool {
	return !ce.deadline.After(t)
}

func (ce *entry) asPage(n offset) *page {
	p := newPage(n, ce.data)
	p.expire = ce.deadline
	return p
}

type entries struct {
	ents map[group]map[offset]*entry
}

func newEntries() *entries {
	return &entries{
		ents: make(map[group]map[offset]*entry),
	}
}

func (e *entries) count() int {
	tot := 0
	for k := range e.ents {
		tot += len(e.ents[k])
	}
	return tot
}

func (e *entries) get(cg group, n offset) (*entry, bool) {
	pages, ok := e.ents[cg]
	if !ok {
		return nil, false
	}
	ce, ok := pages[n]
	return ce, ok
}

func (e *entries) put(cg group, n offset, ce *entry) {
	_, ok := e.ents[cg]
	if !ok {
		e.ents[cg] = make(map[offset]*entry)
	}
	e.ents[cg][n] = ce
}

func (e *entries) has(cg group, n offset, t time.Time) bool {
	_, ok := e.ents[cg]
	if !ok {
		return false
	}
	ent, ok := e.ents[cg][n]
	if !ok {
		return false
	}
	return !ent.invalid(t)
}

// remove assumes the entry exists
func (e *entries) remove(cg group, n offset) {
	delete(e.ents[cg], n)
	if len(e.ents[cg]) == 0 {
		delete(e.ents, cg)
	}
}

func (e *entries) purge(cg group, st *stats) {
	st.mem(-e.sizeof(cg))
	delete(e.ents, cg)
}

func (e *entries) sizeof(cg group) int {
	var s int
	for n := range e.ents[cg] {
		s += len(e.ents[cg][n].data)
	}
	return s
}

func (e *entries) oldestDeadline(cg group) time.Time {
	var t time.Time
	for off := range e.ents[cg] {
		entry := e.ents[cg][off]
		if t.IsZero() || t.After(entry.deadline) {
			t = entry.deadline
		}
	}
	return t
}

func (e *entries) gc(t time.Time, st *stats) {
	for cg, ents := range e.ents {
		for n := range ents {
			if ents[n].invalid(t) {
				st.mem(-len(ents[n].data))
				e.remove(cg, n)
			}
		}
	}
}

type cacheFunc func() error

type cache struct {
	entries *entries
	waits   *waiters
	fetcher *fetcher
	config  *config
	stat    *stats
	events  chan cacheFunc
	debug   func(string, ...interface{})
}

func newCache(f *fetcher, logs *logbuf, cf *config) *cache {
	c := &cache{
		fetcher: f,
		config:  cf,
		events:  make(chan cacheFunc),
		entries: newEntries(),
		waits:   newWaiters(),
		stat:    newStats(),
		debug:   logs.debug,
	}
	go c.gc(cf.gcpause)
	go c.run()
	return c
}

func (c *cache) run() {
	for f := range c.events {
		if err := f(); err != nil {
			log.Print("cache: ", err)
		}
	}
}

func (c *cache) gc(d time.Duration) {
	done := make(chan struct{})
	for {
		time.Sleep(d)
		c.events <- func() error {
			c.debug("running garbage collector cycle, memory is %d", c.stat.Mem)
			c.entries.gc(time.Now(), c.stat)
			c.debug("garbage collection done, memory is %d", c.stat.Mem)
			done <- struct{}{}
			return nil
		}
		<-done
	}
}

type timeGroup struct {
	t  time.Time
	cg group
}

type timeGroups struct {
	entries []timeGroup
}

type byTime []timeGroup

func (a byTime) Len() int           { return len(a) }
func (a byTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byTime) Less(i, j int) bool { return a[i].t.After(a[j].t) }

func makeTimeGroups(e *entries) timeGroups {
	var i int
	tg := make([]timeGroup, len(e.ents))
	for cg := range e.ents {
		tg[i] = timeGroup{
			t:  e.oldestDeadline(cg),
			cg: cg,
		}
		i++
	}
	sort.Sort(byTime(tg))
	return timeGroups{tg}
}

func (tg *timeGroups) purgeOldest(c *cache) {
	entry, ts := tg.entries[len(tg.entries)-1], tg.entries[0:len(tg.entries)-1]
	c.entries.purge(entry.cg, c.stat)
	tg.entries = ts
}

func (tg *timeGroups) empty() bool {
	return len(tg.entries) == 0
}

func (c *cache) oom(target int64) {
	c.debug("OOM called: using %d, limit is %d", c.stat.Mem, target)
	tg := makeTimeGroups(c.entries)
	for {
		tg.purgeOldest(c)
		// Repeat unless we have no more entries or we are using less memory than target
		if tg.empty() || !c.stat.above(target) {
			break
		}
	}
	c.debug("OOM: mem now %d", c.stat.Mem)
}

// put inserts a page into the cache (after it was fetched).
func (c *cache) put(cg group, p *page, err error) {
	c.events <- func() error {
		ce := newEntry(p.body, c.config.lifetime)
		// If an entry exists, it will be overwritten
		if ent, ok := c.entries.get(cg, p.n); ok {
			c.stat.mem(-len(ent.data))
		}
		c.entries.put(cg, p.n, ce)
		c.stat.mem(len(p.body))
		c.debug("added page %s/%d", cg, p.n)
		if c.stat.above(c.config.maxMemory) {
			go func() {
				c.events <- func() error {
					c.oom(c.config.maxMemory)
					return nil
				}
			}()
		}
		// If there were waiters, signal that the wait is over
		c.waits.done(cg, p.n)
		return err
	}
}

// prefetch requests pages (n-m, n) and (n, n+m) if not already fetched
func (c *cache) prefetch(cg group, n, m int, t time.Time) {
	i := n - m
	if i < 0 {
		i = 0
	}
	for ; i < n+m; i++ {
		if i == n {
			// just fetched
			continue
		}
		off := offset(i * c.config.incr)
		if c.entries.has(cg, off, t) || c.waits.has(cg, off) {
			// already fetched or requested
			continue
		}
		c.waits.wait(cg, off)
		res := newResource(c.config.tmpl, cg, off)
		c.fetcher.request(newJob(res, c))
	}
}

func (c *cache) request(cg group, n int, t time.Time) chan struct{} {
	off := offset(n * c.config.incr)
	wait := c.waits.wait(cg, off)
	res := newResource(c.config.tmpl, cg, off)
	c.fetcher.request(newJob(res, c))
	c.prefetch(cg, n, c.config.npref, t)
	return wait
}

func (c *cache) stats() *stats {
	var st *stats
	wait := make(chan struct{})
	c.events <- func() error {
		c.stat.Entries = c.entries.count()
		c.stat.Waiters = c.waits.count()
		st = c.stat.clone()
		wait <- struct{}{}
		return nil
	}
	<-wait
	return st
}

func (c *cache) get(cg group, n int) *page {
	var (
		page *page
		wait chan struct{}
	)
	cached := true
	requested := make(chan struct{})
	off := offset(n * c.config.incr)
	c.debug("%s/%d: requesting from cache", cg, off)
	for {
		wait = nil
		c.events <- func() error {
			defer func() { requested <- struct{}{} }()
			if c.waits.has(cg, off) {
				c.waits.wait(cg, off)
			}
			var now time.Time
			ce, ok := c.entries.get(cg, off)
			if ok {
				now = time.Now()
			}
			if !ok || ce.invalid(now) {
				c.debug("%s/%d: not cached, requested", cg, off)
				wait = c.request(cg, n, now)
				return nil
			}
			c.debug("%s/%d: found", cg, off)
			c.prefetch(cg, n, c.config.npref, now)
			c.stat.hit(cached)
			page = ce.asPage(off)
			return nil
		}
		<-requested
		// content was already in cache, return it
		if wait == nil {
			page.cached = cached
			return page
		}
		// We needed to request the object, it was not cached
		cached = false
		// content is being fetched, wait and try to get again
		<-wait
	}
}
<|endoftext|>


<commit_before>package httpcache

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/textproto"
	"os"
	pathutil "path"
	"time"

	vfs "gopkgs.com/vfs.v1"
)

const (
	headerPrefix = "header/"
	bodyPrefix   = "body/"
)

// Returned when a resource doesn't exist
var ErrNotFoundInCache = errors.New("Not found in cache")

// Cache provides a storage mechanism for cached Resources
type Cache struct {
	fs    vfs.VFS
	stale map[string]time.Time
}

// NewCache returns a Cache backed off the provided VFS
func NewCache(fs vfs.VFS) *Cache {
	return &Cache{fs: fs, stale: map[string]time.Time{}}
}

// NewMemoryCache returns an ephemeral cache in memory
func NewMemoryCache() *Cache {
	return NewCache(vfs.Memory())
}

// NewDiskCache returns a disk-backed cache
func NewDiskCache(dir string) (*Cache, error) {
	fs, err := vfs.FS(dir)
	if err != nil {
		return nil, err
	}
	chfs, err := vfs.Chroot("/", fs)
	if err != nil {
		return nil, err
	}
	return NewCache(chfs), nil
}

func (c *Cache) vfsWrite(path string, r io.Reader) error {
	if err := vfs.MkdirAll(c.fs, pathutil.Dir(path), 0700); err != nil {
		return err
	}
	f, err := c.fs.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err := io.Copy(f, r); err != nil {
		return err
	}
	return nil
}

// Retrieve the Headers for a given key path
func (c *Cache) Header(key string) (http.Header, error) {
	path := headerPrefix + hashKey(key)
	f, err := c.fs.Open(path)
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}

	return readHeaders(bufio.NewReader(f))
}

// Store a resource against a number of keys
func (c *Cache) Store(res *Resource, keys ...string) error {
	b, err := ioutil.ReadAll(res)
	if err != nil {
		return err
	}

	for _, key := range keys {
		delete(c.stale, key)

		if err := c.storeBody(bytes.NewReader(b), key); err != nil {
			return err
		}

		if err := c.storeHeaders(res.Header(), key); err != nil {
			return err
		}
	}

	return nil
}

func (c *Cache) storeBody(r io.Reader, key string) error {
	if err := c.vfsWrite(bodyPrefix+hashKey(key), r); err != nil {
		return err
	}
	return nil
}

func (c *Cache) storeHeaders(h http.Header, key string) error {
	hb := &bytes.Buffer{}
	writeHeaders(h, hb)

	if err := c.vfsWrite(headerPrefix+hashKey(key), bytes.NewReader(hb.Bytes())); err != nil {
		return err
	}
	return nil
}

// Retrieve returns a cached Resource for the given key
func (c *Cache) Retrieve(key string) (*Resource, error) {
	f, err := c.fs.Open(bodyPrefix + hashKey(key))
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}
	h, err := c.Header(key)
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}
	res := NewResource(http.StatusOK, f, h)
	if staleTime, exists := c.stale[key]; exists {
		if !res.DateAfter(staleTime) {
			log.Printf("stale marker of %s found", staleTime)
			res.MarkStale()
		}
	}
	return res, nil
}

func (c *Cache) Invalidate(keys ...string) {
	log.Printf("invalidating %q", keys)
	for _, key := range keys {
		c.stale[key] = Clock()
	}
}

func (c *Cache) Freshen(res *Resource, keys ...string) error {
	for _, key := range keys {
		if h, err := c.Header(key); err == nil {
			if validateHeaders(h, res.Header()) {
				log.Printf("freshened %s", key)
				if err := c.storeHeaders(res.Header(), key); err != nil {
					return err
				}
			} else {
				log.Printf("freshen failed, invalidating %s", key)
				c.Invalidate(key)
			}
		}
	}
	return nil
}

func hashKey(key string) string {
	h := sha256.New()
	io.WriteString(h, key)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func readHeaders(r *bufio.Reader) (http.Header, error) {
	tp := textproto.NewReader(r)
	mimeHeader, err := tp.ReadMIMEHeader()
	if err != nil {
		return nil, err
	}
	return http.Header(mimeHeader), nil
}

func writeHeaders(h http.Header, w io.Writer) error {
	if err := h.Write(w); err != nil {
		return err
	}
	// ReadMIMEHeader expects a trailing newline
	_, err := w.Write([]byte("\r\n"))
	return err
}
<commit_msg>DiskCache now creates cache dir as needed<commit_after>package httpcache

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/textproto"
	"os"
	pathutil "path"
	"time"

	vfs "gopkgs.com/vfs.v1"
)

const (
	headerPrefix = "header/"
	bodyPrefix   = "body/"
)

// Returned when a resource doesn't exist
var ErrNotFoundInCache = errors.New("Not found in cache")

// Cache provides a storage mechanism for cached Resources
type Cache struct {
	fs    vfs.VFS
	stale map[string]time.Time
}

// NewCache returns a Cache backed off the provided VFS
func NewCache(fs vfs.VFS) *Cache {
	return &Cache{fs: fs, stale: map[string]time.Time{}}
}

// NewMemoryCache returns an ephemeral cache in memory
func NewMemoryCache() *Cache {
	return NewCache(vfs.Memory())
}

// NewDiskCache returns a disk-backed cache
func NewDiskCache(dir string) (*Cache, error) {
	if err := os.MkdirAll(dir, 0777); err != nil {
		return nil, err
	}
	fs, err := vfs.FS(dir)
	if err != nil {
		return nil, err
	}
	chfs, err := vfs.Chroot("/", fs)
	if err != nil {
		return nil, err
	}
	return NewCache(chfs), nil
}

func (c *Cache) vfsWrite(path string, r io.Reader) error {
	if err := vfs.MkdirAll(c.fs, pathutil.Dir(path), 0700); err != nil {
		return err
	}
	f, err := c.fs.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err := io.Copy(f, r); err != nil {
		return err
	}
	return nil
}

// Retrieve the Headers for a given key path
func (c *Cache) Header(key string) (http.Header, error) {
	path := headerPrefix + hashKey(key)
	f, err := c.fs.Open(path)
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}

	return readHeaders(bufio.NewReader(f))
}

// Store a resource against a number of keys
func (c *Cache) Store(res *Resource, keys ...string) error {
	b, err := ioutil.ReadAll(res)
	if err != nil {
		return err
	}

	for _, key := range keys {
		delete(c.stale, key)

		if err := c.storeBody(bytes.NewReader(b), key); err != nil {
			return err
		}

		if err := c.storeHeaders(res.Header(), key); err != nil {
			return err
		}
	}

	return nil
}

func (c *Cache) storeBody(r io.Reader, key string) error {
	if err := c.vfsWrite(bodyPrefix+hashKey(key), r); err != nil {
		return err
	}
	return nil
}

func (c *Cache) storeHeaders(h http.Header, key string) error {
	hb := &bytes.Buffer{}
	writeHeaders(h, hb)

	if err := c.vfsWrite(headerPrefix+hashKey(key), bytes.NewReader(hb.Bytes())); err != nil {
		return err
	}
	return nil
}

// Retrieve returns a cached Resource for the given key
func (c *Cache) Retrieve(key string) (*Resource, error) {
	f, err := c.fs.Open(bodyPrefix + hashKey(key))
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}
	h, err := c.Header(key)
	if err != nil {
		if vfs.IsNotExist(err) {
			return nil, ErrNotFoundInCache
		}
		return nil, err
	}
	res := NewResource(http.StatusOK, f, h)
	if staleTime, exists := c.stale[key]; exists {
		if !res.DateAfter(staleTime) {
			log.Printf("stale marker of %s found", staleTime)
			res.MarkStale()
		}
	}
	return res, nil
}

func (c *Cache) Invalidate(keys ...string) {
	log.Printf("invalidating %q", keys)
	for _, key := range keys {
		c.stale[key] = Clock()
	}
}

func (c *Cache) Freshen(res *Resource, keys ...string) error {
	for _, key := range keys {
		if h, err := c.Header(key); err == nil {
			if validateHeaders(h, res.Header()) {
				log.Printf("freshened %s", key)
				if err := c.storeHeaders(res.Header(), key); err != nil {
					return err
				}
			} else {
				log.Printf("freshen failed, invalidating %s", key)
				c.Invalidate(key)
			}
		}
	}
	return nil
}

func hashKey(key string) string {
	h := sha256.New()
	io.WriteString(h, key)
	return fmt.Sprintf("%x", h.Sum(nil))
}

func readHeaders(r *bufio.Reader) (http.Header, error) {
	tp := textproto.NewReader(r)
	mimeHeader, err := tp.ReadMIMEHeader()
	if err != nil {
		return nil, err
	}
	return http.Header(mimeHeader), nil
}

func writeHeaders(h http.Header, w io.Writer) error {
	if err := h.Write(w); err != nil {
		return err
	}
	// ReadMIMEHeader expects a trailing newline
	_, err := w.Write([]byte("\r\n"))
	return err
}
<|endoftext|>


<commit_before>package containerd

import (
	"context"
	"os"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	introspection "github.com/containerd/containerd/api/services/introspection/v1"
	"github.com/containerd/containerd/gc"
	"github.com/containerd/containerd/leases"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/cache/metadata"
	"github.com/moby/buildkit/executor/containerdexecutor"
	"github.com/moby/buildkit/executor/oci"
	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
	"github.com/moby/buildkit/util/leaseutil"
	"github.com/moby/buildkit/util/network/netproviders"
	"github.com/moby/buildkit/util/winlayers"
	"github.com/moby/buildkit/worker/base"
	specs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

// NewWorkerOpt creates a WorkerOpt.
func NewWorkerOpt(root string, address, snapshotterName, ns string, labels map[string]string, dns *oci.DNSConfig, nopt netproviders.Opt, opts ...containerd.ClientOpt) (base.WorkerOpt, error) {
	opts = append(opts,
		containerd.WithDefaultNamespace(ns),
		containerd.WithDefaultRuntime("io.containerd.runtime.v1.linux"),
	)
	client, err := containerd.New(address, opts...)
	if err != nil {
		return base.WorkerOpt{}, errors.Wrapf(err, "failed to connect client to %q . make sure containerd is running", address)
	}
	return newContainerd(root, client, snapshotterName, ns, labels, dns, nopt)
}

func newContainerd(root string, client *containerd.Client, snapshotterName, ns string, labels map[string]string, dns *oci.DNSConfig, nopt netproviders.Opt) (base.WorkerOpt, error) {
	if strings.Contains(snapshotterName, "/") {
		return base.WorkerOpt{}, errors.Errorf("bad snapshotter name: %q", snapshotterName)
	}
	name := "containerd-" + snapshotterName
	root = filepath.Join(root, name)
	if err := os.MkdirAll(root, 0700); err != nil {
		return base.WorkerOpt{}, errors.Wrapf(err, "failed to create %s", root)
	}

	df := client.DiffService()
	// TODO: should use containerd daemon instance ID (containerd/containerd#1862)?
	id, err := base.ID(root)
	if err != nil {
		return base.WorkerOpt{}, err
	}
	xlabels := base.Labels("containerd", snapshotterName)
	for k, v := range labels {
		xlabels[k] = v
	}

	lm := leaseutil.WithNamespace(client.LeasesService(), ns)

	gc := func(ctx context.Context) (gc.Stats, error) {
		l, err := lm.Create(ctx)
		if err != nil {
			return nil, nil
		}
		return nil, lm.Delete(ctx, leases.Lease{ID: l.ID}, leases.SynchronousDelete)
	}

	cs := containerdsnapshot.NewContentStore(client.ContentStore(), ns)

	resp, err := client.IntrospectionService().Plugins(context.TODO(), &introspection.PluginsRequest{Filters: []string{"type==io.containerd.runtime.v1"}})
	if err != nil {
		return base.WorkerOpt{}, errors.Wrap(err, "failed to list runtime plugin")
	}
	if len(resp.Plugins) == 0 {
		return base.WorkerOpt{}, errors.New("failed to find any runtime plugins")
	}

	var platforms []specs.Platform
	for _, plugin := range resp.Plugins {
		for _, p := range plugin.Platforms {
			platforms = append(platforms, specs.Platform{
				OS:           p.OS,
				Architecture: p.Architecture,
				Variant:      p.Variant,
			})
		}
	}

	np, err := netproviders.Providers(nopt)
	if err != nil {
		return base.WorkerOpt{}, err
	}

	snap := containerdsnapshot.NewSnapshotter(snapshotterName, client.SnapshotService(snapshotterName), ns, nil)

	if err := cache.MigrateV2(context.TODO(), filepath.Join(root, "metadata.db"), filepath.Join(root, "metadata_v2.db"), cs, snap, lm); err != nil {
		return base.WorkerOpt{}, err
	}

	md, err := metadata.NewStore(filepath.Join(root, "metadata_v2.db"))
	if err != nil {
		return base.WorkerOpt{}, err
	}

	opt := base.WorkerOpt{
		ID:             id,
		Labels:         xlabels,
		MetadataStore:  md,
		Executor:       containerdexecutor.New(client, root, "", np, dns),
		Snapshotter:    snap,
		ContentStore:   cs,
		Applier:        winlayers.NewFileSystemApplierWithWindows(cs, df),
		Differ:         winlayers.NewWalkingDiffWithWindows(cs, df),
		ImageStore:     client.ImageService(),
		Platforms:      platforms,
		LeaseManager:   lm,
		GarbageCollect: gc,
	}
	return opt, nil
}
<commit_msg>Recognise Runtime V2 containerd plugins<commit_after>package containerd

import (
	"context"
	"os"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	introspection "github.com/containerd/containerd/api/services/introspection/v1"
	"github.com/containerd/containerd/gc"
	"github.com/containerd/containerd/leases"
	"github.com/moby/buildkit/cache"
	"github.com/moby/buildkit/cache/metadata"
	"github.com/moby/buildkit/executor/containerdexecutor"
	"github.com/moby/buildkit/executor/oci"
	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
	"github.com/moby/buildkit/util/leaseutil"
	"github.com/moby/buildkit/util/network/netproviders"
	"github.com/moby/buildkit/util/winlayers"
	"github.com/moby/buildkit/worker/base"
	specs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

// NewWorkerOpt creates a WorkerOpt.
func NewWorkerOpt(root string, address, snapshotterName, ns string, labels map[string]string, dns *oci.DNSConfig, nopt netproviders.Opt, opts ...containerd.ClientOpt) (base.WorkerOpt, error) {
	opts = append(opts,
		containerd.WithDefaultNamespace(ns),
		containerd.WithDefaultRuntime("io.containerd.runtime.v1.linux"),
	)
	client, err := containerd.New(address, opts...)
	if err != nil {
		return base.WorkerOpt{}, errors.Wrapf(err, "failed to connect client to %q . make sure containerd is running", address)
	}
	return newContainerd(root, client, snapshotterName, ns, labels, dns, nopt)
}

func newContainerd(root string, client *containerd.Client, snapshotterName, ns string, labels map[string]string, dns *oci.DNSConfig, nopt netproviders.Opt) (base.WorkerOpt, error) {
	if strings.Contains(snapshotterName, "/") {
		return base.WorkerOpt{}, errors.Errorf("bad snapshotter name: %q", snapshotterName)
	}
	name := "containerd-" + snapshotterName
	root = filepath.Join(root, name)
	if err := os.MkdirAll(root, 0700); err != nil {
		return base.WorkerOpt{}, errors.Wrapf(err, "failed to create %s", root)
	}

	df := client.DiffService()
	// TODO: should use containerd daemon instance ID (containerd/containerd#1862)?
	id, err := base.ID(root)
	if err != nil {
		return base.WorkerOpt{}, err
	}
	xlabels := base.Labels("containerd", snapshotterName)
	for k, v := range labels {
		xlabels[k] = v
	}

	lm := leaseutil.WithNamespace(client.LeasesService(), ns)

	gc := func(ctx context.Context) (gc.Stats, error) {
		l, err := lm.Create(ctx)
		if err != nil {
			return nil, nil
		}
		return nil, lm.Delete(ctx, leases.Lease{ID: l.ID}, leases.SynchronousDelete)
	}

	cs := containerdsnapshot.NewContentStore(client.ContentStore(), ns)

	resp, err := client.IntrospectionService().Plugins(context.TODO(), &introspection.PluginsRequest{Filters: []string{"type==io.containerd.runtime.v1", "type==io.containerd.runtime.v2"}})
	if err != nil {
		return base.WorkerOpt{}, errors.Wrap(err, "failed to list runtime plugin")
	}
	if len(resp.Plugins) == 0 {
		return base.WorkerOpt{}, errors.New("failed to find any runtime plugins")
	}

	var platforms []specs.Platform
	for _, plugin := range resp.Plugins {
		for _, p := range plugin.Platforms {
			platforms = append(platforms, specs.Platform{
				OS:           p.OS,
				Architecture: p.Architecture,
				Variant:      p.Variant,
			})
		}
	}

	np, err := netproviders.Providers(nopt)
	if err != nil {
		return base.WorkerOpt{}, err
	}

	snap := containerdsnapshot.NewSnapshotter(snapshotterName, client.SnapshotService(snapshotterName), ns, nil)

	if err := cache.MigrateV2(context.TODO(), filepath.Join(root, "metadata.db"), filepath.Join(root, "metadata_v2.db"), cs, snap, lm); err != nil {
		return base.WorkerOpt{}, err
	}

	md, err := metadata.NewStore(filepath.Join(root, "metadata_v2.db"))
	if err != nil {
		return base.WorkerOpt{}, err
	}

	opt := base.WorkerOpt{
		ID:             id,
		Labels:         xlabels,
		MetadataStore:  md,
		Executor:       containerdexecutor.New(client, root, "", np, dns),
		Snapshotter:    snap,
		ContentStore:   cs,
		Applier:        winlayers.NewFileSystemApplierWithWindows(cs, df),
		Differ:         winlayers.NewWalkingDiffWithWindows(cs, df),
		ImageStore:     client.ImageService(),
		Platforms:      platforms,
		LeaseManager:   lm,
		GarbageCollect: gc,
	}
	return opt, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"regexp"
	"strings"
	"sync"
)

// ScalewayCache is used not to query the API to resolve full identifiers
type ScalewayCache struct {
	// Images contains names of Scaleway images indexed by identifier
	Images map[string]string `json:"images"`

	// Snapshots contains names of Scaleway snapshots indexed by identifier
	Snapshots map[string]string `json:"snapshots"`

	// Bootscripts contains names of Scaleway bootscripts indexed by identifier
	Bootscripts map[string]string `json:"bootscripts"`

	// Servers contains names of Scaleway C1 servers indexed by identifier
	Servers map[string]string `json:"servers"`

	// Path is the path to the cache file
	Path string `json:"-"`

	// Modified tells if the cache needs to be overwritten or not
	Modified bool `json:"-"`

	// Lock allows ScalewayCache to be used concurrently
	Lock sync.Mutex `json:"-"`
}

const (
	IDENTIFIER_SERVER = iota
	IDENTIFIER_IMAGE
	IDENTIFIER_SNAPSHOT
	IDENTIFIER_BOOTSCRIPT
)

// ScalewayIdentifier is a unique identifier on Scaleway
type ScalewayIdentifier struct {
	// Identifier is a unique identifier on
	Identifier string

	// Type of the identifier
	Type int
}

// NewScalewayCache loads a per-user cache
func NewScalewayCache() (*ScalewayCache, error) {
	u, err := user.Current()
	if err != nil {
		return nil, err
	}
	cache_path := fmt.Sprintf("%s/.scw-cache.db", u.HomeDir)
	_, err = os.Stat(cache_path)
	if os.IsNotExist(err) {
		return &ScalewayCache{
			Images:      make(map[string]string),
			Snapshots:   make(map[string]string),
			Bootscripts: make(map[string]string),
			Servers:     make(map[string]string),
			Path:        cache_path,
		}, nil
	} else if err != nil {
		return nil, err
	}
	file, err := ioutil.ReadFile(cache_path)
	if err != nil {
		return nil, err
	}
	var cache ScalewayCache
	cache.Path = cache_path
	err = json.Unmarshal(file, &cache)
	if err != nil {
		return nil, err
	}
	if cache.Images == nil {
		cache.Images = make(map[string]string)
	}
	if cache.Snapshots == nil {
		cache.Snapshots = make(map[string]string)
	}
	if cache.Servers == nil {
		cache.Servers = make(map[string]string)
	}
	if cache.Bootscripts == nil {
		cache.Bootscripts = make(map[string]string)
	}
	return &cache, nil
}

// Save atomically overwrites the current cache database
func (c *ScalewayCache) Save() error {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	if c.Modified {
		file, err := ioutil.TempFile("", "")
		if err != nil {
			return err
		}
		encoder := json.NewEncoder(file)
		err = encoder.Encode(*c)
		if err != nil {
			return err
		}
		return os.Rename(file.Name(), c.Path)
	}
	return nil
}

// LookupImages attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpImages(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	needle = regexp.MustCompile(`^user/`).ReplaceAllString(needle, "")
	// FIXME: if 'user/' is in needle, only watch for a user image
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Images {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupSnapshots attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpSnapshots(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	needle = regexp.MustCompile(`^user/`).ReplaceAllString(needle, "")
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Snapshots {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupBootscripts attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpBootscripts(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Bootscripts {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupServers attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpServers(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Servers {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupIdentifier attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpIdentifiers(needle string) []ScalewayIdentifier {
	result := []ScalewayIdentifier{}

	for _, identifier := range c.LookUpServers(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_SERVER,
		})
	}

	for _, identifier := range c.LookUpImages(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_IMAGE,
		})
	}

	for _, identifier := range c.LookUpSnapshots(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_SNAPSHOT,
		})
	}

	for _, identifier := range c.LookUpBootscripts(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_BOOTSCRIPT,
		})
	}

	return result
}

// InsertServer registers a server in the cache
func (c *ScalewayCache) InsertServer(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Servers[identifier]
	if !exists || current_name != name {
		c.Servers[identifier] = name
		c.Modified = true
	}
}

// InsertImage registers an image in the cache
func (c *ScalewayCache) InsertImage(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Images[identifier]
	if !exists || current_name != name {
		c.Images[identifier] = name
		c.Modified = true
	}
}

// InsertSnapshot registers an snapshot in the cache
func (c *ScalewayCache) InsertSnapshot(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Snapshots[identifier]
	if !exists || current_name != name {
		c.Snapshots[identifier] = name
		c.Modified = true
	}
}

// InsertBootscript registers an bootscript in the cache
func (c *ScalewayCache) InsertBootscript(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Bootscripts[identifier]
	if !exists || current_name != name {
		c.Bootscripts[identifier] = name
		c.Modified = true
	}
}

// GetNbServers returns the number of servers in the cache
func (c *ScalewayCache) GetNbServers() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Servers)
}

// GetNbImages returns the number of images in the cache
func (c *ScalewayCache) GetNbImages() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Images)
}

// GetNbSnapshots returns the number of snapshots in the cache
func (c *ScalewayCache) GetNbSnapshots() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Snapshots)
}

// GetNbBootscripts returns the number of bootscripts in the cache
func (c *ScalewayCache) GetNbBootscripts() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Bootscripts)
}
<commit_msg>Added helpers to remove entries from the cache.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"regexp"
	"strings"
	"sync"
)

// ScalewayCache is used not to query the API to resolve full identifiers
type ScalewayCache struct {
	// Images contains names of Scaleway images indexed by identifier
	Images map[string]string `json:"images"`

	// Snapshots contains names of Scaleway snapshots indexed by identifier
	Snapshots map[string]string `json:"snapshots"`

	// Bootscripts contains names of Scaleway bootscripts indexed by identifier
	Bootscripts map[string]string `json:"bootscripts"`

	// Servers contains names of Scaleway C1 servers indexed by identifier
	Servers map[string]string `json:"servers"`

	// Path is the path to the cache file
	Path string `json:"-"`

	// Modified tells if the cache needs to be overwritten or not
	Modified bool `json:"-"`

	// Lock allows ScalewayCache to be used concurrently
	Lock sync.Mutex `json:"-"`
}

const (
	IDENTIFIER_SERVER = iota
	IDENTIFIER_IMAGE
	IDENTIFIER_SNAPSHOT
	IDENTIFIER_BOOTSCRIPT
)

// ScalewayIdentifier is a unique identifier on Scaleway
type ScalewayIdentifier struct {
	// Identifier is a unique identifier on
	Identifier string

	// Type of the identifier
	Type int
}

// NewScalewayCache loads a per-user cache
func NewScalewayCache() (*ScalewayCache, error) {
	u, err := user.Current()
	if err != nil {
		return nil, err
	}
	cache_path := fmt.Sprintf("%s/.scw-cache.db", u.HomeDir)
	_, err = os.Stat(cache_path)
	if os.IsNotExist(err) {
		return &ScalewayCache{
			Images:      make(map[string]string),
			Snapshots:   make(map[string]string),
			Bootscripts: make(map[string]string),
			Servers:     make(map[string]string),
			Path:        cache_path,
		}, nil
	} else if err != nil {
		return nil, err
	}
	file, err := ioutil.ReadFile(cache_path)
	if err != nil {
		return nil, err
	}
	var cache ScalewayCache
	cache.Path = cache_path
	err = json.Unmarshal(file, &cache)
	if err != nil {
		return nil, err
	}
	if cache.Images == nil {
		cache.Images = make(map[string]string)
	}
	if cache.Snapshots == nil {
		cache.Snapshots = make(map[string]string)
	}
	if cache.Servers == nil {
		cache.Servers = make(map[string]string)
	}
	if cache.Bootscripts == nil {
		cache.Bootscripts = make(map[string]string)
	}
	return &cache, nil
}

// Save atomically overwrites the current cache database
func (c *ScalewayCache) Save() error {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	if c.Modified {
		file, err := ioutil.TempFile("", "")
		if err != nil {
			return err
		}
		encoder := json.NewEncoder(file)
		err = encoder.Encode(*c)
		if err != nil {
			return err
		}
		return os.Rename(file.Name(), c.Path)
	}
	return nil
}

// LookupImages attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpImages(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	needle = regexp.MustCompile(`^user/`).ReplaceAllString(needle, "")
	// FIXME: if 'user/' is in needle, only watch for a user image
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Images {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupSnapshots attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpSnapshots(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	needle = regexp.MustCompile(`^user/`).ReplaceAllString(needle, "")
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Snapshots {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupBootscripts attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpBootscripts(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Bootscripts {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupServers attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpServers(needle string) []string {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	var res []string
	nameRegex := regexp.MustCompile(`(?i)` + regexp.MustCompile(`[_-]`).ReplaceAllString(needle, ".*"))
	for identifier, name := range c.Servers {
		if strings.HasPrefix(identifier, needle) || nameRegex.MatchString(name) {
			res = append(res, identifier)
		}
	}
	return res
}

// LookupIdentifier attempts to return identifiers matching a pattern
func (c *ScalewayCache) LookUpIdentifiers(needle string) []ScalewayIdentifier {
	result := []ScalewayIdentifier{}

	for _, identifier := range c.LookUpServers(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_SERVER,
		})
	}

	for _, identifier := range c.LookUpImages(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_IMAGE,
		})
	}

	for _, identifier := range c.LookUpSnapshots(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_SNAPSHOT,
		})
	}

	for _, identifier := range c.LookUpBootscripts(needle) {
		result = append(result, ScalewayIdentifier{
			Identifier: identifier,
			Type:       IDENTIFIER_BOOTSCRIPT,
		})
	}

	return result
}

// InsertServer registers a server in the cache
func (c *ScalewayCache) InsertServer(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Servers[identifier]
	if !exists || current_name != name {
		c.Servers[identifier] = name
		c.Modified = true
	}
}

// RemoveServer removes a server from the cache
func (c *ScalewayCache) RemoveServer(identifier string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	delete(c.Servers, identifier)
}

// InsertImage registers an image in the cache
func (c *ScalewayCache) InsertImage(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Images[identifier]
	if !exists || current_name != name {
		c.Images[identifier] = name
		c.Modified = true
	}
}

// RemoveImage removes a server from the cache
func (c *ScalewayCache) RemoveImage(identifier string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	delete(c.Images, identifier)
}

// InsertSnapshot registers an snapshot in the cache
func (c *ScalewayCache) InsertSnapshot(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Snapshots[identifier]
	if !exists || current_name != name {
		c.Snapshots[identifier] = name
		c.Modified = true
	}
}

// RemoveSnapshot removes a server from the cache
func (c *ScalewayCache) RemoveSnapshot(identifier string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	delete(c.Snapshots, identifier)
}

// InsertBootscript registers an bootscript in the cache
func (c *ScalewayCache) InsertBootscript(identifier, name string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	current_name, exists := c.Bootscripts[identifier]
	if !exists || current_name != name {
		c.Bootscripts[identifier] = name
		c.Modified = true
	}
}

// RemoveBootscript removes a bootscript from the cache
func (c *ScalewayCache) RemoveBootscript(identifier string) {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	delete(c.Bootscripts, identifier)
}

// GetNbServers returns the number of servers in the cache
func (c *ScalewayCache) GetNbServers() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Servers)
}

// GetNbImages returns the number of images in the cache
func (c *ScalewayCache) GetNbImages() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Images)
}

// GetNbSnapshots returns the number of snapshots in the cache
func (c *ScalewayCache) GetNbSnapshots() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Snapshots)
}

// GetNbBootscripts returns the number of bootscripts in the cache
func (c *ScalewayCache) GetNbBootscripts() int {
	c.Lock.Lock()
	defer c.Lock.Unlock()

	return len(c.Bootscripts)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/garyburd/redigo/redis"
)

const prefix = "genproxy:"

type Cache struct {
	pool *redis.Pool
}

// set records a key in cache (redis)
func (c *Cache) set(key string, value []byte) error {

	client := c.pool.Get()
	defer client.Close()

	_, err := client.Do("SET", fmt.Sprintf(prefix+key), value)

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Failed to record request...")
	} else {
		log.WithFields(log.Fields{
			"key": fmt.Sprintf(prefix + key),
		}).Info("Request recorded!")
	}

	return err
}

// getRedisPool returns thread safe Redis connection pool
func getRedisPool() *redis.Pool {

	// getting redis connection
	maxConnections := 10
	mc := os.Getenv("MaxConnections")
	if mc != "" {
		maxCons, err := strconv.Atoi(mc)
		if err != nil {
			maxConnections = 10
		} else {
			maxConnections = maxCons
		}
	}
	// getting redis client for state storing
	redisPool := redis.NewPool(func() (redis.Conn, error) {
		c, err := redis.Dial("tcp", AppConfig.redisAddress)

		if err != nil {
			log.WithFields(log.Fields{"Error": err.Error()}).Panic("Failed to create Redis connection pool!")
			return nil, err
		}
		if AppConfig.redisPassword != "" {
			if _, err := c.Do("AUTH", AppConfig.redisPassword); err != nil {
				log.WithFields(log.Fields{
					"Error":        err.Error(),
					"PasswordUsed": AppConfig.redisPassword,
				}).Panic("Failed to authenticate to Redis!")
				c.Close()
				return nil, err
			} else {
				log.Info("Authenticated to Redis successfully! ")
			}
		}

		return c, err
	}, maxConnections)

	return redisPool
}
<commit_msg>cache get functionality<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"

	log "github.com/Sirupsen/logrus"
	"github.com/garyburd/redigo/redis"
)

const prefix = "genproxy:"

type Cache struct {
	pool *redis.Pool
}

// set records a key in cache (redis)
func (c *Cache) set(key string, value []byte) error {

	client := c.pool.Get()
	defer client.Close()

	_, err := client.Do("SET", fmt.Sprintf(prefix+key), value)

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Failed to record request...")
	} else {
		log.WithFields(log.Fields{
			"key": fmt.Sprintf(prefix + key),
		}).Info("Request recorded!")
	}

	return err
}

// get returns key from cache
func (c *Cache) get(key string) (interface{}, error) {

	client := c.pool.Get()
	defer client.Close()

	value, err := client.Do("GET", fmt.Sprintf(prefix+key))

	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
			"key":   fmt.Sprintf(prefix + key),
		}).Error("Failed to GET key...")
	} else {
		log.WithFields(log.Fields{
			"key": fmt.Sprintf(prefix + key),
		}).Info("Key found!")
	}

	return value, err
}

// getRedisPool returns thread safe Redis connection pool
func getRedisPool() *redis.Pool {

	// getting redis connection
	maxConnections := 10
	mc := os.Getenv("MaxConnections")
	if mc != "" {
		maxCons, err := strconv.Atoi(mc)
		if err != nil {
			maxConnections = 10
		} else {
			maxConnections = maxCons
		}
	}
	// getting redis client for state storing
	redisPool := redis.NewPool(func() (redis.Conn, error) {
		c, err := redis.Dial("tcp", AppConfig.redisAddress)

		if err != nil {
			log.WithFields(log.Fields{"Error": err.Error()}).Panic("Failed to create Redis connection pool!")
			return nil, err
		}
		if AppConfig.redisPassword != "" {
			if _, err := c.Do("AUTH", AppConfig.redisPassword); err != nil {
				log.WithFields(log.Fields{
					"Error":        err.Error(),
					"PasswordUsed": AppConfig.redisPassword,
				}).Panic("Failed to authenticate to Redis!")
				c.Close()
				return nil, err
			} else {
				log.Info("Authenticated to Redis successfully! ")
			}
		}

		return c, err
	}, maxConnections)

	return redisPool
}
<|endoftext|>


<commit_before>package cluster

import (
	"io/ioutil"

	"github.com/cosminrentea/gobbler/server/store/filestore"

	"github.com/cosminrentea/gobbler/protocol"
	"github.com/cosminrentea/gobbler/server/store"

	"github.com/hashicorp/go-multierror"
	"github.com/stretchr/testify/assert"

	"errors"
	"net"
	"testing"
	"time"
)

const basePort = 10000

var (
	index = 1
)

func testConfig() (config Config) {
	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	config = Config{ID: uint8(index), Host: "127.0.0.1", Port: basePort + index, Remotes: remotes}
	index++
	return
}

func testConfigAnother() (config Config) {
	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index - 1}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	config = Config{ID: uint8(index), Host: "127.0.0.1", Port: basePort + index, Remotes: remotes}
	index++
	return
}

func TestCluster_StartCheckStop(t *testing.T) {
	a := assert.New(t)

	conf := testConfig()
	node, err := New(&conf)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	err = node.Check()
	a.NoError(err, "Health-check score of a Cluster with a single node should be OK")

	err = node.Stop()
	a.NoError(err, "No error should be raised when Stopping the Cluster")
}

func TestCluster_BroadcastStringAndMessageAndCheck(t *testing.T) {
	a := assert.New(t)

	config1 := testConfig()
	node1, err := New(&config1)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node1.Router = newDummyRouter(t)

	//start the cluster node 1
	defer node1.Stop()
	err = node1.Start()
	a.NoError(err, "No error should be raised when starting node 1 of the Cluster")

	config2 := testConfigAnother()
	node2, err := New(&config2)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node2.Router = newDummyRouter(t)

	//start the cluster node 2
	defer node2.Stop()
	err = node2.Start()
	a.NoError(err, "No error should be raised when starting node 2 of the Cluster")

	// Send a String Message
	str := "TEST"
	err = node1.BroadcastString(&str)
	a.NoError(err, "No error should be raised when sending a string to Cluster")

	// and a protocol message
	pmsg := protocol.Message{
		ID:            1,
		Path:          "/stuff",
		UserID:        "id",
		ApplicationID: "appId",
		Time:          time.Now().Unix(),
		HeaderJSON:    "{}",
		Body:          []byte("test"),
		NodeID:        1}
	err = node1.BroadcastMessage(&pmsg)
	a.NoError(err, "No error should be raised when sending a protocol message to Cluster")

	err = node1.Check()
	a.NoError(err, "Health-check score of a Cluster with 2 nodes should be OK for node 1")

	err = node2.Check()
	a.NoError(err, "Health-check score of a Cluster with 2 nodes should be OK for node 2")
}

func TestCluster_NewShouldReturnErrorWhenPortIsInvalid(t *testing.T) {
	a := assert.New(t)

	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index - 1}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	index++

	config := Config{ID: 1, Host: "localhost", Port: -1, Remotes: remotes}
	_, err := New(&config)
	if a.Error(err, "An error was expected when Creating the Cluster") {
		a.Equal(err, errors.New("Failed to start TCP listener. Err: listen tcp :-1: bind: invalid argument"),
			"Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenNoRemotes(t *testing.T) {
	a := assert.New(t)

	var remotes []*net.TCPAddr
	index++

	config := Config{ID: 1, Host: "localhost", Port: basePort + index - 1, Remotes: remotes}
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		a.Equal(err, errors.New("No remote hosts were successfully contacted when this node wanted to join the cluster"),
			"Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenInvalidRemotes(t *testing.T) {
	a := assert.New(t)

	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: 0}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	index++

	config := Config{ID: 1, Host: "localhost", Port: basePort + index - 1, Remotes: remotes}
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		expected := multierror.Append(errors.New("Failed to join 127.0.0.1: dial tcp 127.0.0.1:0: getsockopt: connection refused"))
		a.Equal(err, expected, "Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenNoMessageHandler(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		expected := errors.New("There should be a valid Router already set-up")
		a.Equal(expected, err, "Error should be precisely defined")
	}
}

func TestCluster_NotifyMsgShouldSimplyReturnWhenDecodingInvalidMessage(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	node.NotifyMsg([]byte{})

	//TODO Cosmin check that HandleMessage is not invoked (i.e. invalid message is not dispatched)
}

func TestCluster_broadcastClusterMessage(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	err = node.broadcastClusterMessage(nil)
	if a.Error(err, "An error is expected from broadcastClusterMessage") {
		expected := errors.New("Could not broadcast a nil cluster-message")
		a.Equal(err, expected, "Error should be precisely defined")
	}
}

type dummyRouter struct {
	store store.MessageStore
}

func newDummyRouter(t *testing.T) *dummyRouter {
	dir, err := ioutil.TempDir("", "guble_cluster_test")
	assert.NoError(t, err)
	return &dummyRouter{store: filestore.New(dir)}
}

func (_ *dummyRouter) HandleMessage(pmsg *protocol.Message) error {
	return nil
}

func (d *dummyRouter) MessageStore() (store.MessageStore, error) {
	return d.store, nil
}
<commit_msg>minor<commit_after>package cluster

import (
	"io/ioutil"

	"github.com/cosminrentea/gobbler/server/store/filestore"

	"github.com/cosminrentea/gobbler/protocol"
	"github.com/cosminrentea/gobbler/server/store"

	"github.com/hashicorp/go-multierror"
	"github.com/stretchr/testify/assert"

	"errors"
	"net"
	"testing"
	"time"
)

const basePort = 11000

var (
	index = 1
)

func testConfig() (config Config) {
	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	config = Config{ID: uint8(index), Host: "127.0.0.1", Port: basePort + index, Remotes: remotes}
	index++
	return
}

func testConfigAnother() (config Config) {
	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index - 1}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	config = Config{ID: uint8(index), Host: "127.0.0.1", Port: basePort + index, Remotes: remotes}
	index++
	return
}

func TestCluster_StartCheckStop(t *testing.T) {
	a := assert.New(t)

	conf := testConfig()
	node, err := New(&conf)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	err = node.Check()
	a.NoError(err, "Health-check score of a Cluster with a single node should be OK")

	err = node.Stop()
	a.NoError(err, "No error should be raised when Stopping the Cluster")
}

func TestCluster_BroadcastStringAndMessageAndCheck(t *testing.T) {
	a := assert.New(t)

	config1 := testConfig()
	node1, err := New(&config1)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node1.Router = newDummyRouter(t)

	//start the cluster node 1
	defer node1.Stop()
	err = node1.Start()
	a.NoError(err, "No error should be raised when starting node 1 of the Cluster")

	config2 := testConfigAnother()
	node2, err := New(&config2)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node2.Router = newDummyRouter(t)

	//start the cluster node 2
	defer node2.Stop()
	err = node2.Start()
	a.NoError(err, "No error should be raised when starting node 2 of the Cluster")

	// Send a String Message
	str := "TEST"
	err = node1.BroadcastString(&str)
	a.NoError(err, "No error should be raised when sending a string to Cluster")

	// and a protocol message
	pmsg := protocol.Message{
		ID:            1,
		Path:          "/stuff",
		UserID:        "id",
		ApplicationID: "appId",
		Time:          time.Now().Unix(),
		HeaderJSON:    "{}",
		Body:          []byte("test"),
		NodeID:        1}
	err = node1.BroadcastMessage(&pmsg)
	a.NoError(err, "No error should be raised when sending a protocol message to Cluster")

	err = node1.Check()
	a.NoError(err, "Health-check score of a Cluster with 2 nodes should be OK for node 1")

	err = node2.Check()
	a.NoError(err, "Health-check score of a Cluster with 2 nodes should be OK for node 2")
}

func TestCluster_NewShouldReturnErrorWhenPortIsInvalid(t *testing.T) {
	a := assert.New(t)

	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: basePort + index - 1}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	index++

	config := Config{ID: 1, Host: "localhost", Port: -1, Remotes: remotes}
	_, err := New(&config)
	if a.Error(err, "An error was expected when Creating the Cluster") {
		a.Equal(err, errors.New("Failed to start TCP listener. Err: listen tcp :-1: bind: invalid argument"),
			"Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenNoRemotes(t *testing.T) {
	a := assert.New(t)

	var remotes []*net.TCPAddr
	index++

	config := Config{ID: 1, Host: "localhost", Port: basePort + index - 1, Remotes: remotes}
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		a.Equal(err, errors.New("No remote hosts were successfully contacted when this node wanted to join the cluster"),
			"Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenInvalidRemotes(t *testing.T) {
	a := assert.New(t)

	remoteAddr := net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: 0}
	var remotes []*net.TCPAddr
	remotes = append(remotes, &remoteAddr)
	index++

	config := Config{ID: 1, Host: "localhost", Port: basePort + index - 1, Remotes: remotes}
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		expected := multierror.Append(errors.New("Failed to join 127.0.0.1: dial tcp 127.0.0.1:0: getsockopt: connection refused"))
		a.Equal(err, expected, "Error should be precisely defined")
	}
}

func TestCluster_StartShouldReturnErrorWhenNoMessageHandler(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	defer node.Stop()
	err = node.Start()
	if a.Error(err, "An error is expected when Starting the Cluster") {
		expected := errors.New("There should be a valid Router already set-up")
		a.Equal(expected, err, "Error should be precisely defined")
	}
}

func TestCluster_NotifyMsgShouldSimplyReturnWhenDecodingInvalidMessage(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	node.NotifyMsg([]byte{})

	//TODO Cosmin check that HandleMessage is not invoked (i.e. invalid message is not dispatched)
}

func TestCluster_broadcastClusterMessage(t *testing.T) {
	a := assert.New(t)

	config := testConfig()
	node, err := New(&config)
	a.NoError(err, "No error should be raised when Creating the Cluster")

	node.Router = newDummyRouter(t)

	defer node.Stop()
	err = node.Start()
	a.NoError(err, "No error should be raised when Starting the Cluster")

	err = node.broadcastClusterMessage(nil)
	if a.Error(err, "An error is expected from broadcastClusterMessage") {
		expected := errors.New("Could not broadcast a nil cluster-message")
		a.Equal(err, expected, "Error should be precisely defined")
	}
}

type dummyRouter struct {
	store store.MessageStore
}

func newDummyRouter(t *testing.T) *dummyRouter {
	dir, err := ioutil.TempDir("", "guble_cluster_test")
	assert.NoError(t, err)
	return &dummyRouter{store: filestore.New(dir)}
}

func (_ *dummyRouter) HandleMessage(pmsg *protocol.Message) error {
	return nil
}

func (d *dummyRouter) MessageStore() (store.MessageStore, error) {
	return d.store, nil
}
<|endoftext|>


<commit_before>package base58

import (
	"encoding/hex"
	"reflect"
	"testing"
)

type encodeTestCase struct {
	in  string
	out string
	err error
}

type decodeTestCase struct {
	in  string
	out string
	err error
}

var (
	encodeTestCasesBitcoin = []encodeTestCase{
		{"", "", nil},
		{"00", "1", nil},
		{"00010966776006953d5567439e5e39f86a0d273beed61967f6", "16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM", nil},
	}
	decodeTestCasesBitcoin = []decodeTestCase{
		{"", "", nil},
		{"1", "00", nil},
		{"16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM", "00010966776006953d5567439e5e39f86a0d273beed61967f6", nil},
	}
)

func TestEncodeToString(t *testing.T) {
	for _, testCase := range encodeTestCasesBitcoin {
		inBytes, err := hex.DecodeString(testCase.in)
		if err != nil {
			t.Fatal("can not decode input string")
		}

		b58 := NewBitcoinBase58()

		out, err := b58.EncodeToString(inBytes)
		if !reflect.DeepEqual(err, testCase.err) {
			t.Errorf(`
invalid error
- got: %v
- expected: %v
`, err, testCase.err)
		}
		if out != testCase.out {
			t.Errorf(`
invalid encoding
- got: %s
- expected: %s
`, out, testCase.out)
		}
	}
}

func TestDecodeString(t *testing.T) {
	for _, testCase := range decodeTestCasesBitcoin {
		b58 := NewBitcoinBase58()

		outBytes, err := b58.DecodeString(testCase.in)
		if !reflect.DeepEqual(err, testCase.err) {
			t.Errorf(`
invalid error
- got: %v
- expected: %v
`, err, testCase.err)
		}

		out := hex.EncodeToString(outBytes)
		if out != testCase.out {
			t.Errorf(`
invalid decoding
- got: %s
- expected: %s
`, out, testCase.out)
		}
	}
}
<commit_msg>add test case for invalid char<commit_after>package base58

import (
	"encoding/hex"
	"reflect"
	"testing"
)

type encodeTestCase struct {
	in  string
	out string
	err error
}

type decodeTestCase struct {
	in  string
	out string
	err error
}

var (
	encodeTestCasesBitcoin = []encodeTestCase{
		{"", "", nil},
		{"00", "1", nil},
		{"00010966776006953d5567439e5e39f86a0d273beed61967f6", "16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM", nil},
	}
	decodeTestCasesBitcoin = []decodeTestCase{
		{"", "", nil},
		{"0", "", ErrInvalidChar},
		{"1", "00", nil},
		{"16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM", "00010966776006953d5567439e5e39f86a0d273beed61967f6", nil},
	}
)

func TestEncodeToString(t *testing.T) {
	for _, testCase := range encodeTestCasesBitcoin {
		inBytes, err := hex.DecodeString(testCase.in)
		if err != nil {
			t.Fatal("can not decode input string")
		}

		b58 := NewBitcoinBase58()

		out, err := b58.EncodeToString(inBytes)
		if !reflect.DeepEqual(err, testCase.err) {
			t.Errorf(`
invalid error
- got: %v
- expected: %v
`, err, testCase.err)
		}
		if out != testCase.out {
			t.Errorf(`
invalid encoding
- got: %s
- expected: %s
`, out, testCase.out)
		}
	}
}

func TestDecodeString(t *testing.T) {
	for _, testCase := range decodeTestCasesBitcoin {
		b58 := NewBitcoinBase58()

		outBytes, err := b58.DecodeString(testCase.in)
		if !reflect.DeepEqual(err, testCase.err) {
			t.Errorf(`
invalid error
- got: %v
- expected: %v
`, err, testCase.err)
		}

		out := hex.EncodeToString(outBytes)
		if out != testCase.out {
			t.Errorf(`
invalid decoding
- got: %s
- expected: %s
`, out, testCase.out)
		}
	}
}
<|endoftext|>


<commit_before>package linda

import (
	"github.com/sirupsen/logrus"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"
)

var (
	brokerConn  Broker
	initMutex   sync.Mutex
	initialized bool
	settings    Settings
)

// SetSettings
// set linda settings
func SetSettings(sets Settings) {
	settings = sets
}

// Init linda with settings
// get instance of broker
func Init() error {
	initMutex.Lock()
	defer initMutex.Unlock()
	if !initialized {
		logrus.Debug("start linda and init...")
		if err := flags(); err != nil {
			return err
		}
		logrus.Debugf("init the scheduler with redis connection string %s", settings.Connection)
		b, err := NewBroker(settings.Connection)
		if err != nil {
			return err
		}
		brokerConn = b
		initialized = true
	}
	return nil
}

// Close linda with close broker
func Close() {
	initMutex.Lock()
	defer initMutex.Unlock()
	if initialized {
		logrus.Debug("close the linda...")
		brokerConn.Close()
		initialized = false
	}
}

// GetBroker will return the initialized broker of linda
func GetBroker() Broker {
	if initialized {
		return brokerConn
	}
	return nil
}

// Run linda, it also call init function self
func Run() error {
	err := Init()
	if err != nil {
		return err
	}
	defer Close()
	logrus.Debug("start the poller to get jobs and migrate expired jobs")
	poller, err := newPoller()
	if err != nil {
		return err
	}
	quit := signals()
	jobs := poller.poll(settings.Queue, settings.Timeout, time.Duration(settings.Interval), quit)
	var monitor sync.WaitGroup
	logrus.Debugf("start %d workers to do the job", settings.Concurrency)
	for i := 0; i < settings.Concurrency; i++ {
		worker, err := newWorker(strconv.Itoa(i))
		if err != nil {
			return err
		}
		worker.work(jobs, &monitor)
	}
	monitor.Wait()
	return nil
}

// Signal Handling
func signals() <-chan bool {
	quit := make(chan bool)
	go func() {
		signals := make(chan os.Signal)
		defer close(signals)
		signal.Notify(signals, syscall.SIGQUIT, syscall.SIGTERM, os.Interrupt)
		defer signalStop(signals)
		<-signals
		quit <- true
	}()
	return quit
}

// Stops signals channel.
func signalStop(c chan<- os.Signal) {
	signal.Stop(c)
}
<commit_msg>go lint<commit_after>package linda

import (
	"github.com/sirupsen/logrus"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"
)

var (
	brokerConn  Broker
	initMutex   sync.Mutex
	initialized bool
	settings    Settings
)

// SetSettings set linda settings
func SetSettings(sets Settings) {
	settings = sets
}

// Init linda with settings
// get instance of broker
func Init() error {
	initMutex.Lock()
	defer initMutex.Unlock()
	if !initialized {
		logrus.Debug("start linda and init...")
		if err := flags(); err != nil {
			return err
		}
		logrus.Debugf("init the scheduler with redis connection string %s", settings.Connection)
		b, err := NewBroker(settings.Connection)
		if err != nil {
			return err
		}
		brokerConn = b
		initialized = true
	}
	return nil
}

// Close linda with close broker
func Close() {
	initMutex.Lock()
	defer initMutex.Unlock()
	if initialized {
		logrus.Debug("close the linda...")
		brokerConn.Close()
		initialized = false
	}
}

// GetBroker will return the initialized broker of linda
func GetBroker() Broker {
	if initialized {
		return brokerConn
	}
	return nil
}

// Run linda, it also call init function self
func Run() error {
	err := Init()
	if err != nil {
		return err
	}
	defer Close()
	logrus.Debug("start the poller to get jobs and migrate expired jobs")
	poller, err := newPoller()
	if err != nil {
		return err
	}
	quit := signals()
	jobs := poller.poll(settings.Queue, settings.Timeout, time.Duration(settings.Interval), quit)
	var monitor sync.WaitGroup
	logrus.Debugf("start %d workers to do the job", settings.Concurrency)
	for i := 0; i < settings.Concurrency; i++ {
		worker, err := newWorker(strconv.Itoa(i))
		if err != nil {
			return err
		}
		worker.work(jobs, &monitor)
	}
	monitor.Wait()
	return nil
}

// Signal Handling
func signals() <-chan bool {
	quit := make(chan bool)
	go func() {
		signals := make(chan os.Signal)
		defer close(signals)
		signal.Notify(signals, syscall.SIGQUIT, syscall.SIGTERM, os.Interrupt)
		defer signalStop(signals)
		<-signals
		quit <- true
	}()
	return quit
}

// Stops signals channel.
func signalStop(c chan<- os.Signal) {
	signal.Stop(c)
}
<|endoftext|>