<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// TODO(shyamjvs): Make this exporter work for master too, currently facing
// gcloud auth error when run from within a pod on the master.

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"github.com/golang/glog"
)

// Initialize the log exporter's configuration related flags.
var (
	nodeName             = flag.String("node-name", "", "Name of the node this log exporter is running on")
	gcsPath              = flag.String("gcs-path", "", "Path to the GCS directory under which to upload logs, for eg: gs://my-logs-bucket/logs")
	cloudProvider        = flag.String("cloud-provider", "", "Cloud provider for this node (gce/gke/aws/kubemark/..)")
	gcloudAuthFilePath   = flag.String("gcloud-auth-file-path", "/etc/service-account/service-account.json", "Path to gcloud service account file, for authenticating gsutil to write to GCS bucket")
	enableHollowNodeLogs = flag.Bool("enable-hollow-node-logs", false, "Enable uploading hollow node logs too. Relevant only for kubemark nodes")
	sleepDuration        = flag.Duration("sleep-duration", 60*time.Second, "Duration to sleep before exiting with success. Useful for making pods schedule with hard anti-affinity when run as a job on a k8s cluster")
)

var (
	localLogPath = "/var/log"

	// Node-type specific logfiles.
	// Currently we only handle nodes, and neglect master.
	nodeLogs = []string{"kube-proxy", "node-problem-detector", "fluentd"}

	// Cloud provider specific logfiles.
	awsLogs      = []string{"cloud-init-output"}
	gceLogs      = []string{"startupscript"}
	kubemarkLogs = []string{"*-hollow-node-*"}

	// System services/kernel related logfiles.
	kernelLog            = "kern"
	initdLogs            = []string{"docker"}
	supervisordLogs      = []string{"kubelet", "supervisor/supervisord", "supervisor/kubelet-stdout", "supervisor/kubelet-stderr", "supervisor/docker-stdout", "supervisor/docker-stderr"}
	systemdServices      = []string{"kern", "kubelet", "docker"}
	systemdSetupServices = []string{"kube-node-installation", "kube-node-configuration"}
	nodeSystemdServices  = []string{"node-problem-detector"}
)

// Check if the config provided through the flags take valid values.
func checkConfigValidity() error {
	glog.Info("Verifying if a valid config has been provided through the flags")
	if *nodeName == "" {
		return fmt.Errorf("Flag --node-name has its value unspecified")
	}
	if *gcsPath == "" {
		return fmt.Errorf("Flag --gcs-path has its value unspecified")
	}
	if _, err := os.Stat(*gcloudAuthFilePath); err != nil {
		return fmt.Errorf("Could not find the gcloud service account file: %v", err)
	} else {
		cmd := exec.Command("gcloud", "auth", "activate-service-account", "--key-file="+*gcloudAuthFilePath)
		if err := cmd.Run(); err != nil {
			return fmt.Errorf("Failed to activate gcloud service account: %v", err)
		}
	}
	return nil
}

// Create logfile for systemd service in outputDir with the given journalctl outputMode.
func createSystemdLogfile(service string, outputMode string, outputDir string) error {
	// Generate the journalctl command.
	journalCmdArgs := []string{fmt.Sprintf("--output=%v", outputMode), "-D", "/var/log/journal"}
	if service == "kern" {
		journalCmdArgs = append(journalCmdArgs, "-k")
	} else {
		journalCmdArgs = append(journalCmdArgs, "-u", fmt.Sprintf("%v.service", service))
	}
	cmd := exec.Command("journalctl", journalCmdArgs...)

	// Run the command and record the output to a file.
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("Journalctl command for '%v' service failed: %v", service, err)
	}
	logfile := filepath.Join(outputDir, service+".log")
	if err := ioutil.WriteFile(logfile, output, 0444); err != nil {
		return fmt.Errorf("Writing to file of journalctl logs for '%v' service failed: %v", service, err)
	}
	return nil
}

// Create logfiles for systemd services in outputDir.
func createSystemdLogfiles(outputDir string) {
	services := append(systemdServices, nodeSystemdServices...)
	for _, service := range services {
		if err := createSystemdLogfile(service, "cat", outputDir); err != nil {
			glog.Warningf("Failed to record journalctl logs: %v", err)
		}
	}
	// Service logs specific to VM setup.
	for _, service := range systemdSetupServices {
		if err := createSystemdLogfile(service, "short-precise", outputDir); err != nil {
			glog.Warningf("Failed to record journalctl logs: %v", err)
		}
	}
}

// Copy logfiles specific to this node based on the cloud-provider, system services, etc
// to a temporary directory. Also create logfiles for systemd services if journalctl is present.
// We do not expect this function to see an error.
func prepareLogfiles(logDir string) {
	glog.Info("Preparing logfiles relevant to this node")
	logfiles := nodeLogs[:]

	switch *cloudProvider {
	case "gce", "gke":
		logfiles = append(logfiles, gceLogs...)
	case "aws":
		logfiles = append(logfiles, awsLogs...)
	default:
		glog.Errorf("Unknown cloud provider '%v' provided, skipping any provider specific logs", *cloudProvider)
	}

	// Grab kubemark logs too, if asked for.
	if *enableHollowNodeLogs {
		logfiles = append(logfiles, kubemarkLogs...)
	}

	// Select system/service specific logs.
	if _, err := os.Stat("/workspace/etc/systemd/journald.conf"); err == nil {
		glog.Info("Journalctl found on host. Collecting systemd logs")
		createSystemdLogfiles(logDir)
	} else {
		glog.Infof("Journalctl not found on host (%v). Collecting supervisord logs instead", err)
		logfiles = append(logfiles, kernelLog)
		logfiles = append(logfiles, initdLogs...)
		logfiles = append(logfiles, supervisordLogs...)
	}

	// Copy all the logfiles that exist, to logDir.
	for _, logfile := range logfiles {
		logfileFullPath := filepath.Join(localLogPath, logfile+".log*") // Append .log* to copy rotated logs too.
		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cp %v %v", logfileFullPath, logDir))
		if err := cmd.Run(); err != nil {
			glog.Warningf("Failed to copy any logfiles with pattern '%v': %v", logfileFullPath, err)
		}
	}
}

func uploadLogfilesToGCS(logDir string) error {
	cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("ls %v/*", logDir))
	if output, err := cmd.Output(); err != nil {
		return fmt.Errorf("Could not list any logfiles: %v", err)
	} else {
		glog.Infof("List of logfiles available: %v", string(output))
	}

	gcsLogPath := *gcsPath + "/" + *nodeName
	glog.Infof("Uploading logfiles to GCS at path '%v'", gcsLogPath)
	var err error
	for uploadAttempt := 0; uploadAttempt < 3; uploadAttempt++ {
		// Upload the files with compression (-z) and parallelism (-m) for speeding
		// up, and set their ACL to make them publicly readable.
		cmd := exec.Command("gsutil", "-m", "-q", "cp", "-a", "public-read", "-c",
			"-z", "log,txt,xml", logDir+"/*", gcsLogPath)
		if err = cmd.Run(); err != nil {
			glog.Errorf("Attempt %v to upload to GCS failed: %v", uploadAttempt, err)
			continue
		}
		return writeSuccessMarkerFile()
	}
	return fmt.Errorf("Multiple attempts of gsutil failed, the final one due to: %v", err)
}

// Write a marker file to GCS named after this node to indicate logexporter's success.
// The directory to which we write this file can then be used as a registry to quickly
// fetch the list of nodes on which logexporter succeeded.
func writeSuccessMarkerFile() error {
	markerFilePath := *gcsPath + "/logexported-nodes-registry/" + *nodeName + ".txt"
	cmd := exec.Command("gsutil", "-q", "cp", "-a", "public-read", "-", markerFilePath)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("Failed to get stdin pipe to write marker file: %v", err)
	}
	io.WriteString(stdin, "")
	stdin.Close()
	if err = cmd.Run(); err != nil {
		return fmt.Errorf("Failed to write marker file to GCS: %v", err)
	}
	return nil
}

func main() {
	flag.Parse()
	if err := checkConfigValidity(); err != nil {
		glog.Fatalf("Bad config provided: %v", err)
	}

	localTmpLogPath, err := ioutil.TempDir("/tmp", "k8s-systemd-logs")
	if err != nil {
		glog.Fatalf("Could not create temporary dir locally for copying logs: %v", err)
	}
	defer os.RemoveAll(localTmpLogPath)

	prepareLogfiles(localTmpLogPath)
	if err := uploadLogfilesToGCS(localTmpLogPath); err != nil {
		glog.Fatalf("Could not upload logs to GCS: %v", err)
	}
	glog.Info("Logs successfully uploaded")

	glog.Infof("Entering sleep for a duration of %v seconds", *sleepDuration)
	time.Sleep(*sleepDuration)
}
<commit_msg>Allow dumping full systemd journal in logexporter.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// TODO(shyamjvs): Make this exporter work for master too, currently facing
// gcloud auth error when run from within a pod on the master.

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"github.com/golang/glog"
)

// Initialize the log exporter's configuration related flags.
var (
	nodeName             = flag.String("node-name", "", "Name of the node this log exporter is running on")
	gcsPath              = flag.String("gcs-path", "", "Path to the GCS directory under which to upload logs, for eg: gs://my-logs-bucket/logs")
	cloudProvider        = flag.String("cloud-provider", "", "Cloud provider for this node (gce/gke/aws/kubemark/..)")
	gcloudAuthFilePath   = flag.String("gcloud-auth-file-path", "/etc/service-account/service-account.json", "Path to gcloud service account file, for authenticating gsutil to write to GCS bucket")
	enableHollowNodeLogs = flag.Bool("enable-hollow-node-logs", false, "Enable uploading hollow node logs too. Relevant only for kubemark nodes")
	sleepDuration        = flag.Duration("sleep-duration", 60*time.Second, "Duration to sleep before exiting with success. Useful for making pods schedule with hard anti-affinity when run as a job on a k8s cluster")
	dumpSystemdJournal   = flag.Bool("dump-systemd-journal", false, "Whether to dump the full systemd journal")
)

var (
	localLogPath = "/var/log"

	// Node-type specific logfiles.
	// Currently we only handle nodes, and neglect master.
	nodeLogs = []string{"kube-proxy", "node-problem-detector", "fluentd"}

	// Cloud provider specific logfiles.
	awsLogs      = []string{"cloud-init-output"}
	gceLogs      = []string{"startupscript"}
	kubemarkLogs = []string{"*-hollow-node-*"}

	// System services/kernel related logfiles.
	kernelLog            = "kern"
	initdLogs            = []string{"docker"}
	supervisordLogs      = []string{"kubelet", "supervisor/supervisord", "supervisor/kubelet-stdout", "supervisor/kubelet-stderr", "supervisor/docker-stdout", "supervisor/docker-stderr"}
	systemdServices      = []string{"kern", "kubelet", "docker"}
	systemdSetupServices = []string{"kube-node-installation", "kube-node-configuration"}
	nodeSystemdServices  = []string{"node-problem-detector"}
)

// Check if the config provided through the flags take valid values.
func checkConfigValidity() error {
	glog.Info("Verifying if a valid config has been provided through the flags")
	if *nodeName == "" {
		return fmt.Errorf("Flag --node-name has its value unspecified")
	}
	if *gcsPath == "" {
		return fmt.Errorf("Flag --gcs-path has its value unspecified")
	}
	if _, err := os.Stat(*gcloudAuthFilePath); err != nil {
		return fmt.Errorf("Could not find the gcloud service account file: %v", err)
	} else {
		cmd := exec.Command("gcloud", "auth", "activate-service-account", "--key-file="+*gcloudAuthFilePath)
		if err := cmd.Run(); err != nil {
			return fmt.Errorf("Failed to activate gcloud service account: %v", err)
		}
	}
	return nil
}

// Create logfile for systemd service in outputDir with the given journalctl outputMode.
func createSystemdLogfile(service string, outputMode string, outputDir string) error {
	// Generate the journalctl command.
	journalCmdArgs := []string{fmt.Sprintf("--output=%v", outputMode), "-D", "/var/log/journal"}
	if service == "kern" {
		journalCmdArgs = append(journalCmdArgs, "-k")
	} else {
		journalCmdArgs = append(journalCmdArgs, "-u", fmt.Sprintf("%v.service", service))
	}
	cmd := exec.Command("journalctl", journalCmdArgs...)

	// Run the command and record the output to a file.
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("Journalctl command for '%v' service failed: %v", service, err)
	}
	logfile := filepath.Join(outputDir, service+".log")
	if err := ioutil.WriteFile(logfile, output, 0444); err != nil {
		return fmt.Errorf("Writing to file of journalctl logs for '%v' service failed: %v", service, err)
	}
	return nil
}

// createFullSystemdLogfile creates logfile for full systemd journal in the outputDir.
func createFullSystemdLogfile(outputDir string) error {
	cmd := exec.Command("journalctl", "--output=short-precise")
	// Run the command and record the output to a file.
	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("Journalctl command failed: %v", err)
	}
	logfile := filepath.Join(outputDir, "systemd.log")
	if err := ioutil.WriteFile(logfile, output, 0444); err != nil {
		return fmt.Errorf("Writing full journalctl logs to file failed: %v", err)
	}
	return nil
}

// Create logfiles for systemd services in outputDir.
func createSystemdLogfiles(outputDir string) {
	services := append(systemdServices, nodeSystemdServices...)
	for _, service := range services {
		if err := createSystemdLogfile(service, "cat", outputDir); err != nil {
			glog.Warningf("Failed to record journalctl logs: %v", err)
		}
	}
	// Service logs specific to VM setup.
	for _, service := range systemdSetupServices {
		if err := createSystemdLogfile(service, "short-precise", outputDir); err != nil {
			glog.Warningf("Failed to record journalctl logs: %v", err)
		}
	}
	if *dumpSystemdJournal {
		if err := createFullSystemdLogfile(outputDir); err != nil {
			glog.Warningf("Failed to record journalctl logs: %v", err)
		}
	}
}

// Copy logfiles specific to this node based on the cloud-provider, system services, etc
// to a temporary directory. Also create logfiles for systemd services if journalctl is present.
// We do not expect this function to see an error.
func prepareLogfiles(logDir string) {
	glog.Info("Preparing logfiles relevant to this node")
	logfiles := nodeLogs[:]

	switch *cloudProvider {
	case "gce", "gke":
		logfiles = append(logfiles, gceLogs...)
	case "aws":
		logfiles = append(logfiles, awsLogs...)
	default:
		glog.Errorf("Unknown cloud provider '%v' provided, skipping any provider specific logs", *cloudProvider)
	}

	// Grab kubemark logs too, if asked for.
	if *enableHollowNodeLogs {
		logfiles = append(logfiles, kubemarkLogs...)
	}

	// Select system/service specific logs.
	if _, err := os.Stat("/workspace/etc/systemd/journald.conf"); err == nil {
		glog.Info("Journalctl found on host. Collecting systemd logs")
		createSystemdLogfiles(logDir)
	} else {
		glog.Infof("Journalctl not found on host (%v). Collecting supervisord logs instead", err)
		logfiles = append(logfiles, kernelLog)
		logfiles = append(logfiles, initdLogs...)
		logfiles = append(logfiles, supervisordLogs...)
	}

	// Copy all the logfiles that exist, to logDir.
	for _, logfile := range logfiles {
		logfileFullPath := filepath.Join(localLogPath, logfile+".log*") // Append .log* to copy rotated logs too.
		cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("cp %v %v", logfileFullPath, logDir))
		if err := cmd.Run(); err != nil {
			glog.Warningf("Failed to copy any logfiles with pattern '%v': %v", logfileFullPath, err)
		}
	}
}

func uploadLogfilesToGCS(logDir string) error {
	cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf("ls %v/*", logDir))
	if output, err := cmd.Output(); err != nil {
		return fmt.Errorf("Could not list any logfiles: %v", err)
	} else {
		glog.Infof("List of logfiles available: %v", string(output))
	}

	gcsLogPath := *gcsPath + "/" + *nodeName
	glog.Infof("Uploading logfiles to GCS at path '%v'", gcsLogPath)
	var err error
	for uploadAttempt := 0; uploadAttempt < 3; uploadAttempt++ {
		// Upload the files with compression (-z) and parallelism (-m) for speeding
		// up, and set their ACL to make them publicly readable.
		cmd := exec.Command("gsutil", "-m", "-q", "cp", "-a", "public-read", "-c",
			"-z", "log,txt,xml", logDir+"/*", gcsLogPath)
		if err = cmd.Run(); err != nil {
			glog.Errorf("Attempt %v to upload to GCS failed: %v", uploadAttempt, err)
			continue
		}
		return writeSuccessMarkerFile()
	}
	return fmt.Errorf("Multiple attempts of gsutil failed, the final one due to: %v", err)
}

// Write a marker file to GCS named after this node to indicate logexporter's success.
// The directory to which we write this file can then be used as a registry to quickly
// fetch the list of nodes on which logexporter succeeded.
func writeSuccessMarkerFile() error {
	markerFilePath := *gcsPath + "/logexported-nodes-registry/" + *nodeName + ".txt"
	cmd := exec.Command("gsutil", "-q", "cp", "-a", "public-read", "-", markerFilePath)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("Failed to get stdin pipe to write marker file: %v", err)
	}
	io.WriteString(stdin, "")
	stdin.Close()
	if err = cmd.Run(); err != nil {
		return fmt.Errorf("Failed to write marker file to GCS: %v", err)
	}
	return nil
}

func main() {
	flag.Parse()
	if err := checkConfigValidity(); err != nil {
		glog.Fatalf("Bad config provided: %v", err)
	}

	localTmpLogPath, err := ioutil.TempDir("/tmp", "k8s-systemd-logs")
	if err != nil {
		glog.Fatalf("Could not create temporary dir locally for copying logs: %v", err)
	}
	defer os.RemoveAll(localTmpLogPath)

	prepareLogfiles(localTmpLogPath)
	if err := uploadLogfilesToGCS(localTmpLogPath); err != nil {
		glog.Fatalf("Could not upload logs to GCS: %v", err)
	}
	glog.Info("Logs successfully uploaded")

	glog.Infof("Entering sleep for a duration of %v seconds", *sleepDuration)
	time.Sleep(*sleepDuration)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"flag"
	"os"
	"runtime"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/density"
	"github.com/phil-mansfield/gotetra/render/io"

	obj "github.com/phil-mansfield/gotetra/los/main/loop_objects"
	util "github.com/phil-mansfield/gotetra/los/main/gtet_util"
	"github.com/phil-mansfield/gotetra/los/main/gtet_util/loop"
)

func main() {
	log.Println("gtet_render")

	p := parseCmd()
	ids, snaps, _, err := util.ParseStdin()
	if err != nil { log.Fatal(err.Error()) }
	if len(ids) == 0 { log.Fatal("No input IDs.") }

	fnames, err := computeGrids(ids, snaps, p)
	if err != nil { log.Fatal(err.Error()) }

	for _, fname := range fnames { fmt.Println(fname) }
}

type Params struct {
	Pixels int
	SubsampleLength int
	Mult float64

	Sphere float64
	Tetra, Linear, Cubic int
}

// parseCmd parses the command line options and returns them as a struct.
func parseCmd() *Params {
	p := &Params{}

	flag.IntVar(&p.Pixels, "RBins", 100,
		"Number of pixels on one side of render.")
	flag.IntVar(&p.SubsampleLength, "SubsampleLength", 1,
		"Grid distance between adjacent points used in interpolation.")
	flag.Float64Var(&p.Mult, "Mult", 3, "Size of box as a multiplier of R200m")
	flag.Float64Var(&p.Sphere, "Sphere", 0.0,
		"Setting to a non-zero value, r, indicates that the particles " +
			"supplied to the profile should be treated as constant density " +
			"spheres of radius r instead of particles.")
	flag.IntVar(&p.Tetra, "Tetra", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from constant density tetrahedra instead of " +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.IntVar(&p.Linear, "Linear", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from tri-linear interpolation instead of from" +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.IntVar(&p.Cubic, "Cubic", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from tri-cubic interpolation instead of from" +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.Parse()

	return p
}

// Determines how the density field is interpolated from the snapshot
// particle locations.
type Method int
const (
	Particle Method = iota
	Sphere
	Tetra
	Linear
	Cubic
)

// This is where the magic happens. Returns a slice of profiles. Each
// profile is represented as a slice of radii with units of comoving Mpc/h
// and densities with units of <rho_m>.
func computeGrids(ids, snaps []int, p *Params) ([]string, error) {
	// Select method.
	method := Particle
	switch {
	case p.Sphere > 0:
		method = Sphere
	case p.Tetra > 0:
		method = Tetra
	case p.Linear > 0:
		method = Linear
	case p.Cubic > 0:
		method = Cubic
	}

	// Set up loop state.
	buf := loop.NewBuffer()
	rs, err := newRenderers(ids, snaps, p, method)
	if err != nil { return nil, err }

	workers := runtime.GOMAXPROCS(runtime.NumCPU())

	snapBins, snapIdxs := binBySnap(snaps, ids)
	for snap := range snapBins {
		runtime.GC()

		// Select out all the halos in a given snapshot.
		idxs := snapIdxs[snap]
		snapRs := make([]loop.Object, len(idxs))
		for i, idx := range idxs { snapRs[i] = rs[idx] }

		// Call loop.Loop().
		switch method {
		case Particle:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		case Tetra:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		case Linear:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, p.Linear, workers,
			)
		case Cubic:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Cubic, p.Cubic, workers,
			)
		case Sphere:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		default:
			panic(":3")
		}
	}

	// Convert Profiles to the needed float slices.
	hd, err := util.ReadSnapHeader(snaps[0])
	dx := hd.TotalWidth / float64(hd.CountWidth) * float64(p.SubsampleLength)

	var pts int
	switch method {
	case Particle, Sphere: pts = 1
	case Tetra: pts = p.Tetra
	case Linear: pts = p.Linear
	case Cubic: pts = p.Cubic
	default:
		panic(":3")
	}

	bufs := processCounts(rs, dx, pts)
	fnames := fileNames(snaps, ids, p, method)
	err = writeBufs(snaps, bufs, rs, p, method, fnames)
	if err != nil { return nil, err }

	return fnames, nil
}

// binBySnap collects the given IDs into groups that share a common snapshot.
// The return format is two maps which accept snap
func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
    snapBins = make(map[int][]int)
    idxBins = make(map[int][]int)
    for i, snap := range snaps {
        id := ids[i]
        snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
    }
    return snapBins, idxBins
}

func newRenderers(
	ids, snaps []int, p *Params, method Method,
) ([]*obj.Renderer, error) {
	snapBins, idxBins := binBySnap(snaps, ids)
	rs := make([]*obj.Renderer, len(ids))

	for _, snap := range snaps {
		snapIDs := snapBins[snap]
        idxs := idxBins[snap]

        vals, err := util.ReadRockstar(
            snap, snapIDs, halo.X, halo.Y, halo.Z, halo.Rad200b,
        )
        if err != nil { return nil, err }

        xs, ys, zs, rads := vals[0], vals[1], vals[2], vals[3]

        for i := range xs {
			origin := [3]float64{ xs[i]-rads[i], ys[i]-rads[i], zs[i]-rads[i] }
			pw := 2 * rads[i] / float64(p.Pixels)
			pixels := [3]int{ p.Pixels, p.Pixels, p.Pixels }

			switch method {
			case Particle, Linear, Cubic:
				rs[idxs[i]] = obj.NewParticleRenderer(origin, pixels, pw)
			case Tetra:
				rs[idxs[i]] = obj.NewTetraRenderer(
					origin, pixels, pw, p.Tetra,
				)
			case Sphere:
				panic("NYI")
			}
        }
    }
    return rs, nil
}

// processCounts takes raw count grids and converts them to density buffers.
func processCounts(rs []*obj.Renderer, dx float64, pts int) []density.Buffer {
	bufs:= make([]density.Buffer, len(rs))
	mp := dx*dx*dx / float64(pts*pts*pts)

	for i, r := range rs {
		rhos := make([]float64, len(r.Counts))
		dV := r.Pw*r.Pw*r.Pw
		for j, n := range r.Counts {
			rhos[j] = float64(n) * mp / dV
		}
		bufs[i] = density.WrapperDensityBuffer(rhos)
	}

	return bufs
}

func fileNames(snaps, ids []int, p *Params, method Method) []string {
	names := make([]string, len(ids))

	var paramStr string
	switch method {
	case Particle:
		paramStr = "ngp"
	case Sphere:
		paramStr = "sph"
	case Tetra:
		paramStr = fmt.Sprintf("tet%d", p.Tetra)
	case Linear:
		paramStr = fmt.Sprintf("lin%d", p.Linear)
	case Cubic:
		paramStr = fmt.Sprintf("cbc%d", p.Cubic)
	}

	for i := range snaps {
		names[i] = fmt.Sprintf("s%d_id%d_%s.gtet",
			snaps[i], snaps[i], ids[i], paramStr)
	}

	return names
}

func writeBufs(
	snaps []int, bufs []density.Buffer, rs []*obj.Renderer,
	p *Params, method Method, fnames []string,
) error {
	var pts int
	switch method {
	case Particle: pts =1
	case Sphere: pts = 1
	case Tetra: pts = p.Tetra
	case Linear: pts = p.Linear
	case Cubic: pts = p.Cubic
	}

	for i := range snaps {
		hd, err := util.ReadSnapHeader(snaps[i])
		if err != nil { return err }

		f, err := os.Create(fnames[i])
		defer f.Close()
		if err != nil { return err }

		// This isn't completely accurate, but that's fine.
		cellOrigin := [3]int{
			int(rs[i].Origin[0] / rs[i].Pw),
			int(rs[i].Origin[1] / rs[i].Pw),
			int(rs[i].Origin[2] / rs[i].Pw),
		}

		loc := io.NewLocationInfo(cellOrigin, rs[i].Pixels, rs[i].Pw)
		cos := io.NewCosmoInfo(
			hd.Cosmo.H100 * 100, hd.Cosmo.OmegaM,
			hd.Cosmo.OmegaL, hd.Cosmo.Z, hd.TotalWidth,
		)
		ren := io.NewRenderInfo(
			pts, int(rs[i].Pw / hd.TotalWidth), p.SubsampleLength, "",
		)

		io.WriteBuffer(bufs[i], cos, ren, loc, f)
	}

	return nil
}
<commit_msg>Small bugfixes to gtet_render<commit_after>package main

import (
	"fmt"
	"log"
	"flag"
	"os"
	"runtime"

	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/density"
	"github.com/phil-mansfield/gotetra/render/io"

	obj "github.com/phil-mansfield/gotetra/los/main/loop_objects"
	util "github.com/phil-mansfield/gotetra/los/main/gtet_util"
	"github.com/phil-mansfield/gotetra/los/main/gtet_util/loop"
)

func main() {
	log.Println("gtet_render")

	p := parseCmd()
	ids, snaps, _, err := util.ParseStdin()
	if err != nil { log.Fatal(err.Error()) }
	if len(ids) == 0 { log.Fatal("No input IDs.") }

	fnames, err := computeGrids(ids, snaps, p)
	if err != nil { log.Fatal(err.Error()) }

	for _, fname := range fnames { fmt.Println(fname) }
}

type Params struct {
	Pixels int
	SubsampleLength int
	Mult float64

	Sphere float64
	Tetra, Linear, Cubic int
}

// parseCmd parses the command line options and returns them as a struct.
func parseCmd() *Params {
	p := &Params{}

	flag.IntVar(&p.Pixels, "Pixels", 100,
		"Number of pixels on one side of render.")
	flag.IntVar(&p.SubsampleLength, "SubsampleLength", 1,
		"Grid distance between adjacent points used in interpolation.")
	flag.Float64Var(&p.Mult, "Mult", 3, "Size of box as a multiplier of R200m")
	flag.Float64Var(&p.Sphere, "Sphere", 0.0,
		"Setting to a non-zero value, r, indicates that the particles " +
			"supplied to the profile should be treated as constant density " +
			"spheres of radius r instead of particles.")
	flag.IntVar(&p.Tetra, "Tetra", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from constant density tetrahedra instead of " +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.IntVar(&p.Linear, "Linear", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from tri-linear interpolation instead of from" +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.IntVar(&p.Cubic, "Cubic", 0,
		"Setting to a positive value, n, indicates that profiles should be " +
			"generated from tri-cubic interpolation instead of from" +
			"particles. n will be the number of particles 'on a side' used " +
			"during interpolation.")
	flag.Parse()

	return p
}

// Determines how the density field is interpolated from the snapshot
// particle locations.
type Method int
const (
	Particle Method = iota
	Sphere
	Tetra
	Linear
	Cubic
)

// This is where the magic happens. Returns a slice of profiles. Each
// profile is represented as a slice of radii with units of comoving Mpc/h
// and densities with units of <rho_m>.
func computeGrids(ids, snaps []int, p *Params) ([]string, error) {
	// Select method.
	method := Particle
	switch {
	case p.Sphere > 0:
		method = Sphere
	case p.Tetra > 0:
		method = Tetra
	case p.Linear > 0:
		method = Linear
	case p.Cubic > 0:
		method = Cubic
	}

	// Set up loop state.
	buf := loop.NewBuffer()
	rs, err := newRenderers(ids, snaps, p, method)
	if err != nil { return nil, err }

	workers := runtime.GOMAXPROCS(runtime.NumCPU())

	snapBins, snapIdxs := binBySnap(snaps, ids)
	for snap := range snapBins {
		runtime.GC()

		// Select out all the halos in a given snapshot.
		idxs := snapIdxs[snap]
		snapRs := make([]loop.Object, len(idxs))
		for i, idx := range idxs { snapRs[i] = rs[idx] }

		// Call loop.Loop().
		switch method {
		case Particle:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		case Tetra:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		case Linear:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, p.Linear, workers,
			)
		case Cubic:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Cubic, p.Cubic, workers,
			)
		case Sphere:
			loop.Loop(
				snap, snapRs, buf, p.SubsampleLength,
				loop.Linear, 1, workers,
			)
		default:
			panic(":3")
		}
	}

	// Convert Profiles to the needed float slices.
	hd, err := util.ReadSnapHeader(snaps[0])
	dx := hd.TotalWidth / float64(hd.CountWidth) * float64(p.SubsampleLength)

	var pts int
	switch method {
	case Particle, Sphere: pts = 1
	case Tetra: pts = p.Tetra
	case Linear: pts = p.Linear
	case Cubic: pts = p.Cubic
	default:
		panic(":3")
	}

	bufs := processCounts(rs, dx, pts)
	fnames := fileNames(snaps, ids, p, method)
	err = writeBufs(snaps, bufs, rs, p, method, fnames)
	if err != nil { return nil, err }

	return fnames, nil
}

// binBySnap collects the given IDs into groups that share a common snapshot.
// The return format is two maps which accept snap
func binBySnap(snaps, ids []int) (snapBins, idxBins map[int][]int) {
    snapBins = make(map[int][]int)
    idxBins = make(map[int][]int)
    for i, snap := range snaps {
        id := ids[i]
        snapBins[snap] = append(snapBins[snap], id)
		idxBins[snap] = append(idxBins[snap], i)
    }
    return snapBins, idxBins
}

func newRenderers(
	ids, snaps []int, p *Params, method Method,
) ([]*obj.Renderer, error) {
	snapBins, idxBins := binBySnap(snaps, ids)
	rs := make([]*obj.Renderer, len(ids))

	for _, snap := range snaps {
		snapIDs := snapBins[snap]
        idxs := idxBins[snap]

        vals, err := util.ReadRockstar(
            snap, snapIDs, halo.X, halo.Y, halo.Z, halo.Rad200b,
        )
        if err != nil { return nil, err }

        xs, ys, zs, rads := vals[0], vals[1], vals[2], vals[3]

        for i := range xs {
			dr := rads[i] * p.Mult
			origin := [3]float64{ xs[i]-dr, ys[i]-dr, zs[i]-dr }
			pw := 2 * dr / float64(p.Pixels)
			pixels := [3]int{ p.Pixels, p.Pixels, p.Pixels }

			switch method {
			case Particle, Linear, Cubic:
				rs[idxs[i]] = obj.NewParticleRenderer(origin, pixels, pw)
			case Tetra:
				rs[idxs[i]] = obj.NewTetraRenderer(
					origin, pixels, pw, p.Tetra,
				)
			case Sphere:
				panic("NYI")
			}
        }
    }
    return rs, nil
}

// processCounts takes raw count grids and converts them to density buffers.
func processCounts(rs []*obj.Renderer, dx float64, pts int) []density.Buffer {
	bufs:= make([]density.Buffer, len(rs))
	mp := dx*dx*dx / float64(pts*pts*pts)

	for i, r := range rs {
		rhos := make([]float64, len(r.Counts))
		dV := r.Pw*r.Pw*r.Pw
		for j, n := range r.Counts {
			rhos[j] = float64(n) * mp / dV
		}
		bufs[i] = density.WrapperDensityBuffer(rhos)
	}

	return bufs
}

func fileNames(snaps, ids []int, p *Params, method Method) []string {
	names := make([]string, len(ids))

	var paramStr string
	switch method {
	case Particle:
		paramStr = "ngp"
	case Sphere:
		paramStr = "sph"
	case Tetra:
		paramStr = fmt.Sprintf("tet%d", p.Tetra)
	case Linear:
		paramStr = fmt.Sprintf("lin%d", p.Linear)
	case Cubic:
		paramStr = fmt.Sprintf("cbc%d", p.Cubic)
	}

	for i := range snaps {
		names[i] = fmt.Sprintf("s%d_id%d_%s.gtet", snaps[i], ids[i], paramStr)
	}

	return names
}

func writeBufs(
	snaps []int, bufs []density.Buffer, rs []*obj.Renderer,
	p *Params, method Method, fnames []string,
) error {
	var pts int
	switch method {
	case Particle: pts =1
	case Sphere: pts = 1
	case Tetra: pts = p.Tetra
	case Linear: pts = p.Linear
	case Cubic: pts = p.Cubic
	}

	for i := range snaps {
		hd, err := util.ReadSnapHeader(snaps[i])
		if err != nil { return err }

		f, err := os.Create(fnames[i])
		defer f.Close()
		if err != nil { return err }

		// This isn't completely accurate, but that's fine.
		cellOrigin := [3]int{
			int(rs[i].Origin[0] / rs[i].Pw),
			int(rs[i].Origin[1] / rs[i].Pw),
			int(rs[i].Origin[2] / rs[i].Pw),
		}

		loc := io.NewLocationInfo(cellOrigin, rs[i].Pixels, rs[i].Pw)
		cos := io.NewCosmoInfo(
			hd.Cosmo.H100 * 100, hd.Cosmo.OmegaM,
			hd.Cosmo.OmegaL, hd.Cosmo.Z, hd.TotalWidth,
		)
		ren := io.NewRenderInfo(
			pts, int(rs[i].Pw / hd.TotalWidth), p.SubsampleLength, "",
		)

		io.WriteBuffer(bufs[i], cos, ren, loc, f)
	}

	return nil
}
<|endoftext|>


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

package storage

import (
	sq "github.com/Masterminds/squirrel"
	"github.com/golang/glog"
	"github.com/kubeflow/pipelines/backend/src/common/util"
)

var (
	defaultDBStatus = sq.Eq{"HaveSamplesLoaded": false}
)

type DBStatusStoreInterface interface {
	HaveSamplesLoaded() (bool, error)
	MarkSampleLoaded() error
}

// Implementation of a DBStatusStoreInterface. This store read/write state of the database.
// For now we store status like whether sample is loaded.
type DBStatusStore struct {
	db *DB
}

func (s *DBStatusStore) InitializeDBStatusTable() error {
	tx, err := s.db.Begin()
	if err != nil {
		return util.NewInternalServerError(err, "Failed to create a new transaction to initialize database status.")
	}

	rows, err := tx.Query("SELECT * FROM db_statuses")
	if err != nil {
		tx.Rollback()
		return util.NewInternalServerError(err, "Failed to load database status.")
	}
	next := rows.Next()
	rows.Close() // "rows" shouldn't be used after this point.

	// The table is not initialized
	if !next {
		sql, args, queryErr := sq.
			Insert("db_statuses").
			SetMap(defaultDBStatus).
			ToSql()

		if queryErr != nil {
			tx.Rollback()
			return util.NewInternalServerError(queryErr, "Error creating query to initialize database status table.")
		}

		_, err = tx.Exec(sql, args...)
		if err != nil {
			tx.Rollback()
			return util.NewInternalServerError(err, "Error initializing the database status table.")
		}
	}
	err = tx.Commit()
	if err != nil {
		glog.Error("Failed to commit transaction to initialize database status table")
		return util.NewInternalServerError(err, "Failed to initializing the database status table.")
	}
	return nil
}

func (s *DBStatusStore) HaveSamplesLoaded() (bool, error) {
	var haveSamplesLoaded bool
	sql, args, err := sq.Select("*").From("db_statuses").ToSql()
	if err != nil {
		return false, util.NewInternalServerError(err, "Error creating query to get load sample status.")
	}
	rows, err := s.db.Query(sql, args...)
	if err != nil {
		return false, util.NewInternalServerError(err, "Error when getting load sample status")
	}
	defer rows.Close()
	if rows.Next() {
		err = rows.Scan(&haveSamplesLoaded)
		if err != nil {
			return false, util.NewInternalServerError(err, "Error when scanning row to load sample status")
		}
		return haveSamplesLoaded, nil
	}
	return false, nil
}

func (s *DBStatusStore) MarkSampleLoaded() error {
	sql, args, err := sq.
		Update("db_statuses").
		SetMap(sq.Eq{"HaveSamplesLoaded": true}).
		ToSql()
	if err != nil {
		return util.NewInternalServerError(err, "Error creating query to mark samples as loaded.")
	}
	_, err = s.db.Exec(sql, args...)
	if err != nil {
		return util.NewInternalServerError(err, "Error mark samples as loaded.")
	}
	return nil
}

// factory function for database status store
func NewDBStatusStore(db *DB) *DBStatusStore {
	s := &DBStatusStore{db: db}
	// Initialize database status table
	s.InitializeDBStatusTable()
	return s
}
<commit_msg>fix(backend): improve forward / backward compatibility on db status table (#4351)<commit_after>// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
	sq "github.com/Masterminds/squirrel"
	"github.com/golang/glog"
	"github.com/kubeflow/pipelines/backend/src/common/util"
)

var (
	defaultDBStatus = sq.Eq{"HaveSamplesLoaded": false}
)

type DBStatusStoreInterface interface {
	HaveSamplesLoaded() (bool, error)
	MarkSampleLoaded() error
}

var (
	dbStatusStoreColumns = []string{
		"HaveSamplesLoaded",
	}
)

// Implementation of a DBStatusStoreInterface. This store read/write state of the database.
// For now we store status like whether sample is loaded.
type DBStatusStore struct {
	db *DB
}

func (s *DBStatusStore) InitializeDBStatusTable() error {
	tx, err := s.db.Begin()
	if err != nil {
		return util.NewInternalServerError(err, "Failed to create a new transaction to initialize database status.")
	}

	rows, err := tx.Query("SELECT * FROM db_statuses")
	if err != nil {
		tx.Rollback()
		return util.NewInternalServerError(err, "Failed to load database status.")
	}
	next := rows.Next()
	rows.Close() // "rows" shouldn't be used after this point.

	// The table is not initialized
	if !next {
		sql, args, queryErr := sq.
			Insert("db_statuses").
			SetMap(defaultDBStatus).
			ToSql()

		if queryErr != nil {
			tx.Rollback()
			return util.NewInternalServerError(queryErr, "Error creating query to initialize database status table.")
		}

		_, err = tx.Exec(sql, args...)
		if err != nil {
			tx.Rollback()
			return util.NewInternalServerError(err, "Error initializing the database status table.")
		}
	}
	err = tx.Commit()
	if err != nil {
		glog.Error("Failed to commit transaction to initialize database status table")
		return util.NewInternalServerError(err, "Failed to initializing the database status table.")
	}
	return nil
}

func (s *DBStatusStore) HaveSamplesLoaded() (bool, error) {
	var haveSamplesLoaded bool
	sql, args, err := sq.Select(dbStatusStoreColumns...).From("db_statuses").ToSql()
	if err != nil {
		return false, util.NewInternalServerError(err, "Error creating query to get load sample status.")
	}
	rows, err := s.db.Query(sql, args...)
	if err != nil {
		return false, util.NewInternalServerError(err, "Error when getting load sample status")
	}
	defer rows.Close()
	if rows.Next() {
		err = rows.Scan(&haveSamplesLoaded)
		if err != nil {
			return false, util.NewInternalServerError(err, "Error when scanning row to load sample status")
		}
		return haveSamplesLoaded, nil
	}
	return false, nil
}

func (s *DBStatusStore) MarkSampleLoaded() error {
	sql, args, err := sq.
		Update("db_statuses").
		SetMap(sq.Eq{"HaveSamplesLoaded": true}).
		ToSql()
	if err != nil {
		return util.NewInternalServerError(err, "Error creating query to mark samples as loaded.")
	}
	_, err = s.db.Exec(sql, args...)
	if err != nil {
		return util.NewInternalServerError(err, "Error mark samples as loaded.")
	}
	return nil
}

// factory function for database status store
func NewDBStatusStore(db *DB) *DBStatusStore {
	s := &DBStatusStore{db: db}
	// Initialize database status table
	s.InitializeDBStatusTable()
	return s
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"fmt"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const emptyCell = "-"

func getHandlerManager() (*grpc.ClientConn, core.HandlerManagerClient) {
	conn, err := grpc.Dial(viper.GetString("ttn-handler"), grpc.WithInsecure())
	if err != nil {
		ctx.Fatalf("Could not connect: %v", err)
	}

	return conn, core.NewHandlerManagerClient(conn)
}

// devicesCmd represents the `devices` command
var devicesCmd = &cobra.Command{
	Use:   "devices",
	Short: "List registered devices on the Handler",
	Long:  `List registered devices on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.List(context.Background(), &core.HandlerListDevicesReq{
			AppEUI: appEUI,
		})

		if err != nil {
			ctx.WithError(err).Fatal("Could not get device list")
		}

		ctx.Infof("Found %d devices", len(res.Devices))

		table := uitable.New()
		table.MaxColWidth = 50

		table.AddRow("DevEUI", "AppKey", "NwkSKey", "AppSKey")
		for _, device := range res.Devices {
			devEUI := fmt.Sprintf("%X", device.DevEUI)

			appKey := emptyCell
			if a := device.GetActivated(); a != nil {
				appKey = fmt.Sprintf("%X", a.AppKey)
			}

			nwkSKey := emptyCell
			appSKey := emptyCell
			if p := device.GetPersonalized(); p != nil {
				nwkSKey = fmt.Sprintf("%X", p.NwkSKey)
				appSKey = fmt.Sprintf("%X", p.AppSKey)
			}

			table.AddRow(devEUI, appKey, nwkSKey, appSKey)
		}
		fmt.Println(table)

	},
}

// devicesRegisterCmd represents the `device register` command
var devicesRegisterCmd = &cobra.Command{
	Use:   "register [DevEUI] [AppKey]",
	Short: "Create or Update OTAA registrations on the Handler",
	Long:  `Create or Update OTAA registrations on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			ctx.Fatal("Insufficient arguments")
		}

		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		devEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid DevEUI: %s", err)
		}

		appKey, err := util.Parse128(args[1])
		if err != nil {
			ctx.Fatalf("Invalid AppKey: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.Upsert(context.Background(), &core.HandlerUpsertDeviceReq{
			AppEUI: appEUI,
			Device: &core.HandlerDevice{
				DevEUI: devEUI,
				Data: &core.HandlerDevice_Activated{
					Activated: &core.HandlerDevice_ActivatedData{
						AppKey: appKey,
					},
				},
			},
		})

		if err != nil || res == nil {
			ctx.WithError(err).Fatal("Could not register device")
		}
		ctx.Info("Ok")

	},
}

// devicesRegisterPersonalizedCmd represents the `device register personalized` command
var devicesRegisterPersonalizedCmd = &cobra.Command{
	Use:   "personalized [DevEUI] [NwkSKey] [AppSKey]",
	Short: "Create or Update ABP registrations on the Handler",
	Long:  `Create or Update ABP registrations on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 3 {
			ctx.Fatal("Insufficient arguments")
		}

		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		devEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid DevEUI: %s", err)
		}

		nwkSKey, err := util.Parse128(args[1])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}

		appSKey, err := util.Parse128(args[2])
		if err != nil {
			ctx.Fatalf("Invalid AppSKey: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.Upsert(context.Background(), &core.HandlerUpsertDeviceReq{
			AppEUI: appEUI,
			Device: &core.HandlerDevice{
				DevEUI: devEUI,
				Data: &core.HandlerDevice_Personalized{
					Personalized: &core.HandlerDevice_PersonalizedData{
						AppSKey: appSKey,
						NwkSKey: nwkSKey,
					},
				},
			},
		})

		if err != nil || res == nil {
			ctx.WithError(err).Fatal("Could not register device")
		}
		ctx.Info("Ok")

	},
}

func init() {
	RootCmd.AddCommand(devicesCmd)

	devicesCmd.AddCommand(devicesRegisterCmd)

	devicesRegisterCmd.AddCommand(devicesRegisterPersonalizedCmd)

	devicesCmd.Flags().String("ttn-handler", "0.0.0.0:1882", "The net address of the TTN Handler")
	viper.BindPFlag("ttn-handler", devicesCmd.Flags().Lookup("ttn-handler"))

	devicesCmd.PersistentFlags().String("app-eui", "0102030405060708", "The app EUI to use")
	viper.BindPFlag("app-eui", devicesCmd.PersistentFlags().Lookup("app-eui"))

	devicesCmd.PersistentFlags().String("app-token", "0102030405060708", "The app Token to use")
	viper.BindPFlag("app-token", devicesCmd.PersistentFlags().Lookup("app-token"))

}
<commit_msg>[feature/manager] Update ttnctl with reworked protos<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package cmd

import (
	"fmt"

	"golang.org/x/net/context"

	"google.golang.org/grpc"

	"github.com/TheThingsNetwork/ttn/core"
	"github.com/TheThingsNetwork/ttn/ttnctl/util"
	"github.com/gosuri/uitable"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const emptyCell = "-"

func getHandlerManager() (*grpc.ClientConn, core.HandlerManagerClient) {
	conn, err := grpc.Dial(viper.GetString("ttn-handler"), grpc.WithInsecure())
	if err != nil {
		ctx.Fatalf("Could not connect: %v", err)
	}

	return conn, core.NewHandlerManagerClient(conn)
}

// devicesCmd represents the `devices` command
var devicesCmd = &cobra.Command{
	Use:   "devices",
	Short: "List registered devices on the Handler",
	Long:  `List registered devices on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.ListDevices(context.Background(), &core.ListDevicesHandlerReq{
			AppEUI: appEUI,
		})

		if err != nil {
			ctx.WithError(err).Fatal("Could not get device list")
		}

		ctx.Infof("Found %d personalized devices (ABP)", len(res.ABP))
		table := uitable.New()
		table.MaxColWidth = 70
		table.AddRow("DevAddr", "NwkSKey", "AppSKey")
		for _, device := range res.ABP {
			devAddr := fmt.Sprintf("%X", device.DevAddr)
			nwkSKey := fmt.Sprintf("%X", device.NwkSKey)
			appSKey := fmt.Sprintf("%X", device.AppSKey)
			table.AddRow(devAddr, nwkSKey, appSKey)
		}
		fmt.Println(table)

		ctx.Infof("Found %d dynamic devices (OTAA)", len(res.OTAA))
		table = uitable.New()
		table.MaxColWidth = 40
		table.AddRow("DevEUI", "DevAddr", "NwkSKey", "AppSKey", "AppKey")
		for _, device := range res.OTAA {
			devEUI := fmt.Sprintf("%X", device.DevEUI)
			devAddr := fmt.Sprintf("%X", device.DevAddr)
			nwkSKey := fmt.Sprintf("%X", device.NwkSKey)
			appSKey := fmt.Sprintf("%X", device.AppSKey)
			appKey := fmt.Sprintf("%X", device.AppKey)
			table.AddRow(devEUI, devAddr, nwkSKey, appSKey, appKey)
		}
		fmt.Println(table)
	},
}

// devicesRegisterCmd represents the `device register` command
var devicesRegisterCmd = &cobra.Command{
	Use:   "register [DevEUI] [AppKey]",
	Short: "Create or Update OTAA registrations on the Handler",
	Long:  `Create or Update OTAA registrations on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 2 {
			ctx.Fatal("Insufficient arguments")
		}

		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		devEUI, err := util.Parse64(args[0])
		if err != nil {
			ctx.Fatalf("Invalid DevEUI: %s", err)
		}

		appKey, err := util.Parse128(args[1])
		if err != nil {
			ctx.Fatalf("Invalid AppKey: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.UpsertOTAA(context.Background(), &core.UpsertOTAAHandlerReq{
			AppEUI: appEUI,
			DevEUI: devEUI,
			AppKey: appKey,
		})

		if err != nil || res == nil {
			ctx.WithError(err).Fatal("Could not register device")
		}
		ctx.Info("Ok")
	},
}

// devicesRegisterPersonalizedCmd represents the `device register personalized` command
var devicesRegisterPersonalizedCmd = &cobra.Command{
	Use:   "personalized [DevAddr] [NwkSKey] [AppSKey]",
	Short: "Create or Update ABP registrations on the Handler",
	Long:  `Create or Update ABP registrations on the Handler`,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) < 3 {
			ctx.Fatal("Insufficient arguments")
		}

		appEUI, err := util.Parse64(viper.GetString("app-eui"))
		if err != nil {
			ctx.Fatalf("Invalid AppEUI: %s", err)
		}

		devAddr, err := util.Parse32(args[0])
		if err != nil {
			ctx.Fatalf("Invalid DevAddr: %s", err)
		}

		nwkSKey, err := util.Parse128(args[1])
		if err != nil {
			ctx.Fatalf("Invalid NwkSKey: %s", err)
		}

		appSKey, err := util.Parse128(args[2])
		if err != nil {
			ctx.Fatalf("Invalid AppSKey: %s", err)
		}

		conn, manager := getHandlerManager()
		defer conn.Close()

		res, err := manager.UpsertABP(context.Background(), &core.UpsertABPHandlerReq{
			AppEUI:  appEUI,
			DevAddr: devAddr,
			AppSKey: appSKey,
			NwkSKey: nwkSKey,
		})

		if err != nil || res == nil {
			ctx.WithError(err).Fatal("Could not register device")
		}
		ctx.Info("Ok")
	},
}

func init() {
	RootCmd.AddCommand(devicesCmd)

	devicesCmd.AddCommand(devicesRegisterCmd)

	devicesRegisterCmd.AddCommand(devicesRegisterPersonalizedCmd)

	devicesCmd.Flags().String("ttn-handler", "0.0.0.0:1882", "The net address of the TTN Handler")
	viper.BindPFlag("ttn-handler", devicesCmd.Flags().Lookup("ttn-handler"))

	devicesCmd.PersistentFlags().String("app-eui", "0102030405060708", "The app EUI to use")
	viper.BindPFlag("app-eui", devicesCmd.PersistentFlags().Lookup("app-eui"))

	devicesCmd.PersistentFlags().String("app-token", "0102030405060708", "The app Token to use")
	viper.BindPFlag("app-token", devicesCmd.PersistentFlags().Lookup("app-token"))
}
<|endoftext|>


<commit_before>package pelicantun

import (
	"fmt"
	"testing"
	"time"

	cv "github.com/glycerine/goconvey/convey"
)

func TestLongPollToGetLowLatency010(t *testing.T) {

	cli, srv, rev, fwd, err := StartTestSystemWithBcast()
	panicOn(err)
	defer srv.Stop()
	defer rev.Stop()
	defer fwd.Stop()
	defer cli.Stop()

	cv.Convey("Given a ForwardProxy and a ReverseProxy communicating over http, in order to acheive low-latency sends from server to client, long-polling with two sockets (one send and one receive) should be used. So a server that has something to send (e.g. broadcast here) should be able to send back to the client immediately.\n", t, func() {
		cli.Start()
		<-srv.FirstClient

		po("got past <-srv.FirstClient\n")

		defer func() { po("deferred frunction running...!\n") }()

		msg := "BREAKING NEWS"
		srv.Bcast(msg)

		// wait for message to get from server to client.
		select {
		case <-cli.MsgRecvd:
			// excllent
			po("excellent, cli got a message!\n")
			po("client saw: '%s'", cli.LastMsgReceived())
		case <-time.After(time.Second * 2):
			po("\n\nWe waited 2 seconds, bailing.\n")
			srv.Stop()
			panic("should have gotten a message from the server at the client by now!")
		}

		cv.So(cli.LastMsgReceived(), cv.ShouldEqual, msg)

	})

	//	fmt.Printf("Given a Forward and Reverse proxy, in order to avoid creating new sockets too often (expensive), we should re-use the existing sockets for up to 5 round trips in 30 seconds.")
}

func SoShouldBeEqual(a, b string) {
	if a != b {
		panic(fmt.Sprintf("'%s' != '%s' but we expected them be equal.", a, b))
	}
}
<commit_msg>ft working at long poll impl now<commit_after>package pelicantun

import (
	"fmt"
	"testing"
	"time"

	cv "github.com/glycerine/goconvey/convey"
)

func TestLongPollToGetLowLatency010(t *testing.T) {

	cli, srv, rev, fwd, err := StartTestSystemWithBcast()
	panicOn(err)
	defer srv.Stop()
	defer rev.Stop()
	defer fwd.Stop()

	cv.Convey("Given a ForwardProxy and a ReverseProxy communicating over http, in order to acheive low-latency sends from server to client, long-polling with two sockets (one send and one receive) should be used. So a server that has something to send (e.g. broadcast here) should be able to send back to the client immediately.\n", t, func() {
		cli.Start()
		defer cli.Stop()
		<-srv.FirstClient

		po("got past <-srv.FirstClient\n")

		defer func() { po("deferred frunction running...!\n") }()

		msg := "BREAKING NEWS"
		srv.Bcast(msg)

		// wait for message to get from server to client.
		select {
		case <-cli.MsgRecvd:
			// excllent
			po("excellent, cli got a message!\n")
			po("client saw: '%s'", cli.LastMsgReceived())
		case <-time.After(time.Second * 2):
			po("\n\nWe waited 2 seconds, bailing.\n")
			srv.Stop()
			// should have gotten a message from the server at the client by now.

			panic("should have gotten a message from the server at the client by now!")
			cv.So(cli.LastMsgReceived(), cv.ShouldEqual, msg)
		}

		cv.So(cli.LastMsgReceived(), cv.ShouldEqual, msg)

	})

	//	fmt.Printf("Given a Forward and Reverse proxy, in order to avoid creating new sockets too often (expensive), we should re-use the existing sockets for up to 5 round trips in 30 seconds.")
}

func SoShouldBeEqual(a, b string) {
	if a != b {
		panic(fmt.Sprintf("'%s' != '%s' but we expected them be equal.", a, b))
	}
}
<|endoftext|>


<commit_before>package apps

import (
	"fmt"
	"time"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa/events"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"
	. "github.com/cloudfoundry/cf-acceptance-tests/helpers/matchers"

	"crypto/tls"
	"strings"

	"encoding/json"
	"os"
	"path/filepath"
)

var _ = Describe("loggregator", func() {
	var appName string
	const hundredthOfOneSecond = 10000 // this app uses millionth of seconds

	BeforeEach(func() {
		appName = generator.PrefixedRandomName("CATS-APP-")

		Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().LoggregatorLoadGenerator, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		app_helpers.SetBackend(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	AfterEach(func() {
		app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

		Expect(cf.Cf("delete", appName, "-f", "-r").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
	})

	Context("cf logs", func() {
		var logs *Session

		BeforeEach(func() {
			logs = cf.Cf("logs", appName)
		})

		AfterEach(func() {
			// logs might be nil if the BeforeEach panics
			if logs != nil {
				logs.Interrupt().Wait(DEFAULT_TIMEOUT)
			}
		})

		It("exercises basic loggregator behavior", func() {
			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Connected, tailing logs for app"))

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Muahaha"))
		})
	})

	Context("cf logs --recent", func() {
		It("makes loggregator buffer and dump log messages", func() {
			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("Muahaha"))
		})
	})

	Context("firehose data", func() {
		It("shows logs and metrics", func() {
			config := helpers.LoadConfig()

			noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
			msgChan := make(chan *events.Envelope, 100000)
			errorChan := make(chan error)
			stopchan := make(chan struct{})
			go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
			defer close(stopchan)

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			timeout := time.After(5 * time.Second)
			messages := make([]*events.Envelope, 0, 100000)

			for {
				select {
				case <-timeout:
					return
				case msg := <-msgChan:
					messages = append(messages, msg)
				}
			}

			Expect(messages).To(ContainElement(EnvelopeContainingMessageLike("Muahaha")), "To enable the logging & metrics firehose feature, please ask your CF administrator to add the 'doppler.firehose' scope to your CF admin user.")
		})
	})

	It("shows container metrics", func() {
		config := helpers.LoadConfig()
		appGuid := strings.TrimSpace(string(cf.Cf("app", appName, "--guid").Wait(DEFAULT_TIMEOUT).Out.Contents()))

		noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
		msgChan := make(chan *events.Envelope, 100000)
		errorChan := make(chan error)
		stopchan := make(chan struct{})
		go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
		defer close(stopchan)

		Eventually(func() bool {
			for {
				select {
				case msg := <-msgChan:
					if cm := msg.GetContainerMetric(); cm != nil {
						if cm.GetApplicationId() == appGuid && cm.GetInstanceIndex() == 0 {
							return true
						}
					}
				default:
					return false
				}
			}
		}, DEFAULT_TIMEOUT).Should(BeTrue())
	})
})

type cfHomeConfig struct {
	AccessToken         string
	LoggregatorEndpoint string
}

func getCfHomeConfig() *cfHomeConfig {
	myCfHomeConfig := &cfHomeConfig{}

	cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
		path := filepath.Join(os.Getenv("CF_HOME"), ".cf", "config.json")

		configFile, err := os.Open(path)
		if err != nil {
			panic(err)
		}

		decoder := json.NewDecoder(configFile)
		err = decoder.Decode(myCfHomeConfig)
		if err != nil {
			panic(err)
		}
	})

	return myCfHomeConfig
}

func getAdminUserAccessToken() string {
	return getCfHomeConfig().AccessToken
}

func getDopplerEndpoint() string {
	return strings.Replace(getCfHomeConfig().LoggregatorEndpoint, "loggregator", "doppler", -1)
}
<commit_msg>Check for errors while slurping<commit_after>package apps

import (
	"fmt"
	"time"

	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/ginkgo"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gbytes"
	. "github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa"
	"github.com/cloudfoundry/cf-acceptance-tests/Godeps/_workspace/src/github.com/cloudfoundry/noaa/events"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/app_helpers"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers/assets"
	. "github.com/cloudfoundry/cf-acceptance-tests/helpers/matchers"

	"crypto/tls"
	"strings"

	"encoding/json"
	"os"
	"path/filepath"
)

var _ = Describe("loggregator", func() {
	var appName string
	const hundredthOfOneSecond = 10000 // this app uses millionth of seconds

	BeforeEach(func() {
		appName = generator.PrefixedRandomName("CATS-APP-")

		Expect(cf.Cf("push", appName, "--no-start", "-b", config.RubyBuildpackName, "-m", DEFAULT_MEMORY_LIMIT, "-p", assets.NewAssets().LoggregatorLoadGenerator, "-d", config.AppsDomain).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		app_helpers.SetBackend(appName)
		Expect(cf.Cf("start", appName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	AfterEach(func() {
		app_helpers.AppReport(appName, DEFAULT_TIMEOUT)

		Expect(cf.Cf("delete", appName, "-f", "-r").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
	})

	Context("cf logs", func() {
		var logs *Session

		BeforeEach(func() {
			logs = cf.Cf("logs", appName)
		})

		AfterEach(func() {
			// logs might be nil if the BeforeEach panics
			if logs != nil {
				logs.Interrupt().Wait(DEFAULT_TIMEOUT)
			}
		})

		It("exercises basic loggregator behavior", func() {
			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Connected, tailing logs for app"))

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(logs, (DEFAULT_TIMEOUT + time.Minute)).Should(Say("Muahaha"))
		})
	})

	Context("cf logs --recent", func() {
		It("makes loggregator buffer and dump log messages", func() {
			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			Eventually(func() *Session {
				appLogsSession := cf.Cf("logs", "--recent", appName)
				Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
				return appLogsSession
			}, DEFAULT_TIMEOUT).Should(Say("Muahaha"))
		})
	})

	Context("firehose data", func() {
		It("shows logs and metrics", func() {
			config := helpers.LoadConfig()

			noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
			msgChan := make(chan *events.Envelope, 100000)
			errorChan := make(chan error)
			stopchan := make(chan struct{})
			go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
			defer close(stopchan)

			Eventually(func() string {
				return helpers.CurlApp(appName, fmt.Sprintf("/log/sleep/%d", hundredthOfOneSecond))
			}, DEFAULT_TIMEOUT).Should(ContainSubstring("Muahaha"))

			timeout := time.After(5 * time.Second)
			messages := make([]*events.Envelope, 0, 100000)

			for {
				select {
				case <-timeout:
					return
				case msg := <-msgChan:
					messages = append(messages, msg)
				}
			}

			Expect(messages).To(ContainElement(EnvelopeContainingMessageLike("Muahaha")), "To enable the logging & metrics firehose feature, please ask your CF administrator to add the 'doppler.firehose' scope to your CF admin user.")
		})

		It("shows container metrics", func() {
			config := helpers.LoadConfig()
			appGuid := strings.TrimSpace(string(cf.Cf("app", appName, "--guid").Wait(DEFAULT_TIMEOUT).Out.Contents()))

			noaaConnection := noaa.NewConsumer(getDopplerEndpoint(), &tls.Config{InsecureSkipVerify: config.SkipSSLValidation}, nil)
			msgChan := make(chan *events.Envelope, 100000)
			errorChan := make(chan error)
			stopchan := make(chan struct{})
			go noaaConnection.Firehose(generator.RandomName(), getAdminUserAccessToken(), msgChan, errorChan, stopchan)
			defer close(stopchan)

			Eventually(func() bool {
				for {
					select {
					case msg := <-msgChan:
						if cm := msg.GetContainerMetric(); cm != nil {
							if cm.GetApplicationId() == appGuid && cm.GetInstanceIndex() == 0 {
								return true
							}
						}
					case e := <-errorChan:
						Expect(e).ToNot(HaveOccurred())
					default:
						return false
					}
				}
			}, DEFAULT_TIMEOUT).Should(BeTrue())
		})
	})
})

type cfHomeConfig struct {
	AccessToken         string
	LoggregatorEndpoint string
}

func getCfHomeConfig() *cfHomeConfig {
	myCfHomeConfig := &cfHomeConfig{}

	cf.AsUser(context.AdminUserContext(), DEFAULT_TIMEOUT, func() {
		path := filepath.Join(os.Getenv("CF_HOME"), ".cf", "config.json")

		configFile, err := os.Open(path)
		if err != nil {
			panic(err)
		}

		decoder := json.NewDecoder(configFile)
		err = decoder.Decode(myCfHomeConfig)
		if err != nil {
			panic(err)
		}
	})

	return myCfHomeConfig
}

func getAdminUserAccessToken() string {
	return getCfHomeConfig().AccessToken
}

func getDopplerEndpoint() string {
	return strings.Replace(getCfHomeConfig().LoggregatorEndpoint, "loggregator", "doppler", -1)
}
<|endoftext|>


<commit_before>package irc

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"time"

	"github.com/kyleterry/tenyks/config"
	"github.com/op/go-logging"
)

var log = logging.MustGetLogger("tenyks")

type Connection struct {
	Name            string
	Config          config.ConnectionConfig
	currentNick     string
	currentServer   string
	nickIndex       int
	connectAttempts uint
	usingSSL        bool
	socket          net.Conn
	In              <-chan string
	Out             chan<- string
	io              *bufio.ReadWriter
	connected       bool
	MessagesRecved  uint
	MessagesSent    uint
	Registry        *HandlerRegistry
	ConnectWait     chan bool
	LastPong        time.Time
	PongIn          chan bool
}

func NewConn(name string, conf config.ConnectionConfig) *Connection {
	registry := NewHandlerRegistry()
	conn := &Connection{
		Name:            name,
		Config:          conf,
		nickIndex:       0,
		connectAttempts: 0,
		usingSSL:        conf.Ssl,
		socket:          nil,
		io:              nil,
		connected:       false,
		Registry:        registry,
		ConnectWait:     make(chan bool, 1),
		PongIn:          make(chan bool, 1),
	}
	conn.addBaseHandlers()
	return conn
}

// Connect is a goroutine that returns a channel that is true if connected
// successfully and false if not.
// It returns a bool channel that when closed or is passed true means success.
func (self *Connection) Connect() chan bool {
	c := make(chan bool, 1)
	go func() {
		retries := 0
		for {
			if retries > self.Config.Retries {
				log.Error("[%s] Max retries reached.",
					self.Name)
				c <- false
				return
			}
			if self.socket != nil {
				break
			}
			server := fmt.Sprintf(
				"%s:%d",
				self.Config.Host,
				self.Config.Port)
			var socket net.Conn
			var err error
			if self.usingSSL {
				socket, err = tls.Dial("tcp", server, nil)
				if err != nil {
					log.Error("[%s] Connection failed... Retrying.",
						self.Name)
					retries += 1
					time.Sleep(time.Second * time.Duration(retries))
					continue
				}
			} else {
				socket, err = net.Dial("tcp", server)
				if err != nil {
					retries += 1
					continue
				}
			}
			self.socket = socket
			self.io = bufio.NewReadWriter(
				bufio.NewReader(self.socket),
				bufio.NewWriter(self.socket))
			self.In = self.recv()
			self.Out = self.send()
			self.connected = true
			c <- true
		}
		c <- true
	}()
	return c
}

// Disconnect will hangup the connection with IRC and reset channels and other
// important bootstrap attributes back to the defaults.
func (self *Connection) Disconnect() {
	if self.connected {
		log.Debug("[%s] Disconnect called", self.Name)
		close(self.Out)
		self.connected = false
		self.socket.Close()
		self.socket = nil
		self.ConnectWait = make(chan bool, 1)
	}
}

// send will kick off a gorouting that will loop forever. It will recieve data
// on a channel and send that to the IRC socket.
// It will return a string channel when called.
func (self *Connection) send() chan<- string {
	c := make(chan string, 1000)
	// goroutine for sending data to the IRC server
	go func() {
		log.Debug("[%s] Starting send loop", self.Name)
		for {
			select {
			case line, ok := <-c:
				if !ok {
					log.Debug("[%s] Stopping send loop", self.Name)
					return
				}
				self.MessagesSent += 1
				self.write(line)
			}
		}
	}()
	return c
}

// write will recieve a string and write it to the IO buffer. It then flushes
// the buffer which in turn will call write() on the socket.
// It might return an error if something goes wrong.
func (self *Connection) write(line string) error {
	_, wrerr := self.io.WriteString(line + "\r\n")
	if wrerr != nil {
		return wrerr
	}
	flerr := self.io.Flush()
	if flerr != nil {
		return flerr
	}
	return nil
}

// recv will kick off a goroutine that will loop forever. It will recieve data
// from a bufio reader and send that to a string channel. Since sockets can
// send nil when a disconnect occurs, it has a minor responsibility of calling
// the Disconnect method when that happens.
// It will return a string channel when called.
func (self *Connection) recv() <-chan string {
	c := make(chan string, 1000)
	// goroutine for receiving data from the IRC server
	go func() {
		for {
			rawLine, err := self.io.ReadString('\n')
			if err != nil {
				self.Disconnect()
				close(c)
				break
			}
			if rawLine != "" {
				//TODO strip newlines
				self.MessagesRecved += 1
				c <- rawLine
			}
		}
	}()
	return c
}

// watchdog will occasionally PING the server and hope to hear back a PONG.
// If no PONG is recieved in reasonable amount of time, it's safe to assume
// we have been disconnected.
func (self *Connection) watchdog() {
	for {
		<-time.After(time.Second * 60)
		dispatch("send_ping", self, nil)
		select {
		case <-self.PongIn:
			continue
		case <-time.After(time.Second * 20):
			self.Disconnect()
			break
		}
	}
}

// IsConnected can be called to detemine if a connection is still connected.
// It returns a bool
func (self *Connection) IsConnected() bool {
	return self.connected
}

// GetCurrentNick will return the nick currently being used in the IRC connection
// It returns a string
func (self *Connection) GetCurrentNick() string {
	return self.currentNick
}

func (self *Connection) String() string {
	msg := "Tenyks Connection for " + self.Name + ".\n"
	if self.connected {
		msg += "Connected to " + self.currentServer + "\n"
	} else {
		msg += "Disconnected\n"
	}
	return msg
}
<commit_msg>More annotations<commit_after>package irc

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net"
	"time"

	"github.com/kyleterry/tenyks/config"
	"github.com/op/go-logging"
)

var log = logging.MustGetLogger("tenyks")

type Connection struct {
	Name            string
	Config          config.ConnectionConfig
	currentNick     string
	currentServer   string
	nickIndex       int
	connectAttempts uint
	usingSSL        bool
	socket          net.Conn
	In              <-chan string
	Out             chan<- string
	io              *bufio.ReadWriter
	connected       bool
	MessagesRecved  uint
	MessagesSent    uint
	Registry        *HandlerRegistry
	ConnectWait     chan bool
	LastPong        time.Time
	PongIn          chan bool
}

// NewConn will create a new instance of an irc.Connection.
// It will return *irc.Connection
func NewConn(name string, conf config.ConnectionConfig) *Connection {
	registry := NewHandlerRegistry()
	conn := &Connection{
		Name:            name,
		Config:          conf,
		nickIndex:       0,
		connectAttempts: 0,
		usingSSL:        conf.Ssl,
		socket:          nil,
		io:              nil,
		connected:       false,
		Registry:        registry,
		ConnectWait:     make(chan bool, 1),
		PongIn:          make(chan bool, 1),
	}
	conn.addBaseHandlers()
	return conn
}

// Connect is a goroutine that returns a channel that is true if connected
// successfully and false if not.
// It returns a bool channel that when closed or is passed true means success.
func (self *Connection) Connect() chan bool {
	c := make(chan bool, 1)
	go func() {
		retries := 0
		for {
			if retries > self.Config.Retries {
				log.Error("[%s] Max retries reached.",
					self.Name)
				c <- false
				return
			}
			if self.socket != nil {
				break
			}
			server := fmt.Sprintf(
				"%s:%d",
				self.Config.Host,
				self.Config.Port)
			var socket net.Conn
			var err error
			if self.usingSSL {
				socket, err = tls.Dial("tcp", server, nil)
				if err != nil {
					log.Error("[%s] Connection failed... Retrying.",
						self.Name)
					retries += 1
					time.Sleep(time.Second * time.Duration(retries))
					continue
				}
			} else {
				socket, err = net.Dial("tcp", server)
				if err != nil {
					retries += 1
					continue
				}
			}
			self.socket = socket
			self.io = bufio.NewReadWriter(
				bufio.NewReader(self.socket),
				bufio.NewWriter(self.socket))
			self.In = self.recv()
			self.Out = self.send()
			self.connected = true
			c <- true
		}
		c <- true
	}()
	return c
}

// Disconnect will hangup the connection with IRC and reset channels and other
// important bootstrap attributes back to the defaults.
func (self *Connection) Disconnect() {
	if self.connected {
		log.Debug("[%s] Disconnect called", self.Name)
		close(self.Out)
		self.connected = false
		self.socket.Close()
		self.socket = nil
		self.ConnectWait = make(chan bool, 1)
	}
}

// send will kick off a gorouting that will loop forever. It will recieve data
// on a channel and send that to the IRC socket.
// It will return a string channel when called.
func (self *Connection) send() chan<- string {
	c := make(chan string, 1000)
	// goroutine for sending data to the IRC server
	go func() {
		log.Debug("[%s] Starting send loop", self.Name)
		for {
			select {
			case line, ok := <-c:
				if !ok {
					log.Debug("[%s] Stopping send loop", self.Name)
					return
				}
				self.MessagesSent += 1
				self.write(line)
			}
		}
	}()
	return c
}

// write will recieve a string and write it to the IO buffer. It then flushes
// the buffer which in turn will call write() on the socket.
// It might return an error if something goes wrong.
func (self *Connection) write(line string) error {
	_, wrerr := self.io.WriteString(line + "\r\n")
	if wrerr != nil {
		return wrerr
	}
	flerr := self.io.Flush()
	if flerr != nil {
		return flerr
	}
	return nil
}

// recv will kick off a goroutine that will loop forever. It will recieve data
// from a bufio reader and send that to a string channel. Since sockets can
// send nil when a disconnect occurs, it has a minor responsibility of calling
// the Disconnect method when that happens.
// It will return a string channel when called.
func (self *Connection) recv() <-chan string {
	c := make(chan string, 1000)
	// goroutine for receiving data from the IRC server
	go func() {
		for {
			rawLine, err := self.io.ReadString('\n')
			if err != nil {
				self.Disconnect()
				close(c)
				break
			}
			if rawLine != "" {
				//TODO strip newlines
				self.MessagesRecved += 1
				c <- rawLine
			}
		}
	}()
	return c
}

// watchdog will occasionally PING the server and hope to hear back a PONG.
// If no PONG is recieved in reasonable amount of time, it's safe to assume
// we have been disconnected.
func (self *Connection) watchdog() {
	for {
		<-time.After(time.Second * 60)
		dispatch("send_ping", self, nil)
		select {
		case <-self.PongIn:
			continue
		case <-time.After(time.Second * 20):
			self.Disconnect()
			break
		}
	}
}

// IsConnected can be called to detemine if a connection is still connected.
// It returns a bool
func (self *Connection) IsConnected() bool {
	return self.connected
}

// GetCurrentNick will return the nick currently being used in the IRC connection
// It returns a string
func (self *Connection) GetCurrentNick() string {
	return self.currentNick
}

func (self *Connection) String() string {
	msg := "Tenyks Connection for " + self.Name + ".\n"
	if self.connected {
		msg += "Connected to " + self.currentServer + "\n"
	} else {
		msg += "Disconnected\n"
	}
	return msg
}
<|endoftext|>


<commit_before>package digispark

import (
	"errors"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/drivers/gpio"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Adaptor = (*Adaptor)(nil)

var _ gpio.DigitalWriter = (*Adaptor)(nil)
var _ gpio.PwmWriter = (*Adaptor)(nil)
var _ gpio.ServoWriter = (*Adaptor)(nil)

type mock struct {
	locationA         uint8
	locationB         uint8
	pwmChannelA       uint8
	pwmChannelB       uint8
	pwmPrescalerValue uint
	pin               uint8
	mode              uint8
	state             uint8
}

func (l *mock) digitalWrite(pin uint8, state uint8) error {
	l.pin = pin
	l.state = state
	return l.error()
}
func (l *mock) pinMode(pin uint8, mode uint8) error {
	l.pin = pin
	l.mode = mode
	return l.error()
}

var pwmInitErrorFunc = func() error { return nil }

func (l *mock) pwmInit() error { return pwmInitErrorFunc() }
func (l *mock) pwmStop() error { return l.error() }
func (l *mock) pwmUpdateCompare(channelA uint8, channelB uint8) error {
	l.pwmChannelA = channelA
	l.pwmChannelB = channelB
	return l.error()
}
func (l *mock) pwmUpdatePrescaler(value uint) error {
	l.pwmPrescalerValue = value
	return l.error()
}
func (l *mock) servoInit() error { return l.error() }
func (l *mock) servoUpdateLocation(locationA uint8, locationB uint8) error {
	l.locationA = locationA
	l.locationB = locationB
	return l.error()
}

var errorFunc = func() error { return nil }

func (l *mock) error() error { return errorFunc() }

func initTestAdaptor() *Adaptor {
	a := NewAdaptor()
	a.connect = func(a *Adaptor) (err error) { return nil }
	a.littleWire = new(mock)
	errorFunc = func() error { return nil }
	pwmInitErrorFunc = func() error { return nil }
	return a
}

func TestAdaptorConnect(t *testing.T) {
	a := NewAdaptor()
	gobottest.Assert(t, a.Connect(), ErrConnection)

	a = initTestAdaptor()
	gobottest.Assert(t, a.Connect(), nil)
}

func TestAdaptorFinalize(t *testing.T) {
	a := initTestAdaptor()
	gobottest.Assert(t, a.Finalize(), nil)
}

func TestAdaptorDigitalWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.DigitalWrite("0", uint8(1))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).pin, uint8(0))
	gobottest.Assert(t, a.littleWire.(*mock).state, uint8(1))

	err = a.DigitalWrite("?", uint8(1))
	gobottest.Refute(t, err, nil)

	errorFunc = func() error { return errors.New("pin mode error") }
	err = a.DigitalWrite("0", uint8(1))
	gobottest.Assert(t, err, errors.New("pin mode error"))
}

func TestAdaptorServoWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.ServoWrite("2", uint8(80))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).locationA, uint8(80))
	gobottest.Assert(t, a.littleWire.(*mock).locationB, uint8(80))

	a = initTestAdaptor()
	errorFunc = func() error { return errors.New("servo error") }
	err = a.ServoWrite("2", uint8(80))
	gobottest.Assert(t, err, errors.New("servo error"))
}

func TestAdaptorPwmWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).pwmChannelA, uint8(100))
	gobottest.Assert(t, a.littleWire.(*mock).pwmChannelB, uint8(100))

	a = initTestAdaptor()
	pwmInitErrorFunc = func() error { return errors.New("pwminit error") }
	err = a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, errors.New("pwminit error"))

	a = initTestAdaptor()
	errorFunc = func() error { return errors.New("pwm error") }
	err = a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, errors.New("pwm error"))
}
<commit_msg>digispark: increase test coverage<commit_after>package digispark

import (
	"errors"
	"strings"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/drivers/gpio"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Adaptor = (*Adaptor)(nil)

var _ gpio.DigitalWriter = (*Adaptor)(nil)
var _ gpio.PwmWriter = (*Adaptor)(nil)
var _ gpio.ServoWriter = (*Adaptor)(nil)

type mock struct {
	locationA         uint8
	locationB         uint8
	pwmChannelA       uint8
	pwmChannelB       uint8
	pwmPrescalerValue uint
	pin               uint8
	mode              uint8
	state             uint8
}

func (l *mock) digitalWrite(pin uint8, state uint8) error {
	l.pin = pin
	l.state = state
	return l.error()
}
func (l *mock) pinMode(pin uint8, mode uint8) error {
	l.pin = pin
	l.mode = mode
	return l.error()
}

var pwmInitErrorFunc = func() error { return nil }

func (l *mock) pwmInit() error { return pwmInitErrorFunc() }
func (l *mock) pwmStop() error { return l.error() }
func (l *mock) pwmUpdateCompare(channelA uint8, channelB uint8) error {
	l.pwmChannelA = channelA
	l.pwmChannelB = channelB
	return l.error()
}
func (l *mock) pwmUpdatePrescaler(value uint) error {
	l.pwmPrescalerValue = value
	return l.error()
}
func (l *mock) servoInit() error { return l.error() }
func (l *mock) servoUpdateLocation(locationA uint8, locationB uint8) error {
	l.locationA = locationA
	l.locationB = locationB
	return l.error()
}

var errorFunc = func() error { return nil }

func (l *mock) error() error { return errorFunc() }

func initTestAdaptor() *Adaptor {
	a := NewAdaptor()
	a.connect = func(a *Adaptor) (err error) { return nil }
	a.littleWire = new(mock)
	errorFunc = func() error { return nil }
	pwmInitErrorFunc = func() error { return nil }
	return a
}

func TestDigisparkAdaptorName(t *testing.T) {
	a := NewAdaptor()
	gobottest.Assert(t, strings.HasPrefix(a.Name(), "Digispark"), true)
	a.SetName("NewName")
	gobottest.Assert(t, a.Name(), "NewName")
}

func TestAdaptorConnect(t *testing.T) {
	a := NewAdaptor()
	gobottest.Assert(t, a.Connect(), ErrConnection)

	a = initTestAdaptor()
	gobottest.Assert(t, a.Connect(), nil)
}

func TestAdaptorFinalize(t *testing.T) {
	a := initTestAdaptor()
	gobottest.Assert(t, a.Finalize(), nil)
}

func TestAdaptorDigitalWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.DigitalWrite("0", uint8(1))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).pin, uint8(0))
	gobottest.Assert(t, a.littleWire.(*mock).state, uint8(1))

	err = a.DigitalWrite("?", uint8(1))
	gobottest.Refute(t, err, nil)

	errorFunc = func() error { return errors.New("pin mode error") }
	err = a.DigitalWrite("0", uint8(1))
	gobottest.Assert(t, err, errors.New("pin mode error"))
}

func TestAdaptorServoWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.ServoWrite("2", uint8(80))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).locationA, uint8(80))
	gobottest.Assert(t, a.littleWire.(*mock).locationB, uint8(80))

	a = initTestAdaptor()
	errorFunc = func() error { return errors.New("servo error") }
	err = a.ServoWrite("2", uint8(80))
	gobottest.Assert(t, err, errors.New("servo error"))
}

func TestAdaptorPwmWrite(t *testing.T) {
	a := initTestAdaptor()
	err := a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, nil)
	gobottest.Assert(t, a.littleWire.(*mock).pwmChannelA, uint8(100))
	gobottest.Assert(t, a.littleWire.(*mock).pwmChannelB, uint8(100))

	a = initTestAdaptor()
	pwmInitErrorFunc = func() error { return errors.New("pwminit error") }
	err = a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, errors.New("pwminit error"))

	a = initTestAdaptor()
	errorFunc = func() error { return errors.New("pwm error") }
	err = a.PwmWrite("1", uint8(100))
	gobottest.Assert(t, err, errors.New("pwm error"))
}
<|endoftext|>