<commit_before>package weed_server

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/chrislusf/raft"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"google.golang.org/grpc/peer"
)

func (ms *MasterServer) SendHeartbeat(stream master_pb.Seaweed_SendHeartbeatServer) error {
	var dn *topology.DataNode
	t := ms.Topo

	defer func() {
		if dn != nil {

			glog.V(0).Infof("unregister disconnected volume server %s:%d", dn.Ip, dn.Port)
			t.UnRegisterDataNode(dn)

			message := &master_pb.VolumeLocation{
				Url:       dn.Url(),
				PublicUrl: dn.PublicUrl,
			}
			for _, v := range dn.GetVolumes() {
				message.DeletedVids = append(message.DeletedVids, uint32(v.Id))
			}
			for _, s := range dn.GetEcShards() {
				message.DeletedVids = append(message.DeletedVids, uint32(s.VolumeId))
			}

			if len(message.DeletedVids) > 0 {
				ms.clientChansLock.RLock()
				for _, ch := range ms.clientChans {
					ch <- message
				}
				ms.clientChansLock.RUnlock()
			}

		}
	}()

	for {
		heartbeat, err := stream.Recv()
		if err != nil {
			return err
		}

		t.Sequence.SetMax(heartbeat.MaxFileKey)

		if dn == nil {
			if heartbeat.Ip == "" {
				if pr, ok := peer.FromContext(stream.Context()); ok {
					if pr.Addr != net.Addr(nil) {
						heartbeat.Ip = pr.Addr.String()[0:strings.LastIndex(pr.Addr.String(), ":")]
						glog.V(0).Infof("remote IP address is detected as %v", heartbeat.Ip)
					}
				}
			}
			dcName, rackName := t.Configuration.Locate(heartbeat.Ip, heartbeat.DataCenter, heartbeat.Rack)
			dc := t.GetOrCreateDataCenter(dcName)
			rack := dc.GetOrCreateRack(rackName)
			dn = rack.GetOrCreateDataNode(heartbeat.Ip,
				int(heartbeat.Port), heartbeat.PublicUrl,
				int64(heartbeat.MaxVolumeCount))
			glog.V(0).Infof("added volume server %v:%d", heartbeat.GetIp(), heartbeat.GetPort())
			if err := stream.Send(&master_pb.HeartbeatResponse{
				VolumeSizeLimit: uint64(ms.option.VolumeSizeLimitMB) * 1024 * 1024,
			}); err != nil {
				return err
			}
		}

		glog.V(4).Infof("master received heartbeat %s", heartbeat.String())
		message := &master_pb.VolumeLocation{
			Url:       dn.Url(),
			PublicUrl: dn.PublicUrl,
		}
		if len(heartbeat.NewVolumes) > 0 || len(heartbeat.DeletedVolumes) > 0 {
			// process delta volume ids if exists for fast volume id updates
			for _, volInfo := range heartbeat.NewVolumes {
				message.NewVids = append(message.NewVids, volInfo.Id)
			}
			for _, volInfo := range heartbeat.DeletedVolumes {
				message.DeletedVids = append(message.DeletedVids, volInfo.Id)
			}
			// update master internal volume layouts
			t.IncrementalSyncDataNodeRegistration(heartbeat.NewVolumes, heartbeat.DeletedVolumes, dn)
		}

		if len(heartbeat.Volumes) > 0 || heartbeat.HasNoVolumes {
			// process heartbeat.Volumes
			newVolumes, deletedVolumes := t.SyncDataNodeRegistration(heartbeat.Volumes, dn)

			for _, v := range newVolumes {
				glog.V(0).Infof("master see new volume %d from %s", uint32(v.Id), dn.Url())
				message.NewVids = append(message.NewVids, uint32(v.Id))
			}
			for _, v := range deletedVolumes {
				glog.V(0).Infof("master see deleted volume %d from %s", uint32(v.Id), dn.Url())
				message.DeletedVids = append(message.DeletedVids, uint32(v.Id))
			}
		}

		if len(heartbeat.NewEcShards) > 0 || len(heartbeat.DeletedEcShards) > 0 {

			// update master internal volume layouts
			t.IncrementalSyncDataNodeEcShards(heartbeat.NewEcShards, heartbeat.DeletedEcShards, dn)

			for _, s := range heartbeat.NewEcShards {
				message.NewVids = append(message.NewVids, s.Id)
			}
			for _, s := range heartbeat.DeletedEcShards {
				if dn.HasVolumesById(needle.VolumeId(s.Id)) {
					continue
				}
				message.DeletedVids = append(message.DeletedVids, s.Id)
			}

		}

		if len(heartbeat.EcShards) > 0 || heartbeat.HasNoEcShards {
			glog.V(1).Infof("master recieved ec shards from %s: %+v", dn.Url(), heartbeat.EcShards)
			newShards, deletedShards := t.SyncDataNodeEcShards(heartbeat.EcShards, dn)

			// broadcast the ec vid changes to master clients
			for _, s := range newShards {
				message.NewVids = append(message.NewVids, uint32(s.VolumeId))
			}
			for _, s := range deletedShards {
				if dn.HasVolumesById(s.VolumeId) {
					continue
				}
				message.DeletedVids = append(message.DeletedVids, uint32(s.VolumeId))
			}

		}

		if len(message.NewVids) > 0 || len(message.DeletedVids) > 0 {
			ms.clientChansLock.RLock()
			for host, ch := range ms.clientChans {
				glog.V(0).Infof("master send to %s: %s", host, message.String())
				ch <- message
			}
			ms.clientChansLock.RUnlock()
		}

		// tell the volume servers about the leader
		newLeader, err := t.Leader()
		if err != nil {
			return err
		}
		if err := stream.Send(&master_pb.HeartbeatResponse{
			Leader:                 newLeader,
			MetricsAddress:         ms.option.MetricsAddress,
			MetricsIntervalSeconds: uint32(ms.option.MetricsIntervalSec),
		}); err != nil {
			return err
		}
	}
}

// KeepConnected keep a stream gRPC call to the master. Used by clients to know the master is up.
// And clients gets the up-to-date list of volume locations
func (ms *MasterServer) KeepConnected(stream master_pb.Seaweed_KeepConnectedServer) error {

	req, err := stream.Recv()
	if err != nil {
		return err
	}

	if !ms.Topo.IsLeader() {
		return ms.informNewLeader(stream)
	}

	// remember client address
	ctx := stream.Context()
	// fmt.Printf("FromContext %+v\n", ctx)
	pr, ok := peer.FromContext(ctx)
	if !ok {
		glog.Error("failed to get peer from ctx")
		return fmt.Errorf("failed to get peer from ctx")
	}
	if pr.Addr == net.Addr(nil) {
		glog.Error("failed to get peer address")
		return fmt.Errorf("failed to get peer address")
	}

	clientName := req.Name + pr.Addr.String()
	glog.V(0).Infof("+ client %v", clientName)

	messageChan := make(chan *master_pb.VolumeLocation)
	stopChan := make(chan bool)

	ms.clientChansLock.Lock()
	ms.clientChans[clientName] = messageChan
	ms.clientChansLock.Unlock()

	defer func() {
		glog.V(0).Infof("- client %v", clientName)
		ms.clientChansLock.Lock()
		delete(ms.clientChans, clientName)
		ms.clientChansLock.Unlock()
	}()

	for _, message := range ms.Topo.ToVolumeLocations() {
		if err := stream.Send(message); err != nil {
			return err
		}
	}

	go func() {
		for {
			_, err := stream.Recv()
			if err != nil {
				glog.V(2).Infof("- client %v: %v", clientName, err)
				stopChan <- true
				break
			}
		}
	}()

	ticker := time.NewTicker(5 * time.Second)
	for {
		select {
		case message := <-messageChan:
			if err := stream.Send(message); err != nil {
				glog.V(0).Infof("=> client %v: %+v", clientName, message)
				return err
			}
		case <-ticker.C:
			if !ms.Topo.IsLeader() {
				return ms.informNewLeader(stream)
			}
		case <-stopChan:
			return nil
		}
	}

	return nil
}

func (ms *MasterServer) informNewLeader(stream master_pb.Seaweed_KeepConnectedServer) error {
	leader, err := ms.Topo.Leader()
	if err != nil {
		glog.Error("topo leader: %v", err)
		return raft.NotLeaderError
	}
	if err := stream.Send(&master_pb.VolumeLocation{
		Leader: leader,
	}); err != nil {
		return err
	}
	return nil
}
<commit_msg>fix test error<commit_after>package weed_server

import (
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/chrislusf/raft"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"google.golang.org/grpc/peer"
)

func (ms *MasterServer) SendHeartbeat(stream master_pb.Seaweed_SendHeartbeatServer) error {
	var dn *topology.DataNode
	t := ms.Topo

	defer func() {
		if dn != nil {

			glog.V(0).Infof("unregister disconnected volume server %s:%d", dn.Ip, dn.Port)
			t.UnRegisterDataNode(dn)

			message := &master_pb.VolumeLocation{
				Url:       dn.Url(),
				PublicUrl: dn.PublicUrl,
			}
			for _, v := range dn.GetVolumes() {
				message.DeletedVids = append(message.DeletedVids, uint32(v.Id))
			}
			for _, s := range dn.GetEcShards() {
				message.DeletedVids = append(message.DeletedVids, uint32(s.VolumeId))
			}

			if len(message.DeletedVids) > 0 {
				ms.clientChansLock.RLock()
				for _, ch := range ms.clientChans {
					ch <- message
				}
				ms.clientChansLock.RUnlock()
			}

		}
	}()

	for {
		heartbeat, err := stream.Recv()
		if err != nil {
			return err
		}

		t.Sequence.SetMax(heartbeat.MaxFileKey)

		if dn == nil {
			if heartbeat.Ip == "" {
				if pr, ok := peer.FromContext(stream.Context()); ok {
					if pr.Addr != net.Addr(nil) {
						heartbeat.Ip = pr.Addr.String()[0:strings.LastIndex(pr.Addr.String(), ":")]
						glog.V(0).Infof("remote IP address is detected as %v", heartbeat.Ip)
					}
				}
			}
			dcName, rackName := t.Configuration.Locate(heartbeat.Ip, heartbeat.DataCenter, heartbeat.Rack)
			dc := t.GetOrCreateDataCenter(dcName)
			rack := dc.GetOrCreateRack(rackName)
			dn = rack.GetOrCreateDataNode(heartbeat.Ip,
				int(heartbeat.Port), heartbeat.PublicUrl,
				int64(heartbeat.MaxVolumeCount))
			glog.V(0).Infof("added volume server %v:%d", heartbeat.GetIp(), heartbeat.GetPort())
			if err := stream.Send(&master_pb.HeartbeatResponse{
				VolumeSizeLimit: uint64(ms.option.VolumeSizeLimitMB) * 1024 * 1024,
			}); err != nil {
				return err
			}
		}

		glog.V(4).Infof("master received heartbeat %s", heartbeat.String())
		message := &master_pb.VolumeLocation{
			Url:       dn.Url(),
			PublicUrl: dn.PublicUrl,
		}
		if len(heartbeat.NewVolumes) > 0 || len(heartbeat.DeletedVolumes) > 0 {
			// process delta volume ids if exists for fast volume id updates
			for _, volInfo := range heartbeat.NewVolumes {
				message.NewVids = append(message.NewVids, volInfo.Id)
			}
			for _, volInfo := range heartbeat.DeletedVolumes {
				message.DeletedVids = append(message.DeletedVids, volInfo.Id)
			}
			// update master internal volume layouts
			t.IncrementalSyncDataNodeRegistration(heartbeat.NewVolumes, heartbeat.DeletedVolumes, dn)
		}

		if len(heartbeat.Volumes) > 0 || heartbeat.HasNoVolumes {
			// process heartbeat.Volumes
			newVolumes, deletedVolumes := t.SyncDataNodeRegistration(heartbeat.Volumes, dn)

			for _, v := range newVolumes {
				glog.V(0).Infof("master see new volume %d from %s", uint32(v.Id), dn.Url())
				message.NewVids = append(message.NewVids, uint32(v.Id))
			}
			for _, v := range deletedVolumes {
				glog.V(0).Infof("master see deleted volume %d from %s", uint32(v.Id), dn.Url())
				message.DeletedVids = append(message.DeletedVids, uint32(v.Id))
			}
		}

		if len(heartbeat.NewEcShards) > 0 || len(heartbeat.DeletedEcShards) > 0 {

			// update master internal volume layouts
			t.IncrementalSyncDataNodeEcShards(heartbeat.NewEcShards, heartbeat.DeletedEcShards, dn)

			for _, s := range heartbeat.NewEcShards {
				message.NewVids = append(message.NewVids, s.Id)
			}
			for _, s := range heartbeat.DeletedEcShards {
				if dn.HasVolumesById(needle.VolumeId(s.Id)) {
					continue
				}
				message.DeletedVids = append(message.DeletedVids, s.Id)
			}

		}

		if len(heartbeat.EcShards) > 0 || heartbeat.HasNoEcShards {
			glog.V(1).Infof("master recieved ec shards from %s: %+v", dn.Url(), heartbeat.EcShards)
			newShards, deletedShards := t.SyncDataNodeEcShards(heartbeat.EcShards, dn)

			// broadcast the ec vid changes to master clients
			for _, s := range newShards {
				message.NewVids = append(message.NewVids, uint32(s.VolumeId))
			}
			for _, s := range deletedShards {
				if dn.HasVolumesById(s.VolumeId) {
					continue
				}
				message.DeletedVids = append(message.DeletedVids, uint32(s.VolumeId))
			}

		}

		if len(message.NewVids) > 0 || len(message.DeletedVids) > 0 {
			ms.clientChansLock.RLock()
			for host, ch := range ms.clientChans {
				glog.V(0).Infof("master send to %s: %s", host, message.String())
				ch <- message
			}
			ms.clientChansLock.RUnlock()
		}

		// tell the volume servers about the leader
		newLeader, err := t.Leader()
		if err != nil {
			return err
		}
		if err := stream.Send(&master_pb.HeartbeatResponse{
			Leader:                 newLeader,
			MetricsAddress:         ms.option.MetricsAddress,
			MetricsIntervalSeconds: uint32(ms.option.MetricsIntervalSec),
		}); err != nil {
			return err
		}
	}
}

// KeepConnected keep a stream gRPC call to the master. Used by clients to know the master is up.
// And clients gets the up-to-date list of volume locations
func (ms *MasterServer) KeepConnected(stream master_pb.Seaweed_KeepConnectedServer) error {

	req, err := stream.Recv()
	if err != nil {
		return err
	}

	if !ms.Topo.IsLeader() {
		return ms.informNewLeader(stream)
	}

	// remember client address
	ctx := stream.Context()
	// fmt.Printf("FromContext %+v\n", ctx)
	pr, ok := peer.FromContext(ctx)
	if !ok {
		glog.Error("failed to get peer from ctx")
		return fmt.Errorf("failed to get peer from ctx")
	}
	if pr.Addr == net.Addr(nil) {
		glog.Error("failed to get peer address")
		return fmt.Errorf("failed to get peer address")
	}

	clientName := req.Name + pr.Addr.String()
	glog.V(0).Infof("+ client %v", clientName)

	messageChan := make(chan *master_pb.VolumeLocation)
	stopChan := make(chan bool)

	ms.clientChansLock.Lock()
	ms.clientChans[clientName] = messageChan
	ms.clientChansLock.Unlock()

	defer func() {
		glog.V(0).Infof("- client %v", clientName)
		ms.clientChansLock.Lock()
		delete(ms.clientChans, clientName)
		ms.clientChansLock.Unlock()
	}()

	for _, message := range ms.Topo.ToVolumeLocations() {
		if err := stream.Send(message); err != nil {
			return err
		}
	}

	go func() {
		for {
			_, err := stream.Recv()
			if err != nil {
				glog.V(2).Infof("- client %v: %v", clientName, err)
				stopChan <- true
				break
			}
		}
	}()

	ticker := time.NewTicker(5 * time.Second)
	for {
		select {
		case message := <-messageChan:
			if err := stream.Send(message); err != nil {
				glog.V(0).Infof("=> client %v: %+v", clientName, message)
				return err
			}
		case <-ticker.C:
			if !ms.Topo.IsLeader() {
				return ms.informNewLeader(stream)
			}
		case <-stopChan:
			return nil
		}
	}

	return nil
}

func (ms *MasterServer) informNewLeader(stream master_pb.Seaweed_KeepConnectedServer) error {
	leader, err := ms.Topo.Leader()
	if err != nil {
		glog.Errorf("topo leader: %v", err)
		return raft.NotLeaderError
	}
	if err := stream.Send(&master_pb.VolumeLocation{
		Leader: leader,
	}); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package shell

import (
	"flag"
	"io"

	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

func init() {
	Commands = append(Commands, &commandVolumeMark{})
}

type commandVolumeMark struct {
}

func (c *commandVolumeMark) Name() string {
	return "volume.mark"
}

func (c *commandVolumeMark) Help() string {
	return `Mark volume readonly from one volume server

	volume.mark -node <volume server host:port> -volumeId <volume id> -readonly true
`
}

func (c *commandVolumeMark) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	volMarkCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeIdInt := volMarkCommand.Int("volumeId", 0, "the volume id")
	nodeStr := volMarkCommand.String("node", "", "the volume server <host>:<port>")
	writable := volMarkCommand.Bool("writable", true, "volume mark writable/readonly")
	if err = volMarkCommand.Parse(args); err != nil {
		return nil
	}

	sourceVolumeServer := *nodeStr

	volumeId := needle.VolumeId(*volumeIdInt)

	return markVolumeWritable(commandEnv.option.GrpcDialOption, volumeId, sourceVolumeServer, *writable)
}
<commit_msg>add readonly<commit_after>package shell

import (
	"flag"
	"fmt"
	"io"

	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

func init() {
	Commands = append(Commands, &commandVolumeMark{})
}

type commandVolumeMark struct {
}

func (c *commandVolumeMark) Name() string {
	return "volume.mark"
}

func (c *commandVolumeMark) Help() string {
	return `Mark volume writable or readonly from one volume server

	volume.mark -node <volume server host:port> -volumeId <volume id> -writable or -readonly
`
}

func (c *commandVolumeMark) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	volMarkCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeIdInt := volMarkCommand.Int("volumeId", 0, "the volume id")
	nodeStr := volMarkCommand.String("node", "", "the volume server <host>:<port>")
	writable := volMarkCommand.Bool("writable", false, "volume mark writable")
	readonly := volMarkCommand.Bool("readonly", false, "volume mark readonly")
	if err = volMarkCommand.Parse(args); err != nil {
		return nil
	}
	markWritable := false
	if (*writable && *readonly) || (!*writable && !*readonly) {
		return fmt.Errorf("use -readonly or -writable")
	} else if *writable {
		markWritable = true
	}

	sourceVolumeServer := *nodeStr

	volumeId := needle.VolumeId(*volumeIdInt)

	return markVolumeWritable(commandEnv.option.GrpcDialOption, volumeId, sourceVolumeServer, markWritable)
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"io"
	"time"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"google.golang.org/grpc"
)

func init() {
	Commands = append(Commands, &commandVolumeTier{})
}

type commandVolumeTier struct {
}

func (c *commandVolumeTier) Name() string {
	return "volume.tier"
}

func (c *commandVolumeTier) Help() string {
	return `move the dat file of a volume to a remote tier

	volume.tier [-collection=""] [-fullPercent=95] [-quietFor=1h]
	volume.tier [-collection=""] -volumeId=<volume_id> -dest=<storage_backend> [-keepLocalDatFile]

	e.g.:
	volume.tier -volumeId=7 -dest=s3
	volume.tier -volumeId=7 -dest=s3.default

	The <storage_backend> is defined in master.toml.
	For example, "s3.default" in [storage.backend.s3.default]

	This command will move the dat file of a volume to a remote tier.

	SeaweedFS enables scalable and fast local access to lots of files,
	and the cloud storage is slower by cost efficient. How to combine them together?

	Usually the data follows 80/20 rule: only 20% of data is frequently accessed.
	We can offload the old volumes to the cloud.

	With this, SeaweedFS can be both fast and scalable, and infinite storage space.
	Just add more local SeaweedFS volume servers to increase the throughput.

	The index file is still local, and the same O(1) disk read is applied to the remote file.

`
}

func (c *commandVolumeTier) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	tierCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := tierCommand.Int("volumeId", 0, "the volume id")
	collection := tierCommand.String("collection", "", "the collection name")
	fullPercentage := tierCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := tierCommand.Duration("quietFor", 24*time.Hour, "select volumes without no writes for this period")
	dest := tierCommand.String("dest", "", "the target tier name")
	keepLocalDatFile := tierCommand.Bool("keepLocalDatFile", false, "whether keep local dat file")
	if err = tierCommand.Parse(args); err != nil {
		return nil
	}

	ctx := context.Background()
	vid := needle.VolumeId(*volumeId)

	// volumeId is provided
	if vid != 0 {
		return doVolumeTier(ctx, commandEnv, writer, *collection, vid, *dest, *keepLocalDatFile)
	}

	// apply to all volumes in the collection
	// reusing collectVolumeIdsForEcEncode for now
	volumeIds, err := collectVolumeIdsForEcEncode(ctx, commandEnv, *collection, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("tiering volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doVolumeTier(ctx, commandEnv, writer, *collection, vid, *dest, *keepLocalDatFile); err != nil {
			return err
		}
	}

	return nil
}

func doVolumeTier(ctx context.Context, commandEnv *CommandEnv, writer io.Writer, collection string, vid needle.VolumeId, dest string, keepLocalDatFile bool) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// mark the volume as readonly
	/*
		err = markVolumeReadonly(ctx, commandEnv.option.GrpcDialOption, needle.VolumeId(vid), locations)
		if err != nil {
			return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
		}
	*/

	// copy the .dat file to remote tier
	err = copyDatToRemoteTier(ctx, commandEnv.option.GrpcDialOption, writer, needle.VolumeId(vid), collection, locations[0].Url, dest, keepLocalDatFile)
	if err != nil {
		return fmt.Errorf("copy dat file for volume %d on %s to %s: %v", vid, locations[0].Url, dest, err)
	}

	return nil
}

func copyDatToRemoteTier(ctx context.Context, grpcDialOption grpc.DialOption, writer io.Writer, volumeId needle.VolumeId, collection string, sourceVolumeServer string, dest string, keepLocalDatFile bool) error {

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		stream, copyErr := volumeServerClient.VolumeTierCopyDatToRemote(ctx, &volume_server_pb.VolumeTierCopyDatToRemoteRequest{
			VolumeId:               uint32(volumeId),
			Collection:             collection,
			DestinationBackendName: dest,
			KeepLocalDatFile:       keepLocalDatFile,
		})

		var lastProcessed int64
		for {
			resp, recvErr := stream.Recv()
			if recvErr != nil {
				if recvErr == io.EOF {
					break
				} else {
					return recvErr
				}
			}

			processingSpeed := float64(resp.Processed-lastProcessed) / 1024.0 / 1024.0

			fmt.Fprintf(writer, "copied %.2f%%, %d bytes, %.2fMB/s\n", resp.ProcessedPercentage, resp.Processed, processingSpeed)

			lastProcessed = resp.Processed
		}

		return copyErr
	})

	return err

}
<commit_msg>adjust command name<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"io"
	"time"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"google.golang.org/grpc"
)

func init() {
	Commands = append(Commands, &commandVolumeTier{})
}

type commandVolumeTier struct {
}

func (c *commandVolumeTier) Name() string {
	return "volume.tier.upload"
}

func (c *commandVolumeTier) Help() string {
	return `move the dat file of a volume to a remote tier

	volume.tier.upload [-collection=""] [-fullPercent=95] [-quietFor=1h]
	volume.tier.upload [-collection=""] -volumeId=<volume_id> -dest=<storage_backend> [-keepLocalDatFile]

	e.g.:
	volume.tier.upload -volumeId=7 -dest=s3
	volume.tier.upload -volumeId=7 -dest=s3.default

	The <storage_backend> is defined in master.toml.
	For example, "s3.default" in [storage.backend.s3.default]

	This command will move the dat file of a volume to a remote tier.

	SeaweedFS enables scalable and fast local access to lots of files,
	and the cloud storage is slower by cost efficient. How to combine them together?

	Usually the data follows 80/20 rule: only 20% of data is frequently accessed.
	We can offload the old volumes to the cloud.

	With this, SeaweedFS can be both fast and scalable, and infinite storage space.
	Just add more local SeaweedFS volume servers to increase the throughput.

	The index file is still local, and the same O(1) disk read is applied to the remote file.

`
}

func (c *commandVolumeTier) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	tierCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := tierCommand.Int("volumeId", 0, "the volume id")
	collection := tierCommand.String("collection", "", "the collection name")
	fullPercentage := tierCommand.Float64("fullPercent", 95, "the volume reaches the percentage of max volume size")
	quietPeriod := tierCommand.Duration("quietFor", 24*time.Hour, "select volumes without no writes for this period")
	dest := tierCommand.String("dest", "", "the target tier name")
	keepLocalDatFile := tierCommand.Bool("keepLocalDatFile", false, "whether keep local dat file")
	if err = tierCommand.Parse(args); err != nil {
		return nil
	}

	ctx := context.Background()
	vid := needle.VolumeId(*volumeId)

	// volumeId is provided
	if vid != 0 {
		return doVolumeTier(ctx, commandEnv, writer, *collection, vid, *dest, *keepLocalDatFile)
	}

	// apply to all volumes in the collection
	// reusing collectVolumeIdsForEcEncode for now
	volumeIds, err := collectVolumeIdsForEcEncode(ctx, commandEnv, *collection, *fullPercentage, *quietPeriod)
	if err != nil {
		return err
	}
	fmt.Printf("tiering volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doVolumeTier(ctx, commandEnv, writer, *collection, vid, *dest, *keepLocalDatFile); err != nil {
			return err
		}
	}

	return nil
}

func doVolumeTier(ctx context.Context, commandEnv *CommandEnv, writer io.Writer, collection string, vid needle.VolumeId, dest string, keepLocalDatFile bool) (err error) {
	// find volume location
	locations, found := commandEnv.MasterClient.GetLocations(uint32(vid))
	if !found {
		return fmt.Errorf("volume %d not found", vid)
	}

	// mark the volume as readonly
	/*
		err = markVolumeReadonly(ctx, commandEnv.option.GrpcDialOption, needle.VolumeId(vid), locations)
		if err != nil {
			return fmt.Errorf("mark volume %d as readonly on %s: %v", vid, locations[0].Url, err)
		}
	*/

	// copy the .dat file to remote tier
	err = copyDatToRemoteTier(ctx, commandEnv.option.GrpcDialOption, writer, needle.VolumeId(vid), collection, locations[0].Url, dest, keepLocalDatFile)
	if err != nil {
		return fmt.Errorf("copy dat file for volume %d on %s to %s: %v", vid, locations[0].Url, dest, err)
	}

	return nil
}

func copyDatToRemoteTier(ctx context.Context, grpcDialOption grpc.DialOption, writer io.Writer, volumeId needle.VolumeId, collection string, sourceVolumeServer string, dest string, keepLocalDatFile bool) error {

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		stream, copyErr := volumeServerClient.VolumeTierCopyDatToRemote(ctx, &volume_server_pb.VolumeTierCopyDatToRemoteRequest{
			VolumeId:               uint32(volumeId),
			Collection:             collection,
			DestinationBackendName: dest,
			KeepLocalDatFile:       keepLocalDatFile,
		})

		var lastProcessed int64
		for {
			resp, recvErr := stream.Recv()
			if recvErr != nil {
				if recvErr == io.EOF {
					break
				} else {
					return recvErr
				}
			}

			processingSpeed := float64(resp.Processed-lastProcessed) / 1024.0 / 1024.0

			fmt.Fprintf(writer, "copied %.2f%%, %d bytes, %.2fMB/s\n", resp.ProcessedPercentage, resp.Processed, processingSpeed)

			lastProcessed = resp.Processed
		}

		return copyErr
	})

	return err

}
<|endoftext|>


<commit_before>package app

import (
	"fmt"
	"github.com/callumj/weave/core"
	"github.com/callumj/weave/remote"
	"github.com/callumj/weave/tools"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

func performExtraction(args []string) {
	if len(args) < 3 {
		panicQuitf("Usage: %v ENCRYPTED_FILE KEY_FILE [OUT_DIRECTORY]\r\n", args[0])
	}

	target := args[1]
	keyfile := args[2]

	var eTag string

	matched, _ := regexp.MatchString("^https?:\\/\\/", target)
	if matched {
		if len(args) == 3 {
			panicQuitf("A OUT_DIRECTORY must be specified if using HTTP\r\n")
		} else {
			resp := remote.DownloadRemoteFile(target, args[3])
			if resp == nil {
				panicQuit()
			} else {
				target = resp.FilePath
				eTag = resp.ETag
			}
		}
	}

	var out string

	var success bool
	if len(args) >= 4 {
		out = strings.Join([]string{args[3], "tmp"}, ".")
		success = core.DecryptFile(target, out, keyfile)
	} else {
		out = strings.Replace(target, ".enc", "", 1)
		out = strings.Join([]string{out, "tmp"}, ".")
		if out == target {
			panicQuitf("Cannot determine the out file, please specify")
		}
		success = core.DecryptFile(target, out, keyfile)
	}

	var ensureDirectory = regexp.MustCompile(`(\.(tmp|tgz|tar|gz))+`)
	directory := ensureDirectory.ReplaceAllString(out, "")

	directory, _ = filepath.Abs(directory)

	if !tools.PathExists(directory) {
		os.Mkdir(directory, 0770)
	}

	preExtraction := core.FetchFile(out, "/pre_extraction.sh")
	if len(preExtraction) != 0 {
		// write the pre file
		preName := fmt.Sprintf("%v/pre_extraction.sh", directory)
		ioutil.WriteFile(preName, []byte(preExtraction), 0770)
		runCallback(preName, directory)
	}

	if success {
		success = core.ExtractArchive(out, directory)
	}

	tools.CleanUpIfNeeded(out)

	if success && len(eTag) != 0 {
		etagfile := fmt.Sprintf("%v/.weave.etag", directory)
		ioutil.WriteFile(etagfile, []byte(eTag), 0775)
	}

	if success {
		callbackPath := fmt.Sprintf("%v/post_extraction.sh", directory)
		runCallback(callbackPath, directory)
	}

	if !success {
		panicQuit()
	}
}

func runCallback(callbackPath, directory string) {
	if tools.PathExists(callbackPath) {
		bashLoc, _ := exec.LookPath("bash")
		if len(bashLoc) == 0 {
			bashLoc = "/bin/bash"
		}
		cmd := exec.Command(bashLoc, callbackPath)
		cmd.Dir = directory
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panicQuitf("Unable to open STDOUT")
		}
		stderr, err := cmd.StderrPipe()
		if err != nil {
			panicQuitf("Unable to open STDERR")
		}

		if err := cmd.Start(); err != nil {
			panicQuitf("Unable to start %v (%v)\r\n", callbackPath, err)
		}

		outLogPath := fmt.Sprintf("%v.stdout.log", callbackPath)
		stdoutLogFile, err := os.Create(outLogPath)
		if err != nil {
			panicQuitf("Unable to open file for STDOUT logging %v\r\n", outLogPath)
		}
		defer stdoutLogFile.Close()

		errLogPath := fmt.Sprintf("%v.stderr.log", callbackPath)
		stderrLogFile, err := os.Create(errLogPath)
		if err != nil {
			panicQuitf("Unable to open file for STDERR logging %v\r\n", errLogPath)
		}
		defer stderrLogFile.Close()

		go io.Copy(stdoutLogFile, stdout)
		go io.Copy(stderrLogFile, stderr)

		if err := cmd.Wait(); err != nil {
			panicQuitf("Failed finalise command (%v)\r\n", err)
		}
	}
}
<commit_msg>Arguments have shifted down 1<commit_after>package app

import (
	"fmt"
	"github.com/callumj/weave/core"
	"github.com/callumj/weave/remote"
	"github.com/callumj/weave/tools"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

func performExtraction(args []string) {
	if len(args) < 2 {
		panicQuitf("Usage: %v ENCRYPTED_FILE KEY_FILE [OUT_DIRECTORY]\r\n", args[0])
	}

	target := args[0]
	keyfile := args[1]

	var eTag string

	matched, _ := regexp.MatchString("^https?:\\/\\/", target)
	if matched {
		if len(args) == 2 {
			panicQuitf("A OUT_DIRECTORY must be specified if using HTTP\r\n")
		} else {
			resp := remote.DownloadRemoteFile(target, args[2])
			if resp == nil {
				panicQuit()
			} else {
				target = resp.FilePath
				eTag = resp.ETag
			}
		}
	}

	var out string

	var success bool
	if len(args) >= 3 {
		out = strings.Join([]string{args[2], "tmp"}, ".")
		success = core.DecryptFile(target, out, keyfile)
	} else {
		out = strings.Replace(target, ".enc", "", 1)
		out = strings.Join([]string{out, "tmp"}, ".")
		if out == target {
			panicQuitf("Cannot determine the out file, please specify")
		}
		success = core.DecryptFile(target, out, keyfile)
	}

	var ensureDirectory = regexp.MustCompile(`(\.(tmp|tgz|tar|gz))+`)
	directory := ensureDirectory.ReplaceAllString(out, "")

	directory, _ = filepath.Abs(directory)

	if !tools.PathExists(directory) {
		os.Mkdir(directory, 0770)
	}

	preExtraction := core.FetchFile(out, "/pre_extraction.sh")
	if len(preExtraction) != 0 {
		// write the pre file
		preName := fmt.Sprintf("%v/pre_extraction.sh", directory)
		ioutil.WriteFile(preName, []byte(preExtraction), 0770)
		runCallback(preName, directory)
	}

	if success {
		success = core.ExtractArchive(out, directory)
	}

	tools.CleanUpIfNeeded(out)

	if success && len(eTag) != 0 {
		etagfile := fmt.Sprintf("%v/.weave.etag", directory)
		ioutil.WriteFile(etagfile, []byte(eTag), 0775)
	}

	if success {
		callbackPath := fmt.Sprintf("%v/post_extraction.sh", directory)
		runCallback(callbackPath, directory)
	}

	if !success {
		panicQuit()
	}
}

func runCallback(callbackPath, directory string) {
	if tools.PathExists(callbackPath) {
		bashLoc, _ := exec.LookPath("bash")
		if len(bashLoc) == 0 {
			bashLoc = "/bin/bash"
		}
		cmd := exec.Command(bashLoc, callbackPath)
		cmd.Dir = directory
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			panicQuitf("Unable to open STDOUT")
		}
		stderr, err := cmd.StderrPipe()
		if err != nil {
			panicQuitf("Unable to open STDERR")
		}

		if err := cmd.Start(); err != nil {
			panicQuitf("Unable to start %v (%v)\r\n", callbackPath, err)
		}

		outLogPath := fmt.Sprintf("%v.stdout.log", callbackPath)
		stdoutLogFile, err := os.Create(outLogPath)
		if err != nil {
			panicQuitf("Unable to open file for STDOUT logging %v\r\n", outLogPath)
		}
		defer stdoutLogFile.Close()

		errLogPath := fmt.Sprintf("%v.stderr.log", callbackPath)
		stderrLogFile, err := os.Create(errLogPath)
		if err != nil {
			panicQuitf("Unable to open file for STDERR logging %v\r\n", errLogPath)
		}
		defer stderrLogFile.Close()

		go io.Copy(stdoutLogFile, stdout)
		go io.Copy(stderrLogFile, stderr)

		if err := cmd.Wait(); err != nil {
			panicQuitf("Failed finalise command (%v)\r\n", err)
		}
	}
}
<|endoftext|>


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

// +build darwin freebsd netbsd openbsd

package fswatch

import (
	"os"
	"syscall"
)

type watcher struct {
	*kqueue
	C     chan *Event
	waker *os.File
}

func newWatcher() (*watcher, error) {
	kq, err := newkqueue()
	if err != nil {
		kq.close()
		return nil, err
	}
	pr, pw, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	if err := kq.add(uintptr(pr.Fd()), syscall.EVFILT_READ, 0, 0); err != nil {
		kq.close()
		pw.Close()
		return nil, err
	}
	return &watcher{
		kqueue: kq,
		C:      make(chan *Event, 1),
		waker:  pw,
	}, nil
}

func (w *watcher) add(path string) error {
	fd, err := syscall.Open(path, syscall.O_NONBLOCK|syscall.O_RDONLY, 0700)
	if err != nil {
		return os.NewSyscallError("open", err)
	}
	return w.kqueue.add(uintptr(fd), syscall.EVFILT_VNODE, 0, syscall.NOTE_DELETE|syscall.NOTE_EXTEND|syscall.NOTE_WRITE)
}

func (w *watcher) run() {
	// nothing
}

func (w *watcher) close() error {
	return w.waker.Close()
}

type event struct {
	mask uint32
}

func (e *event) IsCreate() bool { return e.mask&syscall.NOTE_EXTEND > 0 }
func (e *event) IsRemove() bool { return e.mask&syscall.NOTE_DELETE > 0 }
<commit_msg>added run method<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin freebsd netbsd openbsd

package fswatch

import (
	"os"
	"sync"
	"syscall"
)

type watcher struct {
	*kqueue
	C     chan *Event
	waker *os.File
	sync.RWMutex
	fds   map[string]int
	paths map[int]string
}

func newWatcher() (*watcher, error) {
	kq, err := newkqueue()
	if err != nil {
		kq.close()
		return nil, err
	}
	pr, pw, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	if err := kq.add(uintptr(pr.Fd()), syscall.EVFILT_READ, 0, 0); err != nil {
		kq.close()
		pw.Close()
		return nil, err
	}
	return &watcher{
		kqueue: kq,
		C:      make(chan *Event, 1),
		waker:  pw,
		fds:    make(map[string]int),
		paths:  make(map[int]string),
	}, nil
}

func (w *watcher) add(path string) error {
	fd, err := syscall.Open(path, syscall.O_NONBLOCK|syscall.O_RDONLY, 0700)
	if err != nil {
		return os.NewSyscallError("open", err)
	}
	w.Lock()
	defer w.Unlock()
	if _, exists := w.fds[path]; exists {
		defer syscall.Close(fd)
		return ErrWatchExists
	}
	w.fds[path] = fd
	w.paths[fd] = path
	return w.kqueue.add(uintptr(fd), syscall.EVFILT_VNODE, 0, syscall.NOTE_DELETE|syscall.NOTE_EXTEND|syscall.NOTE_WRITE)
}

func (w *watcher) run() {
	defer w.kqueue.close()
	var buf [18]syscall.Kevent_t
	for done := false; !done; {
		n, err := w.kqueue.wait(buf[:])
		if err != nil {
			println(err.Error())
			return
		}
		for _, e := range buf[:n] {
			if uintptr(e.Ident) == w.waker.Fd() {
				done = true
				continue
			}
			w.C <- &Event{
				Target: w.paths[int(e.Ident)],
				event:  event(e),
			}
		}
	}
	println("done")
}

func (w *watcher) close() error {
	return w.waker.Close()
}

type event syscall.Kevent_t

func (e *event) IsCreate() bool { return e.Flags&syscall.NOTE_EXTEND > 0 }
func (e *event) IsRemove() bool { return e.Flags&syscall.NOTE_DELETE > 0 }
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	storagev1 "google.golang.org/api/storage/v1"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// A request to create an object, accepted by Bucket.CreateObject.
type CreateObjectRequest struct {
	// Attributes with which the object should be created. The Name field must be
	// set; other zero-valued fields are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	Attrs storage.ObjectAttrs

	// A reader from which to obtain the contents of the object. Must be non-nil.
	Contents io.Reader
}

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Create or overwrite an object according to the supplied request. The new
	// object is guaranteed to exist immediately for the purposes of reading (and
	// eventually for listing) after this method returns a nil error. It is
	// guaranteed not to exist before req.Contents returns io.EOF.
	CreateObject(ctx context.Context, req *CreateObjectRequest) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func toRawAcls(in []storage.ACLRule) []*storagev1.ObjectAccessControl {
	out := make([]*storagev1.ObjectAccessControl, len(in))
	for i, rule := range in {
		out[i] = &storagev1.ObjectAccessControl{
			Entity: string(rule.Entity),
			Role:   string(rule.Role),
		}
	}

	return out
}

func fromRawAcls(in []*storagev1.ObjectAccessControl) []storage.ACLRule

func fromRawObject(
	bucketName string,
	in *storagev1.Object) (out *storage.Object, err error) {
	// Convert the easy fields.
	out = &storage.Object{
		Bucket:          bucketName,
		Name:            in.Name,
		ContentType:     in.ContentType,
		ContentLanguage: in.ContentLanguage,
		CacheControl:    in.CacheControl,
		ACL:             fromRawAcls(in.Acl),
		ContentEncoding: in.ContentEncoding,
		Size:            int64(in.Size),
		CRC32C:          in.CRC32C,
		MediaLink:       in.MediaLink,
		Metadata:        in.Metadata,
		Generation:      in.Generation,
		MetaGeneration:  in.MetaGeneration,
		StorageClass:    in.StorageClass,
		Deleted:         in.TimeDeleted,
		Updated:         in.Updated,
	}

	// Handle special cases.
	if in.Owner != nil {
		out.Owner = in.Owner.Entity
	}

	if out.MD5, err = base64.StdEncoding.DecodeString(in.Md5Hash); err != nil {
		err = fmt.Errorf("Decoding Md5Hash field: %v", err)
		return
	}

	return
}

func getRawService(authContext context.Context) *storagev1.Service

func (b *bucket) CreateObject(
	ctx context.Context,
	req *CreateObjectRequest) (o *storage.Object, err error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	objectsService := getRawService(authContext).Objects

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(req.Attrs.Name) {
		err = errors.New("Invalid object name: not valid UTF-8")
		return
	}

	// Set up an object struct based on the supplied attributes.
	inputObj := &storagev1.Object{
		Name:            req.Attrs.Name,
		Bucket:          b.Name(),
		ContentType:     req.Attrs.ContentType,
		ContentLanguage: req.Attrs.ContentLanguage,
		ContentEncoding: req.Attrs.ContentEncoding,
		CacheControl:    req.Attrs.CacheControl,
		Acl:             toRawAcls(req.Attrs.ACL),
		Metadata:        req.Attrs.Metadata,
	}

	// Configure a 'call' object.
	call := objectsService.Insert(b.Name(), inputObj)
	call.Media(req.Contents)
	call.Projection("full")

	// Execute the call.
	rawObject, err := call.Do()
	if err != nil {
		return
	}

	// Convert the returned object.
	o, err = fromRawObject(b.Name(), rawObject)
	if err != nil {
		return
	}

	return
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<commit_msg>Fixed more build errors.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcs

import (
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net/http"
	"unicode/utf8"

	"golang.org/x/net/context"
	storagev1 "google.golang.org/api/storage/v1"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
)

// A request to create an object, accepted by Bucket.CreateObject.
type CreateObjectRequest struct {
	// Attributes with which the object should be created. The Name field must be
	// set; other zero-valued fields are ignored.
	//
	// Object names must:
	//
	// *  be non-empty.
	// *  be no longer than 1024 bytes.
	// *  be valid UTF-8.
	// *  not contain the code point U+000A (line feed).
	// *  not contain the code point U+000D (carriage return).
	//
	// See here for authoritative documentation:
	//     https://cloud.google.com/storage/docs/bucket-naming#objectnames
	Attrs storage.ObjectAttrs

	// A reader from which to obtain the contents of the object. Must be non-nil.
	Contents io.Reader
}

// Bucket represents a GCS bucket, pre-bound with a bucket name and necessary
// authorization information.
//
// Each method that may block accepts a context object that is used for
// deadlines and cancellation. Users need not package authorization information
// into the context object (using cloud.WithContext or similar).
type Bucket interface {
	Name() string

	// List the objects in the bucket that meet the criteria defined by the
	// query, returning a result object that contains the results and potentially
	// a cursor for retrieving the next portion of the larger set of results.
	ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error)

	// Create a reader for the contents of the object with the given name. The
	// caller must arrange for the reader to be closed when it is no longer
	// needed.
	NewReader(ctx context.Context, objectName string) (io.ReadCloser, error)

	// Create or overwrite an object according to the supplied request. The new
	// object is guaranteed to exist immediately for the purposes of reading (and
	// eventually for listing) after this method returns a nil error. It is
	// guaranteed not to exist before req.Contents returns io.EOF.
	CreateObject(ctx context.Context, req *CreateObjectRequest) (*storage.Object, error)

	// Delete the object with the given name.
	DeleteObject(ctx context.Context, name string) error
}

type bucket struct {
	projID string
	client *http.Client
	name   string
}

func (b *bucket) Name() string {
	return b.name
}

func (b *bucket) ListObjects(ctx context.Context, query *storage.Query) (*storage.Objects, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.ListObjects(authContext, b.name, query)
}

func (b *bucket) NewReader(ctx context.Context, objectName string) (io.ReadCloser, error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.NewReader(authContext, b.name, objectName)
}

func toRawAcls(in []storage.ACLRule) []*storagev1.ObjectAccessControl {
	out := make([]*storagev1.ObjectAccessControl, len(in))
	for i, rule := range in {
		out[i] = &storagev1.ObjectAccessControl{
			Entity: string(rule.Entity),
			Role:   string(rule.Role),
		}
	}

	return out
}

func fromRawAcls(in []*storagev1.ObjectAccessControl) []storage.ACLRule

func fromRawObject(
	bucketName string,
	in *storagev1.Object) (out *storage.Object, err error) {
	// Convert the easy fields.
	out = &storage.Object{
		Bucket:          bucketName,
		Name:            in.Name,
		ContentType:     in.ContentType,
		ContentLanguage: in.ContentLanguage,
		CacheControl:    in.CacheControl,
		ACL:             fromRawAcls(in.Acl),
		ContentEncoding: in.ContentEncoding,
		Size:            int64(in.Size),
		MediaLink:       in.MediaLink,
		Metadata:        in.Metadata,
		Generation:      in.Generation,
		MetaGeneration:  in.Metageneration,
		StorageClass:    in.StorageClass,
		Deleted:         in.TimeDeleted,
		Updated:         in.Updated,
	}

	// Handle special cases.
	if in.Owner != nil {
		out.Owner = in.Owner.Entity
	}

	if out.MD5, err = base64.StdEncoding.DecodeString(in.Md5Hash); err != nil {
		err = fmt.Errorf("Decoding Md5Hash field: %v", err)
		return
	}

	crc32cString, err := base64.StdEncoding.DecodeString(in.Crc32c)
	if err != nil {
		err = fmt.Errorf("Decoding Crc32c field: %v", err)
		return
	}

	if len(crc32cString) != 4 {
		err = fmt.Errorf("Wrong length for decoded Crc32c field: %d", len(crc32cString))
		return
	}

	out.CRC32C =
		uint32(crc32cString[0])<<24 |
			uint32(crc32cString[1])<<16 |
			uint32(crc32cString[2])<<8 |
			uint32(crc32cString[3])<<0

	return
}

func getRawService(authContext context.Context) *storagev1.Service

func (b *bucket) CreateObject(
	ctx context.Context,
	req *CreateObjectRequest) (o *storage.Object, err error) {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	objectsService := getRawService(authContext).Objects

	// As of 2015-02, the wrapped storage package doesn't check this for us,
	// causing silently transformed names:
	//     https://github.com/GoogleCloudPlatform/gcloud-golang/issues/111
	if !utf8.ValidString(req.Attrs.Name) {
		err = errors.New("Invalid object name: not valid UTF-8")
		return
	}

	// Set up an object struct based on the supplied attributes.
	inputObj := &storagev1.Object{
		Name:            req.Attrs.Name,
		Bucket:          b.Name(),
		ContentType:     req.Attrs.ContentType,
		ContentLanguage: req.Attrs.ContentLanguage,
		ContentEncoding: req.Attrs.ContentEncoding,
		CacheControl:    req.Attrs.CacheControl,
		Acl:             toRawAcls(req.Attrs.ACL),
		Metadata:        req.Attrs.Metadata,
	}

	// Configure a 'call' object.
	call := objectsService.Insert(b.Name(), inputObj)
	call.Media(req.Contents)
	call.Projection("full")

	// Execute the call.
	rawObject, err := call.Do()
	if err != nil {
		return
	}

	// Convert the returned object.
	o, err = fromRawObject(b.Name(), rawObject)
	if err != nil {
		return
	}

	return
}

func (b *bucket) DeleteObject(ctx context.Context, name string) error {
	authContext := cloud.WithContext(ctx, b.projID, b.client)
	return storage.DeleteObject(authContext, b.name, name)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 IBM Corp.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 k8s

import (
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology/graph"
)

const (
	podToContainerLink = "pod2container"
	policyToPodLink    = "policy2pod"
)

var (
	podToContainerMetadata = graph.Metadata{"RelationType": podToContainerLink}
	policyToPodMetadata    = graph.Metadata{"RelationType": policyToPodLink}
)

// Probe for tracking k8s events
type Probe struct {
	podProbe           *podProbe
	networkPolicyProbe *networkPolicyProbe
	nodeProbe          *nodeProbe
	containerProbe     *containerProbe
	bundle             *probe.ProbeBundle
}

func (p *Probe) makeProbeBundle() *probe.ProbeBundle {
	configProbes := config.GetConfig().GetStringSlice("k8s.probes")
	logging.GetLogger().Infof("K8s probes: %v", configProbes)
	probes := make(map[string]probe.Probe)
	for name, i := range configProbes {
		switch i {
		case "pod":
			probes[i] = p.podProbe
		case "networkpolicy":
			probes[i] = p.networkPolicyProbe
		case "container":
			probes[i] = p.containerProbe
		case "node":
			probes[i] = p.nodeProbe
		default:
			logging.GetLogger().Errorf("skipping unsupported K8s probe %v", name)
		}
	}
	return probe.NewProbeBundle(probes)
}

// Start k8s probe
func (p *Probe) Start() {
	p.bundle.Start()
}

// Stop k8s probe
func (p *Probe) Stop() {
	p.bundle.Stop()
}

// NewProbe create the Probe for tracking k8s events
func NewProbe(g *graph.Graph) (*Probe, error) {
	err := initClientset()
	if err != nil {
		return nil, err
	}

	p := &Probe{}
	p.podProbe = newPodProbe(g)
	p.networkPolicyProbe = newNetworkPolicyProbe(g)
	p.containerProbe = newContainerProbe(g)
	p.nodeProbe = newNodeProbe(g)
	p.bundle = p.makeProbeBundle()
	return p, nil
}
<commit_msg>k8s: remove <kind>Probe objects within k8s.Probe<commit_after>/*
 * Copyright 2017 IBM Corp.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 k8s

import (
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
	"github.com/skydive-project/skydive/probe"
	"github.com/skydive-project/skydive/topology/graph"
)

const (
	podToContainerLink = "pod2container"
	policyToPodLink    = "policy2pod"
)

var (
	podToContainerMetadata = graph.Metadata{"RelationType": podToContainerLink}
	policyToPodMetadata    = graph.Metadata{"RelationType": policyToPodLink}
)

// Probe for tracking k8s events
type Probe struct {
	bundle *probe.ProbeBundle
}

func makeProbeBundle(g *graph.Graph) *probe.ProbeBundle {
	configProbes := config.GetConfig().GetStringSlice("k8s.probes")
	logging.GetLogger().Infof("K8s probes: %v", configProbes)
	probes := make(map[string]probe.Probe)
	for name, i := range configProbes {
		switch i {
		case "pod":
			probes[i] = newPodProbe(g)
		case "networkpolicy":
			probes[i] = newNetworkPolicyProbe(g)
		case "container":
			probes[i] = newContainerProbe(g)
		case "node":
			probes[i] = newNodeProbe(g)
		default:
			logging.GetLogger().Errorf("skipping unsupported K8s probe %v", name)
		}
	}
	return probe.NewProbeBundle(probes)
}

// Start k8s probe
func (p *Probe) Start() {
	p.bundle.Start()
}

// Stop k8s probe
func (p *Probe) Stop() {
	p.bundle.Stop()
}

// NewProbe create the Probe for tracking k8s events
func NewProbe(g *graph.Graph) (*Probe, error) {
	err := initClientset()
	if err != nil {
		return nil, err
	}

	return &Probe{
		bundle: makeProbeBundle(g),
	}, nil
}
<|endoftext|>


<commit_before>// for each target domain in the given prioritized list:
//     for front candidate in the given prioritized list:
//        check whether we can reach the target domain by attempting domain
//        fronting through the front candidate.  Log successfully proxied pairs
//        to stdout.  Log detected errors to stderr.

// XXX: record where we left off, to enable resuming an interrupted scan.

package main

import (
	"bufio"
	"crypto/sha1"
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	"gopkg.in/getlantern/tlsdialer.v1"
)

var workers = runtime.NumCPU()

var headRe *regexp.Regexp = regexp.MustCompile("(?s)(.*)")

const reportFmt = "%-24s %-24s %-12s %-12s\n"

var dialTimeout = 10 * time.Second

// ResponseFeatures contains the characteristics we'll use when comparing
// direct vs domain fronted Responses, to estimate the likelihood that we are
// successfully fronting to a target.
//
// We are also checking whether we get a 200 OK response from the site at all,
// but we don't include this in ResponseFeatures becouse we will not even try
// domain fronting to sites for which we don't get 200 OK in a direct request,
// and we are not reporting sites for which we don't get a 200 OK through
// domain fronting.
type ResponseFeatures struct {
	// Hash of the body's  element, or "" if the body has none
	HeadHash string
	// Hash of the whole body, or "" if the body is the empty string
	BodyHash string
}

type Pairing struct {
	// The domain we're trying to reach.
	Target string
	// A URL where we can hopefully reach Target without redirects.
	TargetURL url.URL
	// The domain we're trying to domain front through.
	Front string
	// A URL where we can hopefully reach Front without redirects.
	FrontURL url.URL

	Features ResponseFeatures
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU() * 10)
	start := time.Now()
	fmt.Println()
	fmt.Printf(reportFmt, "TARGET", "FRONT", "BODIES", "HEADS")
	fmt.Printf(reportFmt, "======", "=====", "======", "=====")
	fmt.Println()
	main_()
	fmt.Println("**************** COMPLETE ******************")
	fmt.Printf("Scan took %.2fs\n", time.Since(start).Seconds())
}

func main_() {

	tasksChan := make(chan Pairing, workers)

	workersWg := sync.WaitGroup{}

	// Size of work queue is O(n^2) on the size of the input lists, so we
	// generate it lazily.
	go feedTasks(tasksChan)

	workersWg.Add(workers)
	for i := 0; i < workers; i++ {
		go work(tasksChan, &workersWg)
	}

	// Task completion handling:
	//
	// The feedTasks goroutine will close tasksChan when it has processed all
	// input.  Each worker reports workersWg.Done() when it learns that
	// tasksChan is closed.  By that time, the worker has sent all its results
	// to the resultsChan.  Therefore, by the time workersWg.Wait() returns,
	// all `Pairing`s have been checked and all the successful ones have been
	// printed out.
	workersWg.Wait()
}

func work(tasksChan <-chan Pairing, workersWg *sync.WaitGroup) {
	for task := range tasksChan {
		features, err := proxiedResponseFeatures(task)
		if err != nil {
			logDebug("proxiedResponseFeatures error:", err)
			continue
		}
		fmt.Printf(
			reportFmt,
			task.Target,
			task.Front,
			matchReport(task.Features.BodyHash, features.BodyHash),
			matchReport(task.Features.HeadHash, features.HeadHash),
		)
	}
	logDebug("One worker done")
	workersWg.Done()
}

func proxiedResponseFeatures(p Pairing) (ResponseFeatures, error) {
	logDebug("canProxy", p)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				return tlsdialer.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					p.Target+":443",
					// We can't use a domain front that requires a properly
					// populated SNI, so let's make those fail.
					false,
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: noRedirect,
	}
	u := p.TargetURL
	u.Host = p.FrontURL.Host
	logDebug("Trying to hit", p.Target, "through", u.String())
	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		logErr("building GET request", err)
		return ResponseFeatures{}, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return ResponseFeatures{}, err
	}
	return responseFeatures(resp)
}

// feedTasks reads input files, enqueues testing tasks in taskChan, and closes
// taskChan when done.
func feedTasks(taskChan chan<- Pairing) {
	//XXX: -fronts command line parameter
	fronts, err := readLines("fronts.txt")
	if err != nil {
		panic(err)
	}
	//XXX: -targets command line parameter
	targets, err := readLines("targets.txt")
	if err != nil {
		panic(err)
	}

	// Domains for which we can't get a hash, even without domain fronting.
	badDomains := make(map[string]bool)

	// No-redirect URLs
	nru := make(map[string]url.URL)

	// No-redirect response features
	nrf := make(map[string]ResponseFeatures)

	nilURL := url.URL{}
	nilRF := ResponseFeatures{}

	// URLnRF just provides caching and some logging on top of followRedirects.
	URLnRF := func(domain string) (url.URL, ResponseFeatures) {
		if _, ok := badDomains[domain]; ok {
			logDebug("Known bad domain", domain)
			return nilURL, nilRF
		}
		u, ok := nru[domain]
		if ok {
			logDebug("Cached domain", domain, u.String())
			return u, nrf[domain]
		} else {
			u, rf, err := followRedirects(domain)
			if err != nil {
				logErr("can't reach "+domain+" *without* a proxy", err)
				badDomains[domain] = true
				return nilURL, nilRF
			} else {
				logDebug("Got new url", u.String(), "fetures", rf)
				nru[domain] = u
				nrf[domain] = rf
				return u, rf
			}
		}
	}

	for _, target := range targets {
		logDebug("Trying", target)
		targetURL, features := URLnRF(target)
		if targetURL == nilURL {
			logDebug("Got null URL for", target)
			continue
		}
		for _, front := range fronts {
			logDebug("Enqueuing pairing:", front, target)
			frontURL, _ := URLnRF(front)
			if frontURL == nilURL {
				// Some CDNs like akamai or Fastly don't like to have their /s
				// requested, but will still domain front.
				frontURL = url.URL{Scheme: "http", Host: front}
			}
			taskChan <- Pairing{
				Target:    target,
				TargetURL: targetURL,
				Front:     front,
				FrontURL:  frontURL,
				Features:  features,
			}
		}
	}
	logDebug("Enqueued all input")
	close(taskChan)
}

// followRedirects makes a request for the given domain without domain fronting,
// returning a URL that will hopefully not redirect, and the relevant features
// of the response.
func followRedirects(domain string) (u url.URL, rf ResponseFeatures, err error) {
	retriesLeft := 10
	// use http scheme to avoid getting double-TLSed
	u = url.URL{Scheme: "http", Host: domain}
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				logDebug("** trying to get", addr)
				return tls.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					domain+":443",
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			retriesLeft--
			if retriesLeft == 0 {
				return errors.New("Too many redirects")
			}
			// use http scheme to avoid getting double-TLSed
			req.URL.Scheme = "http"
			u = *req.URL
			logDebug("(re)setting URL to", u.String())
			return nil
		},
	}
	// use http scheme to avoid getting double-TLSed
	req, err := http.NewRequest("GET", "http://"+domain, nil)
	if err != nil {
		logErr("building GET request", err)
		return url.URL{}, rf, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return url.URL{}, rf, err
	}
	if resp.StatusCode != 200 {
		return url.URL{}, rf, errors.New("Non-200 response:" + resp.Status)
	}
	rf, err = responseFeatures(resp)
	if err != nil {
		logErr("getting response features", err)
		return url.URL{}, rf, err
	}
	return
}

func responseFeatures(resp *http.Response) (ret ResponseFeatures, err error) {
	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return ret, err
	}
	if body == nil {
		ret.BodyHash = ""
		ret.HeadHash = ""
		return ret, nil
	}
	ret.BodyHash = hash(body)
	headMatches := headRe.FindSubmatch(body)
	if headMatches == nil {
		ret.HeadHash = ""
	} else {
		ret.HeadHash = hash(headMatches[1])
	}
	return
}

func hash(s []byte) string {
	h := sha1.New()
	h.Write(s)
	bs := h.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func matchReport(expected, actual string) string {
	if actual == expected {
		if actual == "" {
			return "both empty"
		} else {
			return "match"
		}
	} else {
		if actual == "" {
			return "actual empty"
		} else {
			return "differ"
		}
	}
}

func readLines(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(file)
	ret := make([]string, 0)
	for scanner.Scan() {
		ret = append(ret, scanner.Text())
	}
	return ret, scanner.Err()
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Don't redirect!: " + req.URL.String())
}

func logErr(msg string, err error) {
	if err != nil {
		log.Printf("ERROR %s: %v\n", msg, err)
	}
}

// Poor man's debug level switch.  Comment out the body to suppress noise.
// XXX: learn about the idiomatic way to do this in Go
func logDebug(vs ...interface{}) {
	//log.Println(vs...)
}
<commit_msg>remove obsolete comment bit<commit_after>// for each target domain in the given prioritized list:
//     for front candidate in the given prioritized list:
//        check whether we can reach the target domain by attempting domain
//        fronting through the front candidate.  Log successfully proxied pairs
//        to stdout.  Log detected errors to stderr.

// XXX: record where we left off, to enable resuming an interrupted scan.

package main

import (
	"bufio"
	"crypto/sha1"
	"crypto/tls"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"sync"
	"time"

	"gopkg.in/getlantern/tlsdialer.v1"
)

var workers = runtime.NumCPU()

var headRe *regexp.Regexp = regexp.MustCompile("(?s)(.*)")

const reportFmt = "%-24s %-24s %-12s %-12s\n"

var dialTimeout = 10 * time.Second

// ResponseFeatures contains the characteristics we'll use when comparing
// direct vs domain fronted Responses, to estimate the likelihood that we are
// successfully fronting to a target.
//
// We are also checking whether we get a 200 OK response from the site at all,
// but we don't include this in ResponseFeatures becouse we are not reporting
// sites for which we don't get a 200 OK through domain fronting.
type ResponseFeatures struct {
	// Hash of the body's  element, or "" if the body has none
	HeadHash string
	// Hash of the whole body, or "" if the body is the empty string
	BodyHash string
}

type Pairing struct {
	// The domain we're trying to reach.
	Target string
	// A URL where we can hopefully reach Target without redirects.
	TargetURL url.URL
	// The domain we're trying to domain front through.
	Front string
	// A URL where we can hopefully reach Front without redirects.
	FrontURL url.URL

	Features ResponseFeatures
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU() * 10)
	start := time.Now()
	fmt.Println()
	fmt.Printf(reportFmt, "TARGET", "FRONT", "BODIES", "HEADS")
	fmt.Printf(reportFmt, "======", "=====", "======", "=====")
	fmt.Println()
	main_()
	fmt.Println("**************** COMPLETE ******************")
	fmt.Printf("Scan took %.2fs\n", time.Since(start).Seconds())
}

func main_() {

	tasksChan := make(chan Pairing, workers)

	workersWg := sync.WaitGroup{}

	// Size of work queue is O(n^2) on the size of the input lists, so we
	// generate it lazily.
	go feedTasks(tasksChan)

	workersWg.Add(workers)
	for i := 0; i < workers; i++ {
		go work(tasksChan, &workersWg)
	}

	// Task completion handling:
	//
	// The feedTasks goroutine will close tasksChan when it has processed all
	// input.  Each worker reports workersWg.Done() when it learns that
	// tasksChan is closed.  By that time, the worker has sent all its results
	// to the resultsChan.  Therefore, by the time workersWg.Wait() returns,
	// all `Pairing`s have been checked and all the successful ones have been
	// printed out.
	workersWg.Wait()
}

func work(tasksChan <-chan Pairing, workersWg *sync.WaitGroup) {
	for task := range tasksChan {
		features, err := proxiedResponseFeatures(task)
		if err != nil {
			logDebug("proxiedResponseFeatures error:", err)
			continue
		}
		fmt.Printf(
			reportFmt,
			task.Target,
			task.Front,
			matchReport(task.Features.BodyHash, features.BodyHash),
			matchReport(task.Features.HeadHash, features.HeadHash),
		)
	}
	logDebug("One worker done")
	workersWg.Done()
}

func proxiedResponseFeatures(p Pairing) (ResponseFeatures, error) {
	logDebug("canProxy", p)
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				return tlsdialer.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					p.Target+":443",
					// We can't use a domain front that requires a properly
					// populated SNI, so let's make those fail.
					false,
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: noRedirect,
	}
	u := p.TargetURL
	u.Host = p.FrontURL.Host
	logDebug("Trying to hit", p.Target, "through", u.String())
	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		logErr("building GET request", err)
		return ResponseFeatures{}, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return ResponseFeatures{}, err
	}
	return responseFeatures(resp)
}

// feedTasks reads input files, enqueues testing tasks in taskChan, and closes
// taskChan when done.
func feedTasks(taskChan chan<- Pairing) {
	//XXX: -fronts command line parameter
	fronts, err := readLines("fronts.txt")
	if err != nil {
		panic(err)
	}
	//XXX: -targets command line parameter
	targets, err := readLines("targets.txt")
	if err != nil {
		panic(err)
	}

	// Domains for which we can't get a hash, even without domain fronting.
	badDomains := make(map[string]bool)

	// No-redirect URLs
	nru := make(map[string]url.URL)

	// No-redirect response features
	nrf := make(map[string]ResponseFeatures)

	nilURL := url.URL{}
	nilRF := ResponseFeatures{}

	// URLnRF just provides caching and some logging on top of followRedirects.
	URLnRF := func(domain string) (url.URL, ResponseFeatures) {
		if _, ok := badDomains[domain]; ok {
			logDebug("Known bad domain", domain)
			return nilURL, nilRF
		}
		u, ok := nru[domain]
		if ok {
			logDebug("Cached domain", domain, u.String())
			return u, nrf[domain]
		} else {
			u, rf, err := followRedirects(domain)
			if err != nil {
				logErr("can't reach "+domain+" *without* a proxy", err)
				badDomains[domain] = true
				return nilURL, nilRF
			} else {
				logDebug("Got new url", u.String(), "features", rf)
				nru[domain] = u
				nrf[domain] = rf
				return u, rf
			}
		}
	}

	for _, target := range targets {
		logDebug("Trying", target)
		targetURL, features := URLnRF(target)
		if targetURL == nilURL {
			logDebug("Got null URL for", target)
			continue
		}
		for _, front := range fronts {
			logDebug("Enqueuing pairing:", front, target)
			frontURL, _ := URLnRF(front)
			if frontURL == nilURL {
				// Some CDNs like akamai or Fastly don't like to have their /s
				// requested, but will still domain front.
				frontURL = url.URL{Scheme: "http", Host: front}
			}
			taskChan <- Pairing{
				Target:    target,
				TargetURL: targetURL,
				Front:     front,
				FrontURL:  frontURL,
				Features:  features,
			}
		}
	}
	logDebug("Enqueued all input")
	close(taskChan)
}

// followRedirects makes a request for the given domain without domain fronting,
// returning a URL that will hopefully not redirect, and the relevant features
// of the response.
func followRedirects(domain string) (u url.URL, rf ResponseFeatures, err error) {
	retriesLeft := 10
	// use http scheme to avoid getting double-TLSed
	u = url.URL{Scheme: "http", Host: domain}
	client := &http.Client{
		Transport: &http.Transport{
			Dial: func(network, addr string) (conn net.Conn, err error) {
				logDebug("** trying to get", addr)
				return tls.DialWithDialer(
					&net.Dialer{Timeout: dialTimeout},
					"tcp",
					domain+":443",
					&tls.Config{InsecureSkipVerify: true},
				)
			},
		},
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			retriesLeft--
			if retriesLeft == 0 {
				return errors.New("Too many redirects")
			}
			// use http scheme to avoid getting double-TLSed
			req.URL.Scheme = "http"
			u = *req.URL
			logDebug("(re)setting URL to", u.String())
			return nil
		},
	}
	// use http scheme to avoid getting double-TLSed
	req, err := http.NewRequest("GET", "http://"+domain, nil)
	if err != nil {
		logErr("building GET request", err)
		return url.URL{}, rf, err
	}
	resp, err := client.Do(req)
	if err != nil {
		logErr("performing GET request", err)
		return url.URL{}, rf, err
	}
	if resp.StatusCode != 200 {
		return url.URL{}, rf, errors.New("Non-200 response:" + resp.Status)
	}
	rf, err = responseFeatures(resp)
	if err != nil {
		logErr("getting response features", err)
		return url.URL{}, rf, err
	}
	return
}

func responseFeatures(resp *http.Response) (ret ResponseFeatures, err error) {
	body, err := ioutil.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return ret, err
	}
	if body == nil {
		ret.BodyHash = ""
		ret.HeadHash = ""
		return ret, nil
	}
	ret.BodyHash = hash(body)
	headMatches := headRe.FindSubmatch(body)
	if headMatches == nil {
		ret.HeadHash = ""
	} else {
		ret.HeadHash = hash(headMatches[1])
	}
	return
}

func hash(s []byte) string {
	h := sha1.New()
	h.Write(s)
	bs := h.Sum(nil)
	return fmt.Sprintf("%x", bs)
}

func matchReport(expected, actual string) string {
	if actual == expected {
		if actual == "" {
			return "both empty"
		} else {
			return "match"
		}
	} else {
		if actual == "" {
			return "actual empty"
		} else {
			return "differ"
		}
	}
}

func readLines(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(file)
	ret := make([]string, 0)
	for scanner.Scan() {
		ret = append(ret, scanner.Text())
	}
	return ret, scanner.Err()
}

func noRedirect(req *http.Request, via []*http.Request) error {
	return errors.New("Don't redirect!: " + req.URL.String())
}

func logErr(msg string, err error) {
	if err != nil {
		log.Printf("ERROR %s: %v\n", msg, err)
	}
}

// Poor man's debug level switch.  Comment out the body to suppress noise.
// XXX: learn about the idiomatic way to do this in Go
func logDebug(vs ...interface{}) {
	//log.Println(vs...)
}
<|endoftext|>


<commit_before>package expr

import (
	"fmt"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
)

// Datamap contains all series to feed into the processing chain or generated therein:
// * fetched series, grouped by their expr.Req, such that expr.FuncGet can find the data it needs and feed it into subsequent expr.GraphiteFunc functions
// * additional series generated while handling the request (e.g. function processing, normalization), keyed by an empty expr.Req (such that can't be mistakenly picked up by FuncGet)
// all of these series will need to be returned to the pool once we're done with all processing and have generated our response body by calling Clean()
// eventually we'd like to be able to reuse intermediately computed data. e.g. queries like target=movingAvg(sum(foo), 10)&target=sum(foo) but for now we don't support this
type DataMap map[Req][]models.Series

func NewDataMap() DataMap {
	return make(map[Req][]models.Series)
}

func (dm DataMap) Add(r Req, s ...models.Series) {
	dm[r] = append(dm[r], s...)
}

// Clean returns all contained pointslices back to the pool.
func (dm DataMap) Clean() {
	for _, series := range dm {
		for _, serie := range series {
			pointSlicePool.Put(serie.Datapoints[:0])
		}
	}
}

// CheckForOverlappingPoints runs through all series in the pool and makes sure there
// are no series that are overlapping (otherwise returning them would cause issues)
// This is not efficient and should probably only be called from tests
func (dm DataMap) CheckForOverlappingPoints() error {
	// First flatten series to make the logic easier
	flatseries := make([]models.Series, 0, len(dm))
	for _, series := range dm {
		flatseries = append(flatseries, series...)
	}

	slicesOverlap := func(x, y []schema.Point) bool {
		return cap(x) > 0 && cap(y) > 0 && &x[0:cap(x)][cap(x)-1] ==
			&y[0:cap(y)][cap(y)-1]
	}

	for i := range flatseries {
		for j := i + 1; j < len(flatseries); j++ {
			if slicesOverlap(flatseries[i].Datapoints, flatseries[j].Datapoints) {
				return fmt.Errorf("Found overlapping series at position %d and %d", i, j)
			}
		}
	}

	return nil
}
<commit_msg>Add comment about overlap detection<commit_after>package expr

import (
	"fmt"

	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
)

// Datamap contains all series to feed into the processing chain or generated therein:
// * fetched series, grouped by their expr.Req, such that expr.FuncGet can find the data it needs and feed it into subsequent expr.GraphiteFunc functions
// * additional series generated while handling the request (e.g. function processing, normalization), keyed by an empty expr.Req (such that can't be mistakenly picked up by FuncGet)
// all of these series will need to be returned to the pool once we're done with all processing and have generated our response body by calling Clean()
// eventually we'd like to be able to reuse intermediately computed data. e.g. queries like target=movingAvg(sum(foo), 10)&target=sum(foo) but for now we don't support this
type DataMap map[Req][]models.Series

func NewDataMap() DataMap {
	return make(map[Req][]models.Series)
}

func (dm DataMap) Add(r Req, s ...models.Series) {
	dm[r] = append(dm[r], s...)
}

// Clean returns all contained pointslices back to the pool.
func (dm DataMap) Clean() {
	for _, series := range dm {
		for _, serie := range series {
			pointSlicePool.Put(serie.Datapoints[:0])
		}
	}
}

// CheckForOverlappingPoints runs through all series in the pool and makes sure there
// are no series that are overlapping (otherwise returning them would cause issues)
// This is not efficient and should probably only be called from tests
func (dm DataMap) CheckForOverlappingPoints() error {
	// First flatten series to make the logic easier
	flatseries := make([]models.Series, 0, len(dm))
	for _, series := range dm {
		flatseries = append(flatseries, series...)
	}

	// we leverage the fact here that for any given slice, if you reslice or subslice it - which may result in the
	// first slot's address changing - the last slot of the backing array remains the same and can be used
	// to find slices backed by the same array
	slicesOverlap := func(x, y []schema.Point) bool {
		return cap(x) > 0 && cap(y) > 0 && &x[0:cap(x)][cap(x)-1] ==
			&y[0:cap(y)][cap(y)-1]
	}

	for i := range flatseries {
		for j := i + 1; j < len(flatseries); j++ {
			if slicesOverlap(flatseries[i].Datapoints, flatseries[j].Datapoints) {
				return fmt.Errorf("Found overlapping series at position %d and %d", i, j)
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package bsw

import (
	"errors"
	"net"
	"strings"

	"github.com/miekg/dns"
)

// LookupIP returns hostname from PTR record or error.
func LookupIP(ip, serverAddr string) ([]string, error) {
	fqdn := []string{}
	result, err := net.LookupAddr(ip)
	if err != nil {
		return fqdn, err
	}
	for _, hostname := range result {
		fqdn = append(fqdn, strings.TrimRight(hostname, "."))
	}
	return fqdn, nil
}

// LookupName returns IPv4 address from A record or error.
func LookupName(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.A); ok {
		ip := a.A.String()
		return ip, nil
	}
	return "", errors.New("no A record returned")
}

// LookupCname returns an IPv4 address from CNAME record or error.
func LookupCname(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeCNAME)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.CNAME); ok {
		name := a.Target
		return strings.TrimRight(name, "."), nil
	}
	return "", errors.New("no CNAME record returned")
}

// LookupName6 returns a IPv6 address from AAAA record or error.
func LookupName6(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeAAAA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.AAAA); ok {
		ip := a.AAAA.String()
		return ip, nil
	}
	return "", errors.New("no AAAA record returned")
}

// LookupSRV returns a hostname from SRV record or error.
func LookupSRV(fqdn, dnsServer string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeSRV)
	in, err := dns.Exchange(m, dnsServer+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.SRV); ok {
		return strings.TrimRight(a.Target, "."), nil
	}
	return "", errors.New("no SRV record returned")
}
<commit_msg>Reverse lookup now uses dns library<commit_after>package bsw

import (
	"errors"
	"strings"

	"github.com/miekg/dns"
)

// LookupIP returns hostname from PTR record or error.
func LookupIP(ip, serverAddr string) ([]string, error) {
	names := []string{}
	m := &dns.Msg{}
	ipArpa, err := dns.ReverseAddr(ip)
	if err != nil {
		return names, err
	}
	m.SetQuestion(ipArpa, dns.TypePTR)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return names, err
	}
	if len(in.Answer) < 1 {
		return names, errors.New("no Answer")
	}

	for _, a := range in.Answer {
		ptr := a.(*dns.PTR)
		names = append(names, strings.TrimRight(ptr.Ptr, "."))
	}
	return names, nil
}

// LookupName returns IPv4 address from A record or error.
func LookupName(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.A); ok {
		ip := a.A.String()
		return ip, nil
	}
	return "", errors.New("no A record returned")
}

// LookupCname returns an IPv4 address from CNAME record or error.
func LookupCname(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeCNAME)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.CNAME); ok {
		name := a.Target
		return strings.TrimRight(name, "."), nil
	}
	return "", errors.New("no CNAME record returned")
}

// LookupName6 returns a IPv6 address from AAAA record or error.
func LookupName6(fqdn, serverAddr string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeAAAA)
	in, err := dns.Exchange(m, serverAddr+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.AAAA); ok {
		ip := a.AAAA.String()
		return ip, nil
	}
	return "", errors.New("no AAAA record returned")
}

// LookupSRV returns a hostname from SRV record or error.
func LookupSRV(fqdn, dnsServer string) (string, error) {
	m := &dns.Msg{}
	m.SetQuestion(dns.Fqdn(fqdn), dns.TypeSRV)
	in, err := dns.Exchange(m, dnsServer+":53")
	if err != nil {
		return "", err
	}
	if len(in.Answer) < 1 {
		return "", errors.New("no Answer")
	}
	if a, ok := in.Answer[0].(*dns.SRV); ok {
		return strings.TrimRight(a.Target, "."), nil
	}
	return "", errors.New("no SRV record returned")
}
<|endoftext|>


<commit_before><commit_msg>Refactor setting statuses for commits<commit_after><|endoftext|>


<commit_before><commit_msg>cleanup comments<commit_after><|endoftext|>


<commit_before><commit_msg>add log middleware for debug<commit_after><|endoftext|>


<commit_before><commit_msg>add imagemap<commit_after><|endoftext|>


<commit_before><commit_msg>Remove duplicated "sameFile" logic<commit_after><|endoftext|>


<commit_before><commit_msg>stuff<commit_after><|endoftext|>


<commit_before><commit_msg>Remove _out before building. Add 'clean' command.<commit_after><|endoftext|>


<commit_before>package geom

type geom1 struct {
	geom0
}

func (g *geom1) Coord(i int) []float64 {
	return g.flatCoords[i*g.stride : (i+1)*g.stride]
}

func (g *geom1) Coords() [][]float64 {
	return inflate1(g.flatCoords, 0, len(g.flatCoords), g.stride)
}

func (g *geom1) LastCoord() []float64 {
	if len(g.flatCoords) == 0 {
		return nil
	} else {
		return g.flatCoords[len(g.flatCoords)-g.stride:]
	}
}

func (g *geom1) NumCoords() int {
	return len(g.flatCoords) / g.stride
}

func (g *geom1) PushCoord(coord0 []float64) error {
	if len(coord0) != g.stride {
		return ErrStrideMismatch{Got: len(coord0), Want: g.stride}
	}
	g.flatCoords = append(g.flatCoords, coord0...)
	return nil
}

func (g *geom1) SetCoords(coords1 [][]float64) error {
	var err error
	if g.flatCoords, err = deflate1(nil, coords1, g.stride); err != nil {
		return err
	}
	return nil
}

func (g *geom1) verify() error {
	if g.stride != g.layout.Stride() {
		return errStrideLayoutMismatch
	}
	if g.stride == 0 {
		if len(g.flatCoords) != 0 {
			return errNonEmptyFlatCoords
		}
	} else {
		if len(g.flatCoords)%g.stride != 0 {
			return errLengthStrideMismatch
		}
	}
	return nil
}
<commit_msg>Remove unused functions<commit_after>package geom

type geom1 struct {
	geom0
}

func (g *geom1) Coord(i int) []float64 {
	return g.flatCoords[i*g.stride : (i+1)*g.stride]
}

func (g *geom1) Coords() [][]float64 {
	return inflate1(g.flatCoords, 0, len(g.flatCoords), g.stride)
}

func (g *geom1) NumCoords() int {
	return len(g.flatCoords) / g.stride
}

func (g *geom1) SetCoords(coords1 [][]float64) error {
	var err error
	if g.flatCoords, err = deflate1(nil, coords1, g.stride); err != nil {
		return err
	}
	return nil
}

func (g *geom1) verify() error {
	if g.stride != g.layout.Stride() {
		return errStrideLayoutMismatch
	}
	if g.stride == 0 {
		if len(g.flatCoords) != 0 {
			return errNonEmptyFlatCoords
		}
	} else {
		if len(g.flatCoords)%g.stride != 0 {
			return errLengthStrideMismatch
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package reader

import (
	"fmt"
	"io"

	"github.com/apache/thrift/lib/go/thrift"
	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/source"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/parquet"
)

type ColumnBufferType struct {
	PFile        source.ParquetFile
	ThriftReader *thrift.TBufferedTransport

	Footer        *parquet.FileMetaData
	SchemaHandler *schema.SchemaHandler

	PathStr       string
	RowGroupIndex int64
	ChunkHeader   *parquet.ColumnChunk

	ChunkReadValues int64

	DictPage *layout.Page

	DataTable        *layout.Table
	DataTableNumRows int64
}

func NewColumnBuffer(pFile source.ParquetFile, footer *parquet.FileMetaData, schemaHandler *schema.SchemaHandler, pathStr string) (*ColumnBufferType, error) {
	newPFile, err := pFile.Open("")
	if err != nil {
		return nil, err
	}
	res := &ColumnBufferType{
		PFile:            newPFile,
		Footer:           footer,
		SchemaHandler:    schemaHandler,
		PathStr:          pathStr,
		DataTableNumRows: -1,
	}
	err = res.NextRowGroup()
	return res, err
}

func (self *ColumnBufferType) NextRowGroup() error {
	var err error
	rowGroups := self.Footer.GetRowGroups()
	ln := int64(len(rowGroups))
	if self.RowGroupIndex >= ln {
		self.DataTableNumRows++ //very important, because DataTableNumRows is one smaller than real rows number
		return fmt.Errorf("End of row groups")
	}
	self.RowGroupIndex++

	columnChunks := rowGroups[self.RowGroupIndex-1].GetColumns()
	i := int64(0)
	ln = int64(len(columnChunks))
	for i = 0; i < ln; i++ {
		path := make([]string, 0)
		path = append(path, self.SchemaHandler.GetRootInName())
		path = append(path, columnChunks[i].MetaData.GetPathInSchema()...)

		if self.PathStr == common.PathToStr(path) {
			break
		}
	}
	if i >= ln {
		return fmt.Errorf("[NextRowGroup] Column not found: %v", self.PathStr)
	}

	self.ChunkHeader = columnChunks[i]
	if columnChunks[i].FilePath != nil {
		self.PFile.Close()
		if self.PFile, err = self.PFile.Open(*columnChunks[i].FilePath); err != nil {
			return err
		}
	}
	//offset := columnChunks[i].FileOffset
	offset := columnChunks[i].MetaData.DataPageOffset
	if columnChunks[i].MetaData.DictionaryPageOffset != nil {
		offset = *columnChunks[i].MetaData.DictionaryPageOffset
	}
	size := columnChunks[i].MetaData.GetTotalCompressedSize()
	if self.ThriftReader != nil {
		self.ThriftReader.Close()
	}
	self.ThriftReader = source.ConvertToThriftReader(self.PFile, offset, size)
	self.ChunkReadValues = 0
	self.DictPage = nil
	return nil
}

func (self *ColumnBufferType) ReadPage() error {
	if self.ChunkHeader != nil && self.ChunkHeader.MetaData != nil && self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
		page, numValues, numRows, err := layout.ReadPage(self.ThriftReader, self.SchemaHandler, self.ChunkHeader.MetaData)
		if err != nil {
			//data is nil and rl/dl=0, no pages in file
			if err == io.EOF {
				if self.DataTable == nil {
					index := self.SchemaHandler.MapIndex[self.PathStr]
					self.DataTable = layout.NewEmptyTable()
					self.DataTable.Type = self.SchemaHandler.SchemaElements[index].GetType()
					self.DataTable.Path = common.StrToPath(self.PathStr)

				}
				for self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
					self.DataTable.Values = append(self.DataTable.Values, nil)
					self.DataTable.RepetitionLevels = append(self.DataTable.RepetitionLevels, int32(0))
					self.DataTable.DefinitionLevels = append(self.DataTable.DefinitionLevels, int32(0))
					self.ChunkReadValues++
					self.DataTableNumRows++
				}
			}
			return err
		}
		if page.Header.GetType() == parquet.PageType_DICTIONARY_PAGE {
			self.DictPage = page
			return nil
		}

		page.Decode(self.DictPage)

		if self.DataTable == nil {
			self.DataTable = layout.NewTableFromTable(page.DataTable)
		}
		self.DataTable.Merge(page.DataTable)
		self.ChunkReadValues += numValues

		self.DataTableNumRows += numRows
	} else {
		if err := self.NextRowGroup(); err != nil {
			return err
		}
		return self.ReadPage()
	}

	return nil
}

func (self *ColumnBufferType) ReadPageForSkip() (*layout.Page, error) {
	if self.ChunkHeader != nil && self.ChunkHeader.MetaData != nil && self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
		page, err := layout.ReadPageRawData(self.ThriftReader, self.SchemaHandler, self.ChunkHeader.MetaData)
		if err != nil {
			return nil, err
		}
		numValues, numRows, err := page.GetRLDLFromRawData(self.SchemaHandler)
		if err != nil {
			return nil, err
		}
		if page.Header.GetType() == parquet.PageType_DICTIONARY_PAGE {
			page.GetValueFromRawData(self.SchemaHandler)
			self.DictPage = page
			return page, nil
		}
		if self.DataTable == nil {
			self.DataTable = layout.NewTableFromTable(page.DataTable)
		}
		self.DataTable.Merge(page.DataTable)
		self.ChunkReadValues += numValues
		self.DataTableNumRows += numRows
		return page, nil

	} else {
		if err := self.NextRowGroup(); err != nil {
			return nil, err
		}
		return self.ReadPageForSkip()
	}
}

func (self *ColumnBufferType) SkipRows(num int64) int64 {
	var (
		err  error
		page *layout.Page
	)

	for self.DataTableNumRows < num && err == nil {
		page, err = self.ReadPageForSkip()
	}
	if num > self.DataTableNumRows {
		num = self.DataTableNumRows
	}
	if page != nil {
		if err = page.GetValueFromRawData(self.SchemaHandler); err != nil {
			return 0
		}
		page.Decode(self.DictPage)
		i, j := len(self.DataTable.Values)-1, len(page.DataTable.Values)-1
		for i >= 0 && j >= 0 {
			self.DataTable.Values[i] = page.DataTable.Values[j]
			i, j = i-1, j-1
		}
	}

	self.DataTable.Pop(num)
	self.DataTableNumRows -= num
	if self.DataTableNumRows <= 0 {
		tmp := self.DataTable
		self.DataTable = layout.NewTableFromTable(tmp)
		self.DataTable.Merge(tmp)
	}
	return num
}

func (self *ColumnBufferType) ReadRows(num int64) (*layout.Table, int64) {
	var err error

	for self.DataTableNumRows < num && err == nil {
		err = self.ReadPage()
	}

	if self.DataTableNumRows < 0 {
		self.DataTableNumRows = 0
		self.DataTable = layout.NewEmptyTable()
	}

	if num > self.DataTableNumRows {
		num = self.DataTableNumRows
	}

	res := self.DataTable.Pop(num)
	self.DataTableNumRows -= num

	if self.DataTableNumRows <= 0 { //release previous slice memory
		tmp := self.DataTable
		self.DataTable = layout.NewTableFromTable(tmp)
		self.DataTable.Merge(tmp)
	}
	return res, num

}
<commit_msg>issue176: fix 0 rows reader bug<commit_after>package reader

import (
	"fmt"
	"io"

	"github.com/apache/thrift/lib/go/thrift"
	"github.com/xitongsys/parquet-go/common"
	"github.com/xitongsys/parquet-go/layout"
	"github.com/xitongsys/parquet-go/source"
	"github.com/xitongsys/parquet-go/schema"
	"github.com/xitongsys/parquet-go/parquet"
)

type ColumnBufferType struct {
	PFile        source.ParquetFile
	ThriftReader *thrift.TBufferedTransport

	Footer        *parquet.FileMetaData
	SchemaHandler *schema.SchemaHandler

	PathStr       string
	RowGroupIndex int64
	ChunkHeader   *parquet.ColumnChunk

	ChunkReadValues int64

	DictPage *layout.Page

	DataTable        *layout.Table
	DataTableNumRows int64
}

func NewColumnBuffer(pFile source.ParquetFile, footer *parquet.FileMetaData, schemaHandler *schema.SchemaHandler, pathStr string) (*ColumnBufferType, error) {
	newPFile, err := pFile.Open("")
	if err != nil {
		return nil, err
	}
	res := &ColumnBufferType{
		PFile:            newPFile,
		Footer:           footer,
		SchemaHandler:    schemaHandler,
		PathStr:          pathStr,
		DataTableNumRows: -1,
	}

	if err = res.NextRowGroup(); err == io.EOF {
		err = nil
	}
	return res, err
}

func (self *ColumnBufferType) NextRowGroup() error {
	var err error
	rowGroups := self.Footer.GetRowGroups()
	ln := int64(len(rowGroups))
	if self.RowGroupIndex >= ln {
		self.DataTableNumRows++ //very important, because DataTableNumRows is one smaller than real rows number
		return io.EOF
	}

	self.RowGroupIndex++

	columnChunks := rowGroups[self.RowGroupIndex-1].GetColumns()
	i := int64(0)
	ln = int64(len(columnChunks))
	for i = 0; i < ln; i++ {
		path := make([]string, 0)
		path = append(path, self.SchemaHandler.GetRootInName())
		path = append(path, columnChunks[i].MetaData.GetPathInSchema()...)

		if self.PathStr == common.PathToStr(path) {
			break
		}
	}

	if i >= ln {
		return fmt.Errorf("[NextRowGroup] Column not found: %v", self.PathStr)
	}

	self.ChunkHeader = columnChunks[i]
	if columnChunks[i].FilePath != nil {
		self.PFile.Close()
		if self.PFile, err = self.PFile.Open(*columnChunks[i].FilePath); err != nil {
			return err
		}
	}

	//offset := columnChunks[i].FileOffset
	offset := columnChunks[i].MetaData.DataPageOffset
	if columnChunks[i].MetaData.DictionaryPageOffset != nil {
		offset = *columnChunks[i].MetaData.DictionaryPageOffset
	}

	size := columnChunks[i].MetaData.GetTotalCompressedSize()
	if self.ThriftReader != nil {
		self.ThriftReader.Close()
	}

	self.ThriftReader = source.ConvertToThriftReader(self.PFile, offset, size)
	self.ChunkReadValues = 0
	self.DictPage = nil
	return nil
}

func (self *ColumnBufferType) ReadPage() error {
	if self.ChunkHeader != nil && self.ChunkHeader.MetaData != nil && self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
		page, numValues, numRows, err := layout.ReadPage(self.ThriftReader, self.SchemaHandler, self.ChunkHeader.MetaData)
		if err != nil {
			//data is nil and rl/dl=0, no pages in file
			if err == io.EOF {
				if self.DataTable == nil {
					index := self.SchemaHandler.MapIndex[self.PathStr]
					self.DataTable = layout.NewEmptyTable()
					self.DataTable.Type = self.SchemaHandler.SchemaElements[index].GetType()
					self.DataTable.Path = common.StrToPath(self.PathStr)

				}

				for self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
					self.DataTable.Values = append(self.DataTable.Values, nil)
					self.DataTable.RepetitionLevels = append(self.DataTable.RepetitionLevels, int32(0))
					self.DataTable.DefinitionLevels = append(self.DataTable.DefinitionLevels, int32(0))
					self.ChunkReadValues++
					self.DataTableNumRows++
				}
			}

			return err
		}

		if page.Header.GetType() == parquet.PageType_DICTIONARY_PAGE {
			self.DictPage = page
			return nil
		}

		page.Decode(self.DictPage)

		if self.DataTable == nil {
			self.DataTable = layout.NewTableFromTable(page.DataTable)
		}

		self.DataTable.Merge(page.DataTable)
		self.ChunkReadValues += numValues

		self.DataTableNumRows += numRows
	} else {
		if err := self.NextRowGroup(); err != nil {
			return err
		}

		return self.ReadPage()
	}

	return nil
}

func (self *ColumnBufferType) ReadPageForSkip() (*layout.Page, error) {
	if self.ChunkHeader != nil && self.ChunkHeader.MetaData != nil && self.ChunkReadValues < self.ChunkHeader.MetaData.NumValues {
		page, err := layout.ReadPageRawData(self.ThriftReader, self.SchemaHandler, self.ChunkHeader.MetaData)
		if err != nil {
			return nil, err
		}

		numValues, numRows, err := page.GetRLDLFromRawData(self.SchemaHandler)
		if err != nil {
			return nil, err
		}

		if page.Header.GetType() == parquet.PageType_DICTIONARY_PAGE {
			page.GetValueFromRawData(self.SchemaHandler)
			self.DictPage = page
			return page, nil
		}

		if self.DataTable == nil {
			self.DataTable = layout.NewTableFromTable(page.DataTable)
		}

		self.DataTable.Merge(page.DataTable)
		self.ChunkReadValues += numValues
		self.DataTableNumRows += numRows
		return page, nil

	} else {
		if err := self.NextRowGroup(); err != nil {
			return nil, err
		}

		return self.ReadPageForSkip()
	}
}

func (self *ColumnBufferType) SkipRows(num int64) int64 {
	var (
		err  error
		page *layout.Page
	)

	for self.DataTableNumRows < num && err == nil {
		page, err = self.ReadPageForSkip()
	}

	if num > self.DataTableNumRows {
		num = self.DataTableNumRows
	}

	if page != nil {
		if err = page.GetValueFromRawData(self.SchemaHandler); err != nil {
			return 0
		}

		page.Decode(self.DictPage)
		i, j := len(self.DataTable.Values)-1, len(page.DataTable.Values)-1
		for i >= 0 && j >= 0 {
			self.DataTable.Values[i] = page.DataTable.Values[j]
			i, j = i-1, j-1
		}
	}

	self.DataTable.Pop(num)
	self.DataTableNumRows -= num
	if self.DataTableNumRows <= 0 {
		tmp := self.DataTable
		self.DataTable = layout.NewTableFromTable(tmp)
		self.DataTable.Merge(tmp)
	}

	return num
}

func (self *ColumnBufferType) ReadRows(num int64) (*layout.Table, int64) {
	var err error

	for self.DataTableNumRows < num && err == nil {
		err = self.ReadPage()
	}

	if self.DataTableNumRows < 0 {
		self.DataTableNumRows = 0
		self.DataTable = layout.NewEmptyTable()
	}

	if num > self.DataTableNumRows {
		num = self.DataTableNumRows
	}

	res := self.DataTable.Pop(num)
	self.DataTableNumRows -= num

	if self.DataTableNumRows <= 0 { //release previous slice memory
		tmp := self.DataTable
		self.DataTable = layout.NewTableFromTable(tmp)
		self.DataTable.Merge(tmp)
	}
	return res, num

}
<|endoftext|>