<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/wandoulabs/codis/pkg/models"
	"github.com/wandoulabs/codis/pkg/proxy/router"
	"github.com/wandoulabs/codis/pkg/utils"
	"github.com/wandoulabs/codis/pkg/utils/errors"
	"github.com/wandoulabs/codis/pkg/utils/log"
	"github.com/wandoulabs/codis/pkg/utils/rpc"
)

type Proxy struct {
	mu sync.Mutex

	token string
	xauth string
	model *models.Proxy

	jodis *Jodis

	init, exit struct {
		C chan struct{}
	}
	online bool
	closed bool

	lproxy net.Listener
	ladmin net.Listener

	config *Config
	router *router.Router
}

var ErrClosedProxy = errors.New("use of closed proxy")

func New(config *Config) (*Proxy, error) {
	if !utils.IsValidProduct(config.ProductName) {
		return nil, errors.Errorf("invalid product name = %s", config.ProductName)
	}
	s := &Proxy{config: config}
	s.token = rpc.NewToken()
	s.xauth = rpc.NewXAuth(config.ProductName, config.ProductAuth, s.token)
	s.model = &models.Proxy{
		Token: s.token, StartTime: time.Now().String(),
	}
	s.model.ProductName = config.ProductName
	s.model.Pid = os.Getpid()
	s.model.Pwd, _ = os.Getwd()
	if b, err := exec.Command("uname", "-a").Output(); err != nil {
		log.WarnErrorf(err, "run command uname failed")
	} else {
		s.model.Sys = strings.TrimSpace(string(b))
	}

	s.router = router.NewWithAuth(config.ProductAuth)
	s.init.C = make(chan struct{})
	s.exit.C = make(chan struct{})

	if err := s.setup(); err != nil {
		s.Close()
		return nil, err
	}

	log.Warnf("[%p] create new proxy:\n%s", s, s.model.Encode())

	go s.serveAdmin()
	go s.serveProxy()

	return s, nil
}

func (s *Proxy) setup() error {
	proto := s.config.ProtoType
	if l, err := net.Listen(proto, s.config.ProxyAddr); err != nil {
		return errors.Trace(err)
	} else {
		s.lproxy = l

		x, err := utils.ResolveAddr(proto, l.Addr().String(), s.config.HostProxy)
		if err != nil {
			return err
		}
		s.model.ProtoType = proto
		s.model.ProxyAddr = x
	}

	if l, err := net.Listen("tcp", s.config.AdminAddr); err != nil {
		return errors.Trace(err)
	} else {
		s.ladmin = l

		x, err := utils.ResolveAddr("tcp", l.Addr().String(), s.config.HostAdmin)
		if err != nil {
			return err
		}
		s.model.AdminAddr = x
	}

	return nil
}

func (s *Proxy) Start() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	if s.online {
		return nil
	}
	s.online = true
	close(s.init.C)

	if s.jodis == nil && s.config.JodisAddr != "" {
		s.jodis = NewJodis(s.config.JodisAddr, s.config.JodisTimeout, s.model)
	}
	return nil
}

func (s *Proxy) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	close(s.exit.C)

	if s.jodis != nil {
		s.jodis.Close()
	}
	if s.ladmin != nil {
		s.ladmin.Close()
	}
	if s.lproxy != nil {
		s.lproxy.Close()
	}
	if s.router != nil {
		s.router.Close()
	}
	return nil
}

func (s *Proxy) Token() string {
	return s.token
}

func (s *Proxy) XAuth() string {
	return s.xauth
}

func (s *Proxy) Model() *models.Proxy {
	return s.model
}

func (s *Proxy) Config() *Config {
	return s.config
}

func (s *Proxy) IsOnline() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.online && !s.closed
}

func (s *Proxy) IsClosed() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.closed
}

func (s *Proxy) Slots() []*models.Slot {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.router.GetSlots()
}

func (s *Proxy) FillSlot(i int, addr, from string, locked bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	return s.router.FillSlot(i, addr, from, locked)
}

func (s *Proxy) FillSlots(slots []*models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	for _, slot := range slots {
		i, locked := slot.Id, slot.Locked
		addr := slot.BackendAddr
		from := slot.MigrateFrom
		if err := s.router.FillSlot(i, addr, from, locked); err != nil {
			return err
		}
	}
	return nil
}

func (s *Proxy) serveAdmin() {
	if s.IsClosed() {
		return
	}
	defer s.Close()

	log.Warnf("[%p] admin start service on %s", s, s.ladmin.Addr())

	eh := make(chan error, 1)
	go func(l net.Listener) {
		h := http.NewServeMux()
		h.Handle("/", newApiServer(s))
		hs := &http.Server{Handler: h}
		eh <- hs.Serve(l)
	}(s.ladmin)

	select {
	case <-s.exit.C:
		log.Warnf("[%p] admin shutdown", s)
	case err := <-eh:
		log.ErrorErrorf(err, "[%p] admin exit on error", s)
	}
}

func (s *Proxy) serveProxy() {
	if s.IsClosed() {
		return
	}
	defer s.Close()

	select {
	case <-s.exit.C:
		return
	case <-s.init.C:
	}

	log.Warnf("[%p] proxy start service on %s", s, s.lproxy.Addr())

	ch := make(chan net.Conn, 4096)
	go func() {
		for c := range ch {
			s.newSession(c)
		}
	}()

	eh := make(chan error, 1)
	go func(l net.Listener) {
		defer close(ch)
		for {
			c, err := s.acceptConn(l)
			if err != nil {
				eh <- err
				return
			} else {
				ch <- c
			}
		}
	}(s.lproxy)

	go s.keepAlive()

	if s.jodis != nil {
		go s.jodis.Run()
	}

	select {
	case <-s.exit.C:
		log.Warnf("[%p] proxy shutdown", s)
	case err := <-eh:
		log.ErrorErrorf(err, "[%p] proxy exit on error", s)
	}
}

func (s *Proxy) keepAlive() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	var tick int
	for {
		select {
		case <-s.exit.C:
			return
		case <-ticker.C:
			if maxTick := s.config.BackendPingPeriod; maxTick != 0 {
				if tick++; tick >= maxTick {
					tick = 0
					s.router.KeepAlive()
				}
			}
		}
	}
}

func (s *Proxy) newSession(c net.Conn) {
	x := router.NewSessionSize(c, s.config.ProductAuth,
		s.config.SessionMaxBufSize, s.config.SessionMaxTimeout)
	x.SetKeepAlivePeriod(s.config.SessionKeepAlivePeriod)
	go x.Serve(s.router, s.config.SessionMaxPipeline)
}

func (s *Proxy) acceptConn(l net.Listener) (net.Conn, error) {
	var delay time.Duration
	for {
		c, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				log.WarnErrorf(err, "[%p] proxy accept new connection failed", s)
				delay = delay * 2
				delay = utils.MaxDuration(delay, time.Millisecond*10)
				delay = utils.MinDuration(delay, time.Millisecond*500)
				time.Sleep(delay)
				continue
			}
		}
		return c, err
	}
}
<commit_msg>Update, optimized code<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"net"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"github.com/wandoulabs/codis/pkg/models"
	"github.com/wandoulabs/codis/pkg/proxy/router"
	"github.com/wandoulabs/codis/pkg/utils"
	"github.com/wandoulabs/codis/pkg/utils/errors"
	"github.com/wandoulabs/codis/pkg/utils/log"
	"github.com/wandoulabs/codis/pkg/utils/rpc"
)

type Proxy struct {
	mu sync.Mutex

	token string
	xauth string
	model *models.Proxy

	jodis *Jodis

	init, exit struct {
		C chan struct{}
	}
	online bool
	closed bool

	lproxy net.Listener
	ladmin net.Listener

	config *Config
	router *router.Router
}

var ErrClosedProxy = errors.New("use of closed proxy")

func New(config *Config) (*Proxy, error) {
	if !utils.IsValidProduct(config.ProductName) {
		return nil, errors.Errorf("invalid product name = %s", config.ProductName)
	}
	s := &Proxy{config: config}
	s.token = rpc.NewToken()
	s.xauth = rpc.NewXAuth(config.ProductName, config.ProductAuth, s.token)
	s.model = &models.Proxy{
		Token: s.token, StartTime: time.Now().String(),
	}
	s.model.ProductName = config.ProductName
	s.model.Pid = os.Getpid()
	s.model.Pwd, _ = os.Getwd()
	if b, err := exec.Command("uname", "-a").Output(); err != nil {
		log.WarnErrorf(err, "run command uname failed")
	} else {
		s.model.Sys = strings.TrimSpace(string(b))
	}

	s.router = router.NewWithAuth(config.ProductAuth)
	s.init.C = make(chan struct{})
	s.exit.C = make(chan struct{})

	if err := s.setup(); err != nil {
		s.Close()
		return nil, err
	}

	log.Warnf("[%p] create new proxy:\n%s", s, s.model.Encode())

	go s.serveAdmin()
	go s.serveProxy()

	return s, nil
}

func (s *Proxy) setup() error {
	proto := s.config.ProtoType
	if l, err := net.Listen(proto, s.config.ProxyAddr); err != nil {
		return errors.Trace(err)
	} else {
		s.lproxy = l

		x, err := utils.ResolveAddr(proto, l.Addr().String(), s.config.HostProxy)
		if err != nil {
			return err
		}
		s.model.ProtoType = proto
		s.model.ProxyAddr = x
	}

	if l, err := net.Listen("tcp", s.config.AdminAddr); err != nil {
		return errors.Trace(err)
	} else {
		s.ladmin = l

		x, err := utils.ResolveAddr("tcp", l.Addr().String(), s.config.HostAdmin)
		if err != nil {
			return err
		}
		s.model.AdminAddr = x
	}

	return nil
}

func (s *Proxy) Start() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	if s.online {
		return nil
	}
	s.online = true
	close(s.init.C)

	if s.jodis == nil && s.config.JodisAddr != "" {
		s.jodis = NewJodis(s.config.JodisAddr, s.config.JodisTimeout, s.model)
	}
	return nil
}

func (s *Proxy) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	close(s.exit.C)

	if s.jodis != nil {
		s.jodis.Close()
	}
	if s.ladmin != nil {
		s.ladmin.Close()
	}
	if s.lproxy != nil {
		s.lproxy.Close()
	}
	if s.router != nil {
		s.router.Close()
	}
	return nil
}

func (s *Proxy) Token() string {
	return s.token
}

func (s *Proxy) XAuth() string {
	return s.xauth
}

func (s *Proxy) Model() *models.Proxy {
	return s.model
}

func (s *Proxy) Config() *Config {
	return s.config
}

func (s *Proxy) IsOnline() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.online && !s.closed
}

func (s *Proxy) IsClosed() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.closed
}

func (s *Proxy) Slots() []*models.Slot {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.router.GetSlots()
}

func (s *Proxy) FillSlot(i int, addr, from string, locked bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	return s.router.FillSlot(i, addr, from, locked)
}

func (s *Proxy) FillSlots(slots []*models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedProxy
	}
	for _, slot := range slots {
		i, locked := slot.Id, slot.Locked
		addr := slot.BackendAddr
		from := slot.MigrateFrom
		if err := s.router.FillSlot(i, addr, from, locked); err != nil {
			return err
		}
	}
	return nil
}

func (s *Proxy) serveAdmin() {
	if s.IsClosed() {
		return
	}
	defer s.Close()

	log.Warnf("[%p] admin start service on %s", s, s.ladmin.Addr())

	eh := make(chan error, 1)
	go func(l net.Listener) {
		h := http.NewServeMux()
		h.Handle("/", newApiServer(s))
		hs := &http.Server{Handler: h}
		eh <- hs.Serve(l)
	}(s.ladmin)

	select {
	case <-s.exit.C:
		log.Warnf("[%p] admin shutdown", s)
	case err := <-eh:
		log.ErrorErrorf(err, "[%p] admin exit on error", s)
	}
}

func (s *Proxy) serveProxy() {
	if s.IsClosed() {
		return
	}
	defer s.Close()

	select {
	case <-s.exit.C:
		return
	case <-s.init.C:
	}

	log.Warnf("[%p] proxy start service on %s", s, s.lproxy.Addr())

	ch := make(chan net.Conn, 4096)
	go func() {
		for c := range ch {
			s.newSession(c)
		}
	}()

	eh := make(chan error, 1)
	go func(l net.Listener) (err error) {
		defer func() {
			eh <- err
			close(ch)
		}()
		for {
			c, err := s.acceptConn(l)
			if err != nil {
				return err
			}
			ch <- c
		}
	}(s.lproxy)

	if seconds := s.config.BackendPingPeriod; seconds > 0 {
		go s.keepAlive(seconds)
	}

	if s.jodis != nil {
		go s.jodis.Run()
	}

	select {
	case <-s.exit.C:
		log.Warnf("[%p] proxy shutdown", s)
	case err := <-eh:
		log.ErrorErrorf(err, "[%p] proxy exit on error", s)
	}
}

func (s *Proxy) keepAlive(seconds int) {
	var ticker = time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		for i := 0; i < seconds; i++ {
			select {
			case <-s.exit.C:
				return
			case <-ticker.C:
			}
		}
		s.router.KeepAlive()
	}
}

func (s *Proxy) newSession(c net.Conn) {
	x := router.NewSessionSize(c, s.config.ProductAuth,
		s.config.SessionMaxBufSize, s.config.SessionMaxTimeout)
	x.SetKeepAlivePeriod(s.config.SessionKeepAlivePeriod)
	go x.Serve(s.router, s.config.SessionMaxPipeline)
}

func (s *Proxy) acceptConn(l net.Listener) (net.Conn, error) {
	var delay time.Duration
	for {
		c, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				log.WarnErrorf(err, "[%p] proxy accept new connection failed", s)
				delay = delay * 2
				delay = utils.MaxDuration(delay, time.Millisecond*10)
				delay = utils.MinDuration(delay, time.Millisecond*500)
				time.Sleep(delay)
				continue
			}
		}
		return c, err
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Sorint.lab
//
// 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 store

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"

	"github.com/docker/libkv"
	kvstore "github.com/docker/libkv/store"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
)

func init() {
	etcd.Register()
	consul.Register()
}

// Backend represents a KV Store Backend
type Backend string

const (
	CONSUL Backend = "consul"
	ETCD   Backend = "etcd"
)

const (
	keepersInfoDir         = "/keepers/info/"
	clusterDataFile        = "clusterdata"
	leaderSentinelInfoFile = "/sentinels/leaderinfo"
	sentinelsInfoDir       = "/sentinels/info/"
	proxiesInfoDir         = "/proxies/info/"
)

const (
	DefaultEtcdEndpoints   = "127.0.0.1:2379"
	DefaultConsulEndpoints = "127.0.0.1:8500"
)

const (
	//TODO(sgotti) fix this in libkv?
	// consul min ttl is 10s and libkv divides this by 2
	MinTTL = 20 * time.Second
)

type StoreManager struct {
	clusterPath string
	store       kvstore.Store
}

func NewStore(backend Backend, addrsStr string) (kvstore.Store, error) {

	var kvbackend kvstore.Backend
	switch backend {
	case CONSUL:
		kvbackend = kvstore.CONSUL
	case ETCD:
		kvbackend = kvstore.ETCD
	default:
		return nil, fmt.Errorf("Unknown store backend: %q", backend)
	}

	if addrsStr == "" {
		switch backend {
		case CONSUL:
			addrsStr = DefaultConsulEndpoints
		case ETCD:
			addrsStr = DefaultEtcdEndpoints
		}
	}
	addrs := strings.Split(addrsStr, ",")

	store, err := libkv.NewStore(kvbackend, addrs, &kvstore.Config{ConnectionTimeout: 10 * time.Second})
	if err != nil {
		return nil, err
	}
	return store, nil
}

func NewStoreManager(kvStore kvstore.Store, path string) *StoreManager {
	return &StoreManager{
		clusterPath: path,
		store:       kvStore,
	}
}

func (e *StoreManager) AtomicPutClusterData(cd *cluster.ClusterData, previous *kvstore.KVPair) (*kvstore.KVPair, error) {
	cdj, err := json.Marshal(cd)
	if err != nil {
		return nil, err
	}
	path := filepath.Join(e.clusterPath, clusterDataFile)
	_, pair, err := e.store.AtomicPut(path, cdj, previous, nil)
	return pair, err
}

func (e *StoreManager) PutClusterData(cd *cluster.ClusterData) error {
	cdj, err := json.Marshal(cd)
	if err != nil {
		return err
	}
	path := filepath.Join(e.clusterPath, clusterDataFile)
	return e.store.Put(path, cdj, nil)
}

func (e *StoreManager) GetClusterData() (*cluster.ClusterData, *kvstore.KVPair, error) {
	var cd *cluster.ClusterData
	path := filepath.Join(e.clusterPath, clusterDataFile)
	pair, err := e.store.Get(path)
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, nil, err
		}
		return nil, nil, nil
	}
	if err := json.Unmarshal(pair.Value, &cd); err != nil {
		return nil, nil, err
	}
	return cd, pair, nil
}

func (e *StoreManager) SetKeeperInfo(id string, ms *cluster.KeeperInfo, ttl time.Duration) error {
	msj, err := json.Marshal(ms)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, keepersInfoDir, id), msj, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetKeeperInfo(id string) (*cluster.KeeperInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty keeper id")
	}
	var keeper cluster.KeeperInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, keepersInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	if err := json.Unmarshal(pair.Value, &keeper); err != nil {
		return nil, false, err
	}
	return &keeper, true, nil
}

func (e *StoreManager) GetKeepersInfo() (cluster.KeepersInfo, error) {
	keepers := cluster.KeepersInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, keepersInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return keepers, nil
	}
	for _, pair := range pairs {
		var ki cluster.KeeperInfo
		err = json.Unmarshal(pair.Value, &ki)
		if err != nil {
			return nil, err
		}
		keepers[ki.UID] = &ki
	}
	return keepers, nil
}

func (e *StoreManager) SetSentinelInfo(si *cluster.SentinelInfo, ttl time.Duration) error {
	sij, err := json.Marshal(si)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, sentinelsInfoDir, si.UID), sij, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetSentinelInfo(id string) (*cluster.SentinelInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty sentinel id")
	}
	var si cluster.SentinelInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, sentinelsInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	err = json.Unmarshal(pair.Value, &si)
	if err != nil {
		return nil, false, err
	}
	return &si, true, nil
}

func (e *StoreManager) GetSentinelsInfo() (cluster.SentinelsInfo, error) {
	ssi := cluster.SentinelsInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, sentinelsInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return ssi, nil
	}
	for _, pair := range pairs {
		var si cluster.SentinelInfo
		err = json.Unmarshal(pair.Value, &si)
		if err != nil {
			return nil, err
		}
		ssi = append(ssi, &si)
	}
	return ssi, nil
}

func (e *StoreManager) GetLeaderSentinelId() (string, error) {
	pair, err := e.store.Get(filepath.Join(e.clusterPath, common.SentinelLeaderKey))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return "", err
		}
		return "", nil
	}
	return string(pair.Value), nil
}

func (e *StoreManager) SetProxyInfo(pi *cluster.ProxyInfo, ttl time.Duration) error {
	pij, err := json.Marshal(pi)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, proxiesInfoDir, pi.UID), pij, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetProxyInfo(id string) (*cluster.ProxyInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty proxy id")
	}
	var pi cluster.ProxyInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, proxiesInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	err = json.Unmarshal(pair.Value, &pi)
	if err != nil {
		return nil, false, err
	}
	return &pi, true, nil
}

func (e *StoreManager) GetProxiesInfo() (cluster.ProxiesInfo, error) {
	psi := cluster.ProxiesInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, proxiesInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return psi, nil
	}
	for _, pair := range pairs {
		var pi cluster.ProxyInfo
		err = json.Unmarshal(pair.Value, &pi)
		if err != nil {
			return nil, err
		}
		psi = append(psi, &pi)
	}
	return psi, nil
}
<commit_msg>store: remove prev value before AtomicPut<commit_after>// Copyright 2015 Sorint.lab
//
// 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 store

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/sorintlab/stolon/common"
	"github.com/sorintlab/stolon/pkg/cluster"

	"github.com/docker/libkv"
	kvstore "github.com/docker/libkv/store"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
)

func init() {
	etcd.Register()
	consul.Register()
}

// Backend represents a KV Store Backend
type Backend string

const (
	CONSUL Backend = "consul"
	ETCD   Backend = "etcd"
)

const (
	keepersInfoDir         = "/keepers/info/"
	clusterDataFile        = "clusterdata"
	leaderSentinelInfoFile = "/sentinels/leaderinfo"
	sentinelsInfoDir       = "/sentinels/info/"
	proxiesInfoDir         = "/proxies/info/"
)

const (
	DefaultEtcdEndpoints   = "127.0.0.1:2379"
	DefaultConsulEndpoints = "127.0.0.1:8500"
)

const (
	//TODO(sgotti) fix this in libkv?
	// consul min ttl is 10s and libkv divides this by 2
	MinTTL = 20 * time.Second
)

type StoreManager struct {
	clusterPath string
	store       kvstore.Store
}

func NewStore(backend Backend, addrsStr string) (kvstore.Store, error) {

	var kvbackend kvstore.Backend
	switch backend {
	case CONSUL:
		kvbackend = kvstore.CONSUL
	case ETCD:
		kvbackend = kvstore.ETCD
	default:
		return nil, fmt.Errorf("Unknown store backend: %q", backend)
	}

	if addrsStr == "" {
		switch backend {
		case CONSUL:
			addrsStr = DefaultConsulEndpoints
		case ETCD:
			addrsStr = DefaultEtcdEndpoints
		}
	}
	addrs := strings.Split(addrsStr, ",")

	store, err := libkv.NewStore(kvbackend, addrs, &kvstore.Config{ConnectionTimeout: 10 * time.Second})
	if err != nil {
		return nil, err
	}
	return store, nil
}

func NewStoreManager(kvStore kvstore.Store, path string) *StoreManager {
	return &StoreManager{
		clusterPath: path,
		store:       kvStore,
	}
}

func (e *StoreManager) AtomicPutClusterData(cd *cluster.ClusterData, previous *kvstore.KVPair) (*kvstore.KVPair, error) {
	cdj, err := json.Marshal(cd)
	if err != nil {
		return nil, err
	}
	path := filepath.Join(e.clusterPath, clusterDataFile)
	// Skip prev Value since LastIndex is enough for a CAS and it gives
	// problem with etcd v2 api with big prev values.
	var prev *kvstore.KVPair
	if previous != nil {
		prev = &kvstore.KVPair{
			Key:       previous.Key,
			LastIndex: previous.LastIndex,
		}
	}
	_, pair, err := e.store.AtomicPut(path, cdj, prev, nil)
	return pair, err
}

func (e *StoreManager) PutClusterData(cd *cluster.ClusterData) error {
	cdj, err := json.Marshal(cd)
	if err != nil {
		return err
	}
	path := filepath.Join(e.clusterPath, clusterDataFile)
	return e.store.Put(path, cdj, nil)
}

func (e *StoreManager) GetClusterData() (*cluster.ClusterData, *kvstore.KVPair, error) {
	var cd *cluster.ClusterData
	path := filepath.Join(e.clusterPath, clusterDataFile)
	pair, err := e.store.Get(path)
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, nil, err
		}
		return nil, nil, nil
	}
	if err := json.Unmarshal(pair.Value, &cd); err != nil {
		return nil, nil, err
	}
	return cd, pair, nil
}

func (e *StoreManager) SetKeeperInfo(id string, ms *cluster.KeeperInfo, ttl time.Duration) error {
	msj, err := json.Marshal(ms)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, keepersInfoDir, id), msj, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetKeeperInfo(id string) (*cluster.KeeperInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty keeper id")
	}
	var keeper cluster.KeeperInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, keepersInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	if err := json.Unmarshal(pair.Value, &keeper); err != nil {
		return nil, false, err
	}
	return &keeper, true, nil
}

func (e *StoreManager) GetKeepersInfo() (cluster.KeepersInfo, error) {
	keepers := cluster.KeepersInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, keepersInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return keepers, nil
	}
	for _, pair := range pairs {
		var ki cluster.KeeperInfo
		err = json.Unmarshal(pair.Value, &ki)
		if err != nil {
			return nil, err
		}
		keepers[ki.UID] = &ki
	}
	return keepers, nil
}

func (e *StoreManager) SetSentinelInfo(si *cluster.SentinelInfo, ttl time.Duration) error {
	sij, err := json.Marshal(si)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, sentinelsInfoDir, si.UID), sij, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetSentinelInfo(id string) (*cluster.SentinelInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty sentinel id")
	}
	var si cluster.SentinelInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, sentinelsInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	err = json.Unmarshal(pair.Value, &si)
	if err != nil {
		return nil, false, err
	}
	return &si, true, nil
}

func (e *StoreManager) GetSentinelsInfo() (cluster.SentinelsInfo, error) {
	ssi := cluster.SentinelsInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, sentinelsInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return ssi, nil
	}
	for _, pair := range pairs {
		var si cluster.SentinelInfo
		err = json.Unmarshal(pair.Value, &si)
		if err != nil {
			return nil, err
		}
		ssi = append(ssi, &si)
	}
	return ssi, nil
}

func (e *StoreManager) GetLeaderSentinelId() (string, error) {
	pair, err := e.store.Get(filepath.Join(e.clusterPath, common.SentinelLeaderKey))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return "", err
		}
		return "", nil
	}
	return string(pair.Value), nil
}

func (e *StoreManager) SetProxyInfo(pi *cluster.ProxyInfo, ttl time.Duration) error {
	pij, err := json.Marshal(pi)
	if err != nil {
		return err
	}
	if ttl < MinTTL {
		ttl = MinTTL
	}
	return e.store.Put(filepath.Join(e.clusterPath, proxiesInfoDir, pi.UID), pij, &kvstore.WriteOptions{TTL: ttl})
}

func (e *StoreManager) GetProxyInfo(id string) (*cluster.ProxyInfo, bool, error) {
	if id == "" {
		return nil, false, fmt.Errorf("empty proxy id")
	}
	var pi cluster.ProxyInfo
	pair, err := e.store.Get(filepath.Join(e.clusterPath, proxiesInfoDir, id))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, false, err
		}
		return nil, false, nil
	}
	err = json.Unmarshal(pair.Value, &pi)
	if err != nil {
		return nil, false, err
	}
	return &pi, true, nil
}

func (e *StoreManager) GetProxiesInfo() (cluster.ProxiesInfo, error) {
	psi := cluster.ProxiesInfo{}
	pairs, err := e.store.List(filepath.Join(e.clusterPath, proxiesInfoDir))
	if err != nil {
		if err != kvstore.ErrKeyNotFound {
			return nil, err
		}
		return psi, nil
	}
	for _, pair := range pairs {
		var pi cluster.ProxyInfo
		err = json.Unmarshal(pair.Value, &pi)
		if err != nil {
			return nil, err
		}
		psi = append(psi, &pi)
	}
	return psi, nil
}
<|endoftext|>


<commit_before>package libvirt

/*
#cgo LDFLAGS: -lvirt -ldl
#include <libvirt/libvirt.h>
#include <libvirt/virterror.h>
#include <stdlib.h>
*/
import "C"

import (
	"errors"
	"unsafe"
)

type VirDomain struct {
	ptr C.virDomainPtr
}

type VirDomainInfo struct {
	ptr C.virDomainInfo
}

func (d *VirDomain) Create() error {
	result := C.virDomainCreate(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Destroy() error {
	result := C.virDomainDestroy(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Shutdown() error {
	result := C.virDomainShutdown(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Reboot(flags uint) error {
	result := C.virDomainReboot(d.ptr, C.uint(flags))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) IsActive() (bool, error) {
	result := C.virDomainIsActive(d.ptr)
	if result == -1 {
		return false, errors.New(GetLastError())
	}
	if result == 1 {
		return true, nil
	}
	return false, nil
}

func (d *VirDomain) SetAutostart(autostart bool) error {
	var cAutostart C.int
	switch autostart {
	case true:
		cAutostart = 1
	default:
		cAutostart = 0
	}
	result := C.virDomainSetAutostart(d.ptr, cAutostart)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) GetAutostart() (bool, error) {
	var out C.int
	result := C.virDomainGetAutostart(d.ptr, (*C.int)(unsafe.Pointer(&out)))
	if result == -1 {
		return false, errors.New(GetLastError())
	}
	switch out {
	case 1:
		return true, nil
	default:
		return false, nil
	}
}

func (d *VirDomain) GetName() (string, error) {
	name := C.virDomainGetName(d.ptr)
	if name == nil {
		return "", errors.New(GetLastError())
	}
	return C.GoString(name), nil
}

func (d *VirDomain) GetState() ([]int, error) {
	var cState C.int
	var cReason C.int
	result := C.virDomainGetState(d.ptr,
		(*C.int)(unsafe.Pointer(&cState)),
		(*C.int)(unsafe.Pointer(&cReason)),
		0)
	if int(result) == -1 {
		return []int{}, errors.New(GetLastError())
	}
	return []int{int(cState), int(cReason)}, nil
}

func (d *VirDomain) GetUUID() ([]byte, error) {
	var cUuid [C.VIR_UUID_BUFLEN](byte)
	cuidPtr := unsafe.Pointer(&cUuid)
	result := C.virDomainGetUUID(d.ptr, (*C.uchar)(cuidPtr))
	if result != 0 {
		return []byte{}, errors.New(GetLastError())
	}
	return C.GoBytes(cuidPtr, C.VIR_UUID_BUFLEN), nil
}

func (d *VirDomain) GetUUIDString() (string, error) {
	var cUuid [C.VIR_UUID_STRING_BUFLEN](C.char)
	cuidPtr := unsafe.Pointer(&cUuid)
	result := C.virDomainGetUUIDString(d.ptr, (*C.char)(cuidPtr))
	if result != 0 {
		return "", errors.New(GetLastError())
	}
	return C.GoString((*C.char)(cuidPtr)), nil
}

func (d *VirDomain) GetInfo() (VirDomainInfo, error) {
	di := VirDomainInfo{}
	var ptr C.virDomainInfo
	result := C.virDomainGetInfo(d.ptr, (*C.virDomainInfo)(unsafe.Pointer(&ptr)))
	if result == -1 {
		return di, errors.New(GetLastError())
	}
	di.ptr = ptr
	return di, nil
}

func (d *VirDomain) GetXMLDesc(flags uint32) (string, error) {
	result := C.virDomainGetXMLDesc(d.ptr, C.uint(flags))
	if result == nil {
		return "", errors.New(GetLastError())
	}
	xml := C.GoString(result)
	C.free(unsafe.Pointer(result))
	return xml, nil
}

func (i *VirDomainInfo) GetState() uint8 {
	return uint8(i.ptr.state)
}

func (i *VirDomainInfo) GetMaxMem() uint64 {
	return uint64(i.ptr.maxMem)
}

func (i *VirDomainInfo) GetMemory() uint64 {
	return uint64(i.ptr.memory)
}

func (i *VirDomainInfo) GetNrVirtCpu() uint16 {
	return uint16(i.ptr.nrVirtCpu)
}

func (i *VirDomainInfo) GetCpuTime() uint64 {
	return uint64(i.ptr.cpuTime)
}

func (d *VirDomain) GetMetadata(tipus int, uri string, flags uint32) (string, error) {

	var cUri *C.char
	if uri != "" {
		cUri = C.CString(uri)
		defer C.free(unsafe.Pointer(cUri))
	}

	result := C.virDomainGetMetadata(d.ptr, C.int(tipus), cUri, C.uint(flags))
	if result == nil {
		return "", errors.New(GetLastError())

	}
	defer C.free(unsafe.Pointer(result))
	return C.GoString(result), nil
}

func (d *VirDomain) SetMetadata(metaDataType int, metaDataCont, uriKey, uri string, flags uint32) error {
	var cMetaDataCont *C.char
	var cUriKey *C.char
	var cUri *C.char

	cMetaDataCont = C.CString(metaDataCont)
	defer C.free(unsafe.Pointer(cMetaDataCont))

	if metaDataType == VIR_DOMAIN_METADATA_ELEMENT {
		cUriKey = C.CString(uriKey)
		defer C.free(unsafe.Pointer(cUriKey))
		cUri = C.CString(uri)
		defer C.free(unsafe.Pointer(cUri))
	}
	result := C.virDomainSetMetadata(d.ptr, C.int(metaDataType), cMetaDataCont, cUriKey, cUri, C.uint(flags))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Undefine() error {
	result := C.virDomainUndefine(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) SetMemory(memory uint64) error {
	result := C.virDomainSetMemory(d.ptr, C.ulong(memory))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}
<commit_msg>Add VirDomain.SetMemoryFlags<commit_after>package libvirt

/*
#cgo LDFLAGS: -lvirt -ldl
#include <libvirt/libvirt.h>
#include <libvirt/virterror.h>
#include <stdlib.h>
*/
import "C"

import (
	"errors"
	"unsafe"
)

type VirDomain struct {
	ptr C.virDomainPtr
}

type VirDomainInfo struct {
	ptr C.virDomainInfo
}

func (d *VirDomain) Create() error {
	result := C.virDomainCreate(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Destroy() error {
	result := C.virDomainDestroy(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Shutdown() error {
	result := C.virDomainShutdown(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Reboot(flags uint) error {
	result := C.virDomainReboot(d.ptr, C.uint(flags))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) IsActive() (bool, error) {
	result := C.virDomainIsActive(d.ptr)
	if result == -1 {
		return false, errors.New(GetLastError())
	}
	if result == 1 {
		return true, nil
	}
	return false, nil
}

func (d *VirDomain) SetAutostart(autostart bool) error {
	var cAutostart C.int
	switch autostart {
	case true:
		cAutostart = 1
	default:
		cAutostart = 0
	}
	result := C.virDomainSetAutostart(d.ptr, cAutostart)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) GetAutostart() (bool, error) {
	var out C.int
	result := C.virDomainGetAutostart(d.ptr, (*C.int)(unsafe.Pointer(&out)))
	if result == -1 {
		return false, errors.New(GetLastError())
	}
	switch out {
	case 1:
		return true, nil
	default:
		return false, nil
	}
}

func (d *VirDomain) GetName() (string, error) {
	name := C.virDomainGetName(d.ptr)
	if name == nil {
		return "", errors.New(GetLastError())
	}
	return C.GoString(name), nil
}

func (d *VirDomain) GetState() ([]int, error) {
	var cState C.int
	var cReason C.int
	result := C.virDomainGetState(d.ptr,
		(*C.int)(unsafe.Pointer(&cState)),
		(*C.int)(unsafe.Pointer(&cReason)),
		0)
	if int(result) == -1 {
		return []int{}, errors.New(GetLastError())
	}
	return []int{int(cState), int(cReason)}, nil
}

func (d *VirDomain) GetUUID() ([]byte, error) {
	var cUuid [C.VIR_UUID_BUFLEN](byte)
	cuidPtr := unsafe.Pointer(&cUuid)
	result := C.virDomainGetUUID(d.ptr, (*C.uchar)(cuidPtr))
	if result != 0 {
		return []byte{}, errors.New(GetLastError())
	}
	return C.GoBytes(cuidPtr, C.VIR_UUID_BUFLEN), nil
}

func (d *VirDomain) GetUUIDString() (string, error) {
	var cUuid [C.VIR_UUID_STRING_BUFLEN](C.char)
	cuidPtr := unsafe.Pointer(&cUuid)
	result := C.virDomainGetUUIDString(d.ptr, (*C.char)(cuidPtr))
	if result != 0 {
		return "", errors.New(GetLastError())
	}
	return C.GoString((*C.char)(cuidPtr)), nil
}

func (d *VirDomain) GetInfo() (VirDomainInfo, error) {
	di := VirDomainInfo{}
	var ptr C.virDomainInfo
	result := C.virDomainGetInfo(d.ptr, (*C.virDomainInfo)(unsafe.Pointer(&ptr)))
	if result == -1 {
		return di, errors.New(GetLastError())
	}
	di.ptr = ptr
	return di, nil
}

func (d *VirDomain) GetXMLDesc(flags uint32) (string, error) {
	result := C.virDomainGetXMLDesc(d.ptr, C.uint(flags))
	if result == nil {
		return "", errors.New(GetLastError())
	}
	xml := C.GoString(result)
	C.free(unsafe.Pointer(result))
	return xml, nil
}

func (i *VirDomainInfo) GetState() uint8 {
	return uint8(i.ptr.state)
}

func (i *VirDomainInfo) GetMaxMem() uint64 {
	return uint64(i.ptr.maxMem)
}

func (i *VirDomainInfo) GetMemory() uint64 {
	return uint64(i.ptr.memory)
}

func (i *VirDomainInfo) GetNrVirtCpu() uint16 {
	return uint16(i.ptr.nrVirtCpu)
}

func (i *VirDomainInfo) GetCpuTime() uint64 {
	return uint64(i.ptr.cpuTime)
}

func (d *VirDomain) GetMetadata(tipus int, uri string, flags uint32) (string, error) {

	var cUri *C.char
	if uri != "" {
		cUri = C.CString(uri)
		defer C.free(unsafe.Pointer(cUri))
	}

	result := C.virDomainGetMetadata(d.ptr, C.int(tipus), cUri, C.uint(flags))
	if result == nil {
		return "", errors.New(GetLastError())

	}
	defer C.free(unsafe.Pointer(result))
	return C.GoString(result), nil
}

func (d *VirDomain) SetMetadata(metaDataType int, metaDataCont, uriKey, uri string, flags uint32) error {
	var cMetaDataCont *C.char
	var cUriKey *C.char
	var cUri *C.char

	cMetaDataCont = C.CString(metaDataCont)
	defer C.free(unsafe.Pointer(cMetaDataCont))

	if metaDataType == VIR_DOMAIN_METADATA_ELEMENT {
		cUriKey = C.CString(uriKey)
		defer C.free(unsafe.Pointer(cUriKey))
		cUri = C.CString(uri)
		defer C.free(unsafe.Pointer(cUri))
	}
	result := C.virDomainSetMetadata(d.ptr, C.int(metaDataType), cMetaDataCont, cUriKey, cUri, C.uint(flags))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) Undefine() error {
	result := C.virDomainUndefine(d.ptr)
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) SetMemory(memory uint64) error {
	result := C.virDomainSetMemory(d.ptr, C.ulong(memory))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}

func (d *VirDomain) SetMemoryFlags(memory uint64, flags uint32) error {
	result := C.virDomainSetMemoryFlags(d.ptr, C.ulong(memory), C.uint(flags))
	if result == -1 {
		return errors.New(GetLastError())
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"srcd.works/go-git.v4"
	"srcd.works/go-git.v4/plumbing"
	"srcd.works/go-git.v4/plumbing/object"
	"srcd.works/go-git.v4/storage/memory"

	"bytes"

	"regexp"

	"github.com/BurntSushi/toml"
)

type gitConfig struct {
	uri string
}
type salesforceConfig struct {
	username    string
	password    string
	endpoint    string
	packagePath string
	apiVersion  string
}

type Downloader interface {
	Download() ([]*File, error)
}

type MetaPackageFile struct {
	Version float64 `toml:"version"`
	Types   []*Type `toml:"types"`
}

type Type struct {
	Name    string   `toml:"name"`
	Members []string `toml:"members"`
}

type SalesforceDownloader struct {
	config *salesforceConfig
	client *ForceClient
	logger Logger
}

func NewSalesforceDownloader(logger Logger, config *salesforceConfig) (*SalesforceDownloader, error) {
	d := &SalesforceDownloader{
		logger: logger,
		config: config,
	}
	err := d.init()
	return d, err
}

func (i *SalesforceDownloader) init() (err error) {
	if i.config.username == "" {
		return errors.New("Username is required")
	}
	if i.config.password == "" {
		return errors.New("Password is required")
	}

	err = i.setClient()
	if err != nil {
		return err
	}
	return nil
}

func (i *SalesforceDownloader) setClient() error {
	i.client = NewForceClient(i.config.endpoint, i.config.apiVersion)
	err := i.client.Login(i.config.username, i.config.password)
	if err != nil {
		return err
	}
	return nil
}

func (i *SalesforceDownloader) Download() ([]*File, error) {
	buf, err := ioutil.ReadFile(i.config.packagePath)
	if err != nil {
		return nil, err
	}
	packages := &MetaPackageFile{}
	err = toml.Unmarshal(buf, packages)
	if err != nil {
		return nil, err
	}
	i.logger.Info("Start Retrieve Request...")
	r, err := i.client.Retrieve(createRetrieveRequest(packages))
	if err != nil {
		return nil, err
	}
	for {
		time.Sleep(2 * time.Second)
		i.logger.Info("Check Retrieve Status...")
		ret_res, err := i.client.CheckRetrieveStatus(r.Result.Id)
		if err != nil {
			return nil, err
		}
		if ret_res.Result.Done {
			zb := make([]byte, len(ret_res.Result.ZipFile))
			_, err = base64.StdEncoding.Decode(zb, ret_res.Result.ZipFile)
			return []*File{&File{Body: zb}}, err
		}
	}
	return nil, nil // Todo: error handling
}

type GitDownloader struct {
	logger Logger
	config *gitConfig
}

func NewGitDownloader(logger Logger, config *gitConfig) (*GitDownloader, error) {
	return &GitDownloader{
		logger: logger,
		config: config,
	}, nil
}

func (d *GitDownloader) Download() ([]*File, error) {
	uri, _, _, branch := extractInstallParameter(d.config.uri)
	d.logger.Infof("Clone repository from %s (branch: %s)", uri, branch)

	r, err := git.Clone(memory.NewStorage(), nil, &git.CloneOptions{
		ReferenceName: plumbing.ReferenceName(fmt.Sprintf("refs/heads/%s", branch)),
		SingleBranch:  true,
		URL:           uri,
	})
	if err != nil {
		return nil, err
	}

	ref, _ := r.Head()
	commit, _ := r.Commit(ref.Hash())

	gfiles, err := commit.Files()
	files := make([]*File, 0)
	err = gfiles.ForEach(func(f *object.File) error {
		reader, err := f.Reader()
		if err != nil {
			return err
		}

		b := new(bytes.Buffer)
		if _, err := b.ReadFrom(reader); err != nil {
			return err
		}
		files = append(files, &File{Name: f.Name, Body: b.Bytes()})
		return nil
	})
	if err != nil {
		return nil, err
	}
	return files, nil
}

func dispatchDownloader(logger Logger, uri string) (Downloader, error) {
	r := regexp.MustCompile(`^https://([^/]+?)/([^/]+?)/([^/@]+?)(/([^@]+))?(\?([^/]+))?$`)
	if r.MatchString(uri) {
		return NewGitDownloader(logger, &gitConfig{uri: uri})
	}
	r = regexp.MustCompile(`^sf://([^/]*?):([^/]*)@([^/]+?)\?path=(.+)$`)
	if r.MatchString(uri) {
		group := r.FindAllStringSubmatch(uri, -1)
		return NewSalesforceDownloader(logger, &salesforceConfig{
			username:    group[0][1],
			password:    group[0][2],
			endpoint:    group[0][3],
			packagePath: group[0][4],
			apiVersion:  "38.0",
		})
	}
	return nil, errors.New("Invalid downloader")
}
<commit_msg>Fix incorrect import<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io/ioutil"
	"regexp"
	"time"

	"gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"
	"gopkg.in/src-d/go-git.v4/plumbing/object"
	"gopkg.in/src-d/go-git.v4/storage/memory"

	"github.com/BurntSushi/toml"
)

type gitConfig struct {
	uri string
}
type salesforceConfig struct {
	username    string
	password    string
	endpoint    string
	packagePath string
	apiVersion  string
}

type Downloader interface {
	Download() ([]*File, error)
}

type MetaPackageFile struct {
	Version float64 `toml:"version"`
	Types   []*Type `toml:"types"`
}

type Type struct {
	Name    string   `toml:"name"`
	Members []string `toml:"members"`
}

type SalesforceDownloader struct {
	config *salesforceConfig
	client *ForceClient
	logger Logger
}

func NewSalesforceDownloader(logger Logger, config *salesforceConfig) (*SalesforceDownloader, error) {
	d := &SalesforceDownloader{
		logger: logger,
		config: config,
	}
	err := d.init()
	return d, err
}

func (i *SalesforceDownloader) init() (err error) {
	if i.config.username == "" {
		return errors.New("Username is required")
	}
	if i.config.password == "" {
		return errors.New("Password is required")
	}

	err = i.setClient()
	if err != nil {
		return err
	}
	return nil
}

func (i *SalesforceDownloader) setClient() error {
	i.client = NewForceClient(i.config.endpoint, i.config.apiVersion)
	err := i.client.Login(i.config.username, i.config.password)
	if err != nil {
		return err
	}
	return nil
}

func (i *SalesforceDownloader) Download() ([]*File, error) {
	buf, err := ioutil.ReadFile(i.config.packagePath)
	if err != nil {
		return nil, err
	}
	packages := &MetaPackageFile{}
	err = toml.Unmarshal(buf, packages)
	if err != nil {
		return nil, err
	}
	i.logger.Info("Start Retrieve Request...")
	r, err := i.client.Retrieve(createRetrieveRequest(packages))
	if err != nil {
		return nil, err
	}
	for {
		time.Sleep(2 * time.Second)
		i.logger.Info("Check Retrieve Status...")
		ret_res, err := i.client.CheckRetrieveStatus(r.Result.Id)
		if err != nil {
			return nil, err
		}
		if ret_res.Result.Done {
			zb := make([]byte, len(ret_res.Result.ZipFile))
			_, err = base64.StdEncoding.Decode(zb, ret_res.Result.ZipFile)
			return []*File{&File{Body: zb}}, err
		}
	}
	return nil, nil // Todo: error handling
}

type GitDownloader struct {
	logger Logger
	config *gitConfig
}

func NewGitDownloader(logger Logger, config *gitConfig) (*GitDownloader, error) {
	return &GitDownloader{
		logger: logger,
		config: config,
	}, nil
}

func (d *GitDownloader) Download() ([]*File, error) {
	uri, _, _, branch := extractInstallParameter(d.config.uri)
	d.logger.Infof("Clone repository from %s (branch: %s)", uri, branch)

	r, err := git.Clone(memory.NewStorage(), nil, &git.CloneOptions{
		ReferenceName: plumbing.ReferenceName(fmt.Sprintf("refs/heads/%s", branch)),
		SingleBranch:  true,
		URL:           uri,
	})
	if err != nil {
		return nil, err
	}

	ref, _ := r.Head()
	commit, _ := r.Commit(ref.Hash())

	gfiles, err := commit.Files()
	files := make([]*File, 0)
	err = gfiles.ForEach(func(f *object.File) error {
		reader, err := f.Reader()
		if err != nil {
			return err
		}

		b := new(bytes.Buffer)
		if _, err := b.ReadFrom(reader); err != nil {
			return err
		}
		files = append(files, &File{Name: f.Name, Body: b.Bytes()})
		return nil
	})
	if err != nil {
		return nil, err
	}
	return files, nil
}

func dispatchDownloader(logger Logger, uri string) (Downloader, error) {
	r := regexp.MustCompile(`^https://([^/]+?)/([^/]+?)/([^/@]+?)(/([^@]+))?(\?([^/]+))?$`)
	if r.MatchString(uri) {
		return NewGitDownloader(logger, &gitConfig{uri: uri})
	}
	r = regexp.MustCompile(`^sf://([^/]*?):([^/]*)@([^/]+?)(\?path=(.+)&version=(.+))?$`)
	if r.MatchString(uri) {
		group := r.FindAllStringSubmatch(uri, -1)
		return NewSalesforceDownloader(logger, &salesforceConfig{
			username:    group[0][1],
			password:    group[0][2],
			endpoint:    group[0][3],
			packagePath: group[0][4],
			apiVersion:  "38.0",
		})
	}
	return nil, errors.New("Invalid downloader")
}
<|endoftext|>


<commit_before>package doublestar

import (
  "path"
  "os"
  "strings"
  "unicode/utf8"
)

var ErrBadPattern = path.ErrBadPattern

func SplitPathOnSeparator(path string, separator rune) []string {
  // if the separator is '\\', then we can just split...
  if separator == '\\' { return strings.Split(path, string(separator)) }

  // otherwise, we need to be careful of situations where the separator was escaped
  cnt := strings.Count(path, string(separator))
  if cnt == 0 { return []string{path} }
  ret := make([]string, cnt + 1)
  pathlen := len(path)
  separatorLen := utf8.RuneLen(separator)
  idx := 0
  for start := 0; start < pathlen; {
    end := indexRuneWithEscaping(path[start:], separator)
    if end == -1 {
      end = pathlen
    } else {
      end += start
    }
    ret[idx] = path[start:end]
    start = end + separatorLen
    idx++
  }
  return ret[:idx]
}

func indexRuneWithEscaping(s string, r rune) int {
  end := strings.IndexRune(s, r)
  if end == -1 { return -1 }
  if end > 0 && s[end - 1] == '\\' {
    start := end + utf8.RuneLen(r)
    end = indexRuneWithEscaping(s[start:], r)
    if end != -1 { end += start }
  }
  return end
}

// Match returns true if name matches the shell file name pattern.
// The pattern syntax is:
//
//	pattern:
//		{ term }
//	term:
//		'*'         matches any sequence of non-path-separators
//              '**'        matches any sequence of characters, including
//                          path separators.
//		'?'         matches any single non-path-separator character
//		'[' [ '^' ] { character-range } ']'
//		            character class (must be non-empty)
//		'{' { term } [ ',' { term } ... ] '}'
//		c           matches character c (c != '*', '?', '\\', '[')
//		'\\' c      matches character c
//
//	character-range:
//		c           matches character c (c != '\\', '-', ']')
//		'\\' c      matches character c
//		lo '-' hi   matches character c for lo <= c <= hi
//
// Match requires pattern to match all of name, not just a substring.
// The path-separator defaults to the '/' character. The only possible
// returned error is ErrBadPattern, when pattern is malformed.
//
func Match(pattern, name string) (bool, error) {
  return matchWithSeparator(pattern, name, '/')
}

// PathMatch is like Match except that it uses your system's path separator.
// For most systems, this will be '/'. However, for Windows, it would be '\\'.
// Note that for systems where the path separator is '\\', escaping is
// disabled.
//
func PathMatch(pattern, name string) (bool, error) {
  return matchWithSeparator(pattern, name, os.PathSeparator)
}

// Match returns true if name matches the shell file name pattern.
// The pattern syntax is:
//
//	pattern:
//		{ term }
//	term:
//		'*'         matches any sequence of non-path-separators
//              '**'        matches any sequence of characters, including
//                          path separators.
//		'?'         matches any single non-path-separator character
//		'[' [ '^' ] { character-range } ']'
//		            character class (must be non-empty)
//		'{' { term } [ ',' { term } ... ] '}'
//		c           matches character c (c != '*', '?', '\\', '[')
//		'\\' c      matches character c
//
//	character-range:
//		c           matches character c (c != '\\', '-', ']')
//		'\\' c      matches character c, unless separator is '\\'
//		lo '-' hi   matches character c for lo <= c <= hi
//
// Match requires pattern to match all of name, not just a substring.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
//
func matchWithSeparator(pattern, name string, separator rune) (bool, error) {
  patternComponents := SplitPathOnSeparator(pattern, separator)
  nameComponents := SplitPathOnSeparator(name, separator)
  return doMatching(patternComponents, nameComponents)
}

// Glob returns the names of all files matching pattern or nil
// if there is no matching file. The syntax of pattern is the same
// as in Match. The pattern may describe hierarchical names such as
// /usr/*/bin/ed (assuming the Separator is '/').
//
// Glob ignores file system errors such as I/O errors reading directories.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
//
// Your system path separator is automatically used. This means on
// systems where the separator is '\\' (Windows), escaping will be
// disabled.
//
func Glob(pattern string) (matches []string, err error) {
  patternComponents := SplitPathOnSeparator(pattern, os.PathSeparator)
  patternLen := len(patternComponents)
  if patternLen == 0 { return false, nil }
  patIdx := 0
  for ; patIdx < patternLen; {
  }
}

func doMatching(patternComponents, nameComponents []string) (matched bool, err error) {
  patternLen, nameLen := len(patternComponents), len(nameComponents)
  if patternLen == 0 && nameLen == 0 { return true, nil }
  if patternLen == 0 || nameLen == 0 { return false, nil }
  patIdx, nameIdx := 0, 0
  for ; patIdx < patternLen && nameIdx < nameLen; {
    if patternComponents[patIdx] == "**" {
      if patIdx++; patIdx >= patternLen { return true, nil }
      for ; nameIdx < nameLen; nameIdx++ {
	if m, _ := doMatching(patternComponents[patIdx:], nameComponents[nameIdx:]); m {
	  return true, nil
	}
      }
      return false, nil
    } else {
      matched, err = matchComponent(patternComponents[patIdx], nameComponents[nameIdx])
      if !matched || err != nil { return }
    }
    patIdx++
    nameIdx++
  }
  return patIdx >= patternLen && nameIdx >= nameLen, nil
}

func matchComponent(pattern, name string) (bool, error) {
  patternLen, nameLen := len(pattern), len(name)
  if patternLen == 0 && nameLen == 0 { return true, nil }
  if patternLen == 0 { return false, nil }
  if nameLen == 0 && pattern != "*" { return false, nil }
  patIdx, nameIdx := 0, 0
  for ; patIdx < patternLen && nameIdx < nameLen; {
    patRune, patAdj := utf8.DecodeRuneInString(pattern[patIdx:])
    nameRune, nameAdj := utf8.DecodeRuneInString(name[nameIdx:])
    if patRune == '\\' {
      patIdx += patAdj
      patRune, patAdj = utf8.DecodeRuneInString(pattern[patIdx:])
      if patRune == utf8.RuneError {
	return false, ErrBadPattern
      } else if patRune == nameRune {
	patIdx += patAdj
	nameIdx += nameAdj
      } else {
	return false, nil
      }
    } else if patRune == '*' {
      if patIdx += patAdj; patIdx >= patternLen { return true, nil }
      for ; nameIdx < nameLen; nameIdx += nameAdj {
	if m, _ := matchComponent(pattern[patIdx:], name[nameIdx:]); m {
	  return true, nil
	}
      }
      return false, nil
    } else if patRune == '[' {
      patIdx += patAdj
      endClass := indexRuneWithEscaping(pattern[patIdx:], ']')
      if endClass == -1 { return false, ErrBadPattern }
      endClass += patIdx
      classRunes := []rune(pattern[patIdx:endClass])
      classRunesLen := len(classRunes)
      if classRunesLen > 0 {
	classIdx := 0
	matchClass := false
	if classRunes[0] == '^' { classIdx++ }
	for classIdx < classRunesLen {
	  low := classRunes[classIdx]
	  if low == '-' { return false, ErrBadPattern }
	  classIdx++
	  if low == '\\' {
	    if classIdx < classRunesLen {
	      low = classRunes[classIdx]
	      classIdx++
	    } else {
	      return false, ErrBadPattern
	    }
	  }
	  high := low
	  if classIdx < classRunesLen && classRunes[classIdx] == '-' {
	    if classIdx++; classIdx >= classRunesLen { return false, ErrBadPattern }
	    high = classRunes[classIdx]
	    if high == '-' { return false, ErrBadPattern }
	    classIdx++
	    if high == '\\' {
	      if classIdx < classRunesLen {
		high = classRunes[classIdx]
		classIdx++
	      } else {
		return false, ErrBadPattern
	      }
	    }
	  }
	  if low <= nameRune && nameRune <= high { matchClass = true }
	}
	if matchClass == (classRunes[0] == '^') { return false, nil }
      } else {
	return false, ErrBadPattern
      }
      patIdx = endClass + 1
      nameIdx += nameAdj
    } else if patRune == '{' {
      patIdx += patAdj
      endOptions := indexRuneWithEscaping(pattern[patIdx:], '}')
      if endOptions == -1 { return false, ErrBadPattern }
      endOptions += patIdx
      options := SplitPathOnSeparator(pattern[patIdx:endOptions], ',')
      patIdx = endOptions + 1
      for _, o := range options {
	m, e := matchComponent(o + pattern[patIdx:], name[nameIdx:])
	if e != nil { return false, e }
	if m { return true, nil }
      }
      return false, nil
    } else if patRune == '?' || patRune == nameRune {
      patIdx += patAdj
      nameIdx += nameAdj
    } else {
      return false, nil
    }
  }
  if patIdx >= patternLen && nameIdx >= nameLen { return true, nil }
  if nameIdx >= nameLen && pattern[patIdx:] == "*" { return true, nil }
  return false, nil
}

<commit_msg>cleaned up<commit_after>package doublestar

import (
  "path"
  "os"
  "strings"
  "unicode/utf8"
)

var ErrBadPattern = path.ErrBadPattern

func SplitPathOnSeparator(path string, separator rune) []string {
  // if the separator is '\\', then we can just split...
  if separator == '\\' { return strings.Split(path, string(separator)) }

  // otherwise, we need to be careful of situations where the separator was escaped
  cnt := strings.Count(path, string(separator))
  if cnt == 0 { return []string{path} }
  ret := make([]string, cnt + 1)
  pathlen := len(path)
  separatorLen := utf8.RuneLen(separator)
  idx := 0
  for start := 0; start < pathlen; {
    end := indexRuneWithEscaping(path[start:], separator)
    if end == -1 {
      end = pathlen
    } else {
      end += start
    }
    ret[idx] = path[start:end]
    start = end + separatorLen
    idx++
  }
  return ret[:idx]
}

func indexRuneWithEscaping(s string, r rune) int {
  end := strings.IndexRune(s, r)
  if end == -1 { return -1 }
  if end > 0 && s[end - 1] == '\\' {
    start := end + utf8.RuneLen(r)
    end = indexRuneWithEscaping(s[start:], r)
    if end != -1 { end += start }
  }
  return end
}

// Match returns true if name matches the shell file name pattern.
// The pattern syntax is:
//
//	pattern:
//		{ term }
//	term:
//		'*'         matches any sequence of non-path-separators
//              '**'        matches any sequence of characters, including
//                          path separators.
//		'?'         matches any single non-path-separator character
//		'[' [ '^' ] { character-range } ']'
//		            character class (must be non-empty)
//		'{' { term } [ ',' { term } ... ] '}'
//		c           matches character c (c != '*', '?', '\\', '[')
//		'\\' c      matches character c
//
//	character-range:
//		c           matches character c (c != '\\', '-', ']')
//		'\\' c      matches character c
//		lo '-' hi   matches character c for lo <= c <= hi
//
// Match requires pattern to match all of name, not just a substring.
// The path-separator defaults to the '/' character. The only possible
// returned error is ErrBadPattern, when pattern is malformed.
//
func Match(pattern, name string) (bool, error) {
  return matchWithSeparator(pattern, name, '/')
}

// PathMatch is like Match except that it uses your system's path separator.
// For most systems, this will be '/'. However, for Windows, it would be '\\'.
// Note that for systems where the path separator is '\\', escaping is
// disabled.
//
func PathMatch(pattern, name string) (bool, error) {
  return matchWithSeparator(pattern, name, os.PathSeparator)
}

// Match returns true if name matches the shell file name pattern.
// The pattern syntax is:
//
//	pattern:
//		{ term }
//	term:
//		'*'         matches any sequence of non-path-separators
//              '**'        matches any sequence of characters, including
//                          path separators.
//		'?'         matches any single non-path-separator character
//		'[' [ '^' ] { character-range } ']'
//		            character class (must be non-empty)
//		'{' { term } [ ',' { term } ... ] '}'
//		c           matches character c (c != '*', '?', '\\', '[')
//		'\\' c      matches character c
//
//	character-range:
//		c           matches character c (c != '\\', '-', ']')
//		'\\' c      matches character c, unless separator is '\\'
//		lo '-' hi   matches character c for lo <= c <= hi
//
// Match requires pattern to match all of name, not just a substring.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
//
func matchWithSeparator(pattern, name string, separator rune) (bool, error) {
  patternComponents := SplitPathOnSeparator(pattern, separator)
  nameComponents := SplitPathOnSeparator(name, separator)
  return doMatching(patternComponents, nameComponents)
}

// Glob returns the names of all files matching pattern or nil
// if there is no matching file. The syntax of pattern is the same
// as in Match. The pattern may describe hierarchical names such as
// /usr/*/bin/ed (assuming the Separator is '/').
//
// Glob ignores file system errors such as I/O errors reading directories.
// The only possible returned error is ErrBadPattern, when pattern
// is malformed.
//
// Your system path separator is automatically used. This means on
// systems where the separator is '\\' (Windows), escaping will be
// disabled.
//
func Glob(pattern string) (matches []string, err error) {
  patternComponents := SplitPathOnSeparator(pattern, os.PathSeparator)
  patternLen := len(patternComponents)
  if patternLen == 0 { return nil, nil }
  patIdx := 0
  for ; patIdx < patternLen; {
  }
  return nil, nil
}

func doMatching(patternComponents, nameComponents []string) (matched bool, err error) {
  patternLen, nameLen := len(patternComponents), len(nameComponents)
  if patternLen == 0 && nameLen == 0 { return true, nil }
  if patternLen == 0 || nameLen == 0 { return false, nil }
  patIdx, nameIdx := 0, 0
  for ; patIdx < patternLen && nameIdx < nameLen; {
    if patternComponents[patIdx] == "**" {
      if patIdx++; patIdx >= patternLen { return true, nil }
      for ; nameIdx < nameLen; nameIdx++ {
	if m, _ := doMatching(patternComponents[patIdx:], nameComponents[nameIdx:]); m {
	  return true, nil
	}
      }
      return false, nil
    } else {
      matched, err = matchComponent(patternComponents[patIdx], nameComponents[nameIdx])
      if !matched || err != nil { return }
    }
    patIdx++
    nameIdx++
  }
  return patIdx >= patternLen && nameIdx >= nameLen, nil
}

func matchComponent(pattern, name string) (bool, error) {
  patternLen, nameLen := len(pattern), len(name)
  if patternLen == 0 && nameLen == 0 { return true, nil }
  if patternLen == 0 { return false, nil }
  if nameLen == 0 && pattern != "*" { return false, nil }
  patIdx, nameIdx := 0, 0
  for ; patIdx < patternLen && nameIdx < nameLen; {
    patRune, patAdj := utf8.DecodeRuneInString(pattern[patIdx:])
    nameRune, nameAdj := utf8.DecodeRuneInString(name[nameIdx:])
    if patRune == '\\' {
      patIdx += patAdj
      patRune, patAdj = utf8.DecodeRuneInString(pattern[patIdx:])
      if patRune == utf8.RuneError {
	return false, ErrBadPattern
      } else if patRune == nameRune {
	patIdx += patAdj
	nameIdx += nameAdj
      } else {
	return false, nil
      }
    } else if patRune == '*' {
      if patIdx += patAdj; patIdx >= patternLen { return true, nil }
      for ; nameIdx < nameLen; nameIdx += nameAdj {
	if m, _ := matchComponent(pattern[patIdx:], name[nameIdx:]); m {
	  return true, nil
	}
      }
      return false, nil
    } else if patRune == '[' {
      patIdx += patAdj
      endClass := indexRuneWithEscaping(pattern[patIdx:], ']')
      if endClass == -1 { return false, ErrBadPattern }
      endClass += patIdx
      classRunes := []rune(pattern[patIdx:endClass])
      classRunesLen := len(classRunes)
      if classRunesLen > 0 {
	classIdx := 0
	matchClass := false
	if classRunes[0] == '^' { classIdx++ }
	for classIdx < classRunesLen {
	  low := classRunes[classIdx]
	  if low == '-' { return false, ErrBadPattern }
	  classIdx++
	  if low == '\\' {
	    if classIdx < classRunesLen {
	      low = classRunes[classIdx]
	      classIdx++
	    } else {
	      return false, ErrBadPattern
	    }
	  }
	  high := low
	  if classIdx < classRunesLen && classRunes[classIdx] == '-' {
	    if classIdx++; classIdx >= classRunesLen { return false, ErrBadPattern }
	    high = classRunes[classIdx]
	    if high == '-' { return false, ErrBadPattern }
	    classIdx++
	    if high == '\\' {
	      if classIdx < classRunesLen {
		high = classRunes[classIdx]
		classIdx++
	      } else {
		return false, ErrBadPattern
	      }
	    }
	  }
	  if low <= nameRune && nameRune <= high { matchClass = true }
	}
	if matchClass == (classRunes[0] == '^') { return false, nil }
      } else {
	return false, ErrBadPattern
      }
      patIdx = endClass + 1
      nameIdx += nameAdj
    } else if patRune == '{' {
      patIdx += patAdj
      endOptions := indexRuneWithEscaping(pattern[patIdx:], '}')
      if endOptions == -1 { return false, ErrBadPattern }
      endOptions += patIdx
      options := SplitPathOnSeparator(pattern[patIdx:endOptions], ',')
      patIdx = endOptions + 1
      for _, o := range options {
	m, e := matchComponent(o + pattern[patIdx:], name[nameIdx:])
	if e != nil { return false, e }
	if m { return true, nil }
      }
      return false, nil
    } else if patRune == '?' || patRune == nameRune {
      patIdx += patAdj
      nameIdx += nameAdj
    } else {
      return false, nil
    }
  }
  if patIdx >= patternLen && nameIdx >= nameLen { return true, nil }
  if nameIdx >= nameLen && pattern[patIdx:] == "*" { return true, nil }
  return false, nil
}

<|endoftext|>


<commit_before>package delayd2

import (
	"database/sql"
	"errors"
	"time"

	"github.com/lib/pq"
)

var ErrMessageDuplicated = errors.New("driver: message is duplicated")

type Driver interface {
	Enqueue(string, int64, string, string) error
	ResetActive() (int64, error)
	MarkActive(time.Time) (int64, error)
	RemoveMessage(string) error
	GetActiveMessages() ([]*QueueMessage, error)
}

type pqDriver struct {
	workerID string
	db       *sql.DB
}

func NewDriver(workerID string, db *sql.DB) Driver {
	return &pqDriver{
		workerID: workerID,
		db:       db,
	}
}

func (d *pqDriver) Enqueue(queueId string, delay int64, relayTo string, payload string) error {
	releaseAt := time.Now().Add(time.Duration(delay) * time.Second)
	_, err := d.db.Exec(`
		INSERT
		INTO
		  queue
		VALUES
			($1, $2, $3, $4, $5)
	;`, queueId, d.workerID, releaseAt, relayTo, payload)

	if perr, ok := err.(*pq.Error); ok && perr.Code.Class() == "23" {
		return ErrMessageDuplicated
	}
	return err
}

func (d *pqDriver) ResetActive() (int64, error) {
	ret, err := d.db.Exec(`
		DELETE
		FROM
		  active
		WHERE
		  worker_id = $1
		;
	`, d.workerID)
	if err != nil {
		return 0, err
	}
	return ret.RowsAffected()
}

func (d *pqDriver) MarkActive(now time.Time) (int64, error) {
	ret, err := d.db.Exec(`
		INSERT
		INTO
		  active
		SELECT queue_id, worker_id
		FROM
		  queue
		WHERE
		      worker_id = $1
		  AND release_at < $2
		  AND NOT EXISTS (
			  SELECT 1 FROM active WHERE queue_id = queue.queue_id
		  )

	;`, d.workerID, now)
	if err != nil {
		return 0, err
	}
	return ret.RowsAffected()
}

func (d *pqDriver) RemoveMessage(queueID string) error {
	_, err := d.db.Exec(`
		DELETE
		FROM
			queue
		WHERE
			queue_id = $1
	;`, queueID)
	return err
}

func (d *pqDriver) GetActiveMessages() ([]*QueueMessage, error) {
	rows, err := d.db.Query(`
		SELECT queue.queue_id, queue.worker_id, queue.release_at, queue.relay_to, queue.payload
		FROM
		  queue INNER JOIN active USING (queue_id)
		WHERE
		  queue.worker_id = $1
		ORDER BY queue.release_at
		;
	`, d.workerID)

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	messages := []*QueueMessage{}
	for rows.Next() {
		var m QueueMessage
		if err := rows.Scan(&m.QueueID, &m.WorkerID, &m.ReleaseAt, &m.RelayTo, &m.Payload); err != nil {
			return nil, err
		}
		messages = append(messages, &m)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return messages, nil
}
<commit_msg>driver: Limiting 1000<commit_after>package delayd2

import (
	"database/sql"
	"errors"
	"time"

	"github.com/lib/pq"
)

var ErrMessageDuplicated = errors.New("driver: message is duplicated")

type Driver interface {
	Enqueue(string, int64, string, string) error
	ResetActive() (int64, error)
	MarkActive(time.Time) (int64, error)
	RemoveMessage(string) error
	GetActiveMessages() ([]*QueueMessage, error)
}

type pqDriver struct {
	workerID string
	db       *sql.DB
}

func NewDriver(workerID string, db *sql.DB) Driver {
	return &pqDriver{
		workerID: workerID,
		db:       db,
	}
}

func (d *pqDriver) Enqueue(queueId string, delay int64, relayTo string, payload string) error {
	releaseAt := time.Now().Add(time.Duration(delay) * time.Second)
	_, err := d.db.Exec(`
		INSERT
		INTO
		  queue
		VALUES
			($1, $2, $3, $4, $5)
	;`, queueId, d.workerID, releaseAt, relayTo, payload)

	if perr, ok := err.(*pq.Error); ok && perr.Code.Class() == "23" {
		return ErrMessageDuplicated
	}
	return err
}

func (d *pqDriver) ResetActive() (int64, error) {
	ret, err := d.db.Exec(`
		DELETE
		FROM
		  active
		WHERE
		  worker_id = $1
		;
	`, d.workerID)
	if err != nil {
		return 0, err
	}
	return ret.RowsAffected()
}

func (d *pqDriver) MarkActive(now time.Time) (int64, error) {
	ret, err := d.db.Exec(`
		INSERT
		INTO
		  active
		SELECT queue_id, worker_id
		FROM
		  queue
		WHERE
		      worker_id = $1
		  AND release_at < $2
		  AND NOT EXISTS (
			  SELECT 1 FROM active WHERE queue_id = queue.queue_id
		  )

	;`, d.workerID, now)
	if err != nil {
		return 0, err
	}
	return ret.RowsAffected()
}

func (d *pqDriver) RemoveMessage(queueID string) error {
	_, err := d.db.Exec(`
		DELETE
		FROM
			queue
		WHERE
			queue_id = $1
	;`, queueID)
	return err
}

func (d *pqDriver) GetActiveMessages() ([]*QueueMessage, error) {
	rows, err := d.db.Query(`
		SELECT queue.queue_id, queue.worker_id, queue.release_at, queue.relay_to, queue.payload
		FROM
		  queue INNER JOIN active USING (queue_id)
		WHERE
		  queue.worker_id = $1
		ORDER BY queue.release_at
		LIMIT 1000
		;
	`, d.workerID)

	if err != nil {
		return nil, err
	}
	defer rows.Close()

	messages := []*QueueMessage{}
	for rows.Next() {
		var m QueueMessage
		if err := rows.Scan(&m.QueueID, &m.WorkerID, &m.ReleaseAt, &m.RelayTo, &m.Payload); err != nil {
			return nil, err
		}
		messages = append(messages, &m)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return messages, nil
}
<|endoftext|>


<commit_before><commit_msg>Added VolumeDriver.Create comments<commit_after><|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"math"

	"github.com/bitly/go-simplejson"
	"github.com/davecgh/go-spew/spew"
	"github.com/ninjasphere/go-ninja"
	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-ninja/devices"
	"github.com/wolfeidau/lifx"
)

//---------------------------------------------------------------[Bulb]----------------------------------------------------------------

func NewLight(bus *ninja.DriverBus, client *lifx.Client, bulb *lifx.Bulb) (*devices.LightDevice, error) { //TODO cut this down!

	log.Infof("Making light with ID: %s Label: %s", bulb.GetLifxAddress(), bulb.GetLabel())

	sigs, _ := simplejson.NewJson([]byte(`{
			"ninja:manufacturer": "Lifx",
			"ninja:productName": "Lifx Bulb",
			"ninja:productType": "Light",
			"ninja:thingType": "light"
	}`))

	deviceBus, err := bus.AnnounceDevice(bulb.GetLifxAddress(), "light", bulb.GetLabel(), sigs)
	if err != nil {
		log.FatalError(err, "Failed to create light device bus")
	}

	light, err := devices.CreateLightDevice(bulb.GetLabel(), deviceBus)
	if err != nil {
		log.FatalError(err, "Failed to create light device")
	}

	light.EnableOnOffChannel()
	light.EnableBrightnessChannel()
	light.EnableColorChannel("temperature", "hue")
	light.EnableTransitionChannel()

	light.ApplyOnOff = func(state bool) error {
		var err error
		if state {
			err = client.LightOn(bulb)
		} else {
			err = client.LightOff(bulb)
		}
		if err != nil {
			return fmt.Errorf("Failed to set on-off state: %s", err)
		}
		return nil
	}

	light.ApplyLightState = func(state *devices.LightDeviceState) error {
		jsonState, _ := json.Marshal(state)
		log.Debugf("Sending light state to lifx bulb: %s", jsonState)

		if state.OnOff != nil {
			err := light.ApplyOnOff(*state.OnOff)
			if err != nil {
				return err
			}
		}

		if state.Color != nil || state.Brightness != nil || state.Transition != nil {
			if state.Color == nil {
				return fmt.Errorf("Color value missing from batch set")
			}

			if state.Brightness == nil {
				return fmt.Errorf("Brightness value missing from batch set")
			}

			if state.Transition == nil {
				return fmt.Errorf("Transition value missing from batch set")
			}

			switch state.Color.Mode {
			case "hue":
				return client.LightColour(
					bulb,
					uint16(*state.Color.Hue*math.MaxUint16),
					uint16(*state.Color.Saturation*math.MaxUint16),
					uint16(*state.Brightness*math.MaxUint16),
					0,
					uint32(*state.Transition),
				)

			case "xy":
				//TODO: Lifx does not support XY color
				return fmt.Errorf("XY color mode is not yet supported in the Lifx driver")

			case "temperature":
				client.LightColour(
					bulb,
					0,
					0,
					uint16(*state.Brightness*math.MaxUint16),
					uint16(*state.Color.Temperature),
					uint32(*state.Transition),
				)

			default:
				return fmt.Errorf("Unknown color mode %s", state.Color.Mode)
			}

		}

		return nil
	}

	bulb.SetStateHandler(buildStateHandler(bulb, light))

	return light, nil
}

func buildStateHandler(bulb *lifx.Bulb, light *devices.LightDevice) lifx.StateHandler {

	return func(bulbState *lifx.BulbState) {

		spew.Dump(bulbState)

		state := &devices.LightDeviceState{}

		onOff := int(bulbState.Power) > 0
		state.OnOff = &onOff

		brightness := float64(bulbState.Brightness) / math.MaxUint16
		state.Brightness = &brightness

		color := &channels.ColorState{}
		if bulbState.Saturation == 0 {
			color.Mode = "temperature"

			temperature := int(bulbState.Kelvin)
			color.Temperature = &temperature

		} else {
			color.Mode = "hue"

			hue := float64(bulbState.Hue) / float64(math.MaxUint16)
			color.Hue = &hue

			saturation := float64(bulbState.Saturation) / float64(math.MaxUint16)
			color.Saturation = &saturation
		}

		state.Color = color

		light.SetLightState(state)
	}
}

//---------------------------------------------------------------[Utils]----------------------------------------------------------------

func isUnique(newbulb *lifx.Bulb) bool {
	ret := true
	for _, bulb := range seenlights {
		if bulb.LifxAddress == newbulb.LifxAddress {
			ret = false
		}
	}
	return ret
}
<commit_msg>Clean up logging<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"math"

	"github.com/bitly/go-simplejson"
	"github.com/ninjasphere/go-ninja"
	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-ninja/devices"
	"github.com/wolfeidau/lifx"
)

//---------------------------------------------------------------[Bulb]----------------------------------------------------------------

func NewLight(bus *ninja.DriverBus, client *lifx.Client, bulb *lifx.Bulb) (*devices.LightDevice, error) { //TODO cut this down!

	log.Infof("Making light with ID: %s Label: %s", bulb.GetLifxAddress(), bulb.GetLabel())

	sigs, _ := simplejson.NewJson([]byte(`{
			"ninja:manufacturer": "Lifx",
			"ninja:productName": "Lifx Bulb",
			"ninja:productType": "Light",
			"ninja:thingType": "light"
	}`))

	deviceBus, err := bus.AnnounceDevice(bulb.GetLifxAddress(), "light", bulb.GetLabel(), sigs)
	if err != nil {
		log.FatalError(err, "Failed to create light device bus")
	}

	light, err := devices.CreateLightDevice(bulb.GetLabel(), deviceBus)
	if err != nil {
		log.FatalError(err, "Failed to create light device")
	}

	light.EnableOnOffChannel()
	light.EnableBrightnessChannel()
	light.EnableColorChannel("temperature", "hue")
	light.EnableTransitionChannel()

	light.ApplyOnOff = func(state bool) error {
		var err error
		if state {
			err = client.LightOn(bulb)
		} else {
			err = client.LightOff(bulb)
		}
		if err != nil {
			return fmt.Errorf("Failed to set on-off state: %s", err)
		}
		return nil
	}

	light.ApplyLightState = func(state *devices.LightDeviceState) error {
		jsonState, _ := json.Marshal(state)
		log.Debugf("Sending light state to lifx bulb: %s", jsonState)

		if state.OnOff != nil {
			err := light.ApplyOnOff(*state.OnOff)
			if err != nil {
				return err
			}
		}

		if state.Color != nil || state.Brightness != nil || state.Transition != nil {
			if state.Color == nil {
				return fmt.Errorf("Color value missing from batch set")
			}

			if state.Brightness == nil {
				return fmt.Errorf("Brightness value missing from batch set")
			}

			if state.Transition == nil {
				return fmt.Errorf("Transition value missing from batch set")
			}

			switch state.Color.Mode {
			case "hue":
				return client.LightColour(
					bulb,
					uint16(*state.Color.Hue*math.MaxUint16),
					uint16(*state.Color.Saturation*math.MaxUint16),
					uint16(*state.Brightness*math.MaxUint16),
					0,
					uint32(*state.Transition),
				)

			case "xy":
				//TODO: Lifx does not support XY color
				return fmt.Errorf("XY color mode is not yet supported in the Lifx driver")

			case "temperature":
				client.LightColour(
					bulb,
					0,
					0,
					uint16(*state.Brightness*math.MaxUint16),
					uint16(*state.Color.Temperature),
					uint32(*state.Transition),
				)

			default:
				return fmt.Errorf("Unknown color mode %s", state.Color.Mode)
			}

		}

		return nil
	}

	bulb.SetStateHandler(buildStateHandler(bulb, light))

	return light, nil
}

func buildStateHandler(bulb *lifx.Bulb, light *devices.LightDevice) lifx.StateHandler {

	return func(bulbState *lifx.BulbState) {

		jsonState, _ := json.Marshal(bulbState)
		log.Debugf("Incoming state: %s", jsonState)

		state := &devices.LightDeviceState{}

		onOff := int(bulbState.Power) > 0
		state.OnOff = &onOff

		brightness := float64(bulbState.Brightness) / math.MaxUint16
		state.Brightness = &brightness

		color := &channels.ColorState{}
		if bulbState.Saturation == 0 {
			color.Mode = "temperature"

			temperature := int(bulbState.Kelvin)
			color.Temperature = &temperature

		} else {
			color.Mode = "hue"

			hue := float64(bulbState.Hue) / float64(math.MaxUint16)
			color.Hue = &hue

			saturation := float64(bulbState.Saturation) / float64(math.MaxUint16)
			color.Saturation = &saturation
		}

		state.Color = color

		light.SetLightState(state)
	}
}

//---------------------------------------------------------------[Utils]----------------------------------------------------------------

func isUnique(newbulb *lifx.Bulb) bool {
	ret := true
	for _, bulb := range seenlights {
		if bulb.LifxAddress == newbulb.LifxAddress {
			ret = false
		}
	}
	return ret
}
<|endoftext|>


<commit_before>package termite

import (
	"crypto"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sync"
)

// Content based addressing cache.
type ContentCache struct {
	dir           string
	hashFunc      crypto.Hash

	mutex         sync.Mutex
	cond          *sync.Cond
	faulting      map[string]bool
	have          map[string]bool
	inMemoryCache *LruCache

	memoryTries int
	memoryHits  int
}

// NewContentCache creates a content cache based in directory d.
// memorySize sets the maximum number of file contents to keep in
// memory.
func NewContentCache(d string) *ContentCache {
	if fi, _ := os.Lstat(d); fi == nil {
		err := os.MkdirAll(d, 0700)
		if err != nil {
			panic(err)
		}
	}

	c := &ContentCache{
		dir:           d,
		have:          ReadHexDatabase(d),
		inMemoryCache: NewLruCache(1024),
		faulting:      make(map[string]bool),
		hashFunc:      crypto.MD5,
	}
	c.cond = sync.NewCond(&c.mutex)
	return c
}

// SetMemoryCacheSize readjusts the size of the in-memory content
// cache.  Not thread safe.
func (me *ContentCache) SetMemoryCacheSize(fileCount int) {
	if fileCount == 0 {
		me.inMemoryCache = nil
		return
	}
	if me.inMemoryCache.Size() != fileCount {
		me.inMemoryCache = NewLruCache(fileCount)
	}
}

func (me *ContentCache) MemoryHitRate() float64 {
	if me.memoryTries == 0 {
		return 0.0
	}
	return float64(me.memoryHits) / float64(me.memoryTries)
}

func HashPath(dir string, hash string) string {
	s := fmt.Sprintf("%x", hash)
	prefix := s[:2]
	name := s[2:]
	dst := filepath.Join(dir, prefix, name)
	prefixDir, _ := filepath.Split(dst)
	if err := os.MkdirAll(prefixDir, 0700); err != nil {
		log.Fatal("MkdirAll error:", err)
	}
	return dst
}

func (me *ContentCache) HasHash(hash string) bool {
	me.mutex.Lock()
	defer me.mutex.Unlock()

	return me.have[hash]
}

func (me *ContentCache) ContentsIfLoaded(hash string) []byte {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for me.faulting[hash] {
		me.cond.Wait()
	}
	me.memoryTries++
	if me.inMemoryCache == nil {
		return nil
	}
	c := me.inMemoryCache.Get(hash)
	if c != nil {
		me.memoryHits++
		return c.([]byte)
	}
	return nil
}

func (me *ContentCache) Path(hash string) string {
	return HashPath(me.dir, hash)
}

func (me *ContentCache) NewHashWriter() *HashWriter {
	return NewHashWriter(me.dir, me.hashFunc)
}

type HashWriter struct {
	hasher hash.Hash
	dest   *os.File
}

func NewHashWriter(dir string, hashfunc crypto.Hash) *HashWriter {
	me := &HashWriter{}
	tmp, err := ioutil.TempFile(dir, ".hashtemp")
	if err != nil {
		log.Panic("NewHashWriter: ", err)
	}

	me.dest = tmp
	me.hasher = hashfunc.New()
	return me
}

func (me *HashWriter) Sum() string {
	return string(me.hasher.Sum())
}

func (me *HashWriter) Write(p []byte) (n int, err error) {
	n, err = me.dest.Write(p)
	me.hasher.Write(p[:n])
	return n, err
}

func (me *HashWriter) WriteClose(p []byte) (err error) {
	_, err = me.Write(p)
	if err != nil {
		return err
	}
	err = me.Close()
	return err
}

func (me *HashWriter) CopyClose(input io.Reader, size int64) error {
	_, err := io.CopyN(me, input, size)
	if err != nil {
		return err
	}
	err = me.Close()
	return err
}

func (me *HashWriter) Close() error {
	me.dest.Chmod(0444)
	err := me.dest.Close()
	if err != nil {
		return err
	}
	src := me.dest.Name()
	dir, _ := filepath.Split(src)
	sum := me.Sum()
	sumpath := HashPath(dir, sum)

	log.Printf("saving hash %x\n", sum)
	err = os.Rename(src, sumpath)
	if err != nil {
		log.Fatal("Rename failed", err)
	}
	return err
}

const _BUFSIZE = 32 * 1024

func (me *ContentCache) DestructiveSavePath(path string) (hash string, err error) {
	var f *os.File
	f, err = os.Open(path)
	if err != nil {
		return "", err
	}
	before, _ := f.Stat()
	defer f.Close()

	h := me.hashFunc.New()

	var content []byte
	if before.Size < _MEMORY_LIMIT {
		content, err = ioutil.ReadAll(f)
		if err != nil {
			return "", err
		}

		h.Write(content)
	} else {
		io.Copy(h, f)
	}

	s := string(h.Sum())
	if me.HasHash(s) {
		os.Remove(path)
		return s, nil
	}

	me.mutex.Lock()
	me.have[s] = true
	if content != nil && me.inMemoryCache != nil {
		me.inMemoryCache.Add(s, content)
	}
	me.mutex.Unlock()

	p := me.Path(s)
	err = os.Rename(path, p)
	if err != nil {
		log.Fatal("Rename failed", err)
	}
	f.Chmod(0444)
	after, _ := f.Stat()
	if after.Mtime_ns != before.Mtime_ns || after.Size != before.Size {
		log.Fatal("File changed during save", OsFileInfo(*before), OsFileInfo(*after))
	}
	return s, nil
}

func (me *ContentCache) SavePath(path string) (hash string) {
	f, err := os.Open(path)
	if err != nil {
		log.Println("SavePath:", err)
		return ""
	}
	defer f.Close()

	fi, _ := f.Stat()
	return me.SaveStream(f, fi.Size)
}

// FaultIn loads the data from disk into the memory cache.
func (me *ContentCache) FaultIn(hash string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for !me.inMemoryCache.Has(hash) && me.faulting[hash] {
		me.cond.Wait()
	}
	if me.inMemoryCache.Has(hash) {
		return
	}

	me.faulting[hash] = true
	me.mutex.Unlock()
	c, err := ioutil.ReadFile(me.Path(hash))
	me.mutex.Lock()
	if err != nil {
		log.Fatal("FaultIn:", err)
	}
	delete(me.faulting, hash)
	me.inMemoryCache.Add(hash, c)
	me.cond.Broadcast()
}

func (me *ContentCache) Save(content []byte) (hash string) {
	return me.saveViaMemory(content)
}

func (me *ContentCache) saveViaMemory(content []byte) (hash string) {
	writer := me.NewHashWriter()
	err := writer.WriteClose(content)
	if err != nil {
		log.Println("saveViaMemory:", err)
		return ""
	}
	hash = writer.Sum()

	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[hash] = true
	if me.inMemoryCache != nil {
		me.inMemoryCache.Add(hash, content)
	}
	return hash
}

const _MEMORY_LIMIT = 128 * 1024

func (me *ContentCache) SaveStream(input io.Reader, size int64) (hash string) {
	if size < _MEMORY_LIMIT {
		b, err := ioutil.ReadAll(input)
		if int64(len(b)) != size {
			log.Panicf("SaveStream: short read: %v %v", len(b), err)
		}

		return me.saveViaMemory(b)
	}

	dup := me.NewHashWriter()
	err := dup.CopyClose(input, size)
	if err != nil {
		return ""
	}
	hash = dup.Sum()

	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[hash] = true
	return hash
}

func (me *ContentCache) FetchFromServer(fetcher func(req *ContentRequest, rep *ContentResponse) error,
	hash string) error {
	if me.HasHash(hash) {
		return nil
	}
	chunkSize := 1 << 18

	output := me.NewHashWriter()
	written := 0
	for {
		req := &ContentRequest{
			Hash:  hash,
			Start: written,
			End:   written + chunkSize,
		}

		rep := &ContentResponse{}
		err := fetcher(req, rep)
		if err != nil {
			log.Println("FileContent error:", err)
			return err
		}

		if len(rep.Chunk) < chunkSize && written == 0 {
			output.Close()
			saved := me.Save(rep.Chunk)
			if saved != hash {
				log.Fatalf("Corruption: savedHash %x != requested hash %x.", saved, hash)
			}
			return nil
		}

		n, err := output.Write(rep.Chunk)
		written += n
		if err != nil {
			return err
		}
		if len(rep.Chunk) < chunkSize {
			break
		}
	}

	output.Close()
	saved := string(output.hasher.Sum())
	if saved != hash {
		log.Fatalf("Corruption: savedHash %x != requested hash %x.", saved, hash)
	}
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[saved] = true
	return nil
}
<commit_msg>Only create HashWriter when we go through disk.<commit_after>package termite

import (
	"crypto"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sync"
)

// Content based addressing cache.
type ContentCache struct {
	dir           string
	hashFunc      crypto.Hash

	mutex         sync.Mutex
	cond          *sync.Cond
	faulting      map[string]bool
	have          map[string]bool
	inMemoryCache *LruCache

	memoryTries int
	memoryHits  int
}

// NewContentCache creates a content cache based in directory d.
// memorySize sets the maximum number of file contents to keep in
// memory.
func NewContentCache(d string) *ContentCache {
	if fi, _ := os.Lstat(d); fi == nil {
		err := os.MkdirAll(d, 0700)
		if err != nil {
			panic(err)
		}
	}

	c := &ContentCache{
		dir:           d,
		have:          ReadHexDatabase(d),
		inMemoryCache: NewLruCache(1024),
		faulting:      make(map[string]bool),
		hashFunc:      crypto.MD5,
	}
	c.cond = sync.NewCond(&c.mutex)
	return c
}

// SetMemoryCacheSize readjusts the size of the in-memory content
// cache.  Not thread safe.
func (me *ContentCache) SetMemoryCacheSize(fileCount int) {
	if fileCount == 0 {
		me.inMemoryCache = nil
		return
	}
	if me.inMemoryCache.Size() != fileCount {
		me.inMemoryCache = NewLruCache(fileCount)
	}
}

func (me *ContentCache) MemoryHitRate() float64 {
	if me.memoryTries == 0 {
		return 0.0
	}
	return float64(me.memoryHits) / float64(me.memoryTries)
}

func HashPath(dir string, hash string) string {
	s := fmt.Sprintf("%x", hash)
	prefix := s[:2]
	name := s[2:]
	dst := filepath.Join(dir, prefix, name)
	prefixDir, _ := filepath.Split(dst)
	if err := os.MkdirAll(prefixDir, 0700); err != nil {
		log.Fatal("MkdirAll error:", err)
	}
	return dst
}

func (me *ContentCache) HasHash(hash string) bool {
	me.mutex.Lock()
	defer me.mutex.Unlock()

	return me.have[hash]
}

func (me *ContentCache) ContentsIfLoaded(hash string) []byte {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for me.faulting[hash] {
		me.cond.Wait()
	}
	me.memoryTries++
	if me.inMemoryCache == nil {
		return nil
	}
	c := me.inMemoryCache.Get(hash)
	if c != nil {
		me.memoryHits++
		return c.([]byte)
	}
	return nil
}

func (me *ContentCache) Path(hash string) string {
	return HashPath(me.dir, hash)
}

func (me *ContentCache) NewHashWriter() *HashWriter {
	return NewHashWriter(me.dir, me.hashFunc)
}

type HashWriter struct {
	hasher hash.Hash
	dest   *os.File
}

func NewHashWriter(dir string, hashfunc crypto.Hash) *HashWriter {
	me := &HashWriter{}
	tmp, err := ioutil.TempFile(dir, ".hashtemp")
	if err != nil {
		log.Panic("NewHashWriter: ", err)
	}

	me.dest = tmp
	me.hasher = hashfunc.New()
	return me
}

func (me *HashWriter) Sum() string {
	return string(me.hasher.Sum())
}

func (me *HashWriter) Write(p []byte) (n int, err error) {
	n, err = me.dest.Write(p)
	me.hasher.Write(p[:n])
	return n, err
}

func (me *HashWriter) WriteClose(p []byte) (err error) {
	_, err = me.Write(p)
	if err != nil {
		return err
	}
	err = me.Close()
	return err
}

func (me *HashWriter) CopyClose(input io.Reader, size int64) error {
	_, err := io.CopyN(me, input, size)
	if err != nil {
		return err
	}
	err = me.Close()
	return err
}

func (me *HashWriter) Close() error {
	me.dest.Chmod(0444)
	err := me.dest.Close()
	if err != nil {
		return err
	}
	src := me.dest.Name()
	dir, _ := filepath.Split(src)
	sum := me.Sum()
	sumpath := HashPath(dir, sum)

	log.Printf("saving hash %x\n", sum)
	err = os.Rename(src, sumpath)
	if err != nil {
		log.Fatal("Rename failed", err)
	}
	return err
}

const _BUFSIZE = 32 * 1024

func (me *ContentCache) DestructiveSavePath(path string) (hash string, err error) {
	var f *os.File
	f, err = os.Open(path)
	if err != nil {
		return "", err
	}
	before, _ := f.Stat()
	defer f.Close()

	h := me.hashFunc.New()

	var content []byte
	if before.Size < _MEMORY_LIMIT {
		content, err = ioutil.ReadAll(f)
		if err != nil {
			return "", err
		}

		h.Write(content)
	} else {
		io.Copy(h, f)
	}

	s := string(h.Sum())
	if me.HasHash(s) {
		os.Remove(path)
		return s, nil
	}

	me.mutex.Lock()
	me.have[s] = true
	if content != nil && me.inMemoryCache != nil {
		me.inMemoryCache.Add(s, content)
	}
	me.mutex.Unlock()

	p := me.Path(s)
	err = os.Rename(path, p)
	if err != nil {
		log.Fatal("Rename failed", err)
	}
	f.Chmod(0444)
	after, _ := f.Stat()
	if after.Mtime_ns != before.Mtime_ns || after.Size != before.Size {
		log.Fatal("File changed during save", OsFileInfo(*before), OsFileInfo(*after))
	}
	return s, nil
}

func (me *ContentCache) SavePath(path string) (hash string) {
	f, err := os.Open(path)
	if err != nil {
		log.Println("SavePath:", err)
		return ""
	}
	defer f.Close()

	fi, _ := f.Stat()
	return me.SaveStream(f, fi.Size)
}

// FaultIn loads the data from disk into the memory cache.
func (me *ContentCache) FaultIn(hash string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for !me.inMemoryCache.Has(hash) && me.faulting[hash] {
		me.cond.Wait()
	}
	if me.inMemoryCache.Has(hash) {
		return
	}

	me.faulting[hash] = true
	me.mutex.Unlock()
	c, err := ioutil.ReadFile(me.Path(hash))
	me.mutex.Lock()
	if err != nil {
		log.Fatal("FaultIn:", err)
	}
	delete(me.faulting, hash)
	me.inMemoryCache.Add(hash, c)
	me.cond.Broadcast()
}

func (me *ContentCache) Save(content []byte) (hash string) {
	return me.saveViaMemory(content)
}

func (me *ContentCache) saveViaMemory(content []byte) (hash string) {
	writer := me.NewHashWriter()
	err := writer.WriteClose(content)
	if err != nil {
		log.Println("saveViaMemory:", err)
		return ""
	}
	hash = writer.Sum()

	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[hash] = true
	if me.inMemoryCache != nil {
		me.inMemoryCache.Add(hash, content)
	}
	return hash
}

const _MEMORY_LIMIT = 128 * 1024

func (me *ContentCache) SaveStream(input io.Reader, size int64) (hash string) {
	if size < _MEMORY_LIMIT {
		b, err := ioutil.ReadAll(input)
		if int64(len(b)) != size {
			log.Panicf("SaveStream: short read: %v %v", len(b), err)
		}

		return me.saveViaMemory(b)
	}

	dup := me.NewHashWriter()
	err := dup.CopyClose(input, size)
	if err != nil {
		return ""
	}
	hash = dup.Sum()

	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[hash] = true
	return hash
}

func (me *ContentCache) FetchFromServer(fetcher func(req *ContentRequest, rep *ContentResponse) error,
	hash string) error {
	if me.HasHash(hash) {
		return nil
	}
	chunkSize := 1 << 18

	var output *HashWriter
	written := 0
	for {
		req := &ContentRequest{
			Hash:  hash,
			Start: written,
			End:   written + chunkSize,
		}

		rep := &ContentResponse{}
		err := fetcher(req, rep)
		if err != nil {
			log.Println("FileContent error:", err)
			return err
		}

		if len(rep.Chunk) < chunkSize && written == 0 {
			saved := me.Save(rep.Chunk)
			if saved != hash {
				log.Fatalf("Corruption: savedHash %x != requested hash %x.", saved, hash)
			}
			return nil
		} else if output == nil {
			output = me.NewHashWriter()
 		}

		n, err := output.Write(rep.Chunk)
		written += n
		if err != nil {
			return err
		}
		if len(rep.Chunk) < chunkSize {
			break
		}
	}

	output.Close()
	saved := string(output.hasher.Sum())
	if saved != hash {
		log.Fatalf("Corruption: savedHash %x != requested hash %x.", saved, hash)
	}
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.have[saved] = true
	return nil
}
<|endoftext|>