<commit_before>package filesys

import (
	"context"
	"os"
	"sort"
	"time"

	"github.com/seaweedfs/fuse"
	"github.com/seaweedfs/fuse/fs"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

const blockSize = 512

var _ = fs.Node(&File{})
var _ = fs.NodeIdentifier(&File{})
var _ = fs.NodeOpener(&File{})
var _ = fs.NodeFsyncer(&File{})
var _ = fs.NodeSetattrer(&File{})
var _ = fs.NodeGetxattrer(&File{})
var _ = fs.NodeSetxattrer(&File{})
var _ = fs.NodeRemovexattrer(&File{})
var _ = fs.NodeListxattrer(&File{})
var _ = fs.NodeForgetter(&File{})

type File struct {
	Name          string
	dir           *Dir
	wfs           *WFS
	entry         *filer_pb.Entry
	isOpen        int
	dirtyMetadata bool
	id            uint64
}

func (file *File) fullpath() util.FullPath {
	return util.NewFullPath(file.dir.FullPath(), file.Name)
}

func (file *File) Id() uint64 {
	return file.id
}

func (file *File) Attr(ctx context.Context, attr *fuse.Attr) (err error) {

	glog.V(4).Infof("file Attr %s, open:%v existing:%v", file.fullpath(), file.isOpen, attr)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if entry == nil {
		return fuse.ENOENT
	}

	attr.Inode = file.Id()
	attr.Valid = time.Second
	attr.Mode = os.FileMode(entry.Attributes.FileMode)
	attr.Size = filer.FileSize(entry)
	if file.isOpen > 0 {
		attr.Size = entry.Attributes.FileSize
		glog.V(4).Infof("file Attr %s, open:%v, size: %d", file.fullpath(), file.isOpen, attr.Size)
	}
	attr.Crtime = time.Unix(entry.Attributes.Crtime, 0)
	attr.Mtime = time.Unix(entry.Attributes.Mtime, 0)
	attr.Gid = entry.Attributes.Gid
	attr.Uid = entry.Attributes.Uid
	attr.Blocks = attr.Size/blockSize + 1
	attr.BlockSize = uint32(file.wfs.option.ChunkSizeLimit)
	if entry.HardLinkCounter > 0 {
		attr.Nlink = uint32(entry.HardLinkCounter)
	}

	return nil

}

func (file *File) Getxattr(ctx context.Context, req *fuse.GetxattrRequest, resp *fuse.GetxattrResponse) error {

	glog.V(4).Infof("file Getxattr %s", file.fullpath())

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	return getxattr(entry, req, resp)
}

func (file *File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {

	glog.V(4).Infof("file %v open %+v", file.fullpath(), req)

	handle := file.wfs.AcquireHandle(file, req.Uid, req.Gid, req.Flags&fuse.OpenWriteOnly > 0)

	resp.Handle = fuse.HandleID(handle.handle)

	glog.V(4).Infof("%v file open handle id = %d", file.fullpath(), handle.handle)

	return handle, nil

}

func (file *File) Setattr(ctx context.Context, req *fuse.SetattrRequest, resp *fuse.SetattrResponse) error {

	glog.V(4).Infof("%v file setattr %+v", file.fullpath(), req)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}
	if file.isOpen > 0 {
		file.wfs.handlesLock.Lock()
		fileHandle := file.wfs.handles[file.Id()]
		file.wfs.handlesLock.Unlock()

		if fileHandle != nil {
			fileHandle.Lock()
			defer fileHandle.Unlock()
		}
	}

	if req.Valid.Size() {

		glog.V(4).Infof("%v file setattr set size=%v chunks=%d", file.fullpath(), req.Size, len(entry.Chunks))
		if req.Size < filer.FileSize(entry) {
			// fmt.Printf("truncate %v \n", fullPath)
			var chunks []*filer_pb.FileChunk
			var truncatedChunks []*filer_pb.FileChunk
			for _, chunk := range entry.Chunks {
				int64Size := int64(chunk.Size)
				if chunk.Offset+int64Size > int64(req.Size) {
					// this chunk is truncated
					int64Size = int64(req.Size) - chunk.Offset
					if int64Size > 0 {
						chunks = append(chunks, chunk)
						glog.V(4).Infof("truncated chunk %+v from %d to %d\n", chunk.GetFileIdString(), chunk.Size, int64Size)
						chunk.Size = uint64(int64Size)
					} else {
						glog.V(4).Infof("truncated whole chunk %+v\n", chunk.GetFileIdString())
						truncatedChunks = append(truncatedChunks, chunk)
					}
				}
			}
			entry.Chunks = chunks
		}
		entry.Attributes.FileSize = req.Size
		file.dirtyMetadata = true
	}

	if req.Valid.Mode() {
		entry.Attributes.FileMode = uint32(req.Mode)
		file.dirtyMetadata = true
	}

	if req.Valid.Uid() {
		entry.Attributes.Uid = req.Uid
		file.dirtyMetadata = true
	}

	if req.Valid.Gid() {
		entry.Attributes.Gid = req.Gid
		file.dirtyMetadata = true
	}

	if req.Valid.Crtime() {
		entry.Attributes.Crtime = req.Crtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Mtime() {
		entry.Attributes.Mtime = req.Mtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Handle() {
		// fmt.Printf("file handle => %d\n", req.Handle)
	}

	if file.isOpen > 0 {
		return nil
	}

	if !file.dirtyMetadata {
		return nil
	}

	return file.saveEntry(entry)

}

func (file *File) Setxattr(ctx context.Context, req *fuse.SetxattrRequest) error {

	glog.V(4).Infof("file Setxattr %s: %s", file.fullpath(), req.Name)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := setxattr(entry, req); err != nil {
		return err
	}

	return file.saveEntry(entry)

}

func (file *File) Removexattr(ctx context.Context, req *fuse.RemovexattrRequest) error {

	glog.V(4).Infof("file Removexattr %s: %s", file.fullpath(), req.Name)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := removexattr(entry, req); err != nil {
		return err
	}

	return file.saveEntry(entry)

}

func (file *File) Listxattr(ctx context.Context, req *fuse.ListxattrRequest, resp *fuse.ListxattrResponse) error {

	glog.V(4).Infof("file Listxattr %s", file.fullpath())

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := listxattr(entry, req, resp); err != nil {
		return err
	}

	return nil

}

func (file *File) Fsync(ctx context.Context, req *fuse.FsyncRequest) error {
	// fsync works at OS level
	// write the file chunks to the filerGrpcAddress
	glog.V(4).Infof("%s/%s fsync file %+v", file.dir.FullPath(), file.Name, req)

	return nil
}

func (file *File) Forget() {
	t := util.NewFullPath(file.dir.FullPath(), file.Name)
	glog.V(4).Infof("Forget file %s", t)
	file.wfs.ReleaseHandle(t, fuse.HandleID(t.AsInode()))
}

func (file *File) maybeLoadEntry(ctx context.Context) (entry *filer_pb.Entry, err error) {

	file.wfs.handlesLock.Lock()
	handle, found := file.wfs.handles[file.Id()]
	file.wfs.handlesLock.Unlock()
	entry = file.entry
	if found {
		glog.V(4).Infof("maybeLoadEntry found opened file %s/%s", file.dir.FullPath(), file.Name)
		entry = handle.f.entry
	}

	if entry != nil {
		if len(entry.HardLinkId) == 0 {
			// only always reload hard link
			return entry, nil
		}
	}
	entry, err = file.wfs.maybeLoadEntry(file.dir.FullPath(), file.Name)
	if err != nil {
		glog.V(3).Infof("maybeLoadEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
		return entry, err
	}
	if entry != nil {
		// file.entry = entry
	} else {
		glog.Warningf("maybeLoadEntry not found entry %s/%s: %v", file.dir.FullPath(), file.Name, err)
	}
	return entry, nil
}

func lessThan(a, b *filer_pb.FileChunk) bool {
	if a.Mtime == b.Mtime {
		return a.Fid.FileKey < b.Fid.FileKey
	}
	return a.Mtime < b.Mtime
}

func (file *File) addChunks(chunks []*filer_pb.FileChunk) {

	// find the earliest incoming chunk
	newChunks := chunks
	earliestChunk := newChunks[0]
	for i := 1; i < len(newChunks); i++ {
		if lessThan(earliestChunk, newChunks[i]) {
			earliestChunk = newChunks[i]
		}
	}

	entry := file.getEntry()
	if entry == nil {
		return
	}

	// pick out-of-order chunks from existing chunks
	for _, chunk := range entry.Chunks {
		if lessThan(earliestChunk, chunk) {
			chunks = append(chunks, chunk)
		}
	}

	// sort incoming chunks
	sort.Slice(chunks, func(i, j int) bool {
		return lessThan(chunks[i], chunks[j])
	})

	glog.V(4).Infof("%s existing %d chunks adds %d more", file.fullpath(), len(entry.Chunks), len(chunks))

	entry.Chunks = append(entry.Chunks, newChunks...)
}

func (file *File) saveEntry(entry *filer_pb.Entry) error {
	return file.wfs.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		file.wfs.mapPbIdFromLocalToFiler(entry)
		defer file.wfs.mapPbIdFromFilerToLocal(entry)

		request := &filer_pb.CreateEntryRequest{
			Directory:  file.dir.FullPath(),
			Entry:      entry,
			Signatures: []int32{file.wfs.signature},
		}

		glog.V(4).Infof("save file entry: %v", request)
		_, err := client.CreateEntry(context.Background(), request)
		if err != nil {
			glog.Errorf("UpdateEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return fuse.EIO
		}

		file.wfs.metaCache.InsertEntry(context.Background(), filer.FromPbEntry(request.Directory, request.Entry))

		return nil
	})
}

func (file *File) getEntry() *filer_pb.Entry {
	return file.entry
}
<commit_msg>reduce logs<commit_after>package filesys

import (
	"context"
	"os"
	"sort"
	"time"

	"github.com/seaweedfs/fuse"
	"github.com/seaweedfs/fuse/fs"

	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

const blockSize = 512

var _ = fs.Node(&File{})
var _ = fs.NodeIdentifier(&File{})
var _ = fs.NodeOpener(&File{})
var _ = fs.NodeFsyncer(&File{})
var _ = fs.NodeSetattrer(&File{})
var _ = fs.NodeGetxattrer(&File{})
var _ = fs.NodeSetxattrer(&File{})
var _ = fs.NodeRemovexattrer(&File{})
var _ = fs.NodeListxattrer(&File{})
var _ = fs.NodeForgetter(&File{})

type File struct {
	Name          string
	dir           *Dir
	wfs           *WFS
	entry         *filer_pb.Entry
	isOpen        int
	dirtyMetadata bool
	id            uint64
}

func (file *File) fullpath() util.FullPath {
	return util.NewFullPath(file.dir.FullPath(), file.Name)
}

func (file *File) Id() uint64 {
	return file.id
}

func (file *File) Attr(ctx context.Context, attr *fuse.Attr) (err error) {

	glog.V(4).Infof("file Attr %s, open:%v existing:%v", file.fullpath(), file.isOpen, attr)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if entry == nil {
		return fuse.ENOENT
	}

	attr.Inode = file.Id()
	attr.Valid = time.Second
	attr.Mode = os.FileMode(entry.Attributes.FileMode)
	attr.Size = filer.FileSize(entry)
	if file.isOpen > 0 {
		attr.Size = entry.Attributes.FileSize
		glog.V(4).Infof("file Attr %s, open:%v, size: %d", file.fullpath(), file.isOpen, attr.Size)
	}
	attr.Crtime = time.Unix(entry.Attributes.Crtime, 0)
	attr.Mtime = time.Unix(entry.Attributes.Mtime, 0)
	attr.Gid = entry.Attributes.Gid
	attr.Uid = entry.Attributes.Uid
	attr.Blocks = attr.Size/blockSize + 1
	attr.BlockSize = uint32(file.wfs.option.ChunkSizeLimit)
	if entry.HardLinkCounter > 0 {
		attr.Nlink = uint32(entry.HardLinkCounter)
	}

	return nil

}

func (file *File) Getxattr(ctx context.Context, req *fuse.GetxattrRequest, resp *fuse.GetxattrResponse) error {

	// glog.V(4).Infof("file Getxattr %s", file.fullpath())

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	return getxattr(entry, req, resp)
}

func (file *File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {

	glog.V(4).Infof("file %v open %+v", file.fullpath(), req)

	handle := file.wfs.AcquireHandle(file, req.Uid, req.Gid, req.Flags&fuse.OpenWriteOnly > 0)

	resp.Handle = fuse.HandleID(handle.handle)

	glog.V(4).Infof("%v file open handle id = %d", file.fullpath(), handle.handle)

	return handle, nil

}

func (file *File) Setattr(ctx context.Context, req *fuse.SetattrRequest, resp *fuse.SetattrResponse) error {

	glog.V(4).Infof("%v file setattr %+v", file.fullpath(), req)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}
	if file.isOpen > 0 {
		file.wfs.handlesLock.Lock()
		fileHandle := file.wfs.handles[file.Id()]
		file.wfs.handlesLock.Unlock()

		if fileHandle != nil {
			fileHandle.Lock()
			defer fileHandle.Unlock()
		}
	}

	if req.Valid.Size() {

		glog.V(4).Infof("%v file setattr set size=%v chunks=%d", file.fullpath(), req.Size, len(entry.Chunks))
		if req.Size < filer.FileSize(entry) {
			// fmt.Printf("truncate %v \n", fullPath)
			var chunks []*filer_pb.FileChunk
			var truncatedChunks []*filer_pb.FileChunk
			for _, chunk := range entry.Chunks {
				int64Size := int64(chunk.Size)
				if chunk.Offset+int64Size > int64(req.Size) {
					// this chunk is truncated
					int64Size = int64(req.Size) - chunk.Offset
					if int64Size > 0 {
						chunks = append(chunks, chunk)
						glog.V(4).Infof("truncated chunk %+v from %d to %d\n", chunk.GetFileIdString(), chunk.Size, int64Size)
						chunk.Size = uint64(int64Size)
					} else {
						glog.V(4).Infof("truncated whole chunk %+v\n", chunk.GetFileIdString())
						truncatedChunks = append(truncatedChunks, chunk)
					}
				}
			}
			entry.Chunks = chunks
		}
		entry.Attributes.FileSize = req.Size
		file.dirtyMetadata = true
	}

	if req.Valid.Mode() {
		entry.Attributes.FileMode = uint32(req.Mode)
		file.dirtyMetadata = true
	}

	if req.Valid.Uid() {
		entry.Attributes.Uid = req.Uid
		file.dirtyMetadata = true
	}

	if req.Valid.Gid() {
		entry.Attributes.Gid = req.Gid
		file.dirtyMetadata = true
	}

	if req.Valid.Crtime() {
		entry.Attributes.Crtime = req.Crtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Mtime() {
		entry.Attributes.Mtime = req.Mtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Handle() {
		// fmt.Printf("file handle => %d\n", req.Handle)
	}

	if file.isOpen > 0 {
		return nil
	}

	if !file.dirtyMetadata {
		return nil
	}

	return file.saveEntry(entry)

}

func (file *File) Setxattr(ctx context.Context, req *fuse.SetxattrRequest) error {

	glog.V(4).Infof("file Setxattr %s: %s", file.fullpath(), req.Name)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := setxattr(entry, req); err != nil {
		return err
	}

	return file.saveEntry(entry)

}

func (file *File) Removexattr(ctx context.Context, req *fuse.RemovexattrRequest) error {

	glog.V(4).Infof("file Removexattr %s: %s", file.fullpath(), req.Name)

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := removexattr(entry, req); err != nil {
		return err
	}

	return file.saveEntry(entry)

}

func (file *File) Listxattr(ctx context.Context, req *fuse.ListxattrRequest, resp *fuse.ListxattrResponse) error {

	glog.V(4).Infof("file Listxattr %s", file.fullpath())

	entry, err := file.maybeLoadEntry(ctx)
	if err != nil {
		return err
	}

	if err := listxattr(entry, req, resp); err != nil {
		return err
	}

	return nil

}

func (file *File) Fsync(ctx context.Context, req *fuse.FsyncRequest) error {
	// fsync works at OS level
	// write the file chunks to the filerGrpcAddress
	glog.V(4).Infof("%s/%s fsync file %+v", file.dir.FullPath(), file.Name, req)

	return nil
}

func (file *File) Forget() {
	t := util.NewFullPath(file.dir.FullPath(), file.Name)
	glog.V(4).Infof("Forget file %s", t)
	file.wfs.ReleaseHandle(t, fuse.HandleID(t.AsInode()))
}

func (file *File) maybeLoadEntry(ctx context.Context) (entry *filer_pb.Entry, err error) {

	file.wfs.handlesLock.Lock()
	handle, found := file.wfs.handles[file.Id()]
	file.wfs.handlesLock.Unlock()
	entry = file.entry
	if found {
		// glog.V(4).Infof("maybeLoadEntry found opened file %s/%s", file.dir.FullPath(), file.Name)
		entry = handle.f.entry
	}

	if entry != nil {
		if len(entry.HardLinkId) == 0 {
			// only always reload hard link
			return entry, nil
		}
	}
	entry, err = file.wfs.maybeLoadEntry(file.dir.FullPath(), file.Name)
	if err != nil {
		glog.V(3).Infof("maybeLoadEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
		return entry, err
	}
	if entry != nil {
		// file.entry = entry
	} else {
		glog.Warningf("maybeLoadEntry not found entry %s/%s: %v", file.dir.FullPath(), file.Name, err)
	}
	return entry, nil
}

func lessThan(a, b *filer_pb.FileChunk) bool {
	if a.Mtime == b.Mtime {
		return a.Fid.FileKey < b.Fid.FileKey
	}
	return a.Mtime < b.Mtime
}

func (file *File) addChunks(chunks []*filer_pb.FileChunk) {

	// find the earliest incoming chunk
	newChunks := chunks
	earliestChunk := newChunks[0]
	for i := 1; i < len(newChunks); i++ {
		if lessThan(earliestChunk, newChunks[i]) {
			earliestChunk = newChunks[i]
		}
	}

	entry := file.getEntry()
	if entry == nil {
		return
	}

	// pick out-of-order chunks from existing chunks
	for _, chunk := range entry.Chunks {
		if lessThan(earliestChunk, chunk) {
			chunks = append(chunks, chunk)
		}
	}

	// sort incoming chunks
	sort.Slice(chunks, func(i, j int) bool {
		return lessThan(chunks[i], chunks[j])
	})

	glog.V(4).Infof("%s existing %d chunks adds %d more", file.fullpath(), len(entry.Chunks), len(chunks))

	entry.Chunks = append(entry.Chunks, newChunks...)
}

func (file *File) saveEntry(entry *filer_pb.Entry) error {
	return file.wfs.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		file.wfs.mapPbIdFromLocalToFiler(entry)
		defer file.wfs.mapPbIdFromFilerToLocal(entry)

		request := &filer_pb.CreateEntryRequest{
			Directory:  file.dir.FullPath(),
			Entry:      entry,
			Signatures: []int32{file.wfs.signature},
		}

		glog.V(4).Infof("save file entry: %v", request)
		_, err := client.CreateEntry(context.Background(), request)
		if err != nil {
			glog.Errorf("UpdateEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return fuse.EIO
		}

		file.wfs.metaCache.InsertEntry(context.Background(), filer.FromPbEntry(request.Directory, request.Entry))

		return nil
	})
}

func (file *File) getEntry() *filer_pb.Entry {
	return file.entry
}
<|endoftext|>


<commit_before>package particle

import (
	"math"
	"time"

	"bitbucket.org/oakmoundstudio/oak/alg"
	"bitbucket.org/oakmoundstudio/oak/event"
	"bitbucket.org/oakmoundstudio/oak/physics"
	"bitbucket.org/oakmoundstudio/oak/render"
	"bitbucket.org/oakmoundstudio/oak/timing"
)

const (
	RECYCLED = -1000
)

// A Source is used to store and control a set of particles.
type Source struct {
	render.Layered
	Generator  Generator
	particles  [BLOCK_SIZE]Particle
	nextPID    int
	recycled   []int
	cID        event.CID
	pIDBlock   int
	stackLevel int
	paused     bool
}

func NewSource(g Generator, stackLevel int) *Source {
	ps := new(Source)
	ps.Generator = g
	ps.stackLevel = stackLevel
	ps.Init()
	return ps
}

func (ps *Source) Init() event.CID {
	cID := event.NextID(ps)
	cID.Bind(rotateParticles, "EnterFrame")
	ps.cID = cID
	ps.pIDBlock = Allocate(ps.cID)
	if ps.Generator.GetBaseGenerator().Duration != Inf {
		go func(ps_p *Source, duration alg.IntRange) {
			timing.DoAfter(time.Duration(duration.Poll())*time.Millisecond, func() {
				ps_p.Stop()
			})
		}(ps, ps.Generator.GetBaseGenerator().Duration)
	}
	return cID
}

func (ps *Source) CycleParticles() bool {
	pg := ps.Generator.GetBaseGenerator()
	cycled := false
	for i := 0; i < ps.nextPID; i++ {
		p := ps.particles[i]
		bp := p.GetBaseParticle()

		// Ignore dead particles
		if bp.Life > 0 {
			cycled = true
			bp.Life--
			// Apply rotational acceleration
			if pg.Rotation != nil {
				bp.Vel = bp.Vel.Rotate(pg.Rotation.Poll())
			}

			if pg.SpeedDecay.X != 0 {
				bp.Vel.X *= pg.SpeedDecay.X
				if math.Abs(bp.Vel.X) < 0.2 {
					bp.Vel.X = 0
				}
			}
			if pg.SpeedDecay.Y != 0 {
				bp.Vel.Y *= pg.SpeedDecay.Y
				if math.Abs(bp.Vel.Y) < 0.2 {
					bp.Vel.Y = 0
				}
			}

			bp.Vel = bp.Vel.Add(pg.Gravity)
			bp.Pos = bp.Pos.Add(bp.Vel)
			bp.SetLayer(ps.Layer(bp.Pos))
			p.Cycle(ps.Generator)

		} else if bp.Life != RECYCLED {
			p.UnDraw()
			cycled = true
			if pg.EndFunc != nil {
				pg.EndFunc(p)
			}
			// This relies on Life going down by 1 per frame
			bp.Life = RECYCLED
			ps.recycled = append(ps.recycled, bp.pID%BLOCK_SIZE)
		}

	}
	return cycled
}

// Shorthand
func (ps *Source) Layer(v physics.Vector) int {
	return ps.Generator.GetBaseGenerator().LayerFunc(v)
}

func (ps *Source) AddParticles() {
	pg := ps.Generator.GetBaseGenerator()
	// Regularly create particles (up until max particles)
	newParticleCount := int(pg.NewPerFrame.Poll())
	//fmt.Println("New particle count:", newParticleCount)
	ri := 0
	for ri < len(ps.recycled) && ri < newParticleCount {

		j := ps.recycled[ri]
		bp := ps.particles[j].GetBaseParticle()
		angle := pg.Angle.Poll() * math.Pi / 180.0
		speed := pg.Speed.Poll()
		startLife := pg.LifeSpan.Poll()

		bp.Pos = physics.NewVector(
			pg.X+floatFromSpread(pg.Spread.X),
			pg.Y+floatFromSpread(pg.Spread.Y))
		bp.Vel = physics.NewVector(
			speed*math.Cos(angle)*-1,
			speed*math.Sin(angle)*-1)
		bp.Life = startLife
		bp.totalLife = startLife
		ps.particles[ps.recycled[ri]] = ps.Generator.GenerateParticle(bp)

		ps.particles[ps.recycled[ri]].SetLayer(ps.Layer(bp.Pos))
		render.Draw(ps.particles[ps.recycled[ri]], ps.stackLevel)
		ri++
	}
	newParticleCount -= len(ps.recycled)
	if ri > 0 {
		ps.recycled = ps.recycled[ri:]
	}

	if ps.nextPID >= BLOCK_SIZE {
		return
	}
	if newParticleCount > 0 {
		//fmt.Println("Making non-recycled particles")
	}
	for i := 0; i < newParticleCount; i++ {
		angle := pg.Angle.Poll() * math.Pi / 180.0
		speed := pg.Speed.Poll()
		startLife := pg.LifeSpan.Poll()

		bp := &BaseParticle{
			Src: ps,
			Pos: physics.NewVector(
				pg.X+floatFromSpread(pg.Spread.X),
				pg.Y+floatFromSpread(pg.Spread.Y)),
			Vel: physics.NewVector(
				speed*math.Cos(angle)*-1,
				speed*math.Sin(angle)*-1),
			Life:      startLife,
			totalLife: startLife,
			pID:       ps.nextPID + ps.pIDBlock*BLOCK_SIZE,
		}

		p := ps.Generator.GenerateParticle(bp)
		ps.particles[ps.nextPID] = p
		ps.nextPID++
		if ps.nextPID >= BLOCK_SIZE {
			return
		}
		p.SetLayer(ps.Layer(bp.Pos))
		render.Draw(p, ps.stackLevel)
	}

}

// rotateParticles updates particles over time as long
// as a Source is active.
func rotateParticles(id int, nothing interface{}) int {
	ps := event.GetEntity(id).(*Source)
	//fmt.Print("rotated")
	if !ps.paused {
		ps.CycleParticles()
		ps.AddParticles()
	}
	return 0
}

// clearParticles is used after a Source has been stopped
// to continue moving old particles for as long as they exist.
func clearParticles(id int, nothing interface{}) int {
	ps := event.GetEntity(id).(*Source)
	if !ps.paused {
		if ps.CycleParticles() {
		} else {
			//fmt.Print("Killed a Source")
			event.DestroyEntity(id)
			Deallocate(ps.pIDBlock)
			return event.UNBIND_EVENT
		}
	}
	return 0
}

// Stop manually stops a Source, if its duration is infinite
// or if it should be stopped before expriring naturally.
func (ps *Source) Stop() {
	ps.cID.UnbindAllAndRebind([]event.Bindable{clearParticles}, []string{"EnterFrame"})
}

// Pausing a Source just stops the repetition
// of its rotation function, which moves, destroys,
// ages and generates particles. Existing particles will
// stay in place.
func (ps *Source) Pause() {
	ps.paused = true
}

// Unpausing a Source rebinds it's rotate function.
func (ps *Source) UnPause() {
	ps.paused = false
}

// Placeholder
func (ps *Source) String() string {
	return "ParticleSource"
}

func (ps *Source) ShiftX(x float64) {
	ps.Generator.ShiftX(x)
}

func (ps *Source) ShiftY(y float64) {
	ps.Generator.ShiftY(y)
}

func (ps *Source) SetPos(x, y float64) {
	ps.Generator.SetPos(x, y)
}
<commit_msg>Added ability to tell sources what to do when they stop<commit_after>package particle

import (
	"math"
	"time"

	"bitbucket.org/oakmoundstudio/oak/alg"
	"bitbucket.org/oakmoundstudio/oak/event"
	"bitbucket.org/oakmoundstudio/oak/physics"
	"bitbucket.org/oakmoundstudio/oak/render"
	"bitbucket.org/oakmoundstudio/oak/timing"
)

const (
	RECYCLED = -1000
)

// A Source is used to store and control a set of particles.
type Source struct {
	render.Layered
	Generator  Generator
	particles  [BLOCK_SIZE]Particle
	nextPID    int
	recycled   []int
	CID        event.CID
	pIDBlock   int
	stackLevel int
	paused     bool
	EndFunc    func()
}

func NewSource(g Generator, stackLevel int) *Source {
	ps := new(Source)
	ps.Generator = g
	ps.stackLevel = stackLevel
	ps.Init()
	return ps
}

func (ps *Source) Init() event.CID {
	CID := event.NextID(ps)
	CID.Bind(rotateParticles, "EnterFrame")
	ps.CID = CID
	ps.pIDBlock = Allocate(ps.CID)
	if ps.Generator.GetBaseGenerator().Duration != Inf {
		go func(ps_p *Source, duration alg.IntRange) {
			timing.DoAfter(time.Duration(duration.Poll())*time.Millisecond, func() {
				ps_p.Stop()
			})
		}(ps, ps.Generator.GetBaseGenerator().Duration)
	}
	return CID
}

func (ps *Source) CycleParticles() bool {
	pg := ps.Generator.GetBaseGenerator()
	cycled := false
	for i := 0; i < ps.nextPID; i++ {
		p := ps.particles[i]
		bp := p.GetBaseParticle()

		// Ignore dead particles
		if bp.Life > 0 {
			cycled = true
			bp.Life--
			// Apply rotational acceleration
			if pg.Rotation != nil {
				bp.Vel = bp.Vel.Rotate(pg.Rotation.Poll())
			}

			if pg.SpeedDecay.X != 0 {
				bp.Vel.X *= pg.SpeedDecay.X
				if math.Abs(bp.Vel.X) < 0.2 {
					bp.Vel.X = 0
				}
			}
			if pg.SpeedDecay.Y != 0 {
				bp.Vel.Y *= pg.SpeedDecay.Y
				if math.Abs(bp.Vel.Y) < 0.2 {
					bp.Vel.Y = 0
				}
			}

			bp.Vel = bp.Vel.Add(pg.Gravity)
			bp.Pos = bp.Pos.Add(bp.Vel)
			bp.SetLayer(ps.Layer(bp.Pos))
			p.Cycle(ps.Generator)

		} else if bp.Life != RECYCLED {
			p.UnDraw()
			cycled = true
			if pg.EndFunc != nil {
				pg.EndFunc(p)
			}
			// This relies on Life going down by 1 per frame
			bp.Life = RECYCLED
			ps.recycled = append(ps.recycled, bp.pID%BLOCK_SIZE)
		}

	}
	return cycled
}

// Shorthand
func (ps *Source) Layer(v physics.Vector) int {
	return ps.Generator.GetBaseGenerator().LayerFunc(v)
}

func (ps *Source) AddParticles() {
	pg := ps.Generator.GetBaseGenerator()
	// Regularly create particles (up until max particles)
	newParticleCount := int(pg.NewPerFrame.Poll())
	//fmt.Println("New particle count:", newParticleCount)
	ri := 0
	for ri < len(ps.recycled) && ri < newParticleCount {

		j := ps.recycled[ri]
		bp := ps.particles[j].GetBaseParticle()
		angle := pg.Angle.Poll() * math.Pi / 180.0
		speed := pg.Speed.Poll()
		startLife := pg.LifeSpan.Poll()

		bp.Pos = physics.NewVector(
			pg.X+floatFromSpread(pg.Spread.X),
			pg.Y+floatFromSpread(pg.Spread.Y))
		bp.Vel = physics.NewVector(
			speed*math.Cos(angle)*-1,
			speed*math.Sin(angle)*-1)
		bp.Life = startLife
		bp.totalLife = startLife
		ps.particles[ps.recycled[ri]] = ps.Generator.GenerateParticle(bp)

		ps.particles[ps.recycled[ri]].SetLayer(ps.Layer(bp.Pos))
		render.Draw(ps.particles[ps.recycled[ri]], ps.stackLevel)
		ri++
	}
	newParticleCount -= len(ps.recycled)
	if ri > 0 {
		ps.recycled = ps.recycled[ri:]
	}

	if ps.nextPID >= BLOCK_SIZE {
		return
	}
	if newParticleCount > 0 {
		//fmt.Println("Making non-recycled particles")
	}
	for i := 0; i < newParticleCount; i++ {
		angle := pg.Angle.Poll() * math.Pi / 180.0
		speed := pg.Speed.Poll()
		startLife := pg.LifeSpan.Poll()

		bp := &BaseParticle{
			Src: ps,
			Pos: physics.NewVector(
				pg.X+floatFromSpread(pg.Spread.X),
				pg.Y+floatFromSpread(pg.Spread.Y)),
			Vel: physics.NewVector(
				speed*math.Cos(angle)*-1,
				speed*math.Sin(angle)*-1),
			Life:      startLife,
			totalLife: startLife,
			pID:       ps.nextPID + ps.pIDBlock*BLOCK_SIZE,
		}

		p := ps.Generator.GenerateParticle(bp)
		ps.particles[ps.nextPID] = p
		ps.nextPID++
		if ps.nextPID >= BLOCK_SIZE {
			return
		}
		p.SetLayer(ps.Layer(bp.Pos))
		render.Draw(p, ps.stackLevel)
	}

}

// rotateParticles updates particles over time as long
// as a Source is active.
func rotateParticles(id int, nothing interface{}) int {
	ps := event.GetEntity(id).(*Source)
	//fmt.Print("rotated")
	if !ps.paused {
		ps.CycleParticles()
		ps.AddParticles()
	}
	return 0
}

// clearParticles is used after a Source has been stopped
// to continue moving old particles for as long as they exist.
func clearParticles(id int, nothing interface{}) int {
	ps := event.GetEntity(id).(*Source)
	if !ps.paused {
		if ps.CycleParticles() {
		} else {
			//fmt.Print("Killed a Source")
			if ps.EndFunc != nil {
				ps.EndFunc()
			}
			event.DestroyEntity(id)
			Deallocate(ps.pIDBlock)
			return event.UNBIND_EVENT
		}
	}
	return 0
}

// Stop manually stops a Source, if its duration is infinite
// or if it should be stopped before expriring naturally.
func (ps *Source) Stop() {
	ps.CID.UnbindAllAndRebind([]event.Bindable{clearParticles}, []string{"EnterFrame"})
}

// Pausing a Source just stops the repetition
// of its rotation function, which moves, destroys,
// ages and generates particles. Existing particles will
// stay in place.
func (ps *Source) Pause() {
	ps.paused = true
}

// Unpausing a Source rebinds it's rotate function.
func (ps *Source) UnPause() {
	ps.paused = false
}

// Placeholder
func (ps *Source) String() string {
	return "ParticleSource"
}

func (ps *Source) ShiftX(x float64) {
	ps.Generator.ShiftX(x)
}

func (ps *Source) ShiftY(y float64) {
	ps.Generator.ShiftY(y)
}

func (ps *Source) SetPos(x, y float64) {
	ps.Generator.SetPos(x, y)
}
<|endoftext|>


<commit_before>package reader

import (
	"testing"

	"os"

	"github.com/stretchr/testify/assert"
)

func TestNewInstanceHavingInputFile(t *testing.T) {
	// given
	file := new(os.File)
	c := newFileReaderMockedConfiguration(file)

	// when
	sut := newFileReader(c)

	// then
	assert.Equal(t, file, sut.inputFile, "Expected file is different than the specified in configuration.")
}

func TestPipe(t *testing.T) {

}
<commit_msg>File reader test<commit_after>package reader

import (
	"testing"

	"os"

	"path/filepath"

	"github.com/stretchr/testify/assert"
)

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

	// given
	file := new(os.File)
	c := newFileReaderMockedConfiguration(file)

	// when
	sut := newFileReader(c)

	// then
	assert.Equal(file, sut.inputFile, "Expected file is different than the specified in configuration.")
}

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

	// given
	file, error := openFuzzSetFile("fuzz_01.txt")
	assert.Nil(error, "Fuzz set file must be available, error: %v", error)

	c := newFileReaderMockedConfiguration(file)
	sut := newFileReader(c)

	var in = make(chan string, 3) // number of lines in fuzz file
	var done = make(chan bool, 1)

	// when
	sut.pipe(in, done)

	// then
	assert.True(<-done, "Reading file should be finished.")
	assert.Equal("001", <-in, "Invalid line read from file.")
	assert.Equal("002", <-in, "Invalid line read from file.")
	assert.Equal("003", <-in, "Invalid line read from file.")
	assert.Len(in, 0, "Read buffer should empty now.")

	assert.Equal(^(uintptr(0)), file.Fd(), "File descriptor should be closed now.")
}

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

	// given
	file, error := openFuzzSetFile("fuzz_01.txt")
	assert.Nil(error, "Fuzz set file must be available, error: %v", error)

	c := newFileReaderMockedConfiguration(file)
	sut := newFileReader(c)

	// when
	sut.closeFile()

	// then
	assert.Equal(^(uintptr(0)), file.Fd(), "File descriptor should be closed now.")
}

func openFuzzSetFile(fileName string) (*os.File, error) {
	fuzzSetFile := filepath.Join("../../../input-data/", fileName)
	absoluteFilePath, _ := filepath.Abs(fuzzSetFile)

	file, error := os.OpenFile(absoluteFilePath, os.O_RDONLY, 0666)
	if nil != error {
		return nil, error
	}

	_, error = file.Stat()
	if nil != error {
		return nil, error
	}

	return file, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package data

import (
	"path/filepath"
	"testing"

	"github.com/caixw/typing/vars"
	"github.com/issue9/assert"
)

func TestPost_sanitize(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	post, err := loadPost(p, filepath.Clean("./testdata/data/posts/post1/meta.yaml"))
	a.NotError(err).NotNil(post)

	a.NotError(post.sanitize())
	a.Equal(len(post.Tags), 0) // 未调用 sanitize2 初始化
	a.Equal(post.Modified, 0)
	a.Equal(post.Template, "post") // 默认模板
}

func TestLoadPost(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	post, err := loadPost(p, filepath.Clean("./testdata/data/posts/post1/meta.yaml"))
	a.NotError(err).NotNil(post)
	a.Equal(len(post.Tags), 0) // 未调用 Data.sanitize2 初始化
	a.Equal(post.Modified, 0)
	a.Equal(post.Template, "") // 未调用 Post.sanitize
	a.Equal(post.Content, "
a1
\n")

	post, err = loadPost(p, filepath.Clean("./testdata/data/posts/folder/post2/meta.yaml"))
	a.NotError(err).NotNil(post)
	a.Equal(post.Slug, "folder/post2")
	a.Equal(post.Template, "t1") // 模板
}

func TestLoadPosts(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	posts, err := loadPosts(p)
	a.NotError(err).NotNil(posts)
	a.Equal(len(posts), 2)
}
<commit_msg>修正测试错误<commit_after>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package data

import (
	"path/filepath"
	"testing"

	"github.com/caixw/typing/vars"
	"github.com/issue9/assert"
)

func TestPost_sanitize(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	post, err := loadPost(p, filepath.Clean("./testdata/data/posts/post1/meta.yaml"))
	a.NotError(err).NotNil(post)

	a.NotError(post.sanitize())
	a.Equal(len(post.Tags), 0)     // 未调用 sanitize2 初始化
	a.Equal(post.Template, "post") // 默认模板
}

func TestLoadPost(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	post, err := loadPost(p, filepath.Clean("./testdata/data/posts/post1/meta.yaml"))
	a.NotError(err).NotNil(post)
	a.Equal(len(post.Tags), 0) // 未调用 Data.sanitize2 初始化
	a.True(post.Modified.IsZero())
	a.Equal(post.Template, "") // 未调用 Post.sanitize
	a.Equal(post.Content, "
a1
\n")

	post, err = loadPost(p, filepath.Clean("./testdata/data/posts/folder/post2/meta.yaml"))
	a.NotError(err).NotNil(post)
	a.Equal(post.Slug, "folder/post2")
	a.Equal(post.Template, "t1") // 模板
}

func TestLoadPosts(t *testing.T) {
	a := assert.New(t)
	p := vars.NewPath("./testdata")

	posts, err := loadPosts(p)
	a.NotError(err).NotNil(posts)
	a.Equal(len(posts), 2)
}
<|endoftext|>


<commit_before>//
// tun.go -- tun interface with cgo for linux / bsd
//

package samtun

/*

#include <linux/if.h>
#include <linux/if_tun.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>

int tundev_open(char * ifname) {
  int fd = open("/dev/net/tun", O_RDWR);
  if (fd > 0) {
    struct ifreq ifr;
    memset(&ifr, 0, sizeof(ifr));
    ifr.ifr_flags = IFF_TUN | IFF_NO_PI;
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    if ( ioctl(fd , TUNSETIFF, (void*) &ifr) < 0) {
      close(fd);
      return -1;
    }
  }
  return fd;
}

int tundev_up(char * ifname, char * addr, char * dstaddr, int mtu) {

  struct ifreq ifr;
  memset(&ifr, 0, sizeof(ifr));
  strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
  int fd = socket(AF_INET6, SOCK_DGRAM, IPPROTO_IP);
  if ( fd > 0 ) {
    if ( ioctl(fd, SIOCGIFINDEX, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    struct sockaddr_in dst;
    memset(&dst, 0, sizeof(dst));
    dst.sin_addr = inet_addr(dstaddr);
    struct sockaddr_in src;
    memset(&src, 0, sizeof(src));
    src.sin_addr = inet_addr(addr);
    memcpy(&ifr.ifr_addr, src, sizeof(src));
    memcpy(&ifr.ifr_dstaddr, dst, sizeof(dst));
    if ( ioctl(fd, SIOCSIFADDR, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    ifr.ifr_mtu = mtu;
    if ( ioctl(fd, SIOCSIFMTU, &ifr) < 0) {
      close(fd);
      return -1;
    }
    if ( ioctl(fd, SIOCGIFFLAGS, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    ifr.ifr_flags |= IFF_UP | IFF_RUNNING;
    if ( ioctl(fd, SIOCSIFFLAGS, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    close(fd);
    return 0;
  }
  return -1;
}

void tundev_close(int fd) {
  close(fd);
}

*/
import "C"

type tunDev C.int

func newTun(ifname, addr, dstaddr string) (t tunDev, err error) {
  t = C.tundev_open(C.CString(ifname))
  if t == -1 {
    err = errors.New("cannot open tun interface")
  } else {
    if C.tundev_up(C.CString(ifname), C.CString(addr), C.CString(dstaddr), C.int(mtu)) < 0 {
      err = errors.New("cannot put up interface")
    }
  }
  return
}

func (d tunDev) Close() {
  C.tundev_close(d)
}
<commit_msg>fix more compiler errors<commit_after>//
// tun.go -- tun interface with cgo for linux / bsd
//

package samtun

/*

#include <linux/if.h>
#include <linux/if_tun.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>

int tundev_open(char * ifname) {
  int fd = open("/dev/net/tun", O_RDWR);
  if (fd > 0) {
    struct ifreq ifr;
    memset(&ifr, 0, sizeof(ifr));
    ifr.ifr_flags = IFF_TUN | IFF_NO_PI;
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    if ( ioctl(fd , TUNSETIFF, (void*) &ifr) < 0) {
      close(fd);
      return -1;
    }
  }
  return fd;
}

int tundev_up(char * ifname, char * addr, char * dstaddr, int mtu) {

  struct ifreq ifr;
  memset(&ifr, 0, sizeof(ifr));
  strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
  int fd = socket(AF_INET6, SOCK_DGRAM, IPPROTO_IP);
  if ( fd > 0 ) {
    if ( ioctl(fd, SIOCGIFINDEX, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    struct sockaddr_in dst;
    memset(&dst, 0, sizeof(dst));
    dst.sin_addr = inet_addr(dstaddr);
    struct sockaddr_in src;
    memset(&src, 0, sizeof(src));
    src.sin_addr = inet_addr(addr);
    memcpy(&ifr.ifr_addr, &src, sizeof(src));
    memcpy(&ifr.ifr_dstaddr, &dst, sizeof(dst));
    if ( ioctl(fd, SIOCSIFADDR, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    ifr.ifr_mtu = mtu;
    if ( ioctl(fd, SIOCSIFMTU, &ifr) < 0) {
      close(fd);
      return -1;
    }
    if ( ioctl(fd, SIOCGIFFLAGS, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    ifr.ifr_flags |= IFF_UP | IFF_RUNNING;
    if ( ioctl(fd, SIOCSIFFLAGS, &ifr) < 0 ) {
      close(fd);
      return -1;
    }
    close(fd);
    return 0;
  }
  return -1;
}

void tundev_close(int fd) {
  close(fd);
}

*/
import "C"

type tunDev C.int

func newTun(ifname, addr, dstaddr string) (t tunDev, err error) {
  t = C.tundev_open(C.CString(ifname))
  if t == -1 {
    err = errors.New("cannot open tun interface")
  } else {
    if C.tundev_up(C.CString(ifname), C.CString(addr), C.CString(dstaddr), C.int(mtu)) < 0 {
      err = errors.New("cannot put up interface")
    }
  }
  return
}

func (d tunDev) Close() {
  C.tundev_close(d)
}
<|endoftext|>


<commit_before>package pt

import (
	"bufio"
	"code.google.com/p/mahonia"
	"fmt"
	"github.com/monochromegane/the_platinum_searcher/util"
	"os"
	"path/filepath"
	"strings"
)

type Searcher struct {
	Root, Pattern string
	Option        *SearchOption
}

type SearchOption struct {
	NoColor bool `long:"nocolor" description:"Don't print color codes in results (Disabled by default)"`
	NoGroup bool `long:"nogroup" description:"Don't print file name at header (Disabled by default)"`
}

type GrepArgument struct {
	Path, Pattern, Encode string
}

type Match struct {
	LineNum int
	Match   string
}

type PrintArgument struct {
	Pattern string
	Path    string
	Matches []*Match
}

func (self *Searcher) Search() {
	grep := make(chan *GrepArgument, 2)
	match := make(chan *PrintArgument, 2)
	done := make(chan bool)
	go self.find(grep)
	go self.grep(grep, match)
	go self.print(match, done)
	<-done
}

func (self *Searcher) find(grep chan *GrepArgument) {
	filepath.Walk(self.Root, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if len(info.Name()) > 1 && strings.Index(info.Name(), ".") == 0 {
				return filepath.SkipDir
			} else {
				return nil
			}
		}
		fileType := pt.IdentifyFileType(path)
		if fileType == pt.BINARY {
			return nil
		}
		grep <- &GrepArgument{path, self.Pattern, fileType}
		return nil
	})
	grep <- nil
}

func (self *Searcher) grep(grep chan *GrepArgument, match chan *PrintArgument) {
	for {
		arg := <-grep
		if arg == nil {
			break
		}

		fh, err := os.Open(arg.Path)
		f := bufio.NewReader(fh)
		if err != nil {
			panic(err)
		}
		buf := make([]byte, 1024)
		decoder := mahonia.NewDecoder(arg.Encode)

		m := make([]*Match, 0)

		var lineNum = 1
		for {
			buf, _, err = f.ReadLine()
			if err != nil {
				break
			}

			s := string(buf)
			if decoder != nil && arg.Encode != pt.UTF8 && arg.Encode != pt.ASCII {
				s = decoder.ConvertString(s)
			}
			if strings.Contains(s, arg.Pattern) {
				m = append(m, &Match{lineNum, s})
			}
			lineNum++
		}
		match <- &PrintArgument{arg.Pattern, arg.Path, m}
		fh.Close()

	}
	match <- nil
}

func (self *Searcher) print(match chan *PrintArgument, done chan bool) {
	for {
		arg := <-match
		if arg == nil {
			break
		}
		if len(arg.Matches) == 0 {
			continue
		}

		if !self.Option.NoGroup {
			pt.PrintPath(arg.Path, self.Option.NoColor)
			fmt.Printf("\n")
		}
		for _, v := range arg.Matches {
			if v == nil {
				continue
			}
			if self.Option.NoGroup {
				pt.PrintPath(arg.Path, self.Option.NoColor)
			}
			pt.PrintLineNumber(v.LineNum, self.Option.NoColor)
			pt.PrintMatch(arg.Pattern, v.Match, self.Option.NoColor)
			fmt.Printf("\n")
		}
		if !self.Option.NoGroup {
			fmt.Printf("\n")
		}
	}
	done <- true
}
<commit_msg>Used range for channel loop.<commit_after>package pt

import (
	"bufio"
	"code.google.com/p/mahonia"
	"fmt"
	"github.com/monochromegane/the_platinum_searcher/util"
	"os"
	"path/filepath"
	"strings"
)

type Searcher struct {
	Root, Pattern string
	Option        *SearchOption
}

type SearchOption struct {
	NoColor bool `long:"nocolor" description:"Don't print color codes in results (Disabled by default)"`
	NoGroup bool `long:"nogroup" description:"Don't print file name at header (Disabled by default)"`
}

type GrepArgument struct {
	Path, Pattern, Encode string
}

type Match struct {
	LineNum int
	Match   string
}

type PrintArgument struct {
	Pattern string
	Path    string
	Matches []*Match
}

func (self *Searcher) Search() {
	grep := make(chan *GrepArgument, 2)
	match := make(chan *PrintArgument, 2)
	done := make(chan bool)
	go self.find(grep)
	go self.grep(grep, match)
	go self.print(match, done)
	<-done
}

func (self *Searcher) find(grep chan *GrepArgument) {
	filepath.Walk(self.Root, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if len(info.Name()) > 1 && strings.Index(info.Name(), ".") == 0 {
				return filepath.SkipDir
			} else {
				return nil
			}
		}
		fileType := pt.IdentifyFileType(path)
		if fileType == pt.BINARY {
			return nil
		}
		grep <- &GrepArgument{path, self.Pattern, fileType}
		return nil
	})
	close(grep)
}

func (self *Searcher) grep(grep chan *GrepArgument, match chan *PrintArgument) {
	for arg := range grep {

		fh, err := os.Open(arg.Path)
		f := bufio.NewReader(fh)
		if err != nil {
			panic(err)
		}
		buf := make([]byte, 1024)
		decoder := mahonia.NewDecoder(arg.Encode)

		m := make([]*Match, 0)

		var lineNum = 1
		for {
			buf, _, err = f.ReadLine()
			if err != nil {
				break
			}

			s := string(buf)
			if decoder != nil && arg.Encode != pt.UTF8 && arg.Encode != pt.ASCII {
				s = decoder.ConvertString(s)
			}
			if strings.Contains(s, arg.Pattern) {
				m = append(m, &Match{lineNum, s})
			}
			lineNum++
		}
		match <- &PrintArgument{arg.Pattern, arg.Path, m}
		fh.Close()

	}
	close(match)
}

func (self *Searcher) print(match chan *PrintArgument, done chan bool) {
	for arg := range match {

		if len(arg.Matches) == 0 {
			continue
		}

		if !self.Option.NoGroup {
			pt.PrintPath(arg.Path, self.Option.NoColor)
			fmt.Printf("\n")
		}
		for _, v := range arg.Matches {
			if v == nil {
				continue
			}
			if self.Option.NoGroup {
				pt.PrintPath(arg.Path, self.Option.NoColor)
			}
			pt.PrintLineNumber(v.LineNum, self.Option.NoColor)
			pt.PrintMatch(arg.Pattern, v.Match, self.Option.NoColor)
			fmt.Printf("\n")
		}
		if !self.Option.NoGroup {
			fmt.Printf("\n")
		}
	}
	done <- true
}
<|endoftext|>


<commit_before>package scan

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
)

// Connectivity contains scanners testing basic connectivity to the host
var Connectivity = &Family{
	Description: "Scans for basic connectivity with the host through DNS and TCP/TLS dials",
	Scanners: map[string]*Scanner{
		"DNSLookup": {
			"Host can be resolved through DNS",
			dnsLookupScan,
		},
		"CloudFlareStatus": {
			"Host is on CloudFlare",
			onCloudFlareScan,
		},
		"TCPDial": {
			"Host accepts TCP connection",
			tcpDialScan,
		},
		"TLSDial": {
			"Host can perform TLS handshake",
			tlsDialScan,
		},
	},
}

// dnsLookupScan tests that DNS resolution of the host returns at least one address
func dnsLookupScan(addr, hostname string) (grade Grade, output Output, err error) {
	addrs, err := net.LookupHost(hostname)
	if err != nil {
		return
	}

	if len(addrs) == 0 {
		err = errors.New("no addresses found for host")
	}

	grade, output = Good, addrs
	return
}

var (
	cfNets    []*net.IPNet
	cfNetsErr error
)

func initOnCloudFlareScan() ([]*net.IPNet, error) {
	// Propogate previous errors and don't attempt to re-download.
	if cfNetsErr != nil {
		return nil, cfNetsErr
	}

	// Don't re-download ranges if we already have them.
	if len(cfNets) > 0 {
		return cfNets, nil
	}

	// Download CloudFlare CIDR ranges and parse them.
	v4resp, err := Client.Get("https://www.cloudflare.com/ips-v4")
	if err != nil {
		cfNetsErr = fmt.Errorf("Couldn't download CloudFlare IPs: %v", err)
		return nil, cfNetsErr
	}
	defer v4resp.Body.Close()

	v6resp, err := Client.Get("https://www.cloudflare.com/ips-v6")
	if err != nil {
		cfNetsErr = fmt.Errorf("Couldn't download CloudFlare IPs: %v", err)
		return nil, cfNetsErr
	}
	defer v6resp.Body.Close()

	scanner := bufio.NewScanner(io.MultiReader(v4resp.Body, strings.NewReader("\n"), v6resp.Body))
	for scanner.Scan() {
		_, ipnet, err := net.ParseCIDR(scanner.Text())
		if err != nil {
			cfNetsErr = fmt.Errorf("Couldn't parse CIDR range: %v", err)
			return nil, cfNetsErr
		}
		cfNets = append(cfNets, ipnet)
	}
	if err := scanner.Err(); err != nil {
		cfNetsErr = fmt.Errorf("Couldn't read IP bodies: %v", err)
		return nil, cfNetsErr
	}

	return cfNets, nil
}

func onCloudFlareScan(addr, hostname string) (grade Grade, output Output, err error) {
	var cloudflareNets []*net.IPNet
	if cloudflareNets, err = initOnCloudFlareScan(); err != nil {
		grade = Skipped
		return
	}

	_, addrs, err := dnsLookupScan(addr, hostname)
	if err != nil {
		return
	}

	cfStatus := make(map[string]bool)
	grade = Good
	for _, addr := range addrs.([]string) {
		ip := net.ParseIP(addr)
		for _, cfNet := range cloudflareNets {
			if cfNet.Contains(ip) {
				cfStatus[addr] = true
				break
			}
		}
		if !cfStatus[addr] {
			cfStatus[addr] = false
			grade = Bad
		}
	}

	output = cfStatus
	return
}

// tcpDialScan tests that the host can be connected to through TCP.
func tcpDialScan(addr, hostname string) (grade Grade, output Output, err error) {
	conn, err := Dialer.Dial(Network, addr)
	if err != nil {
		return
	}
	conn.Close()
	grade = Good
	return
}

// tlsDialScan tests that the host can perform a TLS Handshake
// and warns if the server's certificate can't be verified.
func tlsDialScan(addr, hostname string) (grade Grade, output Output, err error) {
	var conn *tls.Conn
	config := defaultTLSConfig(hostname)

	if conn, err = tls.DialWithDialer(Dialer, Network, addr, config); err != nil {
		return
	}
	conn.Close()

	config.InsecureSkipVerify = false
	if conn, err = tls.DialWithDialer(Dialer, Network, addr, config); err != nil {
		grade = Warning
		return
	}
	conn.Close()

	grade = Good
	return
}
<commit_msg>scan/connectivity: Remove newline from IP list for CIDR parsing<commit_after>package scan

import (
	"bufio"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
)

// Connectivity contains scanners testing basic connectivity to the host
var Connectivity = &Family{
	Description: "Scans for basic connectivity with the host through DNS and TCP/TLS dials",
	Scanners: map[string]*Scanner{
		"DNSLookup": {
			"Host can be resolved through DNS",
			dnsLookupScan,
		},
		"CloudFlareStatus": {
			"Host is on CloudFlare",
			onCloudFlareScan,
		},
		"TCPDial": {
			"Host accepts TCP connection",
			tcpDialScan,
		},
		"TLSDial": {
			"Host can perform TLS handshake",
			tlsDialScan,
		},
	},
}

// dnsLookupScan tests that DNS resolution of the host returns at least one address
func dnsLookupScan(addr, hostname string) (grade Grade, output Output, err error) {
	addrs, err := net.LookupHost(hostname)
	if err != nil {
		return
	}

	if len(addrs) == 0 {
		err = errors.New("no addresses found for host")
	}

	grade, output = Good, addrs
	return
}

var (
	cfNets    []*net.IPNet
	cfNetsErr error
)

func initOnCloudFlareScan() ([]*net.IPNet, error) {
	// Propogate previous errors and don't attempt to re-download.
	if cfNetsErr != nil {
		return nil, cfNetsErr
	}

	// Don't re-download ranges if we already have them.
	if len(cfNets) > 0 {
		return cfNets, nil
	}

	// Download CloudFlare CIDR ranges and parse them.
	v4resp, err := Client.Get("https://www.cloudflare.com/ips-v4")
	if err != nil {
		cfNetsErr = fmt.Errorf("Couldn't download CloudFlare IPs: %v", err)
		return nil, cfNetsErr
	}
	defer v4resp.Body.Close()

	v6resp, err := Client.Get("https://www.cloudflare.com/ips-v6")
	if err != nil {
		cfNetsErr = fmt.Errorf("Couldn't download CloudFlare IPs: %v", err)
		return nil, cfNetsErr
	}
	defer v6resp.Body.Close()

	scanner := bufio.NewScanner(io.MultiReader(v4resp.Body, v6resp.Body))
	for scanner.Scan() {
		_, ipnet, err := net.ParseCIDR(scanner.Text())
		if err != nil {
			cfNetsErr = fmt.Errorf("Couldn't parse CIDR range: %v", err)
			return nil, cfNetsErr
		}
		cfNets = append(cfNets, ipnet)
	}
	if err := scanner.Err(); err != nil {
		cfNetsErr = fmt.Errorf("Couldn't read IP bodies: %v", err)
		return nil, cfNetsErr
	}

	return cfNets, nil
}

func onCloudFlareScan(addr, hostname string) (grade Grade, output Output, err error) {
	var cloudflareNets []*net.IPNet
	if cloudflareNets, err = initOnCloudFlareScan(); err != nil {
		grade = Skipped
		return
	}

	_, addrs, err := dnsLookupScan(addr, hostname)
	if err != nil {
		return
	}

	cfStatus := make(map[string]bool)
	grade = Good
	for _, addr := range addrs.([]string) {
		ip := net.ParseIP(addr)
		for _, cfNet := range cloudflareNets {
			if cfNet.Contains(ip) {
				cfStatus[addr] = true
				break
			}
		}
		if !cfStatus[addr] {
			cfStatus[addr] = false
			grade = Bad
		}
	}

	output = cfStatus
	return
}

// tcpDialScan tests that the host can be connected to through TCP.
func tcpDialScan(addr, hostname string) (grade Grade, output Output, err error) {
	conn, err := Dialer.Dial(Network, addr)
	if err != nil {
		return
	}
	conn.Close()
	grade = Good
	return
}

// tlsDialScan tests that the host can perform a TLS Handshake
// and warns if the server's certificate can't be verified.
func tlsDialScan(addr, hostname string) (grade Grade, output Output, err error) {
	var conn *tls.Conn
	config := defaultTLSConfig(hostname)

	if conn, err = tls.DialWithDialer(Dialer, Network, addr, config); err != nil {
		return
	}
	conn.Close()

	config.InsecureSkipVerify = false
	if conn, err = tls.DialWithDialer(Dialer, Network, addr, config); err != nil {
		grade = Warning
		return
	}
	conn.Close()

	grade = Good
	return
}
<|endoftext|>


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

package wireguard

import (
	"fmt"
	"os"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/projectcalico/felix/netlinkshim"
)

const (
	labelHostname      = "hostname"
	labelPublicKey     = "public_key"
	labelInterfaceName = "iface"
	labelListenPort    = "listen_port"
	labelPeerKey       = "peer_key"
	labelPeerEndpoint  = "peer_endpoint"
)

var _ prometheus.Collector = (*Metrics)(nil)

const (
	wireguardMetaFQName   = "wireguard_meta"
	wireguardMetaHelpText = "wireguard interface and runtime metadata"

	wireguardLatestHandshakeIntervalFQName   = "wireguard_latest_handshake_seconds"
	wireguardLatestHandshakeIntervalHelpText = "wireguard interface latest handshake unix timestamp in seconds to a peer"

	wireguardBytesSentFQName   = "wireguard_bytes_sent"
	wireguardBytesSentHelpText = "wireguard interface outgoing bytes count to peer"

	wireguardBytesRcvdFQName   = "wireguard_bytes_rcvd"
	wireguardBytesRcvdHelpText = "wireguard interface incoming bytes count from peer"

	defaultCollectionRatelimit = time.Second
)

func init() {
	prometheus.MustRegister(
		MustNewWireguardMetrics(),
	)
}

type Metrics struct {
	hostname           string
	newWireguardClient func() (netlinkshim.Wireguard, error)
	logCtx             *logrus.Entry

	peerRx, peerTx map[wgtypes.Key]int64

	lastCollectionTime time.Time
	rateLimitInterval  time.Duration
}

func (collector *Metrics) Describe(d chan<- *prometheus.Desc) {
	for _, dev := range collector.getDevices() {
		collector.descsByDevice(d, dev)
	}
}

func (collector *Metrics) descsByDevice(d chan<- *prometheus.Desc, device *wgtypes.Device) {
	if device == nil {
		collector.logCtx.Error("BUG: called descsByDevice with nil device")
		return
	}

	labels := collector.defaultLabelValues("pub", deviceMetaLabelValues(device))
	for fqName, help := range map[string]string{
		wireguardMetaFQName: wireguardMetaHelpText,
	} {
		d <- prometheus.NewDesc(fqName, help, nil, labels)
	}
	for _, peer := range device.Peers {
		collector.descByPeer(d, &peer)
	}
}

func (collector *Metrics) descByPeer(d chan<- *prometheus.Desc, peer *wgtypes.Peer) {
	if peer == nil {
		collector.logCtx.Error("BUG: called descByPeer with nil peer")
		return
	}

	labels := collector.defaultLabelValues("pub", peerServiceLabelValues(peer))
	for fqName, help := range map[string]string{
		wireguardBytesRcvdFQName:               wireguardBytesRcvdHelpText,
		wireguardBytesSentFQName:               wireguardBytesSentHelpText,
		wireguardLatestHandshakeIntervalFQName: wireguardLatestHandshakeIntervalHelpText,
	} {
		d <- prometheus.NewDesc(fqName, help, nil, labels)
	}
}

func (collector *Metrics) Collect(m chan<- prometheus.Metric) {
	collector.refreshStats(m)
}

func MustNewWireguardMetrics() *Metrics {
	wg, err := NewWireguardMetrics()
	if err != nil {
		logrus.Panic(err)
	}
	return wg
}

func NewWireguardMetrics() (*Metrics, error) {
	hostname, err := os.Hostname()
	if err != nil {
		logrus.WithError(err).Error("cannot register wireguard metrics stats")
		return nil, err
	}
	return NewWireguardMetricsWithShims(hostname, netlinkshim.NewRealWireguard, defaultCollectionRatelimit), nil
}

func NewWireguardMetricsWithShims(hostname string, newWireguardClient func() (netlinkshim.Wireguard, error), rateLimitInterval time.Duration) *Metrics {
	logrus.WithField("hostname", hostname).Debug("created wireguard collector for host")
	return &Metrics{
		hostname:           hostname,
		newWireguardClient: newWireguardClient,
		logCtx: logrus.WithFields(logrus.Fields{
			"prometheus_collector": "wireguard",
		}),

		peerRx: map[wgtypes.Key]int64{},
		peerTx: map[wgtypes.Key]int64{},

		rateLimitInterval: rateLimitInterval,
	}
}

func (collector *Metrics) getDevices() []*wgtypes.Device {
	wgClient, err := collector.newWireguardClient()
	if err != nil {
		collector.logCtx.WithError(err).Error("error initializing wireguard client devices")
		return nil
	}

	devices, err := wgClient.Devices()
	if err != nil {
		collector.logCtx.WithError(err).Error("error listing wireguard devices")
		return nil
	}
	return devices
}

func (collector *Metrics) refreshStats(m chan<- prometheus.Metric) {
	if ct := time.Since(collector.lastCollectionTime); ct < collector.rateLimitInterval {
		collector.logCtx.WithFields(logrus.Fields{
			"since":              ct.String(),
			"ratelimit_interval": collector.rateLimitInterval.String(),
		}).Error("refreshStats disallowed due to rate limit")
		return
	}
	devices := collector.getDevices()
	collector.logCtx.WithFields(logrus.Fields{
		"count": len(devices),
		"dev":   devices,
	}).Debug("collect device metrics enumerated devices")

	collector.collectDeviceMetrics(devices, m)
	collector.collectDevicePeerMetrics(devices, m)

	collector.lastCollectionTime = time.Now()
}

func (collector *Metrics) collectDeviceMetrics(devices []*wgtypes.Device, m chan<- prometheus.Metric) {
	collector.logCtx.Debug("collecting wg device metrics")

	for _, device := range devices {
		l := collector.defaultLabelValues(device.PublicKey.String(), deviceMetaLabelValues(device))

		collector.logCtx.WithFields(logrus.Fields{
			"dev":    device.Name,
			"labels": l,
		}).Debug("iterate device")

		m <- prometheus.MustNewConstMetric(
			prometheus.NewDesc(
				wireguardMetaFQName, wireguardMetaHelpText, nil, l,
			),
			prometheus.GaugeValue,
			1,
		)
	}
}

func (collector *Metrics) collectDevicePeerMetrics(devices []*wgtypes.Device, m chan<- prometheus.Metric) {
	collector.logCtx.Debug("collecting wg peer(s) metrics")

	collector.logCtx.WithFields(logrus.Fields{
		"count": len(devices),
	}).Debug("enumerated wireguard devices")

	for _, device := range devices {
		logCtx := collector.logCtx.WithFields(logrus.Fields{
			"key":  device.PublicKey,
			"name": device.Name,
		})
		for _, peer := range device.Peers {
			logCtx.WithFields(logrus.Fields{
				"peer_key":      peer.PublicKey,
				"peer_endpoint": peer.Endpoint,
			}).Debug("collect peer metrics")

			labels := collector.defaultLabelValues(device.PublicKey.String(), peerServiceLabelValues(&peer))

			hs := float64(peer.LastHandshakeTime.Unix())

			collector.logCtx.WithFields(logrus.Fields{
				"rx_bytes_total": peer.ReceiveBytes,
				"tx_bytes_total": peer.TransmitBytes,
				"handshake_ts":   hs,
			}).Debug("collected peer metrics")

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardBytesRcvdFQName, wireguardBytesRcvdHelpText, nil, labels,
				),
				prometheus.CounterValue,
				float64(peer.ReceiveBytes),
			)

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardBytesSentFQName, wireguardBytesSentHelpText, nil, labels,
				),
				prometheus.CounterValue,
				float64(peer.TransmitBytes),
			)

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardLatestHandshakeIntervalFQName, wireguardLatestHandshakeIntervalHelpText, nil, labels,
				),
				prometheus.GaugeValue,
				hs,
			)
		}
	}

}

func (collector *Metrics) defaultLabelValues(publicKey string, extend prometheus.Labels) prometheus.Labels {
	l := prometheus.Labels{labelHostname: collector.hostname}
	if publicKey != "" {
		l[labelPublicKey] = publicKey
	}

	for k, v := range extend {
		l[k] = v
	}
	return l
}

func deviceMetaLabelValues(dev *wgtypes.Device) prometheus.Labels {
	return prometheus.Labels{
		labelInterfaceName: dev.Name,
		labelListenPort:    fmt.Sprintf("%d", dev.ListenPort),
	}
}

func peerServiceLabelValues(peer *wgtypes.Peer) prometheus.Labels {
	return prometheus.Labels{
		labelPeerKey:      peer.PublicKey.String(),
		labelPeerEndpoint: peer.Endpoint.String(),
	}
}
<commit_msg>ease up wireguard metrics log levels (#2945)<commit_after>// Copyright (c) 2021 Tigera, Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package wireguard

import (
	"fmt"
	"os"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/projectcalico/felix/netlinkshim"
)

const (
	labelHostname      = "hostname"
	labelPublicKey     = "public_key"
	labelInterfaceName = "iface"
	labelListenPort    = "listen_port"
	labelPeerKey       = "peer_key"
	labelPeerEndpoint  = "peer_endpoint"
)

var _ prometheus.Collector = (*Metrics)(nil)

const (
	wireguardMetaFQName   = "wireguard_meta"
	wireguardMetaHelpText = "wireguard interface and runtime metadata"

	wireguardLatestHandshakeIntervalFQName   = "wireguard_latest_handshake_seconds"
	wireguardLatestHandshakeIntervalHelpText = "wireguard interface latest handshake unix timestamp in seconds to a peer"

	wireguardBytesSentFQName   = "wireguard_bytes_sent"
	wireguardBytesSentHelpText = "wireguard interface outgoing bytes count to peer"

	wireguardBytesRcvdFQName   = "wireguard_bytes_rcvd"
	wireguardBytesRcvdHelpText = "wireguard interface incoming bytes count from peer"

	defaultCollectionRatelimit = time.Second
)

func init() {
	prometheus.MustRegister(
		MustNewWireguardMetrics(),
	)
}

type Metrics struct {
	hostname           string
	newWireguardClient func() (netlinkshim.Wireguard, error)
	logCtx             *logrus.Entry

	peerRx, peerTx map[wgtypes.Key]int64

	lastCollectionTime time.Time
	rateLimitInterval  time.Duration
}

func (collector *Metrics) Describe(d chan<- *prometheus.Desc) {
	for _, dev := range collector.getDevices() {
		collector.descsByDevice(d, dev)
	}
}

func (collector *Metrics) descsByDevice(d chan<- *prometheus.Desc, device *wgtypes.Device) {
	if device == nil {
		collector.logCtx.Error("BUG: called descsByDevice with nil device")
		return
	}

	labels := collector.defaultLabelValues("pub", deviceMetaLabelValues(device))
	for fqName, help := range map[string]string{
		wireguardMetaFQName: wireguardMetaHelpText,
	} {
		d <- prometheus.NewDesc(fqName, help, nil, labels)
	}
	for _, peer := range device.Peers {
		collector.descByPeer(d, &peer)
	}
}

func (collector *Metrics) descByPeer(d chan<- *prometheus.Desc, peer *wgtypes.Peer) {
	if peer == nil {
		collector.logCtx.Error("BUG: called descByPeer with nil peer")
		return
	}

	labels := collector.defaultLabelValues("pub", peerServiceLabelValues(peer))
	for fqName, help := range map[string]string{
		wireguardBytesRcvdFQName:               wireguardBytesRcvdHelpText,
		wireguardBytesSentFQName:               wireguardBytesSentHelpText,
		wireguardLatestHandshakeIntervalFQName: wireguardLatestHandshakeIntervalHelpText,
	} {
		d <- prometheus.NewDesc(fqName, help, nil, labels)
	}
}

func (collector *Metrics) Collect(m chan<- prometheus.Metric) {
	collector.refreshStats(m)
}

func MustNewWireguardMetrics() *Metrics {
	wg, err := NewWireguardMetrics()
	if err != nil {
		logrus.Panic(err)
	}
	return wg
}

func NewWireguardMetrics() (*Metrics, error) {
	hostname, err := os.Hostname()
	if err != nil {
		return nil, err
	}
	return NewWireguardMetricsWithShims(hostname, netlinkshim.NewRealWireguard, defaultCollectionRatelimit), nil
}

func NewWireguardMetricsWithShims(hostname string, newWireguardClient func() (netlinkshim.Wireguard, error), rateLimitInterval time.Duration) *Metrics {
	logrus.WithField("hostname", hostname).Debug("created wireguard collector for host")
	return &Metrics{
		hostname:           hostname,
		newWireguardClient: newWireguardClient,
		logCtx: logrus.WithFields(logrus.Fields{
			"prometheus_collector": "wireguard",
		}),

		peerRx: map[wgtypes.Key]int64{},
		peerTx: map[wgtypes.Key]int64{},

		rateLimitInterval: rateLimitInterval,
	}
}

func (collector *Metrics) getDevices() []*wgtypes.Device {
	wgClient, err := collector.newWireguardClient()
	if err != nil {
		collector.logCtx.WithError(err).Debug("something went wrong initializing wireguard rpc client")
		return nil
	}

	devices, err := wgClient.Devices()
	if err != nil {
		collector.logCtx.WithError(err).Debug("something went wrong enumerating wireguard devices")
		return nil
	}
	return devices
}

func (collector *Metrics) refreshStats(m chan<- prometheus.Metric) {
	if ct := time.Since(collector.lastCollectionTime); ct < collector.rateLimitInterval {
		collector.logCtx.WithFields(logrus.Fields{
			"since":              ct.String(),
			"ratelimit_interval": collector.rateLimitInterval.String(),
		}).Debug("refreshStats disallowed due to rate limit")
		return
	}
	devices := collector.getDevices()
	collector.logCtx.WithFields(logrus.Fields{
		"count": len(devices),
		"dev":   devices,
	}).Debug("collect device metrics enumerated devices")

	collector.collectDeviceMetrics(devices, m)
	collector.collectDevicePeerMetrics(devices, m)

	collector.lastCollectionTime = time.Now()
}

func (collector *Metrics) collectDeviceMetrics(devices []*wgtypes.Device, m chan<- prometheus.Metric) {
	collector.logCtx.Debug("collecting wg device metrics")

	for _, device := range devices {
		l := collector.defaultLabelValues(device.PublicKey.String(), deviceMetaLabelValues(device))

		collector.logCtx.WithFields(logrus.Fields{
			"dev":    device.Name,
			"labels": l,
		}).Debug("iterate device")

		m <- prometheus.MustNewConstMetric(
			prometheus.NewDesc(
				wireguardMetaFQName, wireguardMetaHelpText, nil, l,
			),
			prometheus.GaugeValue,
			1,
		)
	}
}

func (collector *Metrics) collectDevicePeerMetrics(devices []*wgtypes.Device, m chan<- prometheus.Metric) {
	collector.logCtx.Debug("collecting wg peer(s) metrics")

	collector.logCtx.WithFields(logrus.Fields{
		"count": len(devices),
	}).Debug("enumerated wireguard devices")

	for _, device := range devices {
		logCtx := collector.logCtx.WithFields(logrus.Fields{
			"key":  device.PublicKey,
			"name": device.Name,
		})
		for _, peer := range device.Peers {
			logCtx.WithFields(logrus.Fields{
				"peer_key":      peer.PublicKey,
				"peer_endpoint": peer.Endpoint,
			}).Debug("collect peer metrics")

			labels := collector.defaultLabelValues(device.PublicKey.String(), peerServiceLabelValues(&peer))

			hs := float64(peer.LastHandshakeTime.Unix())

			collector.logCtx.WithFields(logrus.Fields{
				"rx_bytes_total": peer.ReceiveBytes,
				"tx_bytes_total": peer.TransmitBytes,
				"handshake_ts":   hs,
			}).Debug("collected peer metrics")

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardBytesRcvdFQName, wireguardBytesRcvdHelpText, nil, labels,
				),
				prometheus.CounterValue,
				float64(peer.ReceiveBytes),
			)

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardBytesSentFQName, wireguardBytesSentHelpText, nil, labels,
				),
				prometheus.CounterValue,
				float64(peer.TransmitBytes),
			)

			m <- prometheus.MustNewConstMetric(
				prometheus.NewDesc(
					wireguardLatestHandshakeIntervalFQName, wireguardLatestHandshakeIntervalHelpText, nil, labels,
				),
				prometheus.GaugeValue,
				hs,
			)
		}
	}

}

func (collector *Metrics) defaultLabelValues(publicKey string, extend prometheus.Labels) prometheus.Labels {
	l := prometheus.Labels{labelHostname: collector.hostname}
	if publicKey != "" {
		l[labelPublicKey] = publicKey
	}

	for k, v := range extend {
		l[k] = v
	}
	return l
}

func deviceMetaLabelValues(dev *wgtypes.Device) prometheus.Labels {
	return prometheus.Labels{
		labelInterfaceName: dev.Name,
		labelListenPort:    fmt.Sprintf("%d", dev.ListenPort),
	}
}

func peerServiceLabelValues(peer *wgtypes.Peer) prometheus.Labels {
	return prometheus.Labels{
		labelPeerKey:      peer.PublicKey.String(),
		labelPeerEndpoint: peer.Endpoint.String(),
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"container/heap"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/pd/server"
	"github.com/pingcap/pd/server/core"
	"github.com/unrolled/render"
)

// RegionInfo records detail region info for api usage.
type RegionInfo struct {
	ID          uint64              `json:"id"`
	StartKey    string              `json:"start_key"`
	EndKey      string              `json:"end_key"`
	RegionEpoch *metapb.RegionEpoch `json:"epoch,omitempty"`
	Peers       []*metapb.Peer      `json:"peers,omitempty"`

	Leader          *metapb.Peer      `json:"leader,omitempty"`
	DownPeers       []*pdpb.PeerStats `json:"down_peers,omitempty"`
	PendingPeers    []*metapb.Peer    `json:"pending_peers,omitempty"`
	WrittenBytes    uint64            `json:"written_bytes,omitempty"`
	ReadBytes       uint64            `json:"read_bytes,omitempty"`
	WrittenKeys     uint64            `json:"written_keys,omitempty"`
	ReadKeys        uint64            `json:"read_keys,omitempty"`
	ApproximateSize int64             `json:"approximate_size,omitempty"`
	ApproximateKeys int64             `json:"approximate_keys,omitempty"`
}

// NewRegionInfo create a new api RegionInfo.
func NewRegionInfo(r *core.RegionInfo) *RegionInfo {
	return InitRegion(r, &RegionInfo{})
}

// InitRegion init a new api RegionInfo from the core.RegionInfo.
func InitRegion(r *core.RegionInfo, s *RegionInfo) *RegionInfo {
	if r == nil {
		return nil
	}

	s.ID = r.GetID()
	s.StartKey = core.HexRegionKeyStr(r.GetStartKey())
	s.EndKey = core.HexRegionKeyStr(r.GetEndKey())
	s.RegionEpoch = r.GetRegionEpoch()
	s.Peers = r.GetPeers()
	s.Leader = r.GetLeader()
	s.DownPeers = r.GetDownPeers()
	s.PendingPeers = r.GetPendingPeers()
	s.WrittenBytes = r.GetBytesWritten()
	s.WrittenKeys = r.GetKeysWritten()
	s.ReadBytes = r.GetBytesRead()
	s.ReadKeys = r.GetKeysRead()
	s.ApproximateSize = r.GetApproximateSize()
	s.ApproximateKeys = r.GetApproximateKeys()

	return s
}

// RegionsInfo contains some regions with the detailed region info.
type RegionsInfo struct {
	Count   int           `json:"count"`
	Regions []*RegionInfo `json:"regions"`
}

type regionHandler struct {
	svr *server.Server
	rd  *render.Render
}

func newRegionHandler(svr *server.Server, rd *render.Render) *regionHandler {
	return &regionHandler{
		svr: svr,
		rd:  rd,
	}
}

func (h *regionHandler) GetRegionByID(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	regionIDStr := vars["id"]
	regionID, err := strconv.ParseUint(regionIDStr, 10, 64)
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}

	regionInfo := rc.GetRegion(regionID)
	h.rd.JSON(w, http.StatusOK, NewRegionInfo(regionInfo))
}

func (h *regionHandler) GetRegionByKey(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	vars := mux.Vars(r)
	key := vars["key"]
	regionInfo := rc.GetRegionInfoByKey([]byte(key))
	h.rd.JSON(w, http.StatusOK, NewRegionInfo(regionInfo))
}

type regionsHandler struct {
	svr *server.Server
	rd  *render.Render
}

func newRegionsHandler(svr *server.Server, rd *render.Render) *regionsHandler {
	return &regionsHandler{
		svr: svr,
		rd:  rd,
	}
}

func convertToAPIRegions(regions []*core.RegionInfo) *RegionsInfo {
	regionInfos := make([]RegionInfo, len(regions))
	regionInfosRefs := make([]*RegionInfo, len(regions))

	for i := 0; i < len(regions); i++ {
		regionInfosRefs[i] = &regionInfos[i]
	}
	for i, r := range regions {
		regionInfosRefs[i] = InitRegion(r, regionInfosRefs[i])
	}
	return &RegionsInfo{
		Count:   len(regions),
		Regions: regionInfosRefs,
	}
}

func (h *regionsHandler) GetAll(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	regions := rc.GetRegions()
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) ScanRegions(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	startKey := r.URL.Query().Get("key")

	limit := defaultRegionLimit
	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		var err error
		limit, err = strconv.Atoi(limitStr)
		if err != nil {
			h.rd.JSON(w, http.StatusBadRequest, err.Error())
			return
		}
	}
	if limit > maxRegionLimit {
		limit = maxRegionLimit
	}
	regions := rc.ScanRegions([]byte(startKey), nil, limit)
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetRegionCount(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	count := rc.GetRegionCount()
	h.rd.JSON(w, http.StatusOK, &RegionsInfo{Count: count})
}

func (h *regionsHandler) GetStoreRegions(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	id, err := strconv.Atoi(vars["id"])
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}
	regions := rc.GetStoreRegions(uint64(id))
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetMissPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetMissPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetExtraPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetExtraPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetPendingPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetPendingPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetDownPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetDownPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetOfflinePeer(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetOfflinePeer()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetEmptyRegion(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetEmptyRegion()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetRegionSiblings(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	id, err := strconv.Atoi(vars["id"])
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}
	region := rc.GetRegion(uint64(id))
	if region == nil {
		h.rd.JSON(w, http.StatusNotFound, server.ErrRegionNotFound(uint64(id)).Error())
		return
	}

	left, right := rc.GetAdjacentRegions(region)
	regionsInfo := convertToAPIRegions([]*core.RegionInfo{left, right})
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

const (
	defaultRegionLimit = 16
	maxRegionLimit     = 10240
)

func (h *regionsHandler) GetTopWriteFlow(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.GetBytesWritten() < b.GetBytesWritten() })
}

func (h *regionsHandler) GetTopReadFlow(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.GetBytesRead() < b.GetBytesRead() })
}

func (h *regionsHandler) GetTopConfVer(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetMeta().GetRegionEpoch().GetConfVer() < b.GetMeta().GetRegionEpoch().GetConfVer()
	})
}

func (h *regionsHandler) GetTopVersion(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetMeta().GetRegionEpoch().GetVersion() < b.GetMeta().GetRegionEpoch().GetVersion()
	})
}

func (h *regionsHandler) GetTopSize(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetApproximateSize() < b.GetApproximateSize()
	})
}

func (h *regionsHandler) GetTopNRegions(w http.ResponseWriter, r *http.Request, less func(a, b *core.RegionInfo) bool) {
	rc := getCluster(r.Context())
	limit := defaultRegionLimit
	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		var err error
		limit, err = strconv.Atoi(limitStr)
		if err != nil {
			h.rd.JSON(w, http.StatusBadRequest, err.Error())
			return
		}
	}
	if limit > maxRegionLimit {
		limit = maxRegionLimit
	}
	regions := TopNRegions(rc.GetRegions(), less, limit)
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

// RegionHeap implements heap.Interface, used for selecting top n regions.
type RegionHeap struct {
	regions []*core.RegionInfo
	less    func(a, b *core.RegionInfo) bool
}

func (h *RegionHeap) Len() int           { return len(h.regions) }
func (h *RegionHeap) Less(i, j int) bool { return h.less(h.regions[i], h.regions[j]) }
func (h *RegionHeap) Swap(i, j int)      { h.regions[i], h.regions[j] = h.regions[j], h.regions[i] }

// Push pushes an element x onto the heap.
func (h *RegionHeap) Push(x interface{}) {
	h.regions = append(h.regions, x.(*core.RegionInfo))
}

// Pop removes the minimum element (according to Less) from the heap and returns
// it.
func (h *RegionHeap) Pop() interface{} {
	pos := len(h.regions) - 1
	x := h.regions[pos]
	h.regions = h.regions[:pos]
	return x
}

// Min returns the minimum region from the heap.
func (h *RegionHeap) Min() *core.RegionInfo {
	if h.Len() == 0 {
		return nil
	}
	return h.regions[0]
}

// TopNRegions returns top n regions according to the given rule.
func TopNRegions(regions []*core.RegionInfo, less func(a, b *core.RegionInfo) bool, n int) []*core.RegionInfo {
	if n <= 0 {
		return nil
	}

	hp := &RegionHeap{
		regions: make([]*core.RegionInfo, 0, n),
		less:    less,
	}
	for _, r := range regions {
		if hp.Len() < n {
			heap.Push(hp, r)
			continue
		}
		if less(hp.Min(), r) {
			heap.Pop(hp)
			heap.Push(hp, r)
		}
	}

	res := make([]*core.RegionInfo, hp.Len())
	for i := hp.Len() - 1; i >= 0; i-- {
		res[i] = heap.Pop(hp).(*core.RegionInfo)
	}
	return res
}
<commit_msg>pdctl: fix some parameters of region not displaying (#2020)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"container/heap"
	"net/http"
	"strconv"

	"github.com/gorilla/mux"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/pd/server"
	"github.com/pingcap/pd/server/core"
	"github.com/unrolled/render"
)

// RegionInfo records detail region info for api usage.
type RegionInfo struct {
	ID          uint64              `json:"id"`
	StartKey    string              `json:"start_key"`
	EndKey      string              `json:"end_key"`
	RegionEpoch *metapb.RegionEpoch `json:"epoch,omitempty"`
	Peers       []*metapb.Peer      `json:"peers,omitempty"`

	Leader          *metapb.Peer      `json:"leader,omitempty"`
	DownPeers       []*pdpb.PeerStats `json:"down_peers,omitempty"`
	PendingPeers    []*metapb.Peer    `json:"pending_peers,omitempty"`
	WrittenBytes    uint64            `json:"written_bytes"`
	ReadBytes       uint64            `json:"read_bytes"`
	WrittenKeys     uint64            `json:"written_keys"`
	ReadKeys        uint64            `json:"read_keys"`
	ApproximateSize int64             `json:"approximate_size"`
	ApproximateKeys int64             `json:"approximate_keys"`
}

// NewRegionInfo create a new api RegionInfo.
func NewRegionInfo(r *core.RegionInfo) *RegionInfo {
	return InitRegion(r, &RegionInfo{})
}

// InitRegion init a new api RegionInfo from the core.RegionInfo.
func InitRegion(r *core.RegionInfo, s *RegionInfo) *RegionInfo {
	if r == nil {
		return nil
	}

	s.ID = r.GetID()
	s.StartKey = core.HexRegionKeyStr(r.GetStartKey())
	s.EndKey = core.HexRegionKeyStr(r.GetEndKey())
	s.RegionEpoch = r.GetRegionEpoch()
	s.Peers = r.GetPeers()
	s.Leader = r.GetLeader()
	s.DownPeers = r.GetDownPeers()
	s.PendingPeers = r.GetPendingPeers()
	s.WrittenBytes = r.GetBytesWritten()
	s.WrittenKeys = r.GetKeysWritten()
	s.ReadBytes = r.GetBytesRead()
	s.ReadKeys = r.GetKeysRead()
	s.ApproximateSize = r.GetApproximateSize()
	s.ApproximateKeys = r.GetApproximateKeys()

	return s
}

// RegionsInfo contains some regions with the detailed region info.
type RegionsInfo struct {
	Count   int           `json:"count"`
	Regions []*RegionInfo `json:"regions"`
}

type regionHandler struct {
	svr *server.Server
	rd  *render.Render
}

func newRegionHandler(svr *server.Server, rd *render.Render) *regionHandler {
	return &regionHandler{
		svr: svr,
		rd:  rd,
	}
}

func (h *regionHandler) GetRegionByID(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	regionIDStr := vars["id"]
	regionID, err := strconv.ParseUint(regionIDStr, 10, 64)
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}

	regionInfo := rc.GetRegion(regionID)
	h.rd.JSON(w, http.StatusOK, NewRegionInfo(regionInfo))
}

func (h *regionHandler) GetRegionByKey(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	vars := mux.Vars(r)
	key := vars["key"]
	regionInfo := rc.GetRegionInfoByKey([]byte(key))
	h.rd.JSON(w, http.StatusOK, NewRegionInfo(regionInfo))
}

type regionsHandler struct {
	svr *server.Server
	rd  *render.Render
}

func newRegionsHandler(svr *server.Server, rd *render.Render) *regionsHandler {
	return &regionsHandler{
		svr: svr,
		rd:  rd,
	}
}

func convertToAPIRegions(regions []*core.RegionInfo) *RegionsInfo {
	regionInfos := make([]RegionInfo, len(regions))
	regionInfosRefs := make([]*RegionInfo, len(regions))

	for i := 0; i < len(regions); i++ {
		regionInfosRefs[i] = &regionInfos[i]
	}
	for i, r := range regions {
		regionInfosRefs[i] = InitRegion(r, regionInfosRefs[i])
	}
	return &RegionsInfo{
		Count:   len(regions),
		Regions: regionInfosRefs,
	}
}

func (h *regionsHandler) GetAll(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	regions := rc.GetRegions()
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) ScanRegions(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	startKey := r.URL.Query().Get("key")

	limit := defaultRegionLimit
	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		var err error
		limit, err = strconv.Atoi(limitStr)
		if err != nil {
			h.rd.JSON(w, http.StatusBadRequest, err.Error())
			return
		}
	}
	if limit > maxRegionLimit {
		limit = maxRegionLimit
	}
	regions := rc.ScanRegions([]byte(startKey), nil, limit)
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetRegionCount(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())
	count := rc.GetRegionCount()
	h.rd.JSON(w, http.StatusOK, &RegionsInfo{Count: count})
}

func (h *regionsHandler) GetStoreRegions(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	id, err := strconv.Atoi(vars["id"])
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}
	regions := rc.GetStoreRegions(uint64(id))
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetMissPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetMissPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetExtraPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetExtraPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetPendingPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetPendingPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetDownPeerRegions(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetDownPeerRegions()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetOfflinePeer(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetOfflinePeer()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetEmptyRegion(w http.ResponseWriter, r *http.Request) {
	handler := h.svr.GetHandler()
	regions, err := handler.GetEmptyRegion()
	if err != nil {
		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
		return
	}
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

func (h *regionsHandler) GetRegionSiblings(w http.ResponseWriter, r *http.Request) {
	rc := getCluster(r.Context())

	vars := mux.Vars(r)
	id, err := strconv.Atoi(vars["id"])
	if err != nil {
		h.rd.JSON(w, http.StatusBadRequest, err.Error())
		return
	}
	region := rc.GetRegion(uint64(id))
	if region == nil {
		h.rd.JSON(w, http.StatusNotFound, server.ErrRegionNotFound(uint64(id)).Error())
		return
	}

	left, right := rc.GetAdjacentRegions(region)
	regionsInfo := convertToAPIRegions([]*core.RegionInfo{left, right})
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

const (
	defaultRegionLimit = 16
	maxRegionLimit     = 10240
)

func (h *regionsHandler) GetTopWriteFlow(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.GetBytesWritten() < b.GetBytesWritten() })
}

func (h *regionsHandler) GetTopReadFlow(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.GetBytesRead() < b.GetBytesRead() })
}

func (h *regionsHandler) GetTopConfVer(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetMeta().GetRegionEpoch().GetConfVer() < b.GetMeta().GetRegionEpoch().GetConfVer()
	})
}

func (h *regionsHandler) GetTopVersion(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetMeta().GetRegionEpoch().GetVersion() < b.GetMeta().GetRegionEpoch().GetVersion()
	})
}

func (h *regionsHandler) GetTopSize(w http.ResponseWriter, r *http.Request) {
	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool {
		return a.GetApproximateSize() < b.GetApproximateSize()
	})
}

func (h *regionsHandler) GetTopNRegions(w http.ResponseWriter, r *http.Request, less func(a, b *core.RegionInfo) bool) {
	rc := getCluster(r.Context())
	limit := defaultRegionLimit
	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		var err error
		limit, err = strconv.Atoi(limitStr)
		if err != nil {
			h.rd.JSON(w, http.StatusBadRequest, err.Error())
			return
		}
	}
	if limit > maxRegionLimit {
		limit = maxRegionLimit
	}
	regions := TopNRegions(rc.GetRegions(), less, limit)
	regionsInfo := convertToAPIRegions(regions)
	h.rd.JSON(w, http.StatusOK, regionsInfo)
}

// RegionHeap implements heap.Interface, used for selecting top n regions.
type RegionHeap struct {
	regions []*core.RegionInfo
	less    func(a, b *core.RegionInfo) bool
}

func (h *RegionHeap) Len() int           { return len(h.regions) }
func (h *RegionHeap) Less(i, j int) bool { return h.less(h.regions[i], h.regions[j]) }
func (h *RegionHeap) Swap(i, j int)      { h.regions[i], h.regions[j] = h.regions[j], h.regions[i] }

// Push pushes an element x onto the heap.
func (h *RegionHeap) Push(x interface{}) {
	h.regions = append(h.regions, x.(*core.RegionInfo))
}

// Pop removes the minimum element (according to Less) from the heap and returns
// it.
func (h *RegionHeap) Pop() interface{} {
	pos := len(h.regions) - 1
	x := h.regions[pos]
	h.regions = h.regions[:pos]
	return x
}

// Min returns the minimum region from the heap.
func (h *RegionHeap) Min() *core.RegionInfo {
	if h.Len() == 0 {
		return nil
	}
	return h.regions[0]
}

// TopNRegions returns top n regions according to the given rule.
func TopNRegions(regions []*core.RegionInfo, less func(a, b *core.RegionInfo) bool, n int) []*core.RegionInfo {
	if n <= 0 {
		return nil
	}

	hp := &RegionHeap{
		regions: make([]*core.RegionInfo, 0, n),
		less:    less,
	}
	for _, r := range regions {
		if hp.Len() < n {
			heap.Push(hp, r)
			continue
		}
		if less(hp.Min(), r) {
			heap.Pop(hp)
			heap.Push(hp, r)
		}
	}

	res := make([]*core.RegionInfo, hp.Len())
	for i := hp.Len() - 1; i >= 0; i-- {
		res[i] = heap.Pop(hp).(*core.RegionInfo)
	}
	return res
}
<|endoftext|>