<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"io"

	"google.golang.org/grpc"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

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

type commandEcDecode struct {
}

func (c *commandEcDecode) Name() string {
	return "ec.decode"
}

func (c *commandEcDecode) Help() string {
	return `decode a erasure coded volume into a normal volume

	ec.decode [-collection=""] [-volumeId=<volume_id>]

`
}

func (c *commandEcDecode) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
	encodeCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := encodeCommand.Int("volumeId", 0, "the volume id")
	collection := encodeCommand.String("collection", "", "the collection name")
	applyChanges := fixCommand.Bool("force", false, "force the encoding even if the cluster has less than recommended 4 nodes")
	if err = encodeCommand.Parse(args); err != nil {
		return nil
	}

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

	vid := needle.VolumeId(*volumeId)

	// collect topology information
	topologyInfo, _, err := collectTopologyInfo(commandEnv)
	if err != nil {
		return err
	}

	if !applyChanges {
		var nodeCount int
		eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
			nodeCount++
		})
		if nodeCount < erasure_coding.ParityShardsCount {
			glog.V(0).Infof("skip erasure coding with %d nodes, less than recommended %d nodes", nodeCount, erasure_coding.ParityShardsCount)
			return nil
		}
	}

	// volumeId is provided
	if vid != 0 {
		return doEcDecode(commandEnv, topologyInfo, *collection, vid)
	}

	// apply to all volumes in the collection
	volumeIds := collectEcShardIds(topologyInfo, *collection)
	fmt.Printf("ec encode volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doEcDecode(commandEnv, topologyInfo, *collection, vid); err != nil {
			return err
		}
	}

	return nil
}

func doEcDecode(commandEnv *CommandEnv, topoInfo *master_pb.TopologyInfo, collection string, vid needle.VolumeId) (err error) {
	// find volume location
	nodeToEcIndexBits := collectEcNodeShardBits(topoInfo, vid)

	fmt.Printf("ec volume %d shard locations: %+v\n", vid, nodeToEcIndexBits)

	// collect ec shards to the server with most space
	targetNodeLocation, err := collectEcShards(commandEnv, nodeToEcIndexBits, collection, vid)
	if err != nil {
		return fmt.Errorf("collectEcShards for volume %d: %v", vid, err)
	}

	// generate a normal volume
	err = generateNormalVolume(commandEnv.option.GrpcDialOption, vid, collection, targetNodeLocation)
	if err != nil {
		return fmt.Errorf("generate normal volume %d on %s: %v", vid, targetNodeLocation, err)
	}

	// delete the previous ec shards
	err = mountVolumeAndDeleteEcShards(commandEnv.option.GrpcDialOption, collection, targetNodeLocation, nodeToEcIndexBits, vid)
	if err != nil {
		return fmt.Errorf("delete ec shards for volume %d: %v", vid, err)
	}

	return nil
}

func mountVolumeAndDeleteEcShards(grpcDialOption grpc.DialOption, collection string, targetNodeLocation pb.ServerAddress, nodeToEcIndexBits map[pb.ServerAddress]erasure_coding.ShardBits, vid needle.VolumeId) error {

	// mount volume
	if err := operation.WithVolumeServerClient(targetNodeLocation, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, mountErr := volumeServerClient.VolumeMount(context.Background(), &volume_server_pb.VolumeMountRequest{
			VolumeId: uint32(vid),
		})
		return mountErr
	}); err != nil {
		return fmt.Errorf("mountVolumeAndDeleteEcShards mount volume %d on %s: %v", vid, targetNodeLocation, err)
	}

	// unmount ec shards
	for location, ecIndexBits := range nodeToEcIndexBits {
		fmt.Printf("unmount ec volume %d on %s has shards: %+v\n", vid, location, ecIndexBits.ShardIds())
		err := unmountEcShards(grpcDialOption, vid, location, ecIndexBits.ToUint32Slice())
		if err != nil {
			return fmt.Errorf("mountVolumeAndDeleteEcShards unmount ec volume %d on %s: %v", vid, location, err)
		}
	}
	// delete ec shards
	for location, ecIndexBits := range nodeToEcIndexBits {
		fmt.Printf("delete ec volume %d on %s has shards: %+v\n", vid, location, ecIndexBits.ShardIds())
		err := sourceServerDeleteEcShards(grpcDialOption, collection, vid, location, ecIndexBits.ToUint32Slice())
		if err != nil {
			return fmt.Errorf("mountVolumeAndDeleteEcShards delete ec volume %d on %s: %v", vid, location, err)
		}
	}

	return nil
}

func generateNormalVolume(grpcDialOption grpc.DialOption, vid needle.VolumeId, collection string, sourceVolumeServer pb.ServerAddress) error {

	fmt.Printf("generateNormalVolume from ec volume %d on %s\n", vid, sourceVolumeServer)

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, genErr := volumeServerClient.VolumeEcShardsToVolume(context.Background(), &volume_server_pb.VolumeEcShardsToVolumeRequest{
			VolumeId:   uint32(vid),
			Collection: collection,
		})
		return genErr
	})

	return err

}

func collectEcShards(commandEnv *CommandEnv, nodeToEcIndexBits map[pb.ServerAddress]erasure_coding.ShardBits, collection string, vid needle.VolumeId) (targetNodeLocation pb.ServerAddress, err error) {

	maxShardCount := 0
	var exisitngEcIndexBits erasure_coding.ShardBits
	for loc, ecIndexBits := range nodeToEcIndexBits {
		toBeCopiedShardCount := ecIndexBits.MinusParityShards().ShardIdCount()
		if toBeCopiedShardCount > maxShardCount {
			maxShardCount = toBeCopiedShardCount
			targetNodeLocation = loc
			exisitngEcIndexBits = ecIndexBits
		}
	}

	fmt.Printf("collectEcShards: ec volume %d collect shards to %s from: %+v\n", vid, targetNodeLocation, nodeToEcIndexBits)

	var copiedEcIndexBits erasure_coding.ShardBits
	for loc, ecIndexBits := range nodeToEcIndexBits {
		if loc == targetNodeLocation {
			continue
		}

		needToCopyEcIndexBits := ecIndexBits.Minus(exisitngEcIndexBits).MinusParityShards()
		if needToCopyEcIndexBits.ShardIdCount() == 0 {
			continue
		}

		err = operation.WithVolumeServerClient(targetNodeLocation, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {

			fmt.Printf("copy %d.%v %s => %s\n", vid, needToCopyEcIndexBits.ShardIds(), loc, targetNodeLocation)

			_, copyErr := volumeServerClient.VolumeEcShardsCopy(context.Background(), &volume_server_pb.VolumeEcShardsCopyRequest{
				VolumeId:       uint32(vid),
				Collection:     collection,
				ShardIds:       needToCopyEcIndexBits.ToUint32Slice(),
				CopyEcxFile:    false,
				CopyEcjFile:    true,
				CopyVifFile:    true,
				SourceDataNode: string(loc),
			})
			if copyErr != nil {
				return fmt.Errorf("copy %d.%v %s => %s : %v\n", vid, needToCopyEcIndexBits.ShardIds(), loc, targetNodeLocation, copyErr)
			}

			return nil
		})

		if err != nil {
			break
		}

		copiedEcIndexBits = copiedEcIndexBits.Plus(needToCopyEcIndexBits)

	}

	nodeToEcIndexBits[targetNodeLocation] = exisitngEcIndexBits.Plus(copiedEcIndexBits)

	return targetNodeLocation, err

}

func lookupVolumeIds(commandEnv *CommandEnv, volumeIds []string) (volumeIdLocations []*master_pb.LookupVolumeResponse_VolumeIdLocation, err error) {
	var resp *master_pb.LookupVolumeResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.LookupVolume(context.Background(), &master_pb.LookupVolumeRequest{VolumeOrFileIds: volumeIds})
		return err
	})
	if err != nil {
		return nil, err
	}
	return resp.VolumeIdLocations, nil
}

func collectTopologyInfo(commandEnv *CommandEnv) (topoInfo *master_pb.TopologyInfo, volumeSizeLimitMb uint64, err error) {

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return
	}

	return resp.TopologyInfo, resp.VolumeSizeLimitMb, nil

}

func collectEcShardIds(topoInfo *master_pb.TopologyInfo, selectedCollection string) (vids []needle.VolumeId) {

	vidMap := make(map[uint32]bool)
	eachDataNode(topoInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		if diskInfo, found := dn.DiskInfos[string(types.HardDriveType)]; found {
			for _, v := range diskInfo.EcShardInfos {
				if v.Collection == selectedCollection {
					vidMap[v.Id] = true
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}

func collectEcNodeShardBits(topoInfo *master_pb.TopologyInfo, vid needle.VolumeId) map[pb.ServerAddress]erasure_coding.ShardBits {

	nodeToEcIndexBits := make(map[pb.ServerAddress]erasure_coding.ShardBits)
	eachDataNode(topoInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		if diskInfo, found := dn.DiskInfos[string(types.HardDriveType)]; found {
			for _, v := range diskInfo.EcShardInfos {
				if v.Id == uint32(vid) {
					nodeToEcIndexBits[pb.NewServerAddressFromDataNode(dn)] = erasure_coding.ShardBits(v.EcIndexBits)
				}
			}
		}
	})

	return nodeToEcIndexBits
}
<commit_msg>fix compilation<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb"
	"github.com/chrislusf/seaweedfs/weed/storage/types"
	"io"

	"google.golang.org/grpc"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
)

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

type commandEcDecode struct {
}

func (c *commandEcDecode) Name() string {
	return "ec.decode"
}

func (c *commandEcDecode) Help() string {
	return `decode a erasure coded volume into a normal volume

	ec.decode [-collection=""] [-volumeId=<volume_id>]

`
}

func (c *commandEcDecode) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
	encodeCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	volumeId := encodeCommand.Int("volumeId", 0, "the volume id")
	collection := encodeCommand.String("collection", "", "the collection name")
	applyChanges := encodeCommand.Bool("force", false, "force the encoding even if the cluster has less than recommended 4 nodes")
	if err = encodeCommand.Parse(args); err != nil {
		return nil
	}

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

	vid := needle.VolumeId(*volumeId)

	// collect topology information
	topologyInfo, _, err := collectTopologyInfo(commandEnv)
	if err != nil {
		return err
	}

	if !*applyChanges {
		var nodeCount int
		eachDataNode(topologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
			nodeCount++
		})
		if nodeCount < erasure_coding.ParityShardsCount {
			glog.V(0).Infof("skip erasure coding with %d nodes, less than recommended %d nodes", nodeCount, erasure_coding.ParityShardsCount)
			return nil
		}
	}

	// volumeId is provided
	if vid != 0 {
		return doEcDecode(commandEnv, topologyInfo, *collection, vid)
	}

	// apply to all volumes in the collection
	volumeIds := collectEcShardIds(topologyInfo, *collection)
	fmt.Printf("ec encode volumes: %v\n", volumeIds)
	for _, vid := range volumeIds {
		if err = doEcDecode(commandEnv, topologyInfo, *collection, vid); err != nil {
			return err
		}
	}

	return nil
}

func doEcDecode(commandEnv *CommandEnv, topoInfo *master_pb.TopologyInfo, collection string, vid needle.VolumeId) (err error) {
	// find volume location
	nodeToEcIndexBits := collectEcNodeShardBits(topoInfo, vid)

	fmt.Printf("ec volume %d shard locations: %+v\n", vid, nodeToEcIndexBits)

	// collect ec shards to the server with most space
	targetNodeLocation, err := collectEcShards(commandEnv, nodeToEcIndexBits, collection, vid)
	if err != nil {
		return fmt.Errorf("collectEcShards for volume %d: %v", vid, err)
	}

	// generate a normal volume
	err = generateNormalVolume(commandEnv.option.GrpcDialOption, vid, collection, targetNodeLocation)
	if err != nil {
		return fmt.Errorf("generate normal volume %d on %s: %v", vid, targetNodeLocation, err)
	}

	// delete the previous ec shards
	err = mountVolumeAndDeleteEcShards(commandEnv.option.GrpcDialOption, collection, targetNodeLocation, nodeToEcIndexBits, vid)
	if err != nil {
		return fmt.Errorf("delete ec shards for volume %d: %v", vid, err)
	}

	return nil
}

func mountVolumeAndDeleteEcShards(grpcDialOption grpc.DialOption, collection string, targetNodeLocation pb.ServerAddress, nodeToEcIndexBits map[pb.ServerAddress]erasure_coding.ShardBits, vid needle.VolumeId) error {

	// mount volume
	if err := operation.WithVolumeServerClient(targetNodeLocation, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, mountErr := volumeServerClient.VolumeMount(context.Background(), &volume_server_pb.VolumeMountRequest{
			VolumeId: uint32(vid),
		})
		return mountErr
	}); err != nil {
		return fmt.Errorf("mountVolumeAndDeleteEcShards mount volume %d on %s: %v", vid, targetNodeLocation, err)
	}

	// unmount ec shards
	for location, ecIndexBits := range nodeToEcIndexBits {
		fmt.Printf("unmount ec volume %d on %s has shards: %+v\n", vid, location, ecIndexBits.ShardIds())
		err := unmountEcShards(grpcDialOption, vid, location, ecIndexBits.ToUint32Slice())
		if err != nil {
			return fmt.Errorf("mountVolumeAndDeleteEcShards unmount ec volume %d on %s: %v", vid, location, err)
		}
	}
	// delete ec shards
	for location, ecIndexBits := range nodeToEcIndexBits {
		fmt.Printf("delete ec volume %d on %s has shards: %+v\n", vid, location, ecIndexBits.ShardIds())
		err := sourceServerDeleteEcShards(grpcDialOption, collection, vid, location, ecIndexBits.ToUint32Slice())
		if err != nil {
			return fmt.Errorf("mountVolumeAndDeleteEcShards delete ec volume %d on %s: %v", vid, location, err)
		}
	}

	return nil
}

func generateNormalVolume(grpcDialOption grpc.DialOption, vid needle.VolumeId, collection string, sourceVolumeServer pb.ServerAddress) error {

	fmt.Printf("generateNormalVolume from ec volume %d on %s\n", vid, sourceVolumeServer)

	err := operation.WithVolumeServerClient(sourceVolumeServer, grpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
		_, genErr := volumeServerClient.VolumeEcShardsToVolume(context.Background(), &volume_server_pb.VolumeEcShardsToVolumeRequest{
			VolumeId:   uint32(vid),
			Collection: collection,
		})
		return genErr
	})

	return err

}

func collectEcShards(commandEnv *CommandEnv, nodeToEcIndexBits map[pb.ServerAddress]erasure_coding.ShardBits, collection string, vid needle.VolumeId) (targetNodeLocation pb.ServerAddress, err error) {

	maxShardCount := 0
	var exisitngEcIndexBits erasure_coding.ShardBits
	for loc, ecIndexBits := range nodeToEcIndexBits {
		toBeCopiedShardCount := ecIndexBits.MinusParityShards().ShardIdCount()
		if toBeCopiedShardCount > maxShardCount {
			maxShardCount = toBeCopiedShardCount
			targetNodeLocation = loc
			exisitngEcIndexBits = ecIndexBits
		}
	}

	fmt.Printf("collectEcShards: ec volume %d collect shards to %s from: %+v\n", vid, targetNodeLocation, nodeToEcIndexBits)

	var copiedEcIndexBits erasure_coding.ShardBits
	for loc, ecIndexBits := range nodeToEcIndexBits {
		if loc == targetNodeLocation {
			continue
		}

		needToCopyEcIndexBits := ecIndexBits.Minus(exisitngEcIndexBits).MinusParityShards()
		if needToCopyEcIndexBits.ShardIdCount() == 0 {
			continue
		}

		err = operation.WithVolumeServerClient(targetNodeLocation, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {

			fmt.Printf("copy %d.%v %s => %s\n", vid, needToCopyEcIndexBits.ShardIds(), loc, targetNodeLocation)

			_, copyErr := volumeServerClient.VolumeEcShardsCopy(context.Background(), &volume_server_pb.VolumeEcShardsCopyRequest{
				VolumeId:       uint32(vid),
				Collection:     collection,
				ShardIds:       needToCopyEcIndexBits.ToUint32Slice(),
				CopyEcxFile:    false,
				CopyEcjFile:    true,
				CopyVifFile:    true,
				SourceDataNode: string(loc),
			})
			if copyErr != nil {
				return fmt.Errorf("copy %d.%v %s => %s : %v\n", vid, needToCopyEcIndexBits.ShardIds(), loc, targetNodeLocation, copyErr)
			}

			return nil
		})

		if err != nil {
			break
		}

		copiedEcIndexBits = copiedEcIndexBits.Plus(needToCopyEcIndexBits)

	}

	nodeToEcIndexBits[targetNodeLocation] = exisitngEcIndexBits.Plus(copiedEcIndexBits)

	return targetNodeLocation, err

}

func lookupVolumeIds(commandEnv *CommandEnv, volumeIds []string) (volumeIdLocations []*master_pb.LookupVolumeResponse_VolumeIdLocation, err error) {
	var resp *master_pb.LookupVolumeResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.LookupVolume(context.Background(), &master_pb.LookupVolumeRequest{VolumeOrFileIds: volumeIds})
		return err
	})
	if err != nil {
		return nil, err
	}
	return resp.VolumeIdLocations, nil
}

func collectTopologyInfo(commandEnv *CommandEnv) (topoInfo *master_pb.TopologyInfo, volumeSizeLimitMb uint64, err error) {

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return
	}

	return resp.TopologyInfo, resp.VolumeSizeLimitMb, nil

}

func collectEcShardIds(topoInfo *master_pb.TopologyInfo, selectedCollection string) (vids []needle.VolumeId) {

	vidMap := make(map[uint32]bool)
	eachDataNode(topoInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		if diskInfo, found := dn.DiskInfos[string(types.HardDriveType)]; found {
			for _, v := range diskInfo.EcShardInfos {
				if v.Collection == selectedCollection {
					vidMap[v.Id] = true
				}
			}
		}
	})

	for vid := range vidMap {
		vids = append(vids, needle.VolumeId(vid))
	}

	return
}

func collectEcNodeShardBits(topoInfo *master_pb.TopologyInfo, vid needle.VolumeId) map[pb.ServerAddress]erasure_coding.ShardBits {

	nodeToEcIndexBits := make(map[pb.ServerAddress]erasure_coding.ShardBits)
	eachDataNode(topoInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		if diskInfo, found := dn.DiskInfos[string(types.HardDriveType)]; found {
			for _, v := range diskInfo.EcShardInfos {
				if v.Id == uint32(vid) {
					nodeToEcIndexBits[pb.NewServerAddressFromDataNode(dn)] = erasure_coding.ShardBits(v.EcIndexBits)
				}
			}
		}
	})

	return nodeToEcIndexBits
}
<|endoftext|>


<commit_before>package CloudForest

/*
WRFTarget wraps a numerical feature as a target for us weigted random forest.
*/
type WRFTarget struct {
	CatFeature
	Weights []float64
}

/*
NewWRFTarget creates a weighted random forest target and initializes its weights.
*/
func NewWRFTarget(f CatFeature, weights map[string]float64) (abt *WRFTarget) {
	abt = &WRFTarget{f, make([]float64, f.NCats())}

	for i := range abt.Weights {
		abt.Weights[i] = weights[f.NumToCat(i)]
	}

	return
}

/*
SplitImpurity is an weigtedRF version of SplitImpurity.
*/
func (target *WRFTarget) SplitImpurity(l *[]int, r *[]int, m *[]int, allocs *BestSplitAllocs) (impurityDecrease float64) {
	nl := float64(len(*l))
	nr := float64(len(*r))
	nm := 0.0

	impurityDecrease = nl * target.Impurity(l, allocs.LCounter)
	impurityDecrease += nr * target.Impurity(r, allocs.RCounter)
	if m != nil && len(*m) > 0 {
		nm = float64(len(*m))
		impurityDecrease += nm * target.Impurity(m, allocs.Counter)
	}

	impurityDecrease /= nl + nr + nm
	return
}

//UpdateSImpFromAllocs willl be called when splits are being built by moving cases from r to l as in learning from numerical variables.
//Here it just wraps SplitImpurity but it can be implemented to provide further optimization.
func (target *WRFTarget) UpdateSImpFromAllocs(l *[]int, r *[]int, m *[]int, allocs *BestSplitAllocs, movedRtoL *[]int) (impurityDecrease float64) {
	return target.SplitImpurity(l, r, m, allocs)
}

//Impurity is Gini impurity that uses the weights specified in WRFTarget.weights.
func (target *WRFTarget) Impurity(cases *[]int, counter *[]int) (e float64) {
	total := 0.0
	counts := *counter
	for i := range counts {
		counts[i] = 0
	}
	for _, i := range *cases {
		if !target.IsMissing(i) {
			cati := target.Geti(i)
			counts[cati]++
			total += target.Weights[cati]
		}
	}
	e = 1.0
	t := float64(total * total)
	for i, v := range counts {
		w := target.Weights[i]
		e -= float64(v*v) * w * w / t
	}
	return
}

//FindPredicted finds the predicted target as the weighted catagorical Mode.
func (target *WRFTarget) FindPredicted(cases []int) (pred string) {

	counts := make([]int, target.NCats())
	for _, i := range cases {
		if !target.IsMissing(i) {
			counts[target.Geti(i)] += 1
		}

	}
	m := 0
	max := 0.0
	for k, v := range counts {
		val := float64(v) * target.Weights[k]
		if val > max {
			m = k
			max = val
		}
	}

	pred = target.NumToCat(m)

	return

}
<commit_msg>Optimize weighted random forest<commit_after>package CloudForest

/*
WRFTarget wraps a numerical feature as a target for us weigted random forest.
*/
type WRFTarget struct {
	CatFeature
	Weights []float64
}

/*
NewWRFTarget creates a weighted random forest target and initializes its weights.
*/
func NewWRFTarget(f CatFeature, weights map[string]float64) (abt *WRFTarget) {
	abt = &WRFTarget{f, make([]float64, f.NCats())}

	for i := range abt.Weights {
		abt.Weights[i] = weights[f.NumToCat(i)]
	}

	return
}

/*
SplitImpurity is an weigtedRF version of SplitImpurity.
*/
func (target *WRFTarget) SplitImpurity(l *[]int, r *[]int, m *[]int, allocs *BestSplitAllocs) (impurityDecrease float64) {
	nl := float64(len(*l))
	nr := float64(len(*r))
	nm := 0.0

	impurityDecrease = nl * target.Impurity(l, allocs.LCounter)
	impurityDecrease += nr * target.Impurity(r, allocs.RCounter)
	if m != nil && len(*m) > 0 {
		nm = float64(len(*m))
		impurityDecrease += nm * target.Impurity(m, allocs.Counter)
	}

	impurityDecrease /= nl + nr + nm
	return
}

//UpdateSImpFromAllocs willl be called when splits are being built by moving cases from r to l
//to avoid recalulatign the entire split impurity.
func (target *WRFTarget) UpdateSImpFromAllocs(l *[]int, r *[]int, m *[]int, allocs *BestSplitAllocs, movedRtoL *[]int) (impurityDecrease float64) {
	var cat, i int
	lcounter := *allocs.LCounter
	rcounter := *allocs.RCounter
	for _, i = range *movedRtoL {

		//most expensive statement:
		cat = target.Geti(i)
		lcounter[cat]++
		rcounter[cat]--
		//counter[target.Geti(i)]++

	}
	nl := float64(len(*l))
	nr := float64(len(*r))
	nm := 0.0

	impurityDecrease = nl * target.ImpFromCounts(allocs.LCounter)
	impurityDecrease += nr * target.ImpFromCounts(allocs.RCounter)
	if m != nil && len(*m) > 0 {
		nm = float64(len(*m))
		impurityDecrease += nm * target.ImpFromCounts(allocs.Counter)
	}

	impurityDecrease /= nl + nr + nm
	return
}

//Impurity is Gini impurity that uses the weights specified in WRFTarget.weights.
func (target *WRFTarget) Impurity(cases *[]int, counter *[]int) (e float64) {

	target.CountPerCat(cases, counter)

	return target.ImpFromCounts(counter)
}

//ImpFromCounts recalculates gini impurity from class counts for us in intertive updates.
func (target *WRFTarget) ImpFromCounts(counter *[]int) (e float64) {

	total := 0.0
	for i, v := range *counter {
		w := target.Weights[i]
		total += float64(v) * w

		e -= float64(v*v) * w * w
	}

	e /= float64(total * total)
	e++

	return

}

//FindPredicted finds the predicted target as the weighted catagorical Mode.
func (target *WRFTarget) FindPredicted(cases []int) (pred string) {

	counts := make([]int, target.NCats())
	for _, i := range cases {
		if !target.IsMissing(i) {
			counts[target.Geti(i)] += 1
		}

	}
	m := 0
	max := 0.0
	for k, v := range counts {
		val := float64(v) * target.Weights[k]
		if val > max {
			m = k
			max = val
		}
	}

	pred = target.NumToCat(m)

	return

}
<|endoftext|>


<commit_before>/*
RPC with a pool of providers each connected via a web-socket.
*/
package wsrpcpool

// The pool server module

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"golang.org/x/net/websocket"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/rpc"
	"strings"
	"sync"
)

/*
PoolServer is used to listen on a set of web-socket URLs for RPC
providers.
*/
type PoolServer struct {
	Server http.Server
	// Provider name to call channel map
	PoolMap map[string]chan *rpc.Call
	// Default call channel
	DefaultPool chan *rpc.Call
	// Used to signal the pool is listening for incoming connections
	Listening <-chan struct{}
	// Used to signal the pool is listening for incoming connections (pool side)
	listening chan struct{}
	// Path to call channel map
	pathMap map[string]chan *rpc.Call
	// Closer list used in Close()
	cList []io.Closer
	// Pool mutex
	lock *sync.RWMutex
}

var (
	// ErrNoDefaultPool signals that provider isn't found and there is no default pool
	ErrNoDefaultPool = errors.New("No default path is bound")
	// ErrNoCertsParsed signals that no SSL certificates were found in the given file
	ErrNoCertsParsed = errors.New("No certificates parsed")
)

/*
NewPool returns a plain PoolServer instance.
*/
func NewPool() *PoolServer {
	listening := make(chan struct{}, 1)
	return &PoolServer{
		Listening: listening,
		listening: listening,
		cList:     make([]io.Closer, 0),
		lock:      &sync.RWMutex{},
	}
}

/*
NewPoolTLS returns a PoolServer instance equipped with the given
SSL certificate.
*/
func NewPoolTLS(certfile, keyfile string) (*PoolServer, error) {
	pool := NewPool()
	if err := pool.AppendCertificate(certfile, keyfile); err != nil {
		return nil, err
	}
	return pool, nil
}

/*
NewPoolTLSAuth returns a PoolServer instance equipped with the given
SSL certificate and a root CA certificates for client authentication.
*/
func NewPoolTLSAuth(certfile, keyfile string, clientCAs ...string) (*PoolServer, error) {
	pool, err := NewPoolTLS(certfile, keyfile)
	if err != nil {
		return nil, err
	}
	err = pool.AppendClientCAs(clientCAs...)
	if err != nil {
		return nil, err
	}
	return pool, err
}

/*
AppendCertificate appends an SSL certificate to the set of server
certificates loading it from the pair of public certificate and private
key files.
*/
func (pool *PoolServer) AppendCertificate(certfile, keyfile string) error {
	if pool.Server.TLSConfig == nil {
		pool.Server.TLSConfig = &tls.Config{}
	}
	return appendCertificate(pool.Server.TLSConfig, certfile, keyfile)
}

/*
appendCertificate appends an SSL certificate to the given tls.Config
loading it from the pair of public certificate and private key files.
*/
func appendCertificate(tlsConfig *tls.Config, certfile, keyfile string) error {
	cert, err := tls.LoadX509KeyPair(certfile, keyfile)
	if err != nil {
		return err
	}
	tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
	tlsConfig.BuildNameToCertificate()
	return nil
}

/*
AppendClientCAs appends the given SSL root CA certificate files to the
set of client CAs to verify client connections against.
*/
func (pool *PoolServer) AppendClientCAs(clientCAs ...string) error {
	if len(clientCAs) == 0 {
		return nil
	}
	if pool.Server.TLSConfig == nil {
		pool.Server.TLSConfig = &tls.Config{}
	}
	if pool.Server.TLSConfig.ClientCAs == nil {
		pool.Server.TLSConfig.ClientCAs = x509.NewCertPool()
	}
	err := appendCAs(pool.Server.TLSConfig.ClientCAs, clientCAs...)
	if err == nil {
		pool.Server.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
	}
	return err
}

/*
appendCAs appends the given SSL root CA certificate files to the
given CA pool.
*/
func appendCAs(caPool *x509.CertPool, caCerts ...string) error {
	for _, caFile := range caCerts {
		caCert, err := ioutil.ReadFile(caFile)
		if err != nil {
			return err
		}
		if !caPool.AppendCertsFromPEM(caCert) {
			return ErrNoCertsParsed
		}
	}
	return nil
}

/*
invoke passes the given call to the client.
*/
func invoke(client *rpc.Client, call *rpc.Call) *rpc.Call {
	return client.Go(call.ServiceMethod, call.Args, call.Reply, call.Done)
}

/*
connObserver used to observe I/O errors in a websocket connection.
*/
type connObserver struct {
	*websocket.Conn
	ioError chan error
}

/*
reportError sends the given error over the ioError channel
if there is a free slot available and do nothing otherwise
(i.e. non blocking).
*/
func (conn *connObserver) reportError(err error) {
	select {
	case conn.ioError <- err:
	default:
	}
}

/*
Read implements io.Reader.
*/
func (conn *connObserver) Read(p []byte) (n int, err error) {
	n, err = conn.Conn.Read(p)
	if err != nil {
		conn.reportError(err)
	}
	return
}

/*
Write implements io.Writer.
*/
func (conn *connObserver) Write(p []byte) (n int, err error) {
	n, err = conn.Conn.Write(p)
	if err != nil {
		conn.reportError(err)
	}
	return
}

/*
handle returns and invoker() function casted to the
websocket.Handler type in order to get the necessary websocket
handshake behavior. The invoker function is wrapped call
to addCloser() to register the connection with the pool.
*/
func (pool *PoolServer) handle(newClient func(conn io.ReadWriteCloser) *rpc.Client, callIn <-chan *rpc.Call) websocket.Handler {
	_invoker := invoker(newClient, callIn, nil)
	return websocket.Handler(func(ws *websocket.Conn) {
		pool.addCloser(ws)
		_invoker(ws)
	})
}

/*
addCloser adds the given connection or a listener to the
set of opened objects.
*/
func (pool *PoolServer) addCloser(c io.Closer) {
	pool.lock.Lock()
	pool.cList = append(pool.cList, c)
	pool.lock.Unlock()
}

/*
invoker returns a function that the passes calls from the
given channel over a websocket connection. In the case of I/O error
it is written to errOut channel if it is provided and the function
returns. The function also returns if callIn channel is closed.
No error is sent in that case. The errOut, if provied, is anyway
closed on return.
*/
func invoker(newClient func(conn io.ReadWriteCloser) *rpc.Client, callIn <-chan *rpc.Call, errOut chan<- error) func(ws *websocket.Conn) {
	return func(ws *websocket.Conn) {
		conn := &connObserver{ws, make(chan error, 10)}
		client := newClient(conn)
		defer client.Close()
		defer func() {
			if errOut != nil {
				close(errOut)
			}
		}()
	loop:
		for {
			select {
			case c, ok := <-callIn:
				if !ok {
					break loop
				}
				invoke(client, c)
			case err := <-conn.ioError:
				if errOut != nil {
					errOut <- err
				}
				break loop
			}
		}
	}
}

/*
assertMux checks for pool.Server.Handler mux and makes
one if it doesn't yet exist.
*/
func (pool *PoolServer) assertMux() *http.ServeMux {
	if pool.Server.Handler == nil {
		pool.Server.Handler = http.NewServeMux()
	}
	return pool.Server.Handler.(*http.ServeMux)
}

/*
Bind makes all the connections to the given path to be
considered as means to make calls to the named providers.
If no providers are specified the connections at the path
are considered the default providers. The expected RPC
protocol is what is used by rpc.NewClient().
*/
func (pool *PoolServer) Bind(path string, providers ...string) {
	pool.BindProto(path, rpc.NewClient, providers...)
}

/*
Bind associates the given path with the particular RPC protocol
client. If no providers are specified the connections at the path
are considered the default providers.
*/
func (pool *PoolServer) BindProto(path string, newClient func(conn io.ReadWriteCloser) *rpc.Client, providers ...string) {
	pool.lock.Lock()
	mux := pool.assertMux()
	if pool.pathMap == nil {
		pool.pathMap = make(map[string]chan *rpc.Call)
	}
	callIn := pool.pathMap[path]
	if callIn == nil {
		callIn = make(chan *rpc.Call)
		pool.pathMap[path] = callIn
	}
	if len(providers) > 0 {
		if pool.PoolMap == nil {
			pool.PoolMap = make(map[string]chan *rpc.Call)
		}
		for _, name := range providers {
			pool.PoolMap[name] = callIn
		}
	} else {
		pool.DefaultPool = callIn
	}
	mux.Handle(path, pool.handle(newClient, callIn))
	pool.lock.Unlock()
}

/*
handleIn returns the websocket.Handler that the serves
incoming RPC calls over a websocket connection using
the given handler.
*/
func (pool *PoolServer) handleIn(serveConn func(conn io.ReadWriteCloser)) websocket.Handler {
	return websocket.Handler(func(ws *websocket.Conn) {
		pool.addCloser(ws)
		serveConn(ws)
	})
}

/*
BindIn handles incoming RPC calls on the given path.
The expected RPC protocol is what is used by rpc.ServeConn().
*/
func (pool *PoolServer) BindIn(path string) {
	pool.BindProtoIn(path, rpc.ServeConn)
}

/*
BindIn handles all the connections to the given path with the
given RPC protocol handler.
*/
func (pool *PoolServer) BindProtoIn(path string, serveConn func(conn io.ReadWriteCloser)) {
	pool.lock.Lock()
	mux := pool.assertMux()
	mux.Handle(path, pool.handleIn(serveConn))
	pool.lock.Unlock()
}

/*
listnObserver wraps a net.Listener providing a special
channel to signal the server pool.Close() was called.
*/
type listnObserver struct {
	net.Listener
	closed chan struct{}
	wg     sync.WaitGroup
}

/*
Close calls Close() on the embedded Listener and aloso closes
the "closed" channel to signal Close() was called on the pool.
*/
func (lo *listnObserver) Close() error {
	close(lo.closed)
	err := lo.Listener.Close()
	lo.wg.Wait()
	return err
}

/*
listen returns the active listener for the current pool config
and an error if any. It also send a signal over the "listening"
channel.
*/
func (pool *PoolServer) listen(addr string, tlsConfig *tls.Config) (*listnObserver, error) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	if tlsConfig != nil {
		l = tls.NewListener(l, tlsConfig)
	}

	lo := &listnObserver{Listener: l, closed: make(chan struct{})}
	pool.addCloser(lo)

	select {
	case pool.listening <- struct{}{}:
	default:
	}

	return lo, nil
}

/*
use uses the given listener waiting for a signal on
the "stop" channel.
*/
func (pool *PoolServer) use(lo *listnObserver) error {
	lo.wg.Add(1)
	err := pool.Server.Serve(lo.Listener)
	lo.wg.Done()
	select {
	case _, opened := <-lo.closed:
		if !opened {
			err = nil // closed by pool.Close()
		}
	default:
	}
	return err
}

/*
ListenAndUse listens the given (or configured if "" is given) address
([host]:port) with no SSL encryption.
*/
func (pool *PoolServer) ListenAndUse(addr string) error {
	if addr == "" {
		addr = pool.Server.Addr
	}
	l, err := pool.listen(addr, nil)
	if err != nil {
		return err
	}
	return pool.use(l)
}

/*
ListenAndUseTLS listens the listens the given (or configured if ""
is given) address ([host]:port) with SSL encryption on.
*/
func (pool *PoolServer) ListenAndUseTLS(addr string) error {
	if addr == "" {
		addr = pool.Server.Addr
	}
	l, err := pool.listen(addr, pool.Server.TLSConfig)
	if err != nil {
		return err
	}
	return pool.use(l)
}

/*
Close closes the pool listener.
*/
func (pool *PoolServer) Close() error {
	var err error

	pool.lock.Lock()
	for i := range pool.cList {
		switch c := pool.cList[i].(type) {
		case *websocket.Conn:
			c.Close() // skip socket error
		default:
			if _err := c.Close(); _err != nil {
				if err == nil {
					err = _err
				}
			}
		}
		pool.cList[i] = nil
	}
	pool.cList = pool.cList[:0]
	pool.lock.Unlock()

	return err
}

/*
Go invokes the given remote function asynchronously. The name of the
provider (if given as the first part of serviceMethod, i.e. "Provider.Function")
is first searched in the PoolMap and the DefaultPool is used if it isn't there
(or isn't specified). If "done" is nil, a new channel is allocated and passed in
the return value. See net/rpc package for details.
*/
func (pool *PoolServer) Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpc.Call) (*rpc.Call, error) {
	var callIn chan *rpc.Call
	if split := strings.SplitN(serviceMethod, ".", 2); len(split) > 1 {
		callIn = pool.PoolMap[split[0]]
	}
	if callIn == nil {
		callIn = pool.DefaultPool
	}
	if callIn == nil {
		return nil, ErrNoDefaultPool
	}

	call := &rpc.Call{
		ServiceMethod: serviceMethod,
		Args:          args,
		Reply:         reply,
	}
	if done == nil {
		done = make(chan *rpc.Call, 1)
	}
	call.Done = done

	callIn <- call
	return call, nil
}

/*
Call invokes the given remote function and waits for it to complete,
returning its error status. If an I/O error encountered, then the
function re-queues the call.
*/
func (pool *PoolServer) Call(serviceMethod string, args interface{}, reply interface{}) error {
	for {
		if call, err := pool.Go(serviceMethod, args, reply, nil); err == nil {
			call = <-call.Done
			switch call.Error {
			case rpc.ErrShutdown, io.ErrUnexpectedEOF:
			default:
				return call.Error
			}
		} else {
			return err
		}
	}
}
<commit_msg>BindProto => BindWith<commit_after>/*
RPC with a pool of providers each connected via a web-socket.
*/
package wsrpcpool

// The pool server module

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"golang.org/x/net/websocket"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/rpc"
	"strings"
	"sync"
)

/*
PoolServer is used to listen on a set of web-socket URLs for RPC
providers.
*/
type PoolServer struct {
	Server http.Server
	// Provider name to call channel map
	PoolMap map[string]chan *rpc.Call
	// Default call channel
	DefaultPool chan *rpc.Call
	// Used to signal the pool is listening for incoming connections
	Listening <-chan struct{}
	// Used to signal the pool is listening for incoming connections (pool side)
	listening chan struct{}
	// Path to call channel map
	pathMap map[string]chan *rpc.Call
	// Closer list used in Close()
	cList []io.Closer
	// Pool mutex
	lock *sync.RWMutex
}

var (
	// ErrNoDefaultPool signals that provider isn't found and there is no default pool
	ErrNoDefaultPool = errors.New("No default path is bound")
	// ErrNoCertsParsed signals that no SSL certificates were found in the given file
	ErrNoCertsParsed = errors.New("No certificates parsed")
)

/*
NewPool returns a plain PoolServer instance.
*/
func NewPool() *PoolServer {
	listening := make(chan struct{}, 1)
	return &PoolServer{
		Listening: listening,
		listening: listening,
		cList:     make([]io.Closer, 0),
		lock:      &sync.RWMutex{},
	}
}

/*
NewPoolTLS returns a PoolServer instance equipped with the given
SSL certificate.
*/
func NewPoolTLS(certfile, keyfile string) (*PoolServer, error) {
	pool := NewPool()
	if err := pool.AppendCertificate(certfile, keyfile); err != nil {
		return nil, err
	}
	return pool, nil
}

/*
NewPoolTLSAuth returns a PoolServer instance equipped with the given
SSL certificate and a root CA certificates for client authentication.
*/
func NewPoolTLSAuth(certfile, keyfile string, clientCAs ...string) (*PoolServer, error) {
	pool, err := NewPoolTLS(certfile, keyfile)
	if err != nil {
		return nil, err
	}
	err = pool.AppendClientCAs(clientCAs...)
	if err != nil {
		return nil, err
	}
	return pool, err
}

/*
AppendCertificate appends an SSL certificate to the set of server
certificates loading it from the pair of public certificate and private
key files.
*/
func (pool *PoolServer) AppendCertificate(certfile, keyfile string) error {
	if pool.Server.TLSConfig == nil {
		pool.Server.TLSConfig = &tls.Config{}
	}
	return appendCertificate(pool.Server.TLSConfig, certfile, keyfile)
}

/*
appendCertificate appends an SSL certificate to the given tls.Config
loading it from the pair of public certificate and private key files.
*/
func appendCertificate(tlsConfig *tls.Config, certfile, keyfile string) error {
	cert, err := tls.LoadX509KeyPair(certfile, keyfile)
	if err != nil {
		return err
	}
	tlsConfig.Certificates = append(tlsConfig.Certificates, cert)
	tlsConfig.BuildNameToCertificate()
	return nil
}

/*
AppendClientCAs appends the given SSL root CA certificate files to the
set of client CAs to verify client connections against.
*/
func (pool *PoolServer) AppendClientCAs(clientCAs ...string) error {
	if len(clientCAs) == 0 {
		return nil
	}
	if pool.Server.TLSConfig == nil {
		pool.Server.TLSConfig = &tls.Config{}
	}
	if pool.Server.TLSConfig.ClientCAs == nil {
		pool.Server.TLSConfig.ClientCAs = x509.NewCertPool()
	}
	err := appendCAs(pool.Server.TLSConfig.ClientCAs, clientCAs...)
	if err == nil {
		pool.Server.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
	}
	return err
}

/*
appendCAs appends the given SSL root CA certificate files to the
given CA pool.
*/
func appendCAs(caPool *x509.CertPool, caCerts ...string) error {
	for _, caFile := range caCerts {
		caCert, err := ioutil.ReadFile(caFile)
		if err != nil {
			return err
		}
		if !caPool.AppendCertsFromPEM(caCert) {
			return ErrNoCertsParsed
		}
	}
	return nil
}

/*
invoke passes the given call to the client.
*/
func invoke(client *rpc.Client, call *rpc.Call) *rpc.Call {
	return client.Go(call.ServiceMethod, call.Args, call.Reply, call.Done)
}

/*
connObserver used to observe I/O errors in a websocket connection.
*/
type connObserver struct {
	*websocket.Conn
	ioError chan error
}

/*
reportError sends the given error over the ioError channel
if there is a free slot available and do nothing otherwise
(i.e. non blocking).
*/
func (conn *connObserver) reportError(err error) {
	select {
	case conn.ioError <- err:
	default:
	}
}

/*
Read implements io.Reader.
*/
func (conn *connObserver) Read(p []byte) (n int, err error) {
	n, err = conn.Conn.Read(p)
	if err != nil {
		conn.reportError(err)
	}
	return
}

/*
Write implements io.Writer.
*/
func (conn *connObserver) Write(p []byte) (n int, err error) {
	n, err = conn.Conn.Write(p)
	if err != nil {
		conn.reportError(err)
	}
	return
}

/*
handle returns and invoker() function casted to the
websocket.Handler type in order to get the necessary websocket
handshake behavior. The invoker function is wrapped call
to addCloser() to register the connection with the pool.
*/
func (pool *PoolServer) handle(newClient func(conn io.ReadWriteCloser) *rpc.Client, callIn <-chan *rpc.Call) websocket.Handler {
	_invoker := invoker(newClient, callIn, nil)
	return websocket.Handler(func(ws *websocket.Conn) {
		pool.addCloser(ws)
		_invoker(ws)
	})
}

/*
addCloser adds the given connection or a listener to the
set of opened objects.
*/
func (pool *PoolServer) addCloser(c io.Closer) {
	pool.lock.Lock()
	pool.cList = append(pool.cList, c)
	pool.lock.Unlock()
}

/*
invoker returns a function that the passes calls from the
given channel over a websocket connection. In the case of I/O error
it is written to errOut channel if it is provided and the function
returns. The function also returns if callIn channel is closed.
No error is sent in that case. The errOut, if provied, is anyway
closed on return.
*/
func invoker(newClient func(conn io.ReadWriteCloser) *rpc.Client, callIn <-chan *rpc.Call, errOut chan<- error) func(ws *websocket.Conn) {
	return func(ws *websocket.Conn) {
		conn := &connObserver{ws, make(chan error, 10)}
		client := newClient(conn)
		defer client.Close()
		defer func() {
			if errOut != nil {
				close(errOut)
			}
		}()
	loop:
		for {
			select {
			case c, ok := <-callIn:
				if !ok {
					break loop
				}
				invoke(client, c)
			case err := <-conn.ioError:
				if errOut != nil {
					errOut <- err
				}
				break loop
			}
		}
	}
}

/*
assertMux checks for pool.Server.Handler mux and makes
one if it doesn't yet exist.
*/
func (pool *PoolServer) assertMux() *http.ServeMux {
	if pool.Server.Handler == nil {
		pool.Server.Handler = http.NewServeMux()
	}
	return pool.Server.Handler.(*http.ServeMux)
}

/*
Bind makes all the connections to the given path to be
considered as means to make calls to the named providers.
If no providers are specified the connections at the path
are considered the default providers. The expected RPC
protocol is what is used by rpc.NewClient().
*/
func (pool *PoolServer) Bind(path string, providers ...string) {
	pool.BindWith(path, rpc.NewClient, providers...)
}

/*
BindWith associates the given path with the particular RPC protocol
client. If no providers are specified the connections at the path
are considered the default providers.
*/
func (pool *PoolServer) BindWith(path string, newClient func(conn io.ReadWriteCloser) *rpc.Client, providers ...string) {
	pool.lock.Lock()
	mux := pool.assertMux()
	if pool.pathMap == nil {
		pool.pathMap = make(map[string]chan *rpc.Call)
	}
	callIn := pool.pathMap[path]
	if callIn == nil {
		callIn = make(chan *rpc.Call)
		pool.pathMap[path] = callIn
	}
	if len(providers) > 0 {
		if pool.PoolMap == nil {
			pool.PoolMap = make(map[string]chan *rpc.Call)
		}
		for _, name := range providers {
			pool.PoolMap[name] = callIn
		}
	} else {
		pool.DefaultPool = callIn
	}
	mux.Handle(path, pool.handle(newClient, callIn))
	pool.lock.Unlock()
}

/*
handleIn returns the websocket.Handler that the serves
incoming RPC calls over a websocket connection using
the given handler.
*/
func (pool *PoolServer) handleIn(serveConn func(conn io.ReadWriteCloser)) websocket.Handler {
	return websocket.Handler(func(ws *websocket.Conn) {
		pool.addCloser(ws)
		serveConn(ws)
	})
}

/*
BindIn handles incoming RPC calls on the given path.
The expected RPC protocol is what is used by rpc.ServeConn().
*/
func (pool *PoolServer) BindIn(path string) {
	pool.BindInWith(path, rpc.ServeConn)
}

/*
BindInWith handles all the connections to the given path with the
given RPC protocol handler.
*/
func (pool *PoolServer) BindInWith(path string, serveConn func(conn io.ReadWriteCloser)) {
	pool.lock.Lock()
	mux := pool.assertMux()
	mux.Handle(path, pool.handleIn(serveConn))
	pool.lock.Unlock()
}

/*
listnObserver wraps a net.Listener providing a special
channel to signal the server pool.Close() was called.
*/
type listnObserver struct {
	net.Listener
	closed chan struct{}
	wg     sync.WaitGroup
}

/*
Close calls Close() on the embedded Listener and aloso closes
the "closed" channel to signal Close() was called on the pool.
*/
func (lo *listnObserver) Close() error {
	close(lo.closed)
	err := lo.Listener.Close()
	lo.wg.Wait()
	return err
}

/*
listen returns the active listener for the current pool config
and an error if any. It also send a signal over the "listening"
channel.
*/
func (pool *PoolServer) listen(addr string, tlsConfig *tls.Config) (*listnObserver, error) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}

	if tlsConfig != nil {
		l = tls.NewListener(l, tlsConfig)
	}

	lo := &listnObserver{Listener: l, closed: make(chan struct{})}
	pool.addCloser(lo)

	select {
	case pool.listening <- struct{}{}:
	default:
	}

	return lo, nil
}

/*
use uses the given listener waiting for a signal on
the "stop" channel.
*/
func (pool *PoolServer) use(lo *listnObserver) error {
	lo.wg.Add(1)
	err := pool.Server.Serve(lo.Listener)
	lo.wg.Done()
	select {
	case _, opened := <-lo.closed:
		if !opened {
			err = nil // closed by pool.Close()
		}
	default:
	}
	return err
}

/*
ListenAndUse listens the given (or configured if "" is given) address
([host]:port) with no SSL encryption.
*/
func (pool *PoolServer) ListenAndUse(addr string) error {
	if addr == "" {
		addr = pool.Server.Addr
	}
	l, err := pool.listen(addr, nil)
	if err != nil {
		return err
	}
	return pool.use(l)
}

/*
ListenAndUseTLS listens the listens the given (or configured if ""
is given) address ([host]:port) with SSL encryption on.
*/
func (pool *PoolServer) ListenAndUseTLS(addr string) error {
	if addr == "" {
		addr = pool.Server.Addr
	}
	l, err := pool.listen(addr, pool.Server.TLSConfig)
	if err != nil {
		return err
	}
	return pool.use(l)
}

/*
Close closes the pool listener.
*/
func (pool *PoolServer) Close() error {
	var err error

	pool.lock.Lock()
	for i := range pool.cList {
		switch c := pool.cList[i].(type) {
		case *websocket.Conn:
			c.Close() // skip socket error
		default:
			if _err := c.Close(); _err != nil {
				if err == nil {
					err = _err
				}
			}
		}
		pool.cList[i] = nil
	}
	pool.cList = pool.cList[:0]
	pool.lock.Unlock()

	return err
}

/*
Go invokes the given remote function asynchronously. The name of the
provider (if given as the first part of serviceMethod, i.e. "Provider.Function")
is first searched in the PoolMap and the DefaultPool is used if it isn't there
(or isn't specified). If "done" is nil, a new channel is allocated and passed in
the return value. See net/rpc package for details.
*/
func (pool *PoolServer) Go(serviceMethod string, args interface{}, reply interface{}, done chan *rpc.Call) (*rpc.Call, error) {
	var callIn chan *rpc.Call
	if split := strings.SplitN(serviceMethod, ".", 2); len(split) > 1 {
		callIn = pool.PoolMap[split[0]]
	}
	if callIn == nil {
		callIn = pool.DefaultPool
	}
	if callIn == nil {
		return nil, ErrNoDefaultPool
	}

	call := &rpc.Call{
		ServiceMethod: serviceMethod,
		Args:          args,
		Reply:         reply,
	}
	if done == nil {
		done = make(chan *rpc.Call, 1)
	}
	call.Done = done

	callIn <- call
	return call, nil
}

/*
Call invokes the given remote function and waits for it to complete,
returning its error status. If an I/O error encountered, then the
function re-queues the call.
*/
func (pool *PoolServer) Call(serviceMethod string, args interface{}, reply interface{}) error {
	for {
		if call, err := pool.Go(serviceMethod, args, reply, nil); err == nil {
			call = <-call.Done
			switch call.Error {
			case rpc.ErrShutdown, io.ErrUnexpectedEOF:
			default:
				return call.Error
			}
		} else {
			return err
		}
	}
}
<|endoftext|>


<commit_before>/* RPC with a pool of providers each connected via web-socket. */
package wsrpcpool

import (
	"crypto/tls"
	"crypto/x509"
	//"golang.org/x/net/websocket"
	"io/ioutil"
	"net/http"
	//"net/rpc/jsonrpc"
)

/* A Web-socket connection endpoint for RPC provideres. */
type PoolServer struct {
	http.Server
}

/* Plain PoolServer instance. */
func NewPool() *PoolServer {
	return &PoolServer{}
}

/* PoolServer instance equipped with the given SSL certificate. */
func NewPoolTLS(certfile, keyfile string) (*PoolServer, error) {
	pool := NewPool()
	if err := pool.AppendCertificate(certfile, keyfile); err != nil {
		return nil, err
	}
	return pool, nil
}

/* Appends an SSL certificate to the set of server certificates loading
it from the pair of public certificate and private key files. */
func (pool *PoolServer) AppendCertificate(certfile, keyfile string) error {
	if pool.TLSConfig == nil {
		pool.TLSConfig = &tls.Config{}
	}

	cert, err := tls.LoadX509KeyPair(certfile, keyfile)
	if err != nil {
		return err
	}

	pool.TLSConfig.Certificates = append(pool.TLSConfig.Certificates, cert)
	pool.TLSConfig.BuildNameToCertificate()

	return nil
}

/* Appends the given SSL root CA certificate files to the set of client
CAs to verify client connections against. */
func (pool *PoolServer) AppendClientCAs(clientCAs ...string) error {
	if pool.TLSConfig == nil {
		pool.TLSConfig = &tls.Config{}
	}
	if pool.TLSConfig.ClientCAs == nil {
		pool.TLSConfig.ClientCAs = x509.NewCertPool()
	}

	for _, caFile := range clientCAs {
		caCert, err := ioutil.ReadFile(caFile)
		if err != nil {
			return err
		}
		pool.TLSConfig.ClientCAs.AppendCertsFromPEM(caCert)
	}

	pool.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
	return nil
}

/* PoolServer instance equipped with the given SSL certificate
and a root CA certificates for client authentication. */
func NewPoolTLSAuth(certfile, keyfile string, clientCAs ...string) (*PoolServer, error) {
	pool, err := NewPoolTLS(certfile, keyfile)
	if err != nil {
		return nil, err
	}

	err = pool.AppendClientCAs(clientCAs...)
	if err != nil {
		return nil, err
	}

	return pool, err
}
<commit_msg>Implement the pool map<commit_after>/* RPC with a pool of providers each connected via web-socket. */
package wsrpcpool

import (
	"crypto/tls"
	"crypto/x509"
	"golang.org/x/net/websocket"
	"io/ioutil"
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync"
)

/* A Web-socket connection endpoint for RPC provideres. */
type PoolServer struct {
	http.Server
	PoolMap     map[string]chan *rpc.Call
	DefaultPool chan *rpc.Call
}

/* Plain PoolServer instance. */
func NewPool() *PoolServer {
	return &PoolServer{}
}

/* PoolServer instance equipped with the given SSL certificate. */
func NewPoolTLS(certfile, keyfile string) (*PoolServer, error) {
	pool := NewPool()
	if err := pool.AppendCertificate(certfile, keyfile); err != nil {
		return nil, err
	}
	return pool, nil
}

/* Appends an SSL certificate to the set of server certificates loading
it from the pair of public certificate and private key files. */
func (pool *PoolServer) AppendCertificate(certfile, keyfile string) error {
	if pool.TLSConfig == nil {
		pool.TLSConfig = &tls.Config{}
	}

	cert, err := tls.LoadX509KeyPair(certfile, keyfile)
	if err != nil {
		return err
	}

	pool.TLSConfig.Certificates = append(pool.TLSConfig.Certificates, cert)
	pool.TLSConfig.BuildNameToCertificate()

	return nil
}

/* Appends the given SSL root CA certificate files to the set of client
CAs to verify client connections against. */
func (pool *PoolServer) AppendClientCAs(clientCAs ...string) error {
	if pool.TLSConfig == nil {
		pool.TLSConfig = &tls.Config{}
	}
	if pool.TLSConfig.ClientCAs == nil {
		pool.TLSConfig.ClientCAs = x509.NewCertPool()
	}

	for _, caFile := range clientCAs {
		caCert, err := ioutil.ReadFile(caFile)
		if err != nil {
			return err
		}
		pool.TLSConfig.ClientCAs.AppendCertsFromPEM(caCert)
	}

	pool.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
	return nil
}

/* PoolServer instance equipped with the given SSL certificate
and a root CA certificates for client authentication. */
func NewPoolTLSAuth(certfile, keyfile string, clientCAs ...string) (*PoolServer, error) {
	pool, err := NewPoolTLS(certfile, keyfile)
	if err != nil {
		return nil, err
	}

	err = pool.AppendClientCAs(clientCAs...)
	if err != nil {
		return nil, err
	}

	return pool, err
}

/* Passes the given call to the client. */
func invoke(client *rpc.Client, call *rpc.Call) *rpc.Call {
	return client.Go(call.ServiceMethod, call.Args, call.Reply, call.Done)
}

/* Associates the given path with the set of remote objects or make
it the default path if no object names given. */
func (pool *PoolServer) BindPath(path string, names ...string) {
	if pool.Handler == nil {
		pool.Handler = http.NewServeMux()
	}

	mux := pool.Handler.(*http.ServeMux)

	if len(names) > 0 {
		if pool.PoolMap == nil {
			pool.PoolMap = make(map[string]chan *rpc.Call)
		}
		for _, name := range names {
			if pool.PoolMap[name] == nil {
				pool.PoolMap[name] = make(chan *rpc.Call)
			}
		}
		mux.Handle(path, websocket.Handler(func(ws *websocket.Conn) {
			var wg sync.WaitGroup
			client := jsonrpc.NewClient(ws)
			for _, name := range names {
				callIn := pool.PoolMap[name]
				go func() {
					defer wg.Done()
					for c := range callIn {
						invoke(client, c)
					}
				}()
			}
			wg.Add(len(names))
			wg.Wait()
			client.Close()
		}))
	} else {
		if pool.DefaultPool == nil {
			pool.DefaultPool = make(chan *rpc.Call)
		}
		chanIn := pool.DefaultPool
		mux.Handle(path, websocket.Handler(func(ws *websocket.Conn) {
			client := jsonrpc.NewClient(ws)
			for c := range chanIn {
				invoke(client, c)
			}
			client.Close()
		}))
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"strings"

	"github.com/github/hub/github"
	"github.com/github/hub/ui"
	"github.com/github/hub/utils"
)

var (
	cmdGist = &Command{
		Run: printGistHelp,
		Usage: `
gist create [--public] [<FILES>...]
gist show <ID> [<FILENAME>]
`,
		Long: `Create and print GitHub Gists

## Commands:

	* _create_:
		Create a new gist. If no <FILES> are specified, the content is read from
		standard input.

    * _show_:
		Print the contents of a gist. If the gist contains multiple files, the
		operation will error out unless <FILENAME> is specified.

## Options:

    --public
        Make the new gist public (default: false).

	-o, --browse
		Open the new gist in a web browser.

	-c, --copy
		Put the URL of the new gist to clipboard instead of printing it.

## Examples:

	$ echo hello | hub gist create --public

	$ hub gist create <file1> <file2>

    # print a specific file within a gist:
    $ hub gist show <ID> testfile1.txt

## See also:

hub(1), hub-api(1)
`,
	}

	cmdShowGist = &Command{
		Key: "show",
		Run: showGist,
	}

	cmdCreateGist = &Command{
		Key: "create",
		Run: createGist,
		KnownFlags: `
		--public
		-o, --browse
		-c, --copy
`,
	}
)

func init() {
	cmdGist.Use(cmdShowGist)
	cmdGist.Use(cmdCreateGist)
	CmdRunner.Use(cmdGist)
}

func getGist(gh *github.Client, id string, filename string) error {
	gist, err := gh.FetchGist(id)
	if err != nil {
		return err
	}

	if len(gist.Files) > 1 && filename == "" {
		filenames := []string{}
		for name := range gist.Files {
			filenames = append(filenames, name)
		}
		return fmt.Errorf("the gist contains multiple files, you must specify one:\n%s", strings.Join(filenames, "\n"))
	}

	if filename != "" {
		if val, ok := gist.Files[filename]; ok {
			ui.Println(val.Content)
		} else {
			return fmt.Errorf("no such file in gist")
		}
	} else {
		for name := range gist.Files {
			file := gist.Files[name]
			ui.Println(file.Content)
		}
	}
	return nil
}

func printGistHelp(command *Command, args *Args) {
	utils.Check(command.UsageError(""))
}

func createGist(cmd *Command, args *Args) {
	args.NoForward()

	host, err := github.CurrentConfig().DefaultHostNoPrompt()
	utils.Check(err)
	gh := github.NewClient(host.Host)

	filenames := []string{}
	if args.IsParamsEmpty() {
		filenames = append(filenames, "-")
	} else {
		filenames = args.Params
	}

	var gist *github.Gist
	if args.Noop {
		ui.Println("Would create gist")
		gist = &github.Gist{
			HtmlUrl: fmt.Sprintf("https://gist.%s/%s", gh.Host.Host, "ID"),
		}
	} else {
		gist, err = gh.CreateGist(filenames, args.Flag.Bool("--public"))
		utils.Check(err)
	}

	flagIssueBrowse := args.Flag.Bool("--browse")
	flagIssueCopy := args.Flag.Bool("--copy")
	printBrowseOrCopy(args, gist.HtmlUrl, flagIssueBrowse, flagIssueCopy)
}

func showGist(cmd *Command, args *Args) {
	args.NoForward()
	if args.ParamsSize() < 1 {
		utils.Check(cmd.UsageError("you must specify a gist ID"))
	}
	id := args.GetParam(0)
	filename := ""
	if args.ParamsSize() > 1 {
		filename = args.GetParam(1)
	}

	host, err := github.CurrentConfig().DefaultHostNoPrompt()
	utils.Check(err)
	gh := github.NewClient(host.Host)
	err = getGist(gh, id, filename)
	utils.Check(err)
}
<commit_msg>[gist] List shorthand flags in usage synopsis<commit_after>package commands

import (
	"fmt"
	"strings"

	"github.com/github/hub/github"
	"github.com/github/hub/ui"
	"github.com/github/hub/utils"
)

var (
	cmdGist = &Command{
		Run: printGistHelp,
		Usage: `
gist create [-oc] [--public] [<FILES>...]
gist show <ID> [<FILENAME>]
`,
		Long: `Create and print GitHub Gists

## Commands:

	* _create_:
		Create a new gist. If no <FILES> are specified, the content is read from
		standard input.

    * _show_:
		Print the contents of a gist. If the gist contains multiple files, the
		operation will error out unless <FILENAME> is specified.

## Options:

    --public
        Make the new gist public (default: false).

	-o, --browse
		Open the new gist in a web browser.

	-c, --copy
		Put the URL of the new gist to clipboard instead of printing it.

## Examples:

	$ echo hello | hub gist create --public

	$ hub gist create <file1> <file2>

    # print a specific file within a gist:
    $ hub gist show <ID> testfile1.txt

## See also:

hub(1), hub-api(1)
`,
	}

	cmdShowGist = &Command{
		Key: "show",
		Run: showGist,
	}

	cmdCreateGist = &Command{
		Key: "create",
		Run: createGist,
		KnownFlags: `
		--public
		-o, --browse
		-c, --copy
`,
	}
)

func init() {
	cmdGist.Use(cmdShowGist)
	cmdGist.Use(cmdCreateGist)
	CmdRunner.Use(cmdGist)
}

func getGist(gh *github.Client, id string, filename string) error {
	gist, err := gh.FetchGist(id)
	if err != nil {
		return err
	}

	if len(gist.Files) > 1 && filename == "" {
		filenames := []string{}
		for name := range gist.Files {
			filenames = append(filenames, name)
		}
		return fmt.Errorf("the gist contains multiple files, you must specify one:\n%s", strings.Join(filenames, "\n"))
	}

	if filename != "" {
		if val, ok := gist.Files[filename]; ok {
			ui.Println(val.Content)
		} else {
			return fmt.Errorf("no such file in gist")
		}
	} else {
		for name := range gist.Files {
			file := gist.Files[name]
			ui.Println(file.Content)
		}
	}
	return nil
}

func printGistHelp(command *Command, args *Args) {
	utils.Check(command.UsageError(""))
}

func createGist(cmd *Command, args *Args) {
	args.NoForward()

	host, err := github.CurrentConfig().DefaultHostNoPrompt()
	utils.Check(err)
	gh := github.NewClient(host.Host)

	filenames := []string{}
	if args.IsParamsEmpty() {
		filenames = append(filenames, "-")
	} else {
		filenames = args.Params
	}

	var gist *github.Gist
	if args.Noop {
		ui.Println("Would create gist")
		gist = &github.Gist{
			HtmlUrl: fmt.Sprintf("https://gist.%s/%s", gh.Host.Host, "ID"),
		}
	} else {
		gist, err = gh.CreateGist(filenames, args.Flag.Bool("--public"))
		utils.Check(err)
	}

	flagIssueBrowse := args.Flag.Bool("--browse")
	flagIssueCopy := args.Flag.Bool("--copy")
	printBrowseOrCopy(args, gist.HtmlUrl, flagIssueBrowse, flagIssueCopy)
}

func showGist(cmd *Command, args *Args) {
	args.NoForward()
	if args.ParamsSize() < 1 {
		utils.Check(cmd.UsageError("you must specify a gist ID"))
	}
	id := args.GetParam(0)
	filename := ""
	if args.ParamsSize() > 1 {
		filename = args.GetParam(1)
	}

	host, err := github.CurrentConfig().DefaultHostNoPrompt()
	utils.Check(err)
	gh := github.NewClient(host.Host)
	err = getGist(gh, id, filename)
	utils.Check(err)
}
<|endoftext|>


<commit_before>package beertasting

import (
	"appengine"
	"appengine/datastore"
	"appengine/urlfetch"
	"appengine/user"
	"encoding/json"
	"fmt"
	"github.com/ant0ine/go-json-rest/rest"
	"html/template"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"time"
)

func init() {
	http.HandleFunc("/", handler)
	http.HandleFunc("/search", searchHandler)
	http.HandleFunc("/feed", feedHandler)
	http.HandleFunc("/displayFeed", displayFeedHandler)
	http.HandleFunc("/oauth/untappd", oauthUntappdHandler)
	restHandler := rest.ResourceHandler{}
	restHandler.SetRoutes(
		&rest.Route{"GET", "/admin/config", getAdminConfig},
		&rest.Route{"PUT", "/admin/config", putAdminConfig},
	)
	http.Handle("/admin/config", &restHandler)
}

var (
	endpoint = url.URL{Scheme: "http", Host: "api.untappd.com", Path: "v4"}
)

type Config struct {
	ClientId     string
	ClientSecret string
	Whitelist    []string
}

func configKey(c appengine.Context) *datastore.Key {
	return datastore.NewKey(c, "Config", "default", 0, nil)
}

func getConfig(c appengine.Context) (Config, error) {
	var cfg Config
	err := datastore.Get(c, configKey(c), &cfg)
	return cfg, err
}

func httpCallback(c appengine.Context, path string) *url.URL {
	return &url.URL{
		Scheme: "http",
		Host:   appengine.DefaultVersionHostname(c),
		Path:   path,
	}
}

func oauthCallback(c appengine.Context, svc string) *url.URL {
	return httpCallback(c, fmt.Sprintf("oauth/%s", svc))
}

func userLoggedIn(c appengine.Context, curUrl *url.URL, w http.ResponseWriter) (*user.User, bool) {
	u := user.Current(c)
	if u != nil {
		return u, true
	}
	url, err := user.LoginURL(c, curUrl.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil, false
	}
	w.Header().Set("Location", url)
	w.WriteHeader(http.StatusFound)
	return nil, false
}

func handler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	fmt.Fprintf(w, "Welcome, %s", user)
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	var err error
	var config Config
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if len(config.Whitelist) == 0 {
		c.Infof("whitelist not found, using test@example.com")
		config.Whitelist = []string{"test@example.com"}
	}
	found := false
	for _, record := range config.Whitelist {
		if record == user.Email {
			found = true
			break
		}
	}
	if !found {
		http.Error(w, fmt.Sprintf("user %s not in whitelist", user), http.StatusInternalServerError)
		return
	}

	t, err := template.ParseFiles("templates/trial1.html")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	s := struct{ Name, Endpoint, ClientId, ClientSecret string }{
		user.String(),
		endpoint.String(),
		config.ClientId,
		config.ClientSecret,
	}
	if err := t.Execute(w, s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func oauthUntappdHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	_, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	if len(r.FormValue("code")) == 0 {
		http.Error(w, "missing code parameter", http.StatusInternalServerError)
		return
	}
	var config Config
	var err error
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	u := url.URL{Scheme: "https", Host: "untappd.com", Path: "oauth/authorize/"}
	q := u.Query()
	q.Add("client_id", config.ClientId)
	q.Add("client_secret", config.ClientSecret)
	q.Add("response_type", "code")
	q.Add("code", r.FormValue("code"))
	q.Add("redirect_url", oauthCallback(c, "untappd").String())
	u.RawQuery = q.Encode()
	c.Infof("authorize URL: %s", u.String())
	resp, err := urlfetch.Client(c).Get(u.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	buf, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	oauthResponse := struct {
		Response struct {
			AccessToken string `json:"access_token"`
		}
	}{}
	err = json.Unmarshal(buf, &oauthResponse)
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(buf))
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	expire := time.Now().AddDate(0, 0, 1)
	hostname := appengine.DefaultVersionHostname(c)
	raw := fmt.Sprintf("access_token=%s", oauthResponse.Response.AccessToken)
	cookie := http.Cookie{
		Name:       "access_token",
		Value:      oauthResponse.Response.AccessToken,
		Path:       "/",
		Domain:     hostname,
		Expires:    expire,
		RawExpires: expire.Format(time.UnixDate),
		MaxAge:     86400,
		Secure:     false,
		HttpOnly:   false,
		Raw:        raw,
		Unparsed:   []string{raw},
	}
	http.SetCookie(w, &cookie)
	http.Redirect(w, r, "/displayFeed", http.StatusFound)
}

func displayFeedHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	t, err := template.ParseFiles("templates/feed.html")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	endpoint.Path = path.Join(endpoint.Path, "checkin/recent")
	s := struct{ Name, FeedRequest string }{user.String(), endpoint.String()}
	if err := t.Execute(w, s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func feedHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	_, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	var config Config
	var err error
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	var untappdOath url.URL
	untappdOath.Scheme = "https"
	untappdOath.Host = "untappd.com"
	untappdOath.Path = "oauth/authenticate/"
	q := untappdOath.Query()
	q.Add("client_id", config.ClientId)
	q.Add("response_type", "code")
	q.Add("redirect_url", oauthCallback(c, "untappd").String())
	untappdOath.RawQuery = q.Encode()
	http.Redirect(w, r, untappdOath.String(), http.StatusFound)
	return
}

func writeJson(w rest.ResponseWriter, v interface{}) {
	if err := w.WriteJson(v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func datastoreRestGet(c appengine.Context, k *datastore.Key, w rest.ResponseWriter, v interface{}) {
	if err := datastore.Get(c, k, v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	writeJson(w, v)
}

func datastoreRestPut(c appengine.Context, k *datastore.Key, w rest.ResponseWriter, v interface{}) {
	if _, err := datastore.Put(c, k, v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	writeJson(w, v)
}

func putAdminConfigCtx(c appengine.Context, w rest.ResponseWriter, r *rest.Request) {
	var config Config
	err := r.DecodeJsonPayload(&config)
	if err != nil {
		err = fmt.Errorf("DecodeJsonPayload(): %s", err)
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	datastoreRestPut(c, configKey(c), w, &config)
}

func putAdminConfig(w rest.ResponseWriter, r *rest.Request) {
	c := appengine.NewContext(r.Request)
	putAdminConfigCtx(c, w, r)
}

func getAdminConfigCtx(c appengine.Context, w rest.ResponseWriter) {
	var config Config
	datastoreRestGet(c, configKey(c), w, &config)
}

func getAdminConfig(w rest.ResponseWriter, r *rest.Request) {
	c := appengine.NewContext(r.Request)
	getAdminConfigCtx(c, w)
}
<commit_msg>add a basic RESTful API for users<commit_after>package beertasting

import (
	"appengine"
	"appengine/datastore"
	"appengine/urlfetch"
	"appengine/user"
	"encoding/json"
	"fmt"
	"github.com/ant0ine/go-json-rest/rest"
	"html/template"
	"io/ioutil"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"time"
)

func init() {
	http.HandleFunc("/", handler)
	http.HandleFunc("/search", searchHandler)
	http.HandleFunc("/feed", feedHandler)
	http.HandleFunc("/displayFeed", displayFeedHandler)
	http.HandleFunc("/oauth/untappd", oauthUntappdHandler)
	restHandler := rest.ResourceHandler{}
	restHandler.SetRoutes(
		&rest.Route{"GET", "/admin/config", getAdminConfig},
		&rest.Route{"PUT", "/admin/config", putAdminConfig},
		&rest.Route{"GET", "/users", getAllUsers},
		&rest.Route{"POST", "/users", postUser},
		&rest.Route{"GET", "/users/:id", getUser},
		&rest.Route{"DELETE", "/users/:id", deleteUser},
	)
	http.Handle("/admin/config", &restHandler)
	http.Handle("/users", &restHandler)
	http.Handle("/users/", &restHandler)
}

var (
	endpoint = url.URL{Scheme: "http", Host: "api.untappd.com", Path: "v4"}
)

type Config struct {
	ClientId     string
	ClientSecret string
	Whitelist    []string
}

type User struct {
	ID    int64 `datastore:"-"`
	Name  string
	Email string
}

func (user *User) DecodeJsonPayload(r *rest.Request) error {
	err := r.DecodeJsonPayload(user)
	if err != nil {
		return err
	}
	if user.Name == "" {
		return fmt.Errorf("name required")
	}
	if user.Email == "" {
		return fmt.Errorf("email required")
	}
	return nil
}

func configKey(c appengine.Context) *datastore.Key {
	return datastore.NewKey(c, "Config", "default", 0, nil)
}

func getConfig(c appengine.Context) (Config, error) {
	var cfg Config
	err := datastore.Get(c, configKey(c), &cfg)
	return cfg, err
}

func httpCallback(c appengine.Context, path string) *url.URL {
	return &url.URL{
		Scheme: "http",
		Host:   appengine.DefaultVersionHostname(c),
		Path:   path,
	}
}

func oauthCallback(c appengine.Context, svc string) *url.URL {
	return httpCallback(c, fmt.Sprintf("oauth/%s", svc))
}

func userLoggedIn(c appengine.Context, curUrl *url.URL, w http.ResponseWriter) (*user.User, bool) {
	u := user.Current(c)
	if u != nil {
		return u, true
	}
	url, err := user.LoginURL(c, curUrl.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return nil, false
	}
	w.Header().Set("Location", url)
	w.WriteHeader(http.StatusFound)
	return nil, false
}

func handler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	fmt.Fprintf(w, "Welcome, %s", user)
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	var err error
	var config Config
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if len(config.Whitelist) == 0 {
		c.Infof("whitelist not found, using test@example.com")
		config.Whitelist = []string{"test@example.com"}
	}
	found := false
	for _, record := range config.Whitelist {
		if record == user.Email {
			found = true
			break
		}
	}
	if !found {
		http.Error(w, fmt.Sprintf("user %s not in whitelist", user), http.StatusInternalServerError)
		return
	}

	t, err := template.ParseFiles("templates/trial1.html")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	s := struct{ Name, Endpoint, ClientId, ClientSecret string }{
		user.String(),
		endpoint.String(),
		config.ClientId,
		config.ClientSecret,
	}
	if err := t.Execute(w, s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func oauthUntappdHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	_, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	if len(r.FormValue("code")) == 0 {
		http.Error(w, "missing code parameter", http.StatusInternalServerError)
		return
	}
	var config Config
	var err error
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	u := url.URL{Scheme: "https", Host: "untappd.com", Path: "oauth/authorize/"}
	q := u.Query()
	q.Add("client_id", config.ClientId)
	q.Add("client_secret", config.ClientSecret)
	q.Add("response_type", "code")
	q.Add("code", r.FormValue("code"))
	q.Add("redirect_url", oauthCallback(c, "untappd").String())
	u.RawQuery = q.Encode()
	c.Infof("authorize URL: %s", u.String())
	resp, err := urlfetch.Client(c).Get(u.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	buf, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	oauthResponse := struct {
		Response struct {
			AccessToken string `json:"access_token"`
		}
	}{}
	err = json.Unmarshal(buf, &oauthResponse)
	if err != nil {
		err = fmt.Errorf("%s: %s", err.Error(), string(buf))
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	expire := time.Now().AddDate(0, 0, 1)
	hostname := appengine.DefaultVersionHostname(c)
	raw := fmt.Sprintf("access_token=%s", oauthResponse.Response.AccessToken)
	cookie := http.Cookie{
		Name:       "access_token",
		Value:      oauthResponse.Response.AccessToken,
		Path:       "/",
		Domain:     hostname,
		Expires:    expire,
		RawExpires: expire.Format(time.UnixDate),
		MaxAge:     86400,
		Secure:     false,
		HttpOnly:   false,
		Raw:        raw,
		Unparsed:   []string{raw},
	}
	http.SetCookie(w, &cookie)
	http.Redirect(w, r, "/displayFeed", http.StatusFound)
}

func displayFeedHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	user, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	t, err := template.ParseFiles("templates/feed.html")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	endpoint.Path = path.Join(endpoint.Path, "checkin/recent")
	s := struct{ Name, FeedRequest string }{user.String(), endpoint.String()}
	if err := t.Execute(w, s); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func feedHandler(w http.ResponseWriter, r *http.Request) {
	c := appengine.NewContext(r)
	_, ok := userLoggedIn(c, r.URL, w)
	if !ok {
		return
	}
	var config Config
	var err error
	if config, err = getConfig(c); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	var untappdOath url.URL
	untappdOath.Scheme = "https"
	untappdOath.Host = "untappd.com"
	untappdOath.Path = "oauth/authenticate/"
	q := untappdOath.Query()
	q.Add("client_id", config.ClientId)
	q.Add("response_type", "code")
	q.Add("redirect_url", oauthCallback(c, "untappd").String())
	untappdOath.RawQuery = q.Encode()
	http.Redirect(w, r, untappdOath.String(), http.StatusFound)
	return
}

func writeJson(w rest.ResponseWriter, v interface{}) {
	if err := w.WriteJson(v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func datastoreRestGet(c appengine.Context, k *datastore.Key, w rest.ResponseWriter, v interface{}) {
	if err := datastore.Get(c, k, v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	writeJson(w, v)
}

func datastoreRestPut(c appengine.Context, k *datastore.Key, w rest.ResponseWriter, v interface{}) {
	if _, err := datastore.Put(c, k, v); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	writeJson(w, v)
}

func putAdminConfigCtx(c appengine.Context, w rest.ResponseWriter, r *rest.Request) {
	var config Config
	err := r.DecodeJsonPayload(&config)
	if err != nil {
		err = fmt.Errorf("DecodeJsonPayload(): %s", err)
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	datastoreRestPut(c, configKey(c), w, &config)
}

func putAdminConfig(w rest.ResponseWriter, r *rest.Request) {
	c := appengine.NewContext(r.Request)
	putAdminConfigCtx(c, w, r)
}

func getAdminConfigCtx(c appengine.Context, w rest.ResponseWriter) {
	var config Config
	datastoreRestGet(c, configKey(c), w, &config)
}

func getAdminConfig(w rest.ResponseWriter, r *rest.Request) {
	c := appengine.NewContext(r.Request)
	getAdminConfigCtx(c, w)
}

func getUser(w rest.ResponseWriter, r *rest.Request) {
	id, err := strconv.Atoi(r.PathParam("id"))
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	c := appengine.NewContext(r.Request)
	key := datastore.NewKey(c, "User", "", int64(id), nil)
	var user User
	datastoreRestGet(c, key, w, &user)
}

func getAllUsers(w rest.ResponseWriter, r *rest.Request) {
	users := []User{}
	c := appengine.NewContext(r.Request)
	q := datastore.NewQuery("User")
	for t := q.Run(c); ; {
		var u User
		key, err := t.Next(&u)
		if err == datastore.Done {
			break
		}
		if err != nil {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		u.ID = key.IntID()
		users = append(users, u)
	}
	w.WriteJson(&users)
}

func postUser(w rest.ResponseWriter, r *rest.Request) {
	user := User{}
	err := user.DecodeJsonPayload(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	c := appengine.NewContext(r.Request)
	key := datastore.NewIncompleteKey(c, "User", nil)
	newKey, err := datastore.Put(c, key, &user)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	user.ID = newKey.IntID()
	writeJson(w, user)
}

func deleteUser(w rest.ResponseWriter, r *rest.Request) {
	c := appengine.NewContext(r.Request)
	id, err := strconv.Atoi(r.PathParam("id"))
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	key := datastore.NewKey(c, "User", "", int64(id), nil)
	err = datastore.Delete(c, key)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

import (
	"fmt"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/log"
	"os"
	"path"
	"path/filepath"
	"time"
)

// Deployer maintains a git repository tracking a series of charm versions,
// and can install and upgrade charm deployments to the current version.
type Deployer struct {
	path    string
	current *GitDir
}

// NewDeployer creates a new Deployer which stores its state in the supplied
// directory.
func NewDeployer(path string) *Deployer {
	return &Deployer{
		path:    path,
		current: NewGitDir(filepath.Join(path, "current")),
	}
}

// Stage causes subsequent calls to Deploy to deploy the supplied charm.
func (d *Deployer) Stage(bun *charm.Bundle, url *charm.URL) error {
	// Read present state of current.
	if err := os.MkdirAll(d.path, 0755); err != nil {
		return err
	}
	defer d.collectOrphans()
	srcExists, err := d.current.Exists()
	if err != nil {
		return err
	}
	if srcExists {
		prevURL, err := ReadCharmURL(d.current)
		if err != nil {
			return err
		}
		if *url == *prevURL {
			return nil
		}
	}

	// Prepare a fresh repository for the update, using current's history
	// if it exists.
	updatePath, err := d.newDir("update")
	if err != nil {
		return err
	}
	var repo *GitDir
	if srcExists {
		repo, err = d.current.Clone(updatePath)
	} else {
		repo = NewGitDir(updatePath)
		err = repo.Init()
	}
	if err != nil {
		return err
	}

	// Write the desired new state and commit.
	if err = bun.ExpandTo(updatePath); err != nil {
		return err
	}
	if err = WriteCharmURL(repo, url); err != nil {
		return err
	}
	if err = repo.Snapshotf("Imported charm %q from %q.", url, bun.Path); err != nil {
		return err
	}

	// Atomically rename fresh repository to current.
	tmplink := filepath.Join(updatePath, "tmplink")
	if err = os.Symlink(updatePath, tmplink); err != nil {
		return err
	}
	return os.Rename(tmplink, d.current.Path())
}

// Deploy deploys the current charm to the target directory.
func (d *Deployer) Deploy(target *GitDir) (err error) {
	defer func() {
		if err == ErrConflict {
			log.Warningf("worker/uniter/charm: charm deployment completed with conflicts")
		} else if err != nil {
			err = fmt.Errorf("charm deployment failed: %s", err)
			log.Errorf("worker/uniter/charm: %v", err)
		} else {
			log.Infof("worker/uniter/charm: charm deployment succeeded")
		}
	}()
	if exists, err := d.current.Exists(); err != nil {
		return err
	} else if !exists {
		return fmt.Errorf("no charm set")
	}
	if exists, err := target.Exists(); err != nil {
		return err
	} else if !exists {
		return d.install(target)
	}
	return d.upgrade(target)
}

// install creates a new deployment of current, and atomically moves it to
// target.
func (d *Deployer) install(target *GitDir) error {
	defer d.collectOrphans()
	log.Infof("worker/uniter/charm: preparing new charm deployment")
	url, err := ReadCharmURL(d.current)
	if err != nil {
		return err
	}
	installPath, err := d.newDir("install")
	if err != nil {
		return err
	}
	repo := NewGitDir(installPath)
	if err = repo.Init(); err != nil {
		return err
	}
	if err = repo.Pull(d.current); err != nil {
		return err
	}
	if err = repo.Snapshotf("Deployed charm %q.", url); err != nil {
		return err
	}
	log.Infof("worker/uniter/charm: deploying charm")
	return os.Rename(installPath, target.Path())
}

// upgrade pulls from current into target. If target has local changes, but
// no conflicts, it will be snapshotted before any changes are made.
func (d *Deployer) upgrade(target *GitDir) error {
	log.Infof("worker/uniter/charm: preparing charm upgrade")
	url, err := ReadCharmURL(d.current)
	if err != nil {
		return err
	}
	if err := target.Init(); err != nil {
		return err
	}
	if dirty, err := target.Dirty(); err != nil {
		return err
	} else if dirty {
		if conflicted, err := target.Conflicted(); err != nil {
			return err
		} else if !conflicted {
			log.Infof("worker/uniter/charm: snapshotting dirty charm before upgrade")
			if err = target.Snapshotf("Pre-upgrade snapshot."); err != nil {
				return err
			}
		}
	}
	log.Infof("worker/uniter/charm: deploying charm")
	if err := target.Pull(d.current); err != nil {
		return err
	}
	return target.Snapshotf("Upgraded charm to %q.", url)
}

// collectOrphans deletes all repos in path except the one pointed to by current.
// Errors are generally ignored; some are logged.
func (d *Deployer) collectOrphans() {
	current, err := os.Readlink(d.current.Path())
	if err != nil {
		return
	}
	filepath.Walk(d.path, func(repoPath string, fi os.FileInfo, err error) error {
		if err != nil && repoPath != d.path && repoPath != current {
			if err = os.RemoveAll(repoPath); err != nil {
				log.Warningf("worker/uniter/charm: failed to remove orphan repo at %s: %s", repoPath, err)
			}
		}
		return err
	})
}

// newDir creates a new timestamped directory with the given prefix. It
// assumes that the deployer will not need to create more than 10
// directories in any given second.
func (d *Deployer) newDir(prefix string) (string, error) {
	prefix = prefix + time.Now().Format("-%Y%m%d-%H%M%S")
	var err error
	var prefixPath string
	for i := 0; i < 10; i++ {
		prefixPath = path.Join(d.path, fmt.Sprintf("%s-%d", prefix, i))
		if err = os.Mkdir(prefixPath, 0755); err == nil {
			return prefixPath, nil
		} else if !os.IsExist(err) {
			break
		}
	}
	return "", fmt.Errorf("failed to create %q: %v", prefixPath, err)
}
<commit_msg>worker/uniter/charm: use ioutil.TempDir and correct Time.Format directive<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package charm

import (
	"fmt"
	"io/ioutil"
	"launchpad.net/juju-core/charm"
	"launchpad.net/juju-core/log"
	"os"
	"path/filepath"
	"time"
)

// Deployer maintains a git repository tracking a series of charm versions,
// and can install and upgrade charm deployments to the current version.
type Deployer struct {
	path    string
	current *GitDir
}

// NewDeployer creates a new Deployer which stores its state in the supplied
// directory.
func NewDeployer(path string) *Deployer {
	return &Deployer{
		path:    path,
		current: NewGitDir(filepath.Join(path, "current")),
	}
}

// Stage causes subsequent calls to Deploy to deploy the supplied charm.
func (d *Deployer) Stage(bun *charm.Bundle, url *charm.URL) error {
	// Read present state of current.
	if err := os.MkdirAll(d.path, 0755); err != nil {
		return err
	}
	defer d.collectOrphans()
	srcExists, err := d.current.Exists()
	if err != nil {
		return err
	}
	if srcExists {
		prevURL, err := ReadCharmURL(d.current)
		if err != nil {
			return err
		}
		if *url == *prevURL {
			return nil
		}
	}

	// Prepare a fresh repository for the update, using current's history
	// if it exists.
	updatePath, err := d.newDir("update")
	if err != nil {
		return err
	}
	var repo *GitDir
	if srcExists {
		repo, err = d.current.Clone(updatePath)
	} else {
		repo = NewGitDir(updatePath)
		err = repo.Init()
	}
	if err != nil {
		return err
	}

	// Write the desired new state and commit.
	if err = bun.ExpandTo(updatePath); err != nil {
		return err
	}
	if err = WriteCharmURL(repo, url); err != nil {
		return err
	}
	if err = repo.Snapshotf("Imported charm %q from %q.", url, bun.Path); err != nil {
		return err
	}

	// Atomically rename fresh repository to current.
	tmplink := filepath.Join(updatePath, "tmplink")
	if err = os.Symlink(updatePath, tmplink); err != nil {
		return err
	}
	return os.Rename(tmplink, d.current.Path())
}

// Deploy deploys the current charm to the target directory.
func (d *Deployer) Deploy(target *GitDir) (err error) {
	defer func() {
		if err == ErrConflict {
			log.Warningf("worker/uniter/charm: charm deployment completed with conflicts")
		} else if err != nil {
			err = fmt.Errorf("charm deployment failed: %s", err)
			log.Errorf("worker/uniter/charm: %v", err)
		} else {
			log.Infof("worker/uniter/charm: charm deployment succeeded")
		}
	}()
	if exists, err := d.current.Exists(); err != nil {
		return err
	} else if !exists {
		return fmt.Errorf("no charm set")
	}
	if exists, err := target.Exists(); err != nil {
		return err
	} else if !exists {
		return d.install(target)
	}
	return d.upgrade(target)
}

// install creates a new deployment of current, and atomically moves it to
// target.
func (d *Deployer) install(target *GitDir) error {
	defer d.collectOrphans()
	log.Infof("worker/uniter/charm: preparing new charm deployment")
	url, err := ReadCharmURL(d.current)
	if err != nil {
		return err
	}
	installPath, err := d.newDir("install")
	if err != nil {
		return err
	}
	repo := NewGitDir(installPath)
	if err = repo.Init(); err != nil {
		return err
	}
	if err = repo.Pull(d.current); err != nil {
		return err
	}
	if err = repo.Snapshotf("Deployed charm %q.", url); err != nil {
		return err
	}
	log.Infof("worker/uniter/charm: deploying charm")
	return os.Rename(installPath, target.Path())
}

// upgrade pulls from current into target. If target has local changes, but
// no conflicts, it will be snapshotted before any changes are made.
func (d *Deployer) upgrade(target *GitDir) error {
	log.Infof("worker/uniter/charm: preparing charm upgrade")
	url, err := ReadCharmURL(d.current)
	if err != nil {
		return err
	}
	if err := target.Init(); err != nil {
		return err
	}
	if dirty, err := target.Dirty(); err != nil {
		return err
	} else if dirty {
		if conflicted, err := target.Conflicted(); err != nil {
			return err
		} else if !conflicted {
			log.Infof("worker/uniter/charm: snapshotting dirty charm before upgrade")
			if err = target.Snapshotf("Pre-upgrade snapshot."); err != nil {
				return err
			}
		}
	}
	log.Infof("worker/uniter/charm: deploying charm")
	if err := target.Pull(d.current); err != nil {
		return err
	}
	return target.Snapshotf("Upgraded charm to %q.", url)
}

// collectOrphans deletes all repos in path except the one pointed to by current.
// Errors are generally ignored; some are logged.
func (d *Deployer) collectOrphans() {
	current, err := os.Readlink(d.current.Path())
	if err != nil {
		return
	}
	filepath.Walk(d.path, func(repoPath string, fi os.FileInfo, err error) error {
		if err != nil && repoPath != d.path && repoPath != current {
			if err = os.RemoveAll(repoPath); err != nil {
				log.Warningf("worker/uniter/charm: failed to remove orphan repo at %s: %s", repoPath, err)
			}
		}
		return err
	})
}

// newDir creates a new timestamped directory with the given prefix. It
// assumes that the deployer will not need to create more than 10
// directories in any given second.
func (d *Deployer) newDir(prefix string) (string, error) {
	return ioutil.TempDir(d.path, prefix+time.Now().Format("-20060102-150405"))
}
<|endoftext|>


<commit_before>// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
	"encoding/json"
	"testing"
)

func TestQueryStringQuery(t *testing.T) {
	q := NewQueryStringQuery(`this AND that OR thus`)
	q = q.DefaultField("content")
	src, err := q.Source()
	if err != nil {
		t.Fatal(err)
	}
	data, err := json.Marshal(src)
	if err != nil {
		t.Fatalf("marshaling to JSON failed: %v", err)
	}
	got := string(data)
	expected := `{"query_string":{"default_field":"content","query":"this AND that OR thus"}}`
	if got != expected {
		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
	}
}

func TestQueryStringQueryTimeZone(t *testing.T) {
	q := NewQueryStringQuery(`tweet_date:[2015-01-01 TO 2017-12-31]`)
	q = q.TimeZone("Europe/Berlin")
	src, err := q.Source()
	if err != nil {
		t.Fatal(err)
	}
	data, err := json.Marshal(src)
	if err != nil {
		t.Fatalf("marshaling to JSON failed: %v", err)
	}
	got := string(data)
	expected := `{"query_string":{"query":"tweet_date:[2015-01-01 TO 2017-12-31]","time_zone":"Europe/Berlin"}}`
	if got != expected {
		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
	}
}
<commit_msg>Add integration test for QueryStringQuery<commit_after>// Copyright 2012-present Oliver Eilhard. All rights reserved.
// Use of this source code is governed by a MIT-license.
// See http://olivere.mit-license.org/license.txt for details.

package elastic

import (
	"context"
	"encoding/json"
	"testing"
)

func TestQueryStringQuery(t *testing.T) {
	q := NewQueryStringQuery(`this AND that OR thus`)
	q = q.DefaultField("content")
	src, err := q.Source()
	if err != nil {
		t.Fatal(err)
	}
	data, err := json.Marshal(src)
	if err != nil {
		t.Fatalf("marshaling to JSON failed: %v", err)
	}
	got := string(data)
	expected := `{"query_string":{"default_field":"content","query":"this AND that OR thus"}}`
	if got != expected {
		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
	}
}

func TestQueryStringQueryTimeZone(t *testing.T) {
	q := NewQueryStringQuery(`tweet_date:[2015-01-01 TO 2017-12-31]`)
	q = q.TimeZone("Europe/Berlin")
	src, err := q.Source()
	if err != nil {
		t.Fatal(err)
	}
	data, err := json.Marshal(src)
	if err != nil {
		t.Fatalf("marshaling to JSON failed: %v", err)
	}
	got := string(data)
	expected := `{"query_string":{"query":"tweet_date:[2015-01-01 TO 2017-12-31]","time_zone":"Europe/Berlin"}}`
	if got != expected {
		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
	}
}

func TestQueryStringQueryIntegration(t *testing.T) {
	client := setupTestClientAndCreateIndexAndAddDocs(t) //, SetTraceLog(log.New(os.Stdout, "", log.LstdFlags)))

	q := NewQueryStringQuery("Golang")

	// Match all should return all documents
	searchResult, err := client.Search().
		Index(testIndexName).
		Query(q).
		Pretty(true).
		Do(context.TODO())
	if err != nil {
		t.Fatal(err)
	}
	if searchResult.Hits == nil {
		t.Errorf("expected SearchResult.Hits != nil; got nil")
	}
	if got, want := searchResult.Hits.TotalHits, int64(1); got != want {
		t.Errorf("expected SearchResult.Hits.TotalHits = %d; got %d", want, got)
	}
	if got, want := len(searchResult.Hits.Hits), 1; got != want {
		t.Errorf("expected len(SearchResult.Hits.Hits) = %d; got %d", want, got)
	}
}
<|endoftext|>