<commit_before>package weed_server

// https://yusufs.medium.com/creating-distributed-kv-database-by-implementing-raft-consensus-using-golang-d0884eef2e28
// https://github.com/Jille/raft-grpc-example/blob/cd5bcab0218f008e044fbeee4facdd01b06018ad/application.go#L18

import (
	"fmt"
	transport "github.com/Jille/raft-grpc-transport"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/hashicorp/raft"
	boltdb "github.com/hashicorp/raft-boltdb"
	"google.golang.org/grpc"
	"math/rand"
	"os"
	"path/filepath"
	"time"
)

func (s *RaftServer) AddPeersConfiguration() (cfg raft.Configuration) {
	for _, peer := range s.peers {
		cfg.Servers = append(cfg.Servers, raft.Server{
			Suffrage: raft.Voter,
			ID:       raft.ServerID(peer.String()),
			Address:  raft.ServerAddress(peer.ToGrpcAddress()),
		})
	}
	return cfg
}

func (s *RaftServer) UpdatePeers() {
	for {
		select {
		case isLeader := <-s.RaftHashicorp.LeaderCh():
			if isLeader {
				peerLeader := s.serverAddr.String()
				existsPeerName := make(map[string]bool)
				for _, server := range s.RaftHashicorp.GetConfiguration().Configuration().Servers {
					if string(server.ID) == peerLeader {
						continue
					}
					existsPeerName[string(server.ID)] = true
				}
				for _, peer := range s.peers {
					if peer.String() == peerLeader || existsPeerName[peer.String()] {
						continue
					}
					glog.V(0).Infof("adding new peer: %s", peer.String())
					s.RaftHashicorp.AddVoter(
						raft.ServerID(peer.String()), raft.ServerAddress(peer.ToGrpcAddress()), 0, 0)
				}
				for peer, _ := range existsPeerName {
					if _, found := s.peers[peer]; !found {
						glog.V(0).Infof("removing old peer: %s", peer)
						s.RaftHashicorp.RemoveServer(raft.ServerID(peer), 0, 0)
					}
				}
				if _, found := s.peers[peerLeader]; !found {
					glog.V(0).Infof("removing old leader peer: %s", peerLeader)
					s.RaftHashicorp.RemoveServer(raft.ServerID(peerLeader), 0, 0)
				}
			}
			break
		}
	}
}

func NewHashicorpRaftServer(option *RaftServerOption) (*RaftServer, error) {
	s := &RaftServer{
		peers:      option.Peers,
		serverAddr: option.ServerAddr,
		dataDir:    option.DataDir,
		topo:       option.Topo,
	}

	c := raft.DefaultConfig()
	c.LocalID = raft.ServerID(s.serverAddr.String()) // TODO maybee the IP:port address will change
	c.NoSnapshotRestoreOnStart = option.RaftResumeState
	c.HeartbeatTimeout = time.Duration(float64(option.HeartbeatInterval) * (rand.Float64()*0.25 + 1))
	c.ElectionTimeout = option.ElectionTimeout
	if c.LeaderLeaseTimeout > c.HeartbeatTimeout {
		c.LeaderLeaseTimeout = c.HeartbeatTimeout
	}
	if glog.V(4) {
		c.LogLevel = "Debug"
	} else if glog.V(2) {
		c.LogLevel = "Info"
	} else if glog.V(1) {
		c.LogLevel = "Warn"
	} else if glog.V(0) {
		c.LogLevel = "Error"
	}

	baseDir := s.dataDir

	ldb, err := boltdb.NewBoltStore(filepath.Join(baseDir, "logs.dat"))
	if err != nil {
		return nil, fmt.Errorf(`boltdb.NewBoltStore(%q): %v`, filepath.Join(baseDir, "logs.dat"), err)
	}

	sdb, err := boltdb.NewBoltStore(filepath.Join(baseDir, "stable.dat"))
	if err != nil {
		return nil, fmt.Errorf(`boltdb.NewBoltStore(%q): %v`, filepath.Join(baseDir, "stable.dat"), err)
	}

	fss, err := raft.NewFileSnapshotStore(baseDir, 3, os.Stderr)
	if err != nil {
		return nil, fmt.Errorf(`raft.NewFileSnapshotStore(%q, ...): %v`, baseDir, err)
	}

	s.TransportManager = transport.New(raft.ServerAddress(s.serverAddr), []grpc.DialOption{option.GrpcDialOption})

	stateMachine := StateMachine{topo: option.Topo}
	s.RaftHashicorp, err = raft.NewRaft(c, &stateMachine, ldb, sdb, fss, s.TransportManager.Transport())
	if err != nil {
		return nil, fmt.Errorf("raft.NewRaft: %v", err)
	}
	if option.RaftBootstrap || len(s.RaftHashicorp.GetConfiguration().Configuration().Servers) == 0 {
		cfg := s.AddPeersConfiguration()
		glog.V(0).Infoln("Bootstrapping new cluster %+v", cfg)
		f := s.RaftHashicorp.BootstrapCluster(cfg)
		if err := f.Error(); err != nil {
			return nil, fmt.Errorf("raft.Raft.BootstrapCluster: %v", err)
		}
	} else {
		go s.UpdatePeers()
	}

	ticker := time.NewTicker(c.HeartbeatTimeout * 10)
	if glog.V(4) {
		go func() {
			for {
				select {
				case <-ticker.C:
					cfuture := s.RaftHashicorp.GetConfiguration()
					if err = cfuture.Error(); err != nil {
						glog.Fatalf("error getting config: %s", err)
					}
					configuration := cfuture.Configuration()
					glog.V(4).Infof("Showing peers known by %s:\n%+v", s.RaftHashicorp.String(), configuration.Servers)
				}
			}
		}()
	}

	return s, nil
}
<commit_msg>fix glog Info<commit_after>package weed_server

// https://yusufs.medium.com/creating-distributed-kv-database-by-implementing-raft-consensus-using-golang-d0884eef2e28
// https://github.com/Jille/raft-grpc-example/blob/cd5bcab0218f008e044fbeee4facdd01b06018ad/application.go#L18

import (
	"fmt"
	transport "github.com/Jille/raft-grpc-transport"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/hashicorp/raft"
	boltdb "github.com/hashicorp/raft-boltdb"
	"google.golang.org/grpc"
	"math/rand"
	"os"
	"path/filepath"
	"time"
)

func (s *RaftServer) AddPeersConfiguration() (cfg raft.Configuration) {
	for _, peer := range s.peers {
		cfg.Servers = append(cfg.Servers, raft.Server{
			Suffrage: raft.Voter,
			ID:       raft.ServerID(peer.String()),
			Address:  raft.ServerAddress(peer.ToGrpcAddress()),
		})
	}
	return cfg
}

func (s *RaftServer) UpdatePeers() {
	for {
		select {
		case isLeader := <-s.RaftHashicorp.LeaderCh():
			if isLeader {
				peerLeader := s.serverAddr.String()
				existsPeerName := make(map[string]bool)
				for _, server := range s.RaftHashicorp.GetConfiguration().Configuration().Servers {
					if string(server.ID) == peerLeader {
						continue
					}
					existsPeerName[string(server.ID)] = true
				}
				for _, peer := range s.peers {
					if peer.String() == peerLeader || existsPeerName[peer.String()] {
						continue
					}
					glog.V(0).Infof("adding new peer: %s", peer.String())
					s.RaftHashicorp.AddVoter(
						raft.ServerID(peer.String()), raft.ServerAddress(peer.ToGrpcAddress()), 0, 0)
				}
				for peer, _ := range existsPeerName {
					if _, found := s.peers[peer]; !found {
						glog.V(0).Infof("removing old peer: %s", peer)
						s.RaftHashicorp.RemoveServer(raft.ServerID(peer), 0, 0)
					}
				}
				if _, found := s.peers[peerLeader]; !found {
					glog.V(0).Infof("removing old leader peer: %s", peerLeader)
					s.RaftHashicorp.RemoveServer(raft.ServerID(peerLeader), 0, 0)
				}
			}
			break
		}
	}
}

func NewHashicorpRaftServer(option *RaftServerOption) (*RaftServer, error) {
	s := &RaftServer{
		peers:      option.Peers,
		serverAddr: option.ServerAddr,
		dataDir:    option.DataDir,
		topo:       option.Topo,
	}

	c := raft.DefaultConfig()
	c.LocalID = raft.ServerID(s.serverAddr.String()) // TODO maybee the IP:port address will change
	c.NoSnapshotRestoreOnStart = option.RaftResumeState
	c.HeartbeatTimeout = time.Duration(float64(option.HeartbeatInterval) * (rand.Float64()*0.25 + 1))
	c.ElectionTimeout = option.ElectionTimeout
	if c.LeaderLeaseTimeout > c.HeartbeatTimeout {
		c.LeaderLeaseTimeout = c.HeartbeatTimeout
	}
	if glog.V(4) {
		c.LogLevel = "Debug"
	} else if glog.V(2) {
		c.LogLevel = "Info"
	} else if glog.V(1) {
		c.LogLevel = "Warn"
	} else if glog.V(0) {
		c.LogLevel = "Error"
	}

	baseDir := s.dataDir

	ldb, err := boltdb.NewBoltStore(filepath.Join(baseDir, "logs.dat"))
	if err != nil {
		return nil, fmt.Errorf(`boltdb.NewBoltStore(%q): %v`, filepath.Join(baseDir, "logs.dat"), err)
	}

	sdb, err := boltdb.NewBoltStore(filepath.Join(baseDir, "stable.dat"))
	if err != nil {
		return nil, fmt.Errorf(`boltdb.NewBoltStore(%q): %v`, filepath.Join(baseDir, "stable.dat"), err)
	}

	fss, err := raft.NewFileSnapshotStore(baseDir, 3, os.Stderr)
	if err != nil {
		return nil, fmt.Errorf(`raft.NewFileSnapshotStore(%q, ...): %v`, baseDir, err)
	}

	s.TransportManager = transport.New(raft.ServerAddress(s.serverAddr), []grpc.DialOption{option.GrpcDialOption})

	stateMachine := StateMachine{topo: option.Topo}
	s.RaftHashicorp, err = raft.NewRaft(c, &stateMachine, ldb, sdb, fss, s.TransportManager.Transport())
	if err != nil {
		return nil, fmt.Errorf("raft.NewRaft: %v", err)
	}
	if option.RaftBootstrap || len(s.RaftHashicorp.GetConfiguration().Configuration().Servers) == 0 {
		cfg := s.AddPeersConfiguration()
		glog.V(0).Infof("Bootstrapping new cluster %+v", cfg)
		f := s.RaftHashicorp.BootstrapCluster(cfg)
		if err := f.Error(); err != nil {
			return nil, fmt.Errorf("raft.Raft.BootstrapCluster: %v", err)
		}
	} else {
		go s.UpdatePeers()
	}

	ticker := time.NewTicker(c.HeartbeatTimeout * 10)
	if glog.V(4) {
		go func() {
			for {
				select {
				case <-ticker.C:
					cfuture := s.RaftHashicorp.GetConfiguration()
					if err = cfuture.Error(); err != nil {
						glog.Fatalf("error getting config: %s", err)
					}
					configuration := cfuture.Configuration()
					glog.V(4).Infof("Showing peers known by %s:\n%+v", s.RaftHashicorp.String(), configuration.Servers)
				}
			}
		}()
	}

	return s, nil
}
<|endoftext|>


<commit_before>package storage

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type DiskLocation struct {
	Directory              string
	IdxDirectory           string
	MaxVolumeCount         int
	OriginalMaxVolumeCount int
	MinFreeSpacePercent    float32
	volumes                map[needle.VolumeId]*Volume
	volumesLock            sync.RWMutex

	// erasure coding
	ecVolumes     map[needle.VolumeId]*erasure_coding.EcVolume
	ecVolumesLock sync.RWMutex

	isDiskSpaceLow bool
}

func NewDiskLocation(dir string, maxVolumeCount int, minFreeSpacePercent float32, idxDir string) *DiskLocation {
	dir = util.ResolvePath(dir)
	if idxDir == "" {
		idxDir = dir
	} else {
		idxDir = util.ResolvePath(idxDir)
	}
	location := &DiskLocation{
		Directory:              dir,
		IdxDirectory:           idxDir,
		MaxVolumeCount:         maxVolumeCount,
		OriginalMaxVolumeCount: maxVolumeCount,
		MinFreeSpacePercent:    minFreeSpacePercent,
	}
	location.volumes = make(map[needle.VolumeId]*Volume)
	location.ecVolumes = make(map[needle.VolumeId]*erasure_coding.EcVolume)
	go location.CheckDiskSpace()
	return location
}

func volumeIdFromFileName(filename string) (needle.VolumeId, string, error) {
	if strings.HasSuffix(filename, ".idx") || strings.HasSuffix(filename, ".vif") {
		base := filename[:len(filename)-4]
		collection, volumeId, err := parseCollectionVolumeId(base)
		return volumeId, collection, err
	}

	return 0, "", fmt.Errorf("file is not a volume: %s", filename)
}

func parseCollectionVolumeId(base string) (collection string, vid needle.VolumeId, err error) {
	i := strings.LastIndex(base, "_")
	if i > 0 {
		collection, base = base[0:i], base[i+1:]
	}
	vol, err := needle.NewVolumeId(base)
	return collection, vol, err
}

func (l *DiskLocation) loadExistingVolume(fileInfo os.FileInfo, needleMapKind NeedleMapType) bool {
	basename := fileInfo.Name()
	if fileInfo.IsDir() {
		return false
	}
	if !strings.HasSuffix(basename, ".idx") && !strings.HasSuffix(basename, ".vif") {
		return false
	}
	volumeName := basename[:len(basename)-4]

	// check for incomplete volume
	noteFile := l.Directory + "/" + volumeName + ".note"
	if util.FileExists(noteFile) {
		note, _ := ioutil.ReadFile(noteFile)
		glog.Warningf("volume %s was not completed: %s", volumeName, string(note))
		removeVolumeFiles(l.Directory + "/" + volumeName)
		removeVolumeFiles(l.IdxDirectory + "/" + volumeName)
		return false
	}

	// parse out collection, volume id
	vid, collection, err := volumeIdFromFileName(basename)
	if err != nil {
		glog.Warningf("get volume id failed, %s, err : %s", volumeName, err)
		return false
	}

	// avoid loading one volume more than once
	l.volumesLock.RLock()
	_, found := l.volumes[vid]
	l.volumesLock.RUnlock()
	if found {
		glog.V(1).Infof("loaded volume, %v", vid)
		return true
	}

	// load the volume
	v, e := NewVolume(l.Directory, l.IdxDirectory, collection, vid, needleMapKind, nil, nil, 0, 0)
	if e != nil {
		glog.V(0).Infof("new volume %s error %s", volumeName, e)
		return false
	}

	l.SetVolume(vid, v)

	size, _, _ := v.FileStat()
	glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s",
		l.Directory+"/"+volumeName+".dat", v.ReplicaPlacement, v.Version(), size, v.Ttl.String())
	return true
}

func (l *DiskLocation) concurrentLoadingVolumes(needleMapKind NeedleMapType, concurrency int) {

	task_queue := make(chan os.FileInfo, 10*concurrency)
	go func() {
		if fileInfos, err := ioutil.ReadDir(l.Directory); err == nil {
			for _, fi := range fileInfos {
				task_queue <- fi
			}
		}
		close(task_queue)
	}()

	var wg sync.WaitGroup
	for workerNum := 0; workerNum < concurrency; workerNum++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for fi := range task_queue {
				_ = l.loadExistingVolume(fi, needleMapKind)
			}
		}()
	}
	wg.Wait()

}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {

	l.concurrentLoadingVolumes(needleMapKind, 10)
	glog.V(0).Infof("Store started on dir: %s with %d volumes max %d", l.Directory, len(l.volumes), l.MaxVolumeCount)

	l.loadAllEcShards()
	glog.V(0).Infof("Store started on dir: %s with %d ec shards", l.Directory, len(l.ecVolumes))

}

func (l *DiskLocation) DeleteCollectionFromDiskLocation(collection string) (e error) {

	l.volumesLock.Lock()
	delVolsMap := l.unmountVolumeByCollection(collection)
	l.volumesLock.Unlock()

	l.ecVolumesLock.Lock()
	delEcVolsMap := l.unmountEcVolumeByCollection(collection)
	l.ecVolumesLock.Unlock()

	errChain := make(chan error, 2)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		for _, v := range delVolsMap {
			if err := v.Destroy(); err != nil {
				errChain <- err
			}
		}
		wg.Done()
	}()

	go func() {
		for _, v := range delEcVolsMap {
			v.Destroy()
		}
		wg.Done()
	}()

	go func() {
		wg.Wait()
		close(errChain)
	}()

	errBuilder := strings.Builder{}
	for err := range errChain {
		errBuilder.WriteString(err.Error())
		errBuilder.WriteString("; ")
	}
	if errBuilder.Len() > 0 {
		e = fmt.Errorf(errBuilder.String())
	}

	return
}

func (l *DiskLocation) deleteVolumeById(vid needle.VolumeId) (found bool, e error) {
	v, ok := l.volumes[vid]
	if !ok {
		return
	}
	e = v.Destroy()
	if e != nil {
		return
	}
	found = true
	delete(l.volumes, vid)
	return
}

func (l *DiskLocation) LoadVolume(vid needle.VolumeId, needleMapKind NeedleMapType) bool {
	if fileInfo, found := l.LocateVolume(vid); found {
		return l.loadExistingVolume(fileInfo, needleMapKind)
	}
	return false
}

func (l *DiskLocation) DeleteVolume(vid needle.VolumeId) error {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not found, VolumeId: %d", vid)
	}
	_, err := l.deleteVolumeById(vid)
	return err
}

func (l *DiskLocation) UnloadVolume(vid needle.VolumeId) error {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	v, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not loaded, VolumeId: %d", vid)
	}
	v.Close()
	delete(l.volumes, vid)
	return nil
}

func (l *DiskLocation) unmountVolumeByCollection(collectionName string) map[needle.VolumeId]*Volume {
	deltaVols := make(map[needle.VolumeId]*Volume, 0)
	for k, v := range l.volumes {
		if v.Collection == collectionName && !v.isCompacting {
			deltaVols[k] = v
		}
	}

	for k := range deltaVols {
		delete(l.volumes, k)
	}
	return deltaVols
}

func (l *DiskLocation) SetVolume(vid needle.VolumeId, volume *Volume) {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	l.volumes[vid] = volume
	volume.location = l
}

func (l *DiskLocation) FindVolume(vid needle.VolumeId) (*Volume, bool) {
	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	v, ok := l.volumes[vid]
	return v, ok
}

func (l *DiskLocation) VolumesLen() int {
	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	return len(l.volumes)
}

func (l *DiskLocation) Close() {
	l.volumesLock.Lock()
	for _, v := range l.volumes {
		v.Close()
	}
	l.volumesLock.Unlock()

	l.ecVolumesLock.Lock()
	for _, ecVolume := range l.ecVolumes {
		ecVolume.Close()
	}
	l.ecVolumesLock.Unlock()

	return
}

func (l *DiskLocation) LocateVolume(vid needle.VolumeId) (os.FileInfo, bool) {
	if fileInfos, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, fileInfo := range fileInfos {
			volId, _, err := volumeIdFromFileName(fileInfo.Name())
			if vid == volId && err == nil {
				return fileInfo, true
			}
		}
	}

	return nil, false
}

func (l *DiskLocation) UnUsedSpace(volumeSizeLimit uint64) (unUsedSpace uint64) {

	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	for _, vol := range l.volumes {
		if vol.IsReadOnly() {
			continue
		}
		datSize, idxSize, _ := vol.FileStat()
		unUsedSpace += volumeSizeLimit - (datSize + idxSize)
	}

	return
}

func (l *DiskLocation) CheckDiskSpace() {
	for {
		if dir, e := filepath.Abs(l.Directory); e == nil {
			s := stats.NewDiskStatus(dir)
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "all").Set(float64(s.All))
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "used").Set(float64(s.Used))
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "free").Set(float64(s.Free))
			if (s.PercentFree < l.MinFreeSpacePercent) != l.isDiskSpaceLow {
				l.isDiskSpaceLow = !l.isDiskSpaceLow
			}
			if l.isDiskSpaceLow {
				glog.V(0).Infof("dir %s freePercent %.2f%% < min %.2f%%, isLowDiskSpace: %v", dir, s.PercentFree, l.MinFreeSpacePercent, l.isDiskSpaceLow)
			} else {
				glog.V(4).Infof("dir %s freePercent %.2f%% < min %.2f%%, isLowDiskSpace: %v", dir, s.PercentFree, l.MinFreeSpacePercent, l.isDiskSpaceLow)
			}
		}
		time.Sleep(time.Minute)
	}

}
<commit_msg>fix race condition when loading volumes concurrently<commit_after>package storage

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/util"
)

type DiskLocation struct {
	Directory              string
	IdxDirectory           string
	MaxVolumeCount         int
	OriginalMaxVolumeCount int
	MinFreeSpacePercent    float32
	volumes                map[needle.VolumeId]*Volume
	volumesLock            sync.RWMutex

	// erasure coding
	ecVolumes     map[needle.VolumeId]*erasure_coding.EcVolume
	ecVolumesLock sync.RWMutex

	isDiskSpaceLow bool
}

func NewDiskLocation(dir string, maxVolumeCount int, minFreeSpacePercent float32, idxDir string) *DiskLocation {
	dir = util.ResolvePath(dir)
	if idxDir == "" {
		idxDir = dir
	} else {
		idxDir = util.ResolvePath(idxDir)
	}
	location := &DiskLocation{
		Directory:              dir,
		IdxDirectory:           idxDir,
		MaxVolumeCount:         maxVolumeCount,
		OriginalMaxVolumeCount: maxVolumeCount,
		MinFreeSpacePercent:    minFreeSpacePercent,
	}
	location.volumes = make(map[needle.VolumeId]*Volume)
	location.ecVolumes = make(map[needle.VolumeId]*erasure_coding.EcVolume)
	go location.CheckDiskSpace()
	return location
}

func volumeIdFromFileName(filename string) (needle.VolumeId, string, error) {
	if strings.HasSuffix(filename, ".idx") || strings.HasSuffix(filename, ".vif") {
		base := filename[:len(filename)-4]
		collection, volumeId, err := parseCollectionVolumeId(base)
		return volumeId, collection, err
	}

	return 0, "", fmt.Errorf("file is not a volume: %s", filename)
}

func parseCollectionVolumeId(base string) (collection string, vid needle.VolumeId, err error) {
	i := strings.LastIndex(base, "_")
	if i > 0 {
		collection, base = base[0:i], base[i+1:]
	}
	vol, err := needle.NewVolumeId(base)
	return collection, vol, err
}

func (l *DiskLocation) loadExistingVolume(fileInfo os.FileInfo, needleMapKind NeedleMapType) bool {
	basename := fileInfo.Name()
	if fileInfo.IsDir() {
		return false
	}
	if !strings.HasSuffix(basename, ".idx") && !strings.HasSuffix(basename, ".vif") {
		return false
	}
	volumeName := basename[:len(basename)-4]

	// check for incomplete volume
	noteFile := l.Directory + "/" + volumeName + ".note"
	if util.FileExists(noteFile) {
		note, _ := ioutil.ReadFile(noteFile)
		glog.Warningf("volume %s was not completed: %s", volumeName, string(note))
		removeVolumeFiles(l.Directory + "/" + volumeName)
		removeVolumeFiles(l.IdxDirectory + "/" + volumeName)
		return false
	}

	// parse out collection, volume id
	vid, collection, err := volumeIdFromFileName(basename)
	if err != nil {
		glog.Warningf("get volume id failed, %s, err : %s", volumeName, err)
		return false
	}

	// avoid loading one volume more than once
	l.volumesLock.Lock()
	if _, found := l.volumes[vid]; found {
		l.volumesLock.Unlock()
		glog.V(1).Infof("loaded volume, %v", vid)
		return true
	}

	// load the volume
	v, e := NewVolume(l.Directory, l.IdxDirectory, collection, vid, needleMapKind, nil, nil, 0, 0)
	if e != nil {
		l.volumesLock.Unlock()
		glog.V(0).Infof("new volume %s error %s", volumeName, e)
		return false
	}

	l.volumes[vid] = v
	v.location = l
	l.volumesLock.Unlock()

	size, _, _ := v.FileStat()
	glog.V(0).Infof("data file %s, replicaPlacement=%s v=%d size=%d ttl=%s",
		l.Directory+"/"+volumeName+".dat", v.ReplicaPlacement, v.Version(), size, v.Ttl.String())
	return true
}

func (l *DiskLocation) concurrentLoadingVolumes(needleMapKind NeedleMapType, concurrency int) {

	task_queue := make(chan os.FileInfo, 10*concurrency)
	go func() {
		if fileInfos, err := ioutil.ReadDir(l.Directory); err == nil {
			for _, fi := range fileInfos {
				task_queue <- fi
			}
		}
		close(task_queue)
	}()

	var wg sync.WaitGroup
	for workerNum := 0; workerNum < concurrency; workerNum++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for fi := range task_queue {
				_ = l.loadExistingVolume(fi, needleMapKind)
			}
		}()
	}
	wg.Wait()

}

func (l *DiskLocation) loadExistingVolumes(needleMapKind NeedleMapType) {

	l.concurrentLoadingVolumes(needleMapKind, 10)
	glog.V(0).Infof("Store started on dir: %s with %d volumes max %d", l.Directory, len(l.volumes), l.MaxVolumeCount)

	l.loadAllEcShards()
	glog.V(0).Infof("Store started on dir: %s with %d ec shards", l.Directory, len(l.ecVolumes))

}

func (l *DiskLocation) DeleteCollectionFromDiskLocation(collection string) (e error) {

	l.volumesLock.Lock()
	delVolsMap := l.unmountVolumeByCollection(collection)
	l.volumesLock.Unlock()

	l.ecVolumesLock.Lock()
	delEcVolsMap := l.unmountEcVolumeByCollection(collection)
	l.ecVolumesLock.Unlock()

	errChain := make(chan error, 2)
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		for _, v := range delVolsMap {
			if err := v.Destroy(); err != nil {
				errChain <- err
			}
		}
		wg.Done()
	}()

	go func() {
		for _, v := range delEcVolsMap {
			v.Destroy()
		}
		wg.Done()
	}()

	go func() {
		wg.Wait()
		close(errChain)
	}()

	errBuilder := strings.Builder{}
	for err := range errChain {
		errBuilder.WriteString(err.Error())
		errBuilder.WriteString("; ")
	}
	if errBuilder.Len() > 0 {
		e = fmt.Errorf(errBuilder.String())
	}

	return
}

func (l *DiskLocation) deleteVolumeById(vid needle.VolumeId) (found bool, e error) {
	v, ok := l.volumes[vid]
	if !ok {
		return
	}
	e = v.Destroy()
	if e != nil {
		return
	}
	found = true
	delete(l.volumes, vid)
	return
}

func (l *DiskLocation) LoadVolume(vid needle.VolumeId, needleMapKind NeedleMapType) bool {
	if fileInfo, found := l.LocateVolume(vid); found {
		return l.loadExistingVolume(fileInfo, needleMapKind)
	}
	return false
}

func (l *DiskLocation) DeleteVolume(vid needle.VolumeId) error {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	_, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not found, VolumeId: %d", vid)
	}
	_, err := l.deleteVolumeById(vid)
	return err
}

func (l *DiskLocation) UnloadVolume(vid needle.VolumeId) error {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	v, ok := l.volumes[vid]
	if !ok {
		return fmt.Errorf("Volume not loaded, VolumeId: %d", vid)
	}
	v.Close()
	delete(l.volumes, vid)
	return nil
}

func (l *DiskLocation) unmountVolumeByCollection(collectionName string) map[needle.VolumeId]*Volume {
	deltaVols := make(map[needle.VolumeId]*Volume, 0)
	for k, v := range l.volumes {
		if v.Collection == collectionName && !v.isCompacting {
			deltaVols[k] = v
		}
	}

	for k := range deltaVols {
		delete(l.volumes, k)
	}
	return deltaVols
}

func (l *DiskLocation) SetVolume(vid needle.VolumeId, volume *Volume) {
	l.volumesLock.Lock()
	defer l.volumesLock.Unlock()

	l.volumes[vid] = volume
	volume.location = l
}

func (l *DiskLocation) FindVolume(vid needle.VolumeId) (*Volume, bool) {
	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	v, ok := l.volumes[vid]
	return v, ok
}

func (l *DiskLocation) VolumesLen() int {
	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	return len(l.volumes)
}

func (l *DiskLocation) Close() {
	l.volumesLock.Lock()
	for _, v := range l.volumes {
		v.Close()
	}
	l.volumesLock.Unlock()

	l.ecVolumesLock.Lock()
	for _, ecVolume := range l.ecVolumes {
		ecVolume.Close()
	}
	l.ecVolumesLock.Unlock()

	return
}

func (l *DiskLocation) LocateVolume(vid needle.VolumeId) (os.FileInfo, bool) {
	if fileInfos, err := ioutil.ReadDir(l.Directory); err == nil {
		for _, fileInfo := range fileInfos {
			volId, _, err := volumeIdFromFileName(fileInfo.Name())
			if vid == volId && err == nil {
				return fileInfo, true
			}
		}
	}

	return nil, false
}

func (l *DiskLocation) UnUsedSpace(volumeSizeLimit uint64) (unUsedSpace uint64) {

	l.volumesLock.RLock()
	defer l.volumesLock.RUnlock()

	for _, vol := range l.volumes {
		if vol.IsReadOnly() {
			continue
		}
		datSize, idxSize, _ := vol.FileStat()
		unUsedSpace += volumeSizeLimit - (datSize + idxSize)
	}

	return
}

func (l *DiskLocation) CheckDiskSpace() {
	for {
		if dir, e := filepath.Abs(l.Directory); e == nil {
			s := stats.NewDiskStatus(dir)
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "all").Set(float64(s.All))
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "used").Set(float64(s.Used))
			stats.VolumeServerResourceGauge.WithLabelValues(l.Directory, "free").Set(float64(s.Free))
			if (s.PercentFree < l.MinFreeSpacePercent) != l.isDiskSpaceLow {
				l.isDiskSpaceLow = !l.isDiskSpaceLow
			}
			if l.isDiskSpaceLow {
				glog.V(0).Infof("dir %s freePercent %.2f%% < min %.2f%%, isLowDiskSpace: %v", dir, s.PercentFree, l.MinFreeSpacePercent, l.isDiskSpaceLow)
			} else {
				glog.V(4).Infof("dir %s freePercent %.2f%% < min %.2f%%, isLowDiskSpace: %v", dir, s.PercentFree, l.MinFreeSpacePercent, l.isDiskSpaceLow)
			}
		}
		time.Sleep(time.Minute)
	}

}
<|endoftext|>


<commit_before>package disgo

import (
	"errors"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/slf4go/logger"
)

type rateBucket struct {
	remaining int
	limit     int
	reset     time.Time
	mutex     sync.Mutex
}

func (s *Session) rateLimit(endPoint EndPoint, call func() (*http.Response, error)) error {
	// Get the bucket, and if it does not exist, create it.
	bucket, exists := s.rateLimitBuckets[endPoint.bucket]
	if !exists {
		bucket = &rateBucket{remaining: 1}
		s.rateLimitBuckets[endPoint.bucket] = bucket
	}

	// Lock this bucket
	bucket.mutex.Lock()
	defer bucket.mutex.Unlock()

	// Wait for the bucket to expire if we're out of attempts
	now := time.Now()
	if bucket.remaining == 0 && bucket.reset.After(now) {
		logger.Warnf("We are out of slots for %s, waiting...", endPoint.bucket)
		time.Sleep(bucket.reset.Sub(now))
	}

	// Once we're past the bucket lock, lock globally
	s.globalRateLimit.Lock()
	defer s.globalRateLimit.Unlock()

	// Wait for the globalRateLimit lock if we're being globally ratelimited
	now = time.Now()
	if s.globalReset.After(now) {
		logger.Warnf("We are waiting for the globalRateLimit ratelimit...")
		time.Sleep(s.globalReset.Sub(now))
	}

	// Okay, we've exhausted all possible ratelimit timers, let's send
	response, err := call()
	now = time.Now()

	// Read the headers
	var (
		headerRemaining  = response.Header.Get("X-RateLimit-Remaining")
		headerLimit      = response.Header.Get("X-RateLimit-Limit")
		headerReset      = response.Header.Get("X-RateLimit-Reset")
		headerRetryAfter = response.Header.Get("Retry-After")
		headerGlobal     = response.Header.Get("X-RateLimit-Global")
	)

	// Are we being ratelimited because of that last request?
	if response.StatusCode == 429 {
		if headerRetryAfter == "" {
			return errors.New("We are being ratelimited, but Discord didn't send a Retry-After header")
		}

		retryAfter, parseError := strconv.Atoi(headerRetryAfter)
		if parseError != nil {
			return errors.New("We are being ratelimited, but Discord didn't send a valid Retry-After header")
		}

		resetTime := now.Add(time.Duration(retryAfter) * time.Millisecond)
		if headerGlobal == "true" {
			logger.Error("We are being globally ratelimited!")
			s.globalReset = resetTime
		} else {
			logger.Errorf("We are being ratelimited on %s!", endPoint.bucket)
			bucket.reset = resetTime
			bucket.remaining = 0
		}

		// Automatically queue a retry, but this one will wait for the timers to expire
		return s.rateLimit(endPoint, call)
	}

	// Nope, not ratelimited, but let's update our bucket first
	var parseError error
	if headerRemaining != "" {
		bucket.remaining, parseError = strconv.Atoi(headerRemaining)
	}
	if headerLimit != "" {
		bucket.limit, parseError = strconv.Atoi(headerLimit)
	}
	if endPoint.resetTime == -1 {
		if headerReset != "" {
			var unix int64
			unix, parseError = strconv.ParseInt(headerReset, 10, 64)
			if parseError == nil {
				bucket.reset = time.Unix(unix, 0)
			}
		}
	} else {
		bucket.reset = now.Add(time.Duration(endPoint.resetTime) * time.Millisecond)
	}

	// Check for errors
	if err != nil {
		return err // If the call previously errored, return that now (we still wanted to try and read the headers
	}
	if parseError != nil {
		return parseError // Did we have any issues reading the headers?
	}

	// No errors? Awesome.
	return nil
}
<commit_msg>Fixed deadlock when ratelimited, fixed some spelling warnings<commit_after>package disgo

import (
	"errors"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/slf4go/logger"
)

type rateBucket struct {
	remaining int
	limit     int
	reset     time.Time
	mutex     sync.Mutex
}

func (s *Session) rateLimit(endPoint EndPoint, call func() (*http.Response, error)) error {
	// Get the bucket, and if it does not exist, create it.
	bucket, exists := s.rateLimitBuckets[endPoint.bucket]
	if !exists {
		bucket = &rateBucket{remaining: 1}
		s.rateLimitBuckets[endPoint.bucket] = bucket
	}

	// Lock this bucket
	bucket.mutex.Lock()
	defer bucket.mutex.Unlock()

	// Wait for the bucket to expire if we're out of attempts
	now := time.Now()
	if bucket.remaining == 0 && bucket.reset.After(now) {
		logger.Warnf("We are out of slots for %s, waiting...", endPoint.bucket)
		time.Sleep(bucket.reset.Sub(now))
	}

	// Once we're past the bucket lock, lock globally
	s.globalRateLimit.Lock()
	defer s.globalRateLimit.Unlock()

	// Wait for the globalRateLimit lock if we're being globally rate limited
	now = time.Now()
	if s.globalReset.After(now) {
		logger.Warnf("We are waiting for the globalRateLimit...")
		time.Sleep(s.globalReset.Sub(now))
	}

	// Okay, we've exhausted all possible ratelimit timers, let's send
	response, err := call()
	now = time.Now()

	// Read the headers
	var (
		headerRemaining  = response.Header.Get("X-RateLimit-Remaining")
		headerLimit      = response.Header.Get("X-RateLimit-Limit")
		headerReset      = response.Header.Get("X-RateLimit-Reset")
		headerRetryAfter = response.Header.Get("Retry-After")
		headerGlobal     = response.Header.Get("X-RateLimit-Global")
	)

	// Are we being rate limited because of that last request?
	if response.StatusCode == 429 {
		if headerRetryAfter == "" {
			return errors.New("We are being ratelimited, but Discord didn't send a Retry-After header")
		}

		retryAfter, _ := strconv.Atoi(headerRetryAfter)
		resetTime := now.Add(time.Duration(retryAfter) * time.Millisecond)

		if headerGlobal == "true" {
			logger.Error("We are being globally ratelimited!")
			s.globalReset = resetTime
		} else {
			logger.Errorf("We are being ratelimited on %s!", endPoint.bucket)
			bucket.reset = resetTime
			bucket.remaining = 0
		}

		// Automatically queue a retry, but this one will wait for the timers to expire
		bucket.mutex.Unlock() // Unlock the mutexes so the recursive call can lock them
		s.globalRateLimit.Unlock()
		err := s.rateLimit(endPoint, call)
		bucket.mutex.Lock() // Re-lock them to prevent the deferred unlock calls from panicking
		s.globalRateLimit.Lock()
		return err
	}

	var parseError error
	// Nope, not rate limited, but let's update our bucket first
	if headerRemaining != "" {
		bucket.remaining, parseError = strconv.Atoi(headerRemaining)
	}
	if headerLimit != "" {
		bucket.limit, parseError = strconv.Atoi(headerLimit)
	}
	if endPoint.resetTime == -1 {
		if headerReset != "" {
			var unix int64
			unix, parseError = strconv.ParseInt(headerReset, 10, 64)
			if parseError == nil {
				bucket.reset = time.Unix(unix, 0)
			}
		}
	} else {
		bucket.reset = now.Add(time.Duration(endPoint.resetTime) * time.Millisecond)
	}

	// Check for errors
	if err != nil {
		return err // If the call previously errored, return that now (we still wanted to try and read the headers
	}

	return parseError
}
<|endoftext|>


<commit_before>package discordgo

import (
	"net/http"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// customRateLimit holds information for defining a custom rate limit
type customRateLimit struct {
	suffix   string
	requests int
	reset    time.Duration
}

// RateLimiter holds all ratelimit buckets
type RateLimiter struct {
	sync.Mutex
	global           *int64
	buckets          map[string]*Bucket
	globalRateLimit  time.Duration
	customRateLimits []*customRateLimit
}

// NewRatelimiter returns a new RateLimiter
func NewRatelimiter() *RateLimiter {

	return &RateLimiter{
		buckets: make(map[string]*Bucket),
		global:  new(int64),
		customRateLimits: []*customRateLimit{
			&customRateLimit{
				suffix:   "//reactions//",
				requests: 1,
				reset:    200 * time.Millisecond,
			},
		},
	}
}

// GetBucket retrieves or creates a bucket
func (r *RateLimiter) GetBucket(key string) *Bucket {
	r.Lock()
	defer r.Unlock()

	if bucket, ok := r.buckets[key]; ok {
		return bucket
	}

	b := &Bucket{
		Remaining: 1,
		Key:       key,
		global:    r.global,
	}

	// Check if there is a custom ratelimit set for this bucket ID.
	for _, rl := range r.customRateLimits {
		if strings.HasSuffix(b.Key, rl.suffix) {
			b.customRateLimit = rl
			break
		}
	}

	r.buckets[key] = b
	return b
}

// GetWaitTime returns the duration you should wait for a Bucket
func (r *RateLimiter) GetWaitTime(b *Bucket, minRemaining int) time.Duration {
	// If we ran out of calls and the reset time is still ahead of us
	// then we need to take it easy and relax a little
	if b.Remaining < minRemaining && b.reset.After(time.Now()) {
		return b.reset.Sub(time.Now())
	}

	// Check for global ratelimits
	sleepTo := time.Unix(0, atomic.LoadInt64(r.global))
	if now := time.Now(); now.Before(sleepTo) {
		return sleepTo.Sub(now)
	}

	return 0
}

// LockBucket Locks until a request can be made
func (r *RateLimiter) LockBucket(bucketID string) *Bucket {
	return r.LockBucketObject(r.GetBucket(bucketID))
}

// LockBucketObject Locks an already resolved bucket until a request can be made
func (r *RateLimiter) LockBucketObject(b *Bucket) *Bucket {
	b.Lock()

	if wait := r.GetWaitTime(b, 1); wait > 0 {
		time.Sleep(wait)
	}

	b.Remaining--
	return b
}

// Bucket represents a ratelimit bucket, each bucket gets ratelimited individually (-global ratelimits)
type Bucket struct {
	sync.Mutex
	Key       string
	Remaining int
	limit     int
	reset     time.Time
	global    *int64

	lastReset       time.Time
	customRateLimit *customRateLimit
	Userdata        interface{}
}

// Release unlocks the bucket and reads the headers to update the buckets ratelimit info
// and locks up the whole thing in case if there's a global ratelimit.
func (b *Bucket) Release(headers http.Header) error {
	defer b.Unlock()

	// Check if the bucket uses a custom ratelimiter
	if rl := b.customRateLimit; rl != nil {
		if time.Now().Sub(b.lastReset) >= rl.reset {
			b.Remaining = rl.requests - 1
			b.lastReset = time.Now()
		}
		if b.Remaining < 1 {
			b.reset = time.Now().Add(rl.reset)
		}
		return nil
	}

	if headers == nil {
		return nil
	}

	remaining := headers.Get("X-RateLimit-Remaining")
	reset := headers.Get("X-RateLimit-Reset")
	global := headers.Get("X-RateLimit-Global")
	retryAfter := headers.Get("Retry-After")

	// Update global and per bucket reset time if the proper headers are available
	// If global is set, then it will block all buckets until after Retry-After
	// If Retry-After without global is provided it will use that for the new reset
	// time since it's more accurate than X-RateLimit-Reset.
	// If Retry-After after is not proided, it will update the reset time from X-RateLimit-Reset
	if retryAfter != "" {
		parsedAfter, err := strconv.ParseInt(retryAfter, 10, 64)
		if err != nil {
			return err
		}

		resetAt := time.Now().Add(time.Duration(parsedAfter) * time.Millisecond)

		// Lock either this single bucket or all buckets
		if global != "" {
			atomic.StoreInt64(b.global, resetAt.UnixNano())
		} else {
			b.reset = resetAt
		}
	} else if reset != "" {
		// Calculate the reset time by using the date header returned from discord
		discordTime, err := http.ParseTime(headers.Get("Date"))
		if err != nil {
			return err
		}

		unix, err := strconv.ParseInt(reset, 10, 64)
		if err != nil {
			return err
		}

		// Calculate the time until reset and add it to the current local time
		// some extra time is added because without it i still encountered 429's.
		// The added amount is the lowest amount that gave no 429's
		// in 1k requests
		delta := time.Unix(unix, 0).Sub(discordTime) + time.Millisecond*250
		b.reset = time.Now().Add(delta)
	}

	// Udpate remaining if header is present
	if remaining != "" {
		parsedRemaining, err := strconv.ParseInt(remaining, 10, 32)
		if err != nil {
			return err
		}
		b.Remaining = int(parsedRemaining)
	}

	return nil
}
<commit_msg>fix custom ratelimit for reactions<commit_after>package discordgo

import (
	"net/http"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// customRateLimit holds information for defining a custom rate limit
type customRateLimit struct {
	suffix   string
	requests int
	reset    time.Duration
}

// RateLimiter holds all ratelimit buckets
type RateLimiter struct {
	sync.Mutex
	global           *int64
	buckets          map[string]*Bucket
	globalRateLimit  time.Duration
	customRateLimits []*customRateLimit
}

// NewRatelimiter returns a new RateLimiter
func NewRatelimiter() *RateLimiter {

	return &RateLimiter{
		buckets: make(map[string]*Bucket),
		global:  new(int64),
		customRateLimits: []*customRateLimit{
			&customRateLimit{
				suffix:   "/reactions//",
				requests: 1,
				reset:    250 * time.Millisecond,
			},
		},
	}
}

// GetBucket retrieves or creates a bucket
func (r *RateLimiter) GetBucket(key string) *Bucket {
	r.Lock()
	defer r.Unlock()

	if bucket, ok := r.buckets[key]; ok {
		return bucket
	}

	b := &Bucket{
		Remaining: 1,
		Key:       key,
		global:    r.global,
	}

	// Check if there is a custom ratelimit set for this bucket ID.
	for _, rl := range r.customRateLimits {
		if strings.HasSuffix(b.Key, rl.suffix) {
			b.customRateLimit = rl
			break
		}
	}

	r.buckets[key] = b
	return b
}

// GetWaitTime returns the duration you should wait for a Bucket
func (r *RateLimiter) GetWaitTime(b *Bucket, minRemaining int) time.Duration {
	// If we ran out of calls and the reset time is still ahead of us
	// then we need to take it easy and relax a little
	if b.Remaining < minRemaining && b.reset.After(time.Now()) {
		return b.reset.Sub(time.Now())
	}

	// Check for global ratelimits
	sleepTo := time.Unix(0, atomic.LoadInt64(r.global))
	if now := time.Now(); now.Before(sleepTo) {
		return sleepTo.Sub(now)
	}

	return 0
}

// LockBucket Locks until a request can be made
func (r *RateLimiter) LockBucket(bucketID string) *Bucket {
	return r.LockBucketObject(r.GetBucket(bucketID))
}

// LockBucketObject Locks an already resolved bucket until a request can be made
func (r *RateLimiter) LockBucketObject(b *Bucket) *Bucket {
	b.Lock()

	if wait := r.GetWaitTime(b, 1); wait > 0 {
		time.Sleep(wait)
	}

	b.Remaining--
	return b
}

// Bucket represents a ratelimit bucket, each bucket gets ratelimited individually (-global ratelimits)
type Bucket struct {
	sync.Mutex
	Key       string
	Remaining int
	limit     int
	reset     time.Time
	global    *int64

	lastReset       time.Time
	customRateLimit *customRateLimit
	Userdata        interface{}
}

// Release unlocks the bucket and reads the headers to update the buckets ratelimit info
// and locks up the whole thing in case if there's a global ratelimit.
func (b *Bucket) Release(headers http.Header) error {
	defer b.Unlock()

	// Check if the bucket uses a custom ratelimiter
	if rl := b.customRateLimit; rl != nil {
		if time.Now().Sub(b.lastReset) >= rl.reset {
			b.Remaining = rl.requests - 1
			b.lastReset = time.Now()
		}
		if b.Remaining < 1 {
			b.reset = time.Now().Add(rl.reset)
		}
		return nil
	}

	if headers == nil {
		return nil
	}

	remaining := headers.Get("X-RateLimit-Remaining")
	reset := headers.Get("X-RateLimit-Reset")
	global := headers.Get("X-RateLimit-Global")
	retryAfter := headers.Get("Retry-After")

	// Update global and per bucket reset time if the proper headers are available
	// If global is set, then it will block all buckets until after Retry-After
	// If Retry-After without global is provided it will use that for the new reset
	// time since it's more accurate than X-RateLimit-Reset.
	// If Retry-After after is not proided, it will update the reset time from X-RateLimit-Reset
	if retryAfter != "" {
		parsedAfter, err := strconv.ParseInt(retryAfter, 10, 64)
		if err != nil {
			return err
		}

		resetAt := time.Now().Add(time.Duration(parsedAfter) * time.Millisecond)

		// Lock either this single bucket or all buckets
		if global != "" {
			atomic.StoreInt64(b.global, resetAt.UnixNano())
		} else {
			b.reset = resetAt
		}
	} else if reset != "" {
		// Calculate the reset time by using the date header returned from discord
		discordTime, err := http.ParseTime(headers.Get("Date"))
		if err != nil {
			return err
		}

		unix, err := strconv.ParseInt(reset, 10, 64)
		if err != nil {
			return err
		}

		// Calculate the time until reset and add it to the current local time
		// some extra time is added because without it i still encountered 429's.
		// The added amount is the lowest amount that gave no 429's
		// in 1k requests
		delta := time.Unix(unix, 0).Sub(discordTime) + time.Millisecond*250
		b.reset = time.Now().Add(delta)
	}

	// Udpate remaining if header is present
	if remaining != "" {
		parsedRemaining, err := strconv.ParseInt(remaining, 10, 32)
		if err != nil {
			return err
		}
		b.Remaining = int(parsedRemaining)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"os"
)

var (
	ErrBufferShort = errors.New("Buffer insufficient for packet")
)

type PacketScanner struct {
	rd      io.Reader
	buf     []byte
	r, w, o int
	lasterr error
}

const DefaultPacketBufferSize = 4096

func NewPacketScanner(i io.Reader) *PacketScanner {
	return NewPacketScannerSize(i, DefaultPacketBufferSize)
	return &PacketScanner{rd: i, buf: make([]byte, 4)}
}

func NewPacketScannerSize(i io.Reader, bufsiz int) *PacketScanner {
	return &PacketScanner{rd: i, buf: make([]byte, bufsiz)}
}

func (s *PacketScanner) SetInput(i io.Reader) {
	s.rd = i
}

func (s *PacketScanner) Bytes() []byte {
	return s.buf[s.o:s.w]
}

func (s *PacketScanner) Error() error {
	return s.lasterr
}

// Scan() returns complete Minecraft packets
func (s *PacketScanner) Scan() bool {
	s.lasterr = nil
	if s.w != 0 {
		copy(s.buf, s.buf[s.w:s.r])
		s.w, s.r = 0, s.r-s.w
	}
	for {
		var epkt, nread int
		if s.r != 0 {
			npkt0, nl := binary.Uvarint(s.buf[:s.r])
			epkt = nl + int(npkt0)
			if nl != 0 && epkt <= s.r {
				s.o, s.w = nl, epkt
				if PACKETDEBUG {
					d := MakeDumper(os.Stdout)
					d.line("pktscan:", nil)
					d.bytes(s.Bytes())
				}
				return true
			}
			if s.lasterr != nil {
				// we don't have a full packet yet, but got an error
				s.o = s.w
				return false
			}
		}
		if s.r == len(s.buf) {
			nlen := len(s.buf)*3/2 + 1
			if epkt > nlen {
				nlen = epkt*3/2 + 1
			}
			nbuf := make([]byte, nlen)
			copy(nbuf, s.buf)
			s.buf = nbuf
		}
		nread, s.lasterr = s.rd.Read(s.buf[s.r:])
		s.r += nread
	}
}

func (s *PacketScanner) PacketId() (uint, bool) {
	i, ni := binary.Uvarint(s.Bytes()) // packet length
	if ni != 0 {
		return uint(i), true
	}
	return 0, false
}

type PacketReader struct {
	r       *bufio.Reader
	scratch []byte
	spos    int
}

func NewPacketReader(r io.Reader) *PacketReader {
	return &PacketReader{bufio.NewReader(r), make([]byte, binary.MaxVarintLen64), 0}
}

func (pr *PacketReader) Read(packet []byte) (int, error) {
	if len(packet) < binary.MaxVarintLen64 {
		return 0, ErrBufferShort
	}
	var lenp, nskip int
	for {
		var lenp0 uint64
		lenp0, nskip = binary.Uvarint(pr.scratch[:pr.spos])
		if nskip != 0 {
			lenp = int(lenp0)
			break
		}
		n, err := pr.r.Read(pr.scratch[pr.spos:binary.MaxVarintLen64])
		if err != nil {
			return 0, err
		}
		pr.spos += n
	}
	if len(packet) < lenp {
		return 0, ErrBufferShort
	}
	copy(packet, pr.scratch[nskip:pr.spos])
	pos := pr.spos - nskip
	_, err := io.ReadFull(pr.r, packet[pos:lenp])
	pr.spos = 0
	return lenp, err
}

type PacketWriter struct {
	w       *bufio.Writer
	scratch []byte
}

func NewPacketWriter(w io.Writer) *PacketWriter {
	return &PacketWriter{bufio.NewWriter(w), make([]byte, binary.MaxVarintLen64)}
}

func (pw *PacketWriter) Write(packet []byte) (int, error) {
	n := binary.PutUvarint(pw.scratch, uint64(len(packet)))
	_, err := pw.w.Write(pw.scratch[:n])
	if err != nil {
		return 0, err
	}
	nw, err := pw.w.Write(packet)
	if err != nil {
		return nw, err
	}
	return nw, pw.w.Flush()
}
<commit_msg>fix infinite loop when encountering EOF<commit_after>package main

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"os"
)

var (
	ErrBufferShort = errors.New("Buffer insufficient for packet")
)

type PacketScanner struct {
	rd      io.Reader
	buf     []byte
	r, w, o int
	lasterr error
}

const DefaultPacketBufferSize = 4096

func NewPacketScanner(i io.Reader) *PacketScanner {
	return NewPacketScannerSize(i, DefaultPacketBufferSize)
	return &PacketScanner{rd: i, buf: make([]byte, 4)}
}

func NewPacketScannerSize(i io.Reader, bufsiz int) *PacketScanner {
	return &PacketScanner{rd: i, buf: make([]byte, bufsiz)}
}

func (s *PacketScanner) SetInput(i io.Reader) {
	s.rd = i
}

func (s *PacketScanner) Bytes() []byte {
	return s.buf[s.o:s.w]
}

func (s *PacketScanner) Error() error {
	return s.lasterr
}

// Scan() returns complete Minecraft packets
func (s *PacketScanner) Scan() bool {
	s.lasterr = nil
	if s.w != 0 {
		copy(s.buf, s.buf[s.w:s.r])
		s.w, s.r = 0, s.r-s.w
	}
	for {
		var epkt, nread int
		if s.r != 0 {
			npkt0, nl := binary.Uvarint(s.buf[:s.r])
			epkt = nl + int(npkt0)
			if nl != 0 && epkt <= s.r {
				s.o, s.w = nl, epkt
				if PACKETDEBUG {
					d := MakeDumper(os.Stdout)
					d.line("pktscan:", nil)
					d.bytes(s.Bytes())
				}
				return true
			}
		}
		if s.lasterr != nil {
			// we don't have a full packet yet, but got an error
			s.o = s.w
			return false
		}
		if s.r == len(s.buf) {
			nlen := len(s.buf)*3/2 + 1
			if epkt > nlen {
				nlen = epkt*3/2 + 1
			}
			nbuf := make([]byte, nlen)
			copy(nbuf, s.buf)
			s.buf = nbuf
		}
		nread, s.lasterr = s.rd.Read(s.buf[s.r:])
		s.r += nread
	}
}

func (s *PacketScanner) PacketId() (uint, bool) {
	i, ni := binary.Uvarint(s.Bytes()) // packet length
	if ni != 0 {
		return uint(i), true
	}
	return 0, false
}

type PacketReader struct {
	r       *bufio.Reader
	scratch []byte
	spos    int
}

func NewPacketReader(r io.Reader) *PacketReader {
	return &PacketReader{bufio.NewReader(r), make([]byte, binary.MaxVarintLen64), 0}
}

func (pr *PacketReader) Read(packet []byte) (int, error) {
	if len(packet) < binary.MaxVarintLen64 {
		return 0, ErrBufferShort
	}
	var lenp, nskip int
	for {
		var lenp0 uint64
		lenp0, nskip = binary.Uvarint(pr.scratch[:pr.spos])
		if nskip != 0 {
			lenp = int(lenp0)
			break
		}
		n, err := pr.r.Read(pr.scratch[pr.spos:binary.MaxVarintLen64])
		if err != nil {
			return 0, err
		}
		pr.spos += n
	}
	if len(packet) < lenp {
		return 0, ErrBufferShort
	}
	copy(packet, pr.scratch[nskip:pr.spos])
	pos := pr.spos - nskip
	_, err := io.ReadFull(pr.r, packet[pos:lenp])
	pr.spos = 0
	return lenp, err
}

type PacketWriter struct {
	w       *bufio.Writer
	scratch []byte
}

func NewPacketWriter(w io.Writer) *PacketWriter {
	return &PacketWriter{bufio.NewWriter(w), make([]byte, binary.MaxVarintLen64)}
}

func (pw *PacketWriter) Write(packet []byte) (int, error) {
	n := binary.PutUvarint(pw.scratch, uint64(len(packet)))
	_, err := pw.w.Write(pw.scratch[:n])
	if err != nil {
		return 0, err
	}
	nw, err := pw.w.Write(packet)
	if err != nil {
		return nw, err
	}
	return nw, pw.w.Flush()
}
<|endoftext|>


<commit_before>package main

import (
	// "path/filepath"
	// "runtime"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"sort"
	"time"

	"github.com/geordanr/go_xwing/attack/modification"
	"github.com/geordanr/go_xwing/attack/step"
	"github.com/geordanr/go_xwing/serialization"
	"github.com/geordanr/go_xwing/ship"
	"github.com/geordanr/go_xwing/shipstats"
	"github.com/gocraft/web"
)

type jsonErrorMessage struct {
	Error   bool   `json:"error"`
	Message string `json:"message"`
}

type genericList struct {
	Data []string `json:"data"`
}

// map ship name to map of collected stat (e.g. "hull", "shields")
type simResults map[string]map[string]shipStatJSONSchema

// map stat to average, stddev
type shipStatsJSONSchema map[string]shipStatJSONSchema
type shipStatJSONSchema struct {
	Histogram [][]float64 `json:"histogram"`
	Average   float64     `json:"average"`
	Stddev    float64     `json:"stddev"`
}

type Context struct{}

func main() {
	rand.Seed(time.Now().UnixNano())

	shipjson := parseArgs()
	factory, err := serialization.ShipsFromJSONPath(*shipjson)
	shipFactory = factory
	if err != nil {
		panic(err)
	}

	currentRoot, _ := os.Getwd()
	router := web.New(Context{}).
		Middleware(web.LoggerMiddleware).
		Middleware(corsMiddleware).
		Middleware(web.StaticMiddleware(path.Join(currentRoot, "web", "public"), web.StaticOption{IndexFile: "index.html"})).
		// Get("/", (*Context).Root).
		Get("/api/v1/modifications", (*Context).Modifications).
		Get("/api/v1/ships", (*Context).Ships).
		Get("/api/v1/steps", (*Context).Steps).
		Post("/api/v1/sim", (*Context).Simulate).
		Error((*Context).Error)

	port := 8080
	log.Println("Listening on port", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), router)
}

func parseArgs() *string {
	shipjson := flag.String("shipjson", "", "Path to JSON file of ship data")

	flag.Parse()

	if *shipjson == "" {
		fmt.Println("Specify the path to the ship JSON file.")
		os.Exit(1)
	}

	return shipjson
}

var shipFactory map[string]func(string, uint) *ship.Ship

func corsMiddleware(rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc) {
	rw.Header().Set("Access-Control-Allow-Origin", "*")
	next(rw, req)
}

func (*Context) Root(rw web.ResponseWriter, req *web.Request) {
	fmt.Fprintf(rw, "Monte Carlo is the best Carlo")
}

// Simulate receives JSON POST data specifying parameters for simulating a list vs. list combat round.
func (*Context) Simulate(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		panic(err)
	}

	output, err := serialization.FromJSON(body, shipFactory)
	if err != nil {
		panic(err)
	}

	// map from ship ID to stats struct
	cbtStats := map[string]*shipstats.Stats{}

	for {
		state, more := <-output
		if !more {
			break
		}
		for name, cbt := range state.Combatants() {
			s, exists := cbtStats[name]
			if !exists {
				cbtStats[name] = shipstats.New()
				s = cbtStats[name]
			}
			// fmt.Printf("Update ship stats for %s (%p) %s\n", cbt.Name(), cbt, cbt)
			s.Update(cbt.(*ship.Ship))
		}
	}

	results := make(simResults)
	for name, s := range cbtStats {
		results[name] = make(shipStatsJSONSchema)
		for statName, combinedStats := range s.StatMap() {
			results[name][statName] = shipStatJSONSchema{
				Average:   combinedStats.Stats.Average(),
				Stddev:    combinedStats.Stats.Stddev(),
				Histogram: combinedStats.Histogram.NormalizedHighchartsSeries(),
			}
		}
	}

	resultsJson, err := json.Marshal(results)
	if err != nil {
		panic(err)
	}

	fmt.Fprintf(rw, string(resultsJson))
}

func (*Context) Error(rw web.ResponseWriter, req *web.Request, err interface{}) {
	rw.WriteHeader(http.StatusInternalServerError)
	rw.Header().Set("Content-Type", "application/json")
	msg := jsonErrorMessage{
		Error:   true,
		Message: err.(error).Error(),
	}
	bytes, marshalError := json.Marshal(msg)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Ships returns a JSON list of supported ships.
func (*Context) Ships(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	shiplist := make([]string, 0, len(shipFactory))
	for shipname, _ := range shipFactory {
		shiplist = append(shiplist, shipname)
	}
	sort.Strings(shiplist)

	d := genericList{Data: shiplist}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Modifications returns a JSON list of supported modifications.
func (*Context) Modifications(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	modList := make([]string, 0, len(modification.All))
	for modName, _ := range modification.All {
		modList = append(modList, modName)
	}
	sort.Strings(modList)

	d := genericList{Data: modList}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Steps returns a JSON list of attack steps in order.
func (*Context) Steps(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	stepList := make([]string, 0, len(step.All)-1) // not including the start step
	s := step.All["__START__"]
	for s.Next() != "" {
		s = step.All[s.Next()]
		stepList = append(stepList, s.Name())
	}

	d := genericList{Data: stepList}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}
<commit_msg>Add flag to set HTTP port.<commit_after>package main

import (
	// "path/filepath"
	// "runtime"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"sort"
	"time"

	"github.com/geordanr/go_xwing/attack/modification"
	"github.com/geordanr/go_xwing/attack/step"
	"github.com/geordanr/go_xwing/serialization"
	"github.com/geordanr/go_xwing/ship"
	"github.com/geordanr/go_xwing/shipstats"
	"github.com/gocraft/web"
)

type jsonErrorMessage struct {
	Error   bool   `json:"error"`
	Message string `json:"message"`
}

type genericList struct {
	Data []string `json:"data"`
}

// map ship name to map of collected stat (e.g. "hull", "shields")
type simResults map[string]map[string]shipStatJSONSchema

// map stat to average, stddev
type shipStatsJSONSchema map[string]shipStatJSONSchema
type shipStatJSONSchema struct {
	Histogram [][]float64 `json:"histogram"`
	Average   float64     `json:"average"`
	Stddev    float64     `json:"stddev"`
}

type Context struct{}

func main() {
	rand.Seed(time.Now().UnixNano())

	shipjson, port := parseArgs()
	factory, err := serialization.ShipsFromJSONPath(*shipjson)
	shipFactory = factory
	if err != nil {
		panic(err)
	}

	currentRoot, _ := os.Getwd()
	router := web.New(Context{}).
		Middleware(web.LoggerMiddleware).
		Middleware(corsMiddleware).
		Middleware(web.StaticMiddleware(path.Join(currentRoot, "web", "public"), web.StaticOption{IndexFile: "index.html"})).
		// Get("/", (*Context).Root).
		Get("/api/v1/modifications", (*Context).Modifications).
		Get("/api/v1/ships", (*Context).Ships).
		Get("/api/v1/steps", (*Context).Steps).
		Post("/api/v1/sim", (*Context).Simulate).
		Error((*Context).Error)

	log.Println("Listening on port", *port)
	http.ListenAndServe(fmt.Sprintf(":%d", *port), router)
}

func parseArgs() (*string, *int) {
	shipjson := flag.String("shipjson", "", "Path to JSON file of ship data")
	port := flag.Int("port", 80, "Port to listen on")

	flag.Parse()

	if *shipjson == "" {
		fmt.Println("Specify the path to the ship JSON file.")
		os.Exit(1)
	}

	return shipjson, port
}

var shipFactory map[string]func(string, uint) *ship.Ship

func corsMiddleware(rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc) {
	rw.Header().Set("Access-Control-Allow-Origin", "*")
	next(rw, req)
}

func (*Context) Root(rw web.ResponseWriter, req *web.Request) {
	fmt.Fprintf(rw, "Monte Carlo is the best Carlo")
}

// Simulate receives JSON POST data specifying parameters for simulating a list vs. list combat round.
func (*Context) Simulate(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		panic(err)
	}

	output, err := serialization.FromJSON(body, shipFactory)
	if err != nil {
		panic(err)
	}

	// map from ship ID to stats struct
	cbtStats := map[string]*shipstats.Stats{}

	for {
		state, more := <-output
		if !more {
			break
		}
		for name, cbt := range state.Combatants() {
			s, exists := cbtStats[name]
			if !exists {
				cbtStats[name] = shipstats.New()
				s = cbtStats[name]
			}
			// fmt.Printf("Update ship stats for %s (%p) %s\n", cbt.Name(), cbt, cbt)
			s.Update(cbt.(*ship.Ship))
		}
	}

	results := make(simResults)
	for name, s := range cbtStats {
		results[name] = make(shipStatsJSONSchema)
		for statName, combinedStats := range s.StatMap() {
			results[name][statName] = shipStatJSONSchema{
				Average:   combinedStats.Stats.Average(),
				Stddev:    combinedStats.Stats.Stddev(),
				Histogram: combinedStats.Histogram.NormalizedHighchartsSeries(),
			}
		}
	}

	resultsJson, err := json.Marshal(results)
	if err != nil {
		panic(err)
	}

	fmt.Fprintf(rw, string(resultsJson))
}

func (*Context) Error(rw web.ResponseWriter, req *web.Request, err interface{}) {
	rw.WriteHeader(http.StatusInternalServerError)
	rw.Header().Set("Content-Type", "application/json")
	msg := jsonErrorMessage{
		Error:   true,
		Message: err.(error).Error(),
	}
	bytes, marshalError := json.Marshal(msg)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Ships returns a JSON list of supported ships.
func (*Context) Ships(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	shiplist := make([]string, 0, len(shipFactory))
	for shipname, _ := range shipFactory {
		shiplist = append(shiplist, shipname)
	}
	sort.Strings(shiplist)

	d := genericList{Data: shiplist}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Modifications returns a JSON list of supported modifications.
func (*Context) Modifications(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	modList := make([]string, 0, len(modification.All))
	for modName, _ := range modification.All {
		modList = append(modList, modName)
	}
	sort.Strings(modList)

	d := genericList{Data: modList}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}

// Steps returns a JSON list of attack steps in order.
func (*Context) Steps(rw web.ResponseWriter, req *web.Request) {
	rw.Header().Set("Content-Type", "application/json")

	stepList := make([]string, 0, len(step.All)-1) // not including the start step
	s := step.All["__START__"]
	for s.Next() != "" {
		s = step.All[s.Next()]
		stepList = append(stepList, s.Name())
	}

	d := genericList{Data: stepList}

	bytes, marshalError := json.Marshal(d)
	if marshalError != nil {
		fmt.Fprintf(rw, `{"error":true,"message":"Error marshaling error message: %s"}`, marshalError)
	} else {
		fmt.Fprintf(rw, "%s", bytes)
	}
}
<|endoftext|>


<commit_before>package driver

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/mail"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/blang/semver"
	"github.com/concourse/semver-resource/version"
)

var gitRepoDir string
var privateKeyPath string
var netRcPath string

var ErrEncryptedKey = errors.New("private keys with passphrases are not supported")

func init() {
	gitRepoDir = filepath.Join(os.TempDir(), "semver-git-repo")
	privateKeyPath = filepath.Join(os.TempDir(), "private-key")
	netRcPath = filepath.Join(os.Getenv("HOME"), ".netrc")
}

type GitDriver struct {
	InitialVersion semver.Version

	URI        string
	Branch     string
	PrivateKey string
	Username   string
	Password   string
	File       string
	GitUser    string
	Depth      string
}

func (driver *GitDriver) Bump(bump version.Bump) (semver.Version, error) {
	err := driver.setUpAuth()
	if err != nil {
		return semver.Version{}, err
	}

	err = driver.setUserInfo()
	if err != nil {
		return semver.Version{}, err
	}

	var newVersion semver.Version

	for {
		err = driver.setUpRepo()
		if err != nil {
			return semver.Version{}, err
		}

		currentVersion, exists, err := driver.readVersion()
		if err != nil {
			return semver.Version{}, err
		}

		if !exists {
			currentVersion = driver.InitialVersion
		}

		newVersion = bump.Apply(currentVersion)

		wrote, err := driver.writeVersion(newVersion)
		if wrote {
			break
		}
	}

	return newVersion, nil
}

func (driver *GitDriver) Set(newVersion semver.Version) error {
	err := driver.setUpAuth()
	if err != nil {
		return err
	}

	err = driver.setUserInfo()
	if err != nil {
		return err
	}

	for {
		err = driver.setUpRepo()
		if err != nil {
			return err
		}

		wrote, err := driver.writeVersion(newVersion)
		if err != nil {
			return err
		}

		if wrote {
			break
		}
	}

	return nil
}

func (driver *GitDriver) Check(cursor *semver.Version) ([]semver.Version, error) {
	err := driver.setUpAuth()
	if err != nil {
		return nil, err
	}

	err = driver.setUpRepo()
	if err != nil {
		return nil, err
	}

	currentVersion, exists, err := driver.readVersion()
	if err != nil {
		return nil, err
	}

	if !exists {
		return []semver.Version{driver.InitialVersion}, nil
	}

	if cursor == nil || currentVersion.GTE(*cursor) {
		return []semver.Version{currentVersion}, nil
	}

	return []semver.Version{}, nil
}

func (driver *GitDriver) setUpRepo() error {
	_, err := os.Stat(gitRepoDir)
	if err != nil {
		gitClone := exec.Command("git", "clone", driver.URI, "--branch", driver.Branch)
		if len(driver.Depth) > 0 {
			gitClone.Args = append(gitClone.Args, "--depth", driver.Depth)
		}
		gitClone.Args = append(gitClone.Args, gitRepoDir)
		gitClone.Stdout = os.Stderr
		gitClone.Stderr = os.Stderr
		if err := gitClone.Run(); err != nil {
			return err
		}
	} else {
		gitFetch := exec.Command("git", "fetch", "origin", driver.Branch)
		gitFetch.Dir = gitRepoDir
		gitFetch.Stdout = os.Stderr
		gitFetch.Stderr = os.Stderr
		if err := gitFetch.Run(); err != nil {
			return err
		}
	}

	gitCheckout := exec.Command("git", "reset", "--hard", "origin/"+driver.Branch)
	gitCheckout.Dir = gitRepoDir
	gitCheckout.Stdout = os.Stderr
	gitCheckout.Stderr = os.Stderr
	if err := gitCheckout.Run(); err != nil {
		return err
	}

	return nil
}

func (driver *GitDriver) setUpAuth() error {
	_, err := os.Stat(netRcPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	} else {
		err := os.Remove(netRcPath)
		if err != nil {
			return err
		}
	}

	if len(driver.PrivateKey) > 0 {
		err := driver.setUpKey()
		if err != nil {
			return err
		}
	}

	if len(driver.Username) > 0 && len(driver.Password) > 0 {
		err := driver.setUpUsernamePassword()
		if err != nil {
			return err
		}
	}

	return nil
}

func (driver *GitDriver) setUpKey() error {
	if strings.Contains(driver.PrivateKey, "ENCRYPTED") {
		return ErrEncryptedKey
	}

	_, err := os.Stat(privateKeyPath)
	if err != nil {
		if os.IsNotExist(err) {
			err := ioutil.WriteFile(privateKeyPath, []byte(driver.PrivateKey), 0600)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	return os.Setenv("GIT_SSH_COMMAND", "ssh -o StrictHostKeyChecking=no -i "+privateKeyPath)
}

func (driver *GitDriver) setUpUsernamePassword() error {
	_, err := os.Stat(netRcPath)
	if err != nil {
		if os.IsNotExist(err) {
			content := fmt.Sprintf("default login %s password %s", driver.Username, driver.Password)
			err := ioutil.WriteFile(netRcPath, []byte(content), 0600)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	return nil
}

func (driver *GitDriver) setUserInfo() error {
	if len(driver.GitUser) == 0 {
		return nil
	}

	e, err := mail.ParseAddress(driver.GitUser)
	if err != nil {
		return err
	}

	if len(e.Name) > 0 {
		gitName := exec.Command("git", "config", "--global", "user.name", e.Name)
		gitName.Stdout = os.Stderr
		gitName.Stderr = os.Stderr
		if err := gitName.Run(); err != nil {
			return err
		}
	}

	gitEmail := exec.Command("git", "config", "--global", "user.email", e.Address)
	gitEmail.Stdout = os.Stderr
	gitEmail.Stderr = os.Stderr
	if err := gitEmail.Run(); err != nil {
		return err
	}
	return nil
}

func (driver *GitDriver) readVersion() (semver.Version, bool, error) {
	var currentVersionStr string
	versionFile, err := os.Open(filepath.Join(gitRepoDir, driver.File))
	if err != nil {
		if os.IsNotExist(err) {
			return semver.Version{}, false, nil
		}

		return semver.Version{}, false, err
	}

	defer versionFile.Close()

	_, err = fmt.Fscanf(versionFile, "%s", &currentVersionStr)
	if err != nil {
		return semver.Version{}, false, err
	}

	currentVersion, err := semver.Parse(currentVersionStr)
	if err != nil {
		return semver.Version{}, false, err
	}

	return currentVersion, true, nil
}

const nothingToCommitString = "nothing to commit"
const falsePushString = "Everything up-to-date"
const pushRejectedString = "[rejected]"
const pushRemoteRejectedString = "[remote rejected]"

func (driver *GitDriver) writeVersion(newVersion semver.Version) (bool, error) {
	err := ioutil.WriteFile(filepath.Join(gitRepoDir, driver.File), []byte(newVersion.String()+"\n"), 0644)
	if err != nil {
		return false, err
	}

	gitAdd := exec.Command("git", "add", driver.File)
	gitAdd.Dir = gitRepoDir
	gitAdd.Stdout = os.Stderr
	gitAdd.Stderr = os.Stderr
	if err := gitAdd.Run(); err != nil {
		return false, err
	}

	gitCommit := exec.Command("git", "commit", "-m", "bump to "+newVersion.String())
	gitCommit.Dir = gitRepoDir

	commitOutput, err := gitCommit.CombinedOutput()

	if strings.Contains(string(commitOutput), nothingToCommitString) {
		return true, nil
	}

	if err != nil {
		os.Stderr.Write(commitOutput)
		return false, err
	}

	gitPush := exec.Command("git", "push", "origin", "HEAD:"+driver.Branch)
	gitPush.Dir = gitRepoDir

	pushOutput, err := gitPush.CombinedOutput()

	if strings.Contains(string(pushOutput), falsePushString) {
		return false, nil
	}

	if strings.Contains(string(pushOutput), pushRejectedString) {
		return false, nil
	}

	if strings.Contains(string(pushOutput), pushRemoteRejectedString) {
		return false, nil
	}

	if err != nil {
		os.Stderr.Write(pushOutput)
		return false, err
	}

	return true, nil
}
<commit_msg>Add default --single-branch option to speed up cloning repos with a very very long history<commit_after>package driver

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/mail"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/blang/semver"
	"github.com/concourse/semver-resource/version"
)

var gitRepoDir string
var privateKeyPath string
var netRcPath string

var ErrEncryptedKey = errors.New("private keys with passphrases are not supported")

func init() {
	gitRepoDir = filepath.Join(os.TempDir(), "semver-git-repo")
	privateKeyPath = filepath.Join(os.TempDir(), "private-key")
	netRcPath = filepath.Join(os.Getenv("HOME"), ".netrc")
}

type GitDriver struct {
	InitialVersion semver.Version

	URI        string
	Branch     string
	PrivateKey string
	Username   string
	Password   string
	File       string
	GitUser    string
	Depth      string
}

func (driver *GitDriver) Bump(bump version.Bump) (semver.Version, error) {
	err := driver.setUpAuth()
	if err != nil {
		return semver.Version{}, err
	}

	err = driver.setUserInfo()
	if err != nil {
		return semver.Version{}, err
	}

	var newVersion semver.Version

	for {
		err = driver.setUpRepo()
		if err != nil {
			return semver.Version{}, err
		}

		currentVersion, exists, err := driver.readVersion()
		if err != nil {
			return semver.Version{}, err
		}

		if !exists {
			currentVersion = driver.InitialVersion
		}

		newVersion = bump.Apply(currentVersion)

		wrote, err := driver.writeVersion(newVersion)
		if wrote {
			break
		}
	}

	return newVersion, nil
}

func (driver *GitDriver) Set(newVersion semver.Version) error {
	err := driver.setUpAuth()
	if err != nil {
		return err
	}

	err = driver.setUserInfo()
	if err != nil {
		return err
	}

	for {
		err = driver.setUpRepo()
		if err != nil {
			return err
		}

		wrote, err := driver.writeVersion(newVersion)
		if err != nil {
			return err
		}

		if wrote {
			break
		}
	}

	return nil
}

func (driver *GitDriver) Check(cursor *semver.Version) ([]semver.Version, error) {
	err := driver.setUpAuth()
	if err != nil {
		return nil, err
	}

	err = driver.setUpRepo()
	if err != nil {
		return nil, err
	}

	currentVersion, exists, err := driver.readVersion()
	if err != nil {
		return nil, err
	}

	if !exists {
		return []semver.Version{driver.InitialVersion}, nil
	}

	if cursor == nil || currentVersion.GTE(*cursor) {
		return []semver.Version{currentVersion}, nil
	}

	return []semver.Version{}, nil
}

func (driver *GitDriver) setUpRepo() error {
	_, err := os.Stat(gitRepoDir)
	if err != nil {
		gitClone := exec.Command("git", "clone", driver.URI, "--branch", driver.Branch)
		if len(driver.Depth) > 0 {
			gitClone.Args = append(gitClone.Args, "--depth", driver.Depth)
		}
		gitClone.Args = append(gitClone.Args, "--single-branch", gitRepoDir)
		gitClone.Stdout = os.Stderr
		gitClone.Stderr = os.Stderr
		if err := gitClone.Run(); err != nil {
			return err
		}
	} else {
		gitFetch := exec.Command("git", "fetch", "origin", driver.Branch)
		gitFetch.Dir = gitRepoDir
		gitFetch.Stdout = os.Stderr
		gitFetch.Stderr = os.Stderr
		if err := gitFetch.Run(); err != nil {
			return err
		}
	}

	gitCheckout := exec.Command("git", "reset", "--hard", "origin/"+driver.Branch)
	gitCheckout.Dir = gitRepoDir
	gitCheckout.Stdout = os.Stderr
	gitCheckout.Stderr = os.Stderr
	if err := gitCheckout.Run(); err != nil {
		return err
	}

	return nil
}

func (driver *GitDriver) setUpAuth() error {
	_, err := os.Stat(netRcPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	} else {
		err := os.Remove(netRcPath)
		if err != nil {
			return err
		}
	}

	if len(driver.PrivateKey) > 0 {
		err := driver.setUpKey()
		if err != nil {
			return err
		}
	}

	if len(driver.Username) > 0 && len(driver.Password) > 0 {
		err := driver.setUpUsernamePassword()
		if err != nil {
			return err
		}
	}

	return nil
}

func (driver *GitDriver) setUpKey() error {
	if strings.Contains(driver.PrivateKey, "ENCRYPTED") {
		return ErrEncryptedKey
	}

	_, err := os.Stat(privateKeyPath)
	if err != nil {
		if os.IsNotExist(err) {
			err := ioutil.WriteFile(privateKeyPath, []byte(driver.PrivateKey), 0600)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	return os.Setenv("GIT_SSH_COMMAND", "ssh -o StrictHostKeyChecking=no -i "+privateKeyPath)
}

func (driver *GitDriver) setUpUsernamePassword() error {
	_, err := os.Stat(netRcPath)
	if err != nil {
		if os.IsNotExist(err) {
			content := fmt.Sprintf("default login %s password %s", driver.Username, driver.Password)
			err := ioutil.WriteFile(netRcPath, []byte(content), 0600)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	return nil
}

func (driver *GitDriver) setUserInfo() error {
	if len(driver.GitUser) == 0 {
		return nil
	}

	e, err := mail.ParseAddress(driver.GitUser)
	if err != nil {
		return err
	}

	if len(e.Name) > 0 {
		gitName := exec.Command("git", "config", "--global", "user.name", e.Name)
		gitName.Stdout = os.Stderr
		gitName.Stderr = os.Stderr
		if err := gitName.Run(); err != nil {
			return err
		}
	}

	gitEmail := exec.Command("git", "config", "--global", "user.email", e.Address)
	gitEmail.Stdout = os.Stderr
	gitEmail.Stderr = os.Stderr
	if err := gitEmail.Run(); err != nil {
		return err
	}
	return nil
}

func (driver *GitDriver) readVersion() (semver.Version, bool, error) {
	var currentVersionStr string
	versionFile, err := os.Open(filepath.Join(gitRepoDir, driver.File))
	if err != nil {
		if os.IsNotExist(err) {
			return semver.Version{}, false, nil
		}

		return semver.Version{}, false, err
	}

	defer versionFile.Close()

	_, err = fmt.Fscanf(versionFile, "%s", &currentVersionStr)
	if err != nil {
		return semver.Version{}, false, err
	}

	currentVersion, err := semver.Parse(currentVersionStr)
	if err != nil {
		return semver.Version{}, false, err
	}

	return currentVersion, true, nil
}

const nothingToCommitString = "nothing to commit"
const falsePushString = "Everything up-to-date"
const pushRejectedString = "[rejected]"
const pushRemoteRejectedString = "[remote rejected]"

func (driver *GitDriver) writeVersion(newVersion semver.Version) (bool, error) {
	err := ioutil.WriteFile(filepath.Join(gitRepoDir, driver.File), []byte(newVersion.String()+"\n"), 0644)
	if err != nil {
		return false, err
	}

	gitAdd := exec.Command("git", "add", driver.File)
	gitAdd.Dir = gitRepoDir
	gitAdd.Stdout = os.Stderr
	gitAdd.Stderr = os.Stderr
	if err := gitAdd.Run(); err != nil {
		return false, err
	}

	gitCommit := exec.Command("git", "commit", "-m", "bump to "+newVersion.String())
	gitCommit.Dir = gitRepoDir

	commitOutput, err := gitCommit.CombinedOutput()

	if strings.Contains(string(commitOutput), nothingToCommitString) {
		return true, nil
	}

	if err != nil {
		os.Stderr.Write(commitOutput)
		return false, err
	}

	gitPush := exec.Command("git", "push", "origin", "HEAD:"+driver.Branch)
	gitPush.Dir = gitRepoDir

	pushOutput, err := gitPush.CombinedOutput()

	if strings.Contains(string(pushOutput), falsePushString) {
		return false, nil
	}

	if strings.Contains(string(pushOutput), pushRejectedString) {
		return false, nil
	}

	if strings.Contains(string(pushOutput), pushRemoteRejectedString) {
		return false, nil
	}

	if err != nil {
		os.Stderr.Write(pushOutput)
		return false, err
	}

	return true, nil
}
<|endoftext|>


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

package mysql

// #include <stdlib.h>
// #include <mysql.h>
import "C"

import (
	"code.google.com/p/vitess/go/hack"
	"fmt"
	"unsafe"
)

type SqlError struct {
	Num     int
	Message string
	Query   string
}

func NewSqlError(number int, format string, args ...interface{}) *SqlError {
	return &SqlError{Num: number, Message: fmt.Sprintf(format, args...)}
}

func (self *SqlError) Error() string {
	if self.Query == "" {
		return fmt.Sprintf("%v (errno %v)", self.Message, self.Num)
	}
	return fmt.Sprintf("%v (errno %v) during query: %s", self.Message, self.Num, self.Query)
}

func (self *SqlError) Number() int {
	return self.Num
}

func handleError(err *error) {
	if x := recover(); x != nil {
		terr := x.(*SqlError)
		*err = terr
	}
}

type Connection struct {
	handle *C.MYSQL
}

type QueryResult struct {
	Fields       []Field
	RowsAffected uint64
	InsertId     uint64
	Rows         [][]interface{}
}

type Field struct {
	Name string
	Type int64
}

func Connect(info map[string]interface{}) (conn *Connection, err error) {
	defer handleError(&err)

	host := mapCString(info, "host")
	defer cfree(host)
	port := mapint(info, "port")
	uname := mapCString(info, "uname")
	defer cfree(uname)
	pass := mapCString(info, "pass")
	defer cfree(pass)
	dbname := mapCString(info, "dbname")
	defer cfree(dbname)
	unix_socket := mapCString(info, "unix_socket")
	defer cfree(unix_socket)
	charset := mapCString(info, "charset")
	defer cfree(charset)

	conn = &Connection{}
	conn.handle = C.mysql_init(nil)
	if C.mysql_real_connect(conn.handle, host, uname, pass, dbname, C.uint(port), unix_socket, 0) == nil {
		defer conn.Close()
		return nil, conn.lastError(nil)
	}

	if C.mysql_set_character_set(conn.handle, charset) != 0 {
		defer conn.Close()
		return nil, conn.lastError(nil)
	}
	return conn, nil
}

func (self *Connection) ExecuteFetch(query []byte, maxrows int, wantfields bool) (qr *QueryResult, err error) {
	defer handleError(&err)
	self.validate()

	if C.mysql_real_query(self.handle, (*C.char)(unsafe.Pointer(&query[0])), C.ulong(len(query))) != 0 {
		return nil, self.lastError(query)
	}

	result := C.mysql_store_result(self.handle)
	if result == nil {
		if int(C.mysql_field_count(self.handle)) != 0 { // Query was supposed to return data, but it didn't
			return nil, self.lastError(query)
		}
		qr = &QueryResult{}
		qr.RowsAffected = uint64(C.mysql_affected_rows(self.handle))
		qr.InsertId = uint64(C.mysql_insert_id(self.handle))
		return qr, nil
	}
	defer C.mysql_free_result(result)
	qr = &QueryResult{}
	qr.RowsAffected = uint64(C.mysql_affected_rows(self.handle))
	if qr.RowsAffected > uint64(maxrows) {
		return nil, &SqlError{0, fmt.Sprintf("Row count exceeded %d", maxrows), string(query)}
	}
	if wantfields {
		qr.Fields = self.buildFields(result)
	}
	qr.Rows = self.fetchAll(result)
	return qr, nil
}

func (self *Connection) Id() int64 {
	if self.handle == nil {
		return 0
	}
	return int64(C.mysql_thread_id(self.handle))
}

func (self *Connection) Close() {
	if self.handle == nil {
		return
	}
	C.mysql_close(self.handle)
	self.handle = nil
}

func (self *Connection) IsClosed() bool {
	return self.handle == nil
}

func (self *Connection) buildFields(result *C.MYSQL_RES) (fields []Field) {
	nfields := int(C.mysql_num_fields(result))
	cfieldsptr := C.mysql_fetch_fields(result)
	cfields := (*[1 << 30]C.MYSQL_FIELD)(unsafe.Pointer(cfieldsptr))
	arena := hack.NewStringArena(1024) // prealloc a reasonable amount of space
	fields = make([]Field, nfields)
	for i := 0; i < nfields; i++ {
		length := strlen(cfields[i].name)
		fname := (*[1 << 30]byte)(unsafe.Pointer(cfields[i].name))[:length]
		fields[i].Name = arena.NewString(fname)
		fields[i].Type = int64(cfields[i]._type)
	}
	return fields
}

func (self *Connection) fetchAll(result *C.MYSQL_RES) (rows [][]interface{}) {
	rowCount := int(C.mysql_num_rows(result))
	if rowCount == 0 {
		return nil
	}
	rows = make([][]interface{}, rowCount)
	colCount := int(C.mysql_num_fields(result))
	for i := 0; i < rowCount; i++ {
		rows[i] = self.fetchNext(result, colCount)
	}
	return rows
}

func (self *Connection) fetchNext(result *C.MYSQL_RES, colCount int) (row []interface{}) {
	rowPtr := (*[1 << 30]*[1 << 30]byte)(unsafe.Pointer(C.mysql_fetch_row(result)))
	if rowPtr == nil {
		panic(self.lastError(nil))
	}
	row = make([]interface{}, colCount)
	lengths := (*[1 << 30]uint64)(unsafe.Pointer(C.mysql_fetch_lengths(result)))
	totalLength := uint64(0)
	for i := 0; i < colCount; i++ {
		totalLength += (*lengths)[i]
	}
	arena := hack.NewStringArena(int(totalLength))
	for i := 0; i < colCount; i++ {
		colLength := (*lengths)[i]
		colPtr := (*rowPtr)[i]
		if colPtr == nil {
			continue
		}
		row[i] = arena.NewString((*colPtr)[:colLength])
	}
	return row
}

func (self *Connection) lastError(query []byte) error {
	if err := C.mysql_error(self.handle); *err != 0 {
		return &SqlError{Num: int(C.mysql_errno(self.handle)), Message: C.GoString(err), Query: string(query)}
	}
	return &SqlError{0, "Dummy", string(query)}
}

func (self *Connection) validate() {
	if self.handle == nil {
		panic(NewSqlError(2006, "Connection is closed"))
	}
}

func mapCString(info map[string]interface{}, key string) *C.char {
	ival, ok := info[key]
	if !ok {
		panic(NewSqlError(0, "Missing connection parameter %s", key))
	}
	sval, ok := ival.(string)
	if !ok {
		panic(NewSqlError(0, "Expecting string for %s, received %T", key, ival))
	}
	if sval == "" {
		return nil
	}
	return C.CString(sval)
}

func mapint(info map[string]interface{}, key string) int {
	ival, ok := info[key]
	if !ok {
		panic(NewSqlError(0, "Missing connection parameter %s", key))
	}
	intval, ok := ival.(int)
	if !ok {
		panic(NewSqlError(0, "Expecting int for %s, received %T", key, ival))
	}
	return intval
}

func cfree(str *C.char) {
	if str != nil {
		C.free(unsafe.Pointer(str))
	}
}

// Muahahaha
func strlen(str *C.char) int {
	gstr := (*[1 << 30]byte)(unsafe.Pointer(str))
	length := 0
	for gstr[length] != 0 {
		length++
	}
	return length
}
<commit_msg>fix intermittent segfault under mysql error conditions<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mysql

/*
#include <stdlib.h>
#include <mysql.h>

// The vt_mysql_* functions are needed to honor the mysql client library's
// assumption that all threads using mysql functions have first called
// my_thread_init().
//
// Any goroutine may run on any thread at any time. However, a Cgo call
// is guaranteed to run on a single thread for its duration.  Therefore,
// calling my_thread_init() before every actual operation is sufficient.
// Multiple calls to my_thread_init() are guarded interanlly in mysql and
// the call is cheap.

MYSQL* vt_mysql_real_connect(MYSQL *mysql, const char *host, const char *user, const char *passwd, const char *db, unsigned int port, const char *unix_socket, unsigned long client_flag) {
  my_thread_init();
  return mysql_real_connect(mysql, host, user, passwd, db, port, unix_socket, client_flag);
}

int vt_mysql_set_character_set(MYSQL *mysql, const char *csname) {
  my_thread_init();
  return mysql_set_character_set(mysql, csname);
}

int vt_mysql_real_query(MYSQL *mysql, const char *stmt_str, unsigned long length) {
  my_thread_init();
  return mysql_real_query(mysql, stmt_str, length);
}

MYSQL_RES *vt_mysql_store_result(MYSQL *mysql) {
  my_thread_init();
  return mysql_store_result(mysql);
}

unsigned int vt_mysql_field_count(MYSQL *mysql) {
  my_thread_init();
  return mysql_field_count(mysql);
}

my_ulonglong vt_mysql_affected_rows(MYSQL *mysql) {
  my_thread_init();
  return mysql_affected_rows(mysql);
}

my_ulonglong vt_mysql_insert_id(MYSQL *mysql) {
  my_thread_init();
  return mysql_insert_id(mysql);
}

void vt_mysql_free_result(MYSQL_RES *result) {
  my_thread_init();
  mysql_free_result(result);
}

unsigned long vt_mysql_thread_id(MYSQL *mysql) {
  my_thread_init();
  return mysql_thread_id(mysql);
}

void vt_mysql_close(MYSQL *mysql) {
  my_thread_init();
  mysql_close(mysql);
}

MYSQL_FIELD *vt_mysql_fetch_fields(MYSQL_RES *result) {
  my_thread_init();
  return mysql_fetch_fields(result);
}

unsigned int vt_mysql_num_fields(MYSQL_RES *result) {
  my_thread_init();
  return mysql_num_fields(result);
}

my_ulonglong vt_mysql_num_rows(MYSQL_RES *result) {
  my_thread_init();
  return mysql_num_rows(result);
}

MYSQL_ROW vt_mysql_fetch_row(MYSQL_RES *result) {
  my_thread_init();
  return mysql_fetch_row(result);
}

unsigned long *vt_mysql_fetch_lengths(MYSQL_RES *result) {
  my_thread_init();
  return mysql_fetch_lengths(result);
}

unsigned int vt_mysql_errno(MYSQL *mysql) {
  my_thread_init();
  return mysql_errno(mysql);
}

const char *vt_mysql_error(MYSQL *mysql) {
  my_thread_init();
  return mysql_error(mysql);
}
*/
import "C"

import (
	"code.google.com/p/vitess/go/hack"
	"fmt"
	"unsafe"
)

type SqlError struct {
	Num     int
	Message string
	Query   string
}

func NewSqlError(number int, format string, args ...interface{}) *SqlError {
	return &SqlError{Num: number, Message: fmt.Sprintf(format, args...)}
}

func (self *SqlError) Error() string {
	if self.Query == "" {
		return fmt.Sprintf("%v (errno %v)", self.Message, self.Num)
	}
	return fmt.Sprintf("%v (errno %v) during query: %s", self.Message, self.Num, self.Query)
}

func (self *SqlError) Number() int {
	return self.Num
}

func handleError(err *error) {
	if x := recover(); x != nil {
		terr := x.(*SqlError)
		*err = terr
	}
}

type Connection struct {
	handle *C.MYSQL
}

type QueryResult struct {
	Fields       []Field
	RowsAffected uint64
	InsertId     uint64
	Rows         [][]interface{}
}

type Field struct {
	Name string
	Type int64
}

func Connect(info map[string]interface{}) (conn *Connection, err error) {
	defer handleError(&err)

	host := mapCString(info, "host")
	defer cfree(host)
	port := mapint(info, "port")
	uname := mapCString(info, "uname")
	defer cfree(uname)
	pass := mapCString(info, "pass")
	defer cfree(pass)
	dbname := mapCString(info, "dbname")
	defer cfree(dbname)
	unix_socket := mapCString(info, "unix_socket")
	defer cfree(unix_socket)
	charset := mapCString(info, "charset")
	defer cfree(charset)

	conn = &Connection{}
	conn.handle = C.mysql_init(nil)
	if C.vt_mysql_real_connect(conn.handle, host, uname, pass, dbname, C.uint(port), unix_socket, 0) == nil {
		defer conn.Close()
		return nil, conn.lastError(nil)
	}

	if C.vt_mysql_set_character_set(conn.handle, charset) != 0 {
		defer conn.Close()
		return nil, conn.lastError(nil)
	}
	return conn, nil
}

func (self *Connection) ExecuteFetch(query []byte, maxrows int, wantfields bool) (qr *QueryResult, err error) {
	defer handleError(&err)
	self.validate()

	if C.vt_mysql_real_query(self.handle, (*C.char)(unsafe.Pointer(&query[0])), C.ulong(len(query))) != 0 {
		return nil, self.lastError(query)
	}

	result := C.vt_mysql_store_result(self.handle)
	if result == nil {
		if int(C.vt_mysql_field_count(self.handle)) != 0 { // Query was supposed to return data, but it didn't
			return nil, self.lastError(query)
		}
		qr = &QueryResult{}
		qr.RowsAffected = uint64(C.vt_mysql_affected_rows(self.handle))
		qr.InsertId = uint64(C.vt_mysql_insert_id(self.handle))
		return qr, nil
	}
	defer C.vt_mysql_free_result(result)
	qr = &QueryResult{}
	qr.RowsAffected = uint64(C.vt_mysql_affected_rows(self.handle))
	if qr.RowsAffected > uint64(maxrows) {
		return nil, &SqlError{0, fmt.Sprintf("Row count exceeded %d", maxrows), string(query)}
	}
	if wantfields {
		qr.Fields = self.buildFields(result)
	}
	qr.Rows = self.fetchAll(result)
	return qr, nil
}

func (self *Connection) Id() int64 {
	if self.handle == nil {
		return 0
	}
	return int64(C.vt_mysql_thread_id(self.handle))
}

func (self *Connection) Close() {
	if self.handle == nil {
		return
	}
	C.vt_mysql_close(self.handle)
	self.handle = nil
}

func (self *Connection) IsClosed() bool {
	return self.handle == nil
}

func (self *Connection) buildFields(result *C.MYSQL_RES) (fields []Field) {
	nfields := int(C.vt_mysql_num_fields(result))
	cfieldsptr := C.vt_mysql_fetch_fields(result)
	cfields := (*[1 << 30]C.MYSQL_FIELD)(unsafe.Pointer(cfieldsptr))
	arena := hack.NewStringArena(1024) // prealloc a reasonable amount of space
	fields = make([]Field, nfields)
	for i := 0; i < nfields; i++ {
		length := strlen(cfields[i].name)
		fname := (*[1 << 30]byte)(unsafe.Pointer(cfields[i].name))[:length]
		fields[i].Name = arena.NewString(fname)
		fields[i].Type = int64(cfields[i]._type)
	}
	return fields
}

func (self *Connection) fetchAll(result *C.MYSQL_RES) (rows [][]interface{}) {
	rowCount := int(C.vt_mysql_num_rows(result))
	if rowCount == 0 {
		return nil
	}
	rows = make([][]interface{}, rowCount)
	colCount := int(C.vt_mysql_num_fields(result))
	for i := 0; i < rowCount; i++ {
		rows[i] = self.fetchNext(result, colCount)
	}
	return rows
}

func (self *Connection) fetchNext(result *C.MYSQL_RES, colCount int) (row []interface{}) {
	rowPtr := (*[1 << 30]*[1 << 30]byte)(unsafe.Pointer(C.vt_mysql_fetch_row(result)))
	if rowPtr == nil {
		panic(self.lastError(nil))
	}
	row = make([]interface{}, colCount)
	lengths := (*[1 << 30]uint64)(unsafe.Pointer(C.vt_mysql_fetch_lengths(result)))
	totalLength := uint64(0)
	for i := 0; i < colCount; i++ {
		totalLength += (*lengths)[i]
	}
	arena := hack.NewStringArena(int(totalLength))
	for i := 0; i < colCount; i++ {
		colLength := (*lengths)[i]
		colPtr := (*rowPtr)[i]
		if colPtr == nil {
			continue
		}
		row[i] = arena.NewString((*colPtr)[:colLength])
	}
	return row
}

func (self *Connection) lastError(query []byte) error {
	if err := C.vt_mysql_error(self.handle); *err != 0 {
		return &SqlError{Num: int(C.vt_mysql_errno(self.handle)), Message: C.GoString(err), Query: string(query)}
	}
	return &SqlError{0, "Dummy", string(query)}
}

func (self *Connection) validate() {
	if self.handle == nil {
		panic(NewSqlError(2006, "Connection is closed"))
	}
}

func mapCString(info map[string]interface{}, key string) *C.char {
	ival, ok := info[key]
	if !ok {
		panic(NewSqlError(0, "Missing connection parameter %s", key))
	}
	sval, ok := ival.(string)
	if !ok {
		panic(NewSqlError(0, "Expecting string for %s, received %T", key, ival))
	}
	if sval == "" {
		return nil
	}
	return C.CString(sval)
}

func mapint(info map[string]interface{}, key string) int {
	ival, ok := info[key]
	if !ok {
		panic(NewSqlError(0, "Missing connection parameter %s", key))
	}
	intval, ok := ival.(int)
	if !ok {
		panic(NewSqlError(0, "Expecting int for %s, received %T", key, ival))
	}
	return intval
}

func cfree(str *C.char) {
	if str != nil {
		C.free(unsafe.Pointer(str))
	}
}

// Muahahaha
func strlen(str *C.char) int {
	gstr := (*[1 << 30]byte)(unsafe.Pointer(str))
	length := 0
	for gstr[length] != 0 {
		length++
	}
	return length
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 cmd

import (
	"fmt"
	api "github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/packet"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"io"
	"net"
	"os"
	"strconv"
	"time"
)

func showRPKIServer(args []string) error {
	arg := &api.Arguments{}

	stream, err := client.GetRPKI(context.Background(), arg)
	if err != nil {
		fmt.Println(err)
		return err
	}
	format := "%-18s %-6s %-10s %s\n"
	fmt.Printf(format, "Session", "State", "Uptime", "#IPv4/IPv6 records")
	for {
		r, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		s := "Up"
		uptime := int64(time.Now().Sub(time.Unix(r.State.Uptime, 0)).Seconds())

		fmt.Printf(format, fmt.Sprintf(r.Conf.Address), s, fmt.Sprint(formatTimedelta(uptime)), fmt.Sprintf("%d/%d", r.State.ReceivedIpv4, r.State.ReceivedIpv6))
	}
	return nil
}

func showRPKITable(args []string) error {
	family, err := checkAddressFamily(bgp.RouteFamily(0))
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	arg := &api.Arguments{
		Family: uint32(family),
	}
	if len(args) > 0 {
		arg.Name = args[0]
	}
	stream, err := client.GetROA(context.Background(), arg)
	if err != nil {
		fmt.Println(err)
		return err
	}

	var format string
	afi, _ := bgp.RouteFamilyToAfiSafi(family)
	if afi == bgp.AFI_IP {
		format = "%-18s %-6s %-10s %s\n"
	} else {
		format = "%-42s %-6s %-10s %s\n"
	}
	fmt.Printf(format, "Network", "Maxlen", "AS", "Server")
	for {
		r, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		if len(args) > 0 && args[0] != r.Conf.Address {
			continue
		}

		server := net.JoinHostPort(r.Conf.Address, strconv.Itoa(int(r.Conf.RemotePort)))
		fmt.Printf(format, fmt.Sprintf("%s/%d", r.Prefix, r.Prefixlen), fmt.Sprint(r.Maxlen), fmt.Sprint(r.As), server)
	}
	return nil
}

func NewRPKICmd() *cobra.Command {
	rpkiCmd := &cobra.Command{
		Use: CMD_RPKI,
	}

	modRPKI := func(op api.Operation, address string) {
		arg := &api.ModRpkiArguments{
			Operation: op,
			Address:   address,
			Port:      323,
		}
		client.ModRPKI(context.Background(), arg)
	}

	enableCmd := &cobra.Command{
		Use: CMD_ENABLE,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Println("usage: gobgp rpki enable <ip address>")
				os.Exit(1)
			}
			modRPKI(api.Operation_ADD, args[0])
		},
	}
	rpkiCmd.AddCommand(enableCmd)

	serverCmd := &cobra.Command{
		Use: CMD_RPKI_SERVER,
		Run: func(cmd *cobra.Command, args []string) {
			showRPKIServer(args)
		},
	}

	rpkiCmd.AddCommand(serverCmd)

	tableCmd := &cobra.Command{
		Use: CMD_RPKI_TABLE,
		Run: func(cmd *cobra.Command, args []string) {
			showRPKITable(args)
		},
	}
	tableCmd.PersistentFlags().StringVarP(&subOpts.AddressFamily, "address-family", "a", "", "address family")

	rpkiCmd.AddCommand(tableCmd)
	return rpkiCmd
}
<commit_msg>gobgp: fix rpki cli uptime<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// 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 cmd

import (
	"fmt"
	api "github.com/osrg/gobgp/api"
	"github.com/osrg/gobgp/packet"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"io"
	"net"
	"os"
	"strconv"
	"time"
)

func showRPKIServer(args []string) error {
	arg := &api.Arguments{}

	stream, err := client.GetRPKI(context.Background(), arg)
	if err != nil {
		fmt.Println(err)
		return err
	}
	format := "%-18s %-6s %-10s %s\n"
	fmt.Printf(format, "Session", "State", "Uptime", "#IPv4/IPv6 records")
	for {
		r, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		s := "Down"
		uptime := "never"
		if r.State.Uptime != 0 {
			s = "Up"
			uptime = fmt.Sprint(formatTimedelta(int64(time.Now().Sub(time.Unix(r.State.Uptime, 0)).Seconds())))
		}

		fmt.Printf(format, fmt.Sprintf(r.Conf.Address), s, uptime, fmt.Sprintf("%d/%d", r.State.ReceivedIpv4, r.State.ReceivedIpv6))
	}
	return nil
}

func showRPKITable(args []string) error {
	family, err := checkAddressFamily(bgp.RouteFamily(0))
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	arg := &api.Arguments{
		Family: uint32(family),
	}
	if len(args) > 0 {
		arg.Name = args[0]
	}
	stream, err := client.GetROA(context.Background(), arg)
	if err != nil {
		fmt.Println(err)
		return err
	}

	var format string
	afi, _ := bgp.RouteFamilyToAfiSafi(family)
	if afi == bgp.AFI_IP {
		format = "%-18s %-6s %-10s %s\n"
	} else {
		format = "%-42s %-6s %-10s %s\n"
	}
	fmt.Printf(format, "Network", "Maxlen", "AS", "Server")
	for {
		r, err := stream.Recv()
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}
		if len(args) > 0 && args[0] != r.Conf.Address {
			continue
		}

		server := net.JoinHostPort(r.Conf.Address, strconv.Itoa(int(r.Conf.RemotePort)))
		fmt.Printf(format, fmt.Sprintf("%s/%d", r.Prefix, r.Prefixlen), fmt.Sprint(r.Maxlen), fmt.Sprint(r.As), server)
	}
	return nil
}

func NewRPKICmd() *cobra.Command {
	rpkiCmd := &cobra.Command{
		Use: CMD_RPKI,
	}

	modRPKI := func(op api.Operation, address string) {
		arg := &api.ModRpkiArguments{
			Operation: op,
			Address:   address,
			Port:      323,
		}
		client.ModRPKI(context.Background(), arg)
	}

	enableCmd := &cobra.Command{
		Use: CMD_ENABLE,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Println("usage: gobgp rpki enable <ip address>")
				os.Exit(1)
			}
			modRPKI(api.Operation_ADD, args[0])
		},
	}
	rpkiCmd.AddCommand(enableCmd)

	serverCmd := &cobra.Command{
		Use: CMD_RPKI_SERVER,
		Run: func(cmd *cobra.Command, args []string) {
			showRPKIServer(args)
		},
	}

	rpkiCmd.AddCommand(serverCmd)

	tableCmd := &cobra.Command{
		Use: CMD_RPKI_TABLE,
		Run: func(cmd *cobra.Command, args []string) {
			showRPKITable(args)
		},
	}
	tableCmd.PersistentFlags().StringVarP(&subOpts.AddressFamily, "address-family", "a", "", "address family")

	rpkiCmd.AddCommand(tableCmd)
	return rpkiCmd
}
<|endoftext|>