<commit_before>package workspace

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/GetStream/vg/internal/utils"
	"github.com/pkg/errors"
)

func (ws *Workspace) InstallLocalPackagePersistently(pkg string, localPath string) error {
	err := ws.InstallLocalPackage(pkg, localPath)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Persistent = true

	fmt.Fprintf(os.Stderr, "Persisting the local install for %q\n", pkg)
	return ws.SaveSettings(settings)
}

func commandExists(command string) (bool, error) {
	_, err := exec.LookPath("bindfs")
	if err != nil {
		execErr, ok := err.(*exec.Error)
		if !ok || execErr.Err != exec.ErrNotFound {
			return false, errors.WithStack(err)
		}
		// Command doesn't exist
		return false, nil
	}
	return true, nil
}

func (ws *Workspace) InstallLocalPackage(pkg string, localPath string) error {
	pkgDir := filepath.Join(path.Split(pkg))
	target := filepath.Join(ws.Src(), pkgDir)

	err := ws.Uninstall(pkg, os.Stderr)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg] = &localInstall{
		Path: localPath,
	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	hasBindfs, err := commandExists("bindfs")
	if err != nil {
		return err
	}

	if hasBindfs {
		err = ws.installLocalPackageWithBindfs(pkg, localPath, target)
	} else {
		err = ws.installLocalPackageWithSymlink(pkg, localPath, target)
	}

	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Successful = true
	return ws.SaveSettings(settings)

}

func (ws *Workspace) installLocalPackageWithBindfs(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using bindfs\n", src, pkg)

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Bindfs = true

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	err = os.MkdirAll(target, 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	cmd := exec.Command("bindfs", "--no-allow-other", src, target)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	return errors.WithStack(cmd.Run())
}

func (ws *Workspace) installLocalPackageWithSymlink(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using symbolic links\n", src, pkg)

	err := os.MkdirAll(filepath.Dir(target), 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	err = os.RemoveAll(target)
	if err != nil {
		return errors.WithStack(err)
	}

	return errors.WithStack(os.Symlink(src, target))
}

func (ws *Workspace) InstallSavedLocalPackages() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	for pkg, install := range settings.LocalInstalls {
		if install.Persistent {
			err = ws.InstallLocalPackagePersistently(pkg, install.Path)
		} else {
			err = ws.InstallLocalPackage(pkg, install.Path)
		}

		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *Workspace) Uninstall(pkg string, logWriter io.Writer) error {
	pkgDir := utils.PkgToDir(pkg)
	pkgSrc := filepath.Join(ws.Src(), pkgDir)
	_, err := os.Stat(pkgSrc)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.WithStack(err)
	}

	fmt.Fprintf(logWriter, "Uninstalling %q from workspace\n", pkg)
	// Check if locally installed
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	install, localInstalled := settings.LocalInstalls[pkg]

	if localInstalled && install.Bindfs {
		fmt.Fprintf(logWriter, "  Unmounting bindfs mount at %q\n", pkgSrc)
		stderrBuff := &bytes.Buffer{}
		outputBuff := &bytes.Buffer{}

		var cmd *exec.Cmd
		var notMountedOutput string

		hasFusermount, err := commandExists("fusermount")
		if err != nil {
			return err
		}

		if !hasFusermount {
			// Use fusermount if that exists
			cmd = exec.Command("fusermount", "-u", pkgSrc)
			notMountedOutput = fmt.Sprintf("fusermount: entry for %s not found", pkgSrc)
		} else {
			// Otherwise fallback to umount
			cmd = exec.Command("sudo", "umount", pkgSrc)
			notMountedOutput = fmt.Sprintf("umount: %s: not mounted", pkgSrc)
		}

		cmd.Stderr = io.MultiWriter(stderrBuff, outputBuff)
		cmd.Stdout = outputBuff

		err = cmd.Run()
		if err != nil {
			if !strings.HasPrefix(stderrBuff.String(), notMountedOutput) {
				// We don't care if the write to stderr failed
				_, _ = io.Copy(os.Stderr, outputBuff)

				return errors.WithStack(err)
			}
		}

	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	fmt.Fprintf(logWriter, "  Removing sources at %q\n", pkgSrc)
	err = os.RemoveAll(pkgSrc)
	if err != nil {
		return errors.Wrapf(err, "Couldn't remove package src %q", ws.Name())
	}

	if localInstalled && !install.Persistent {
		fmt.Fprintf(logWriter, "  Removing %q from locally installed packages\n", pkg)
		delete(settings.LocalInstalls, pkg)
		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	pkgInstalledDirs, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledDirs {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}

	pkgInstalledFiles, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir+".a"))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledFiles {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}
	return nil
}

func (ws *Workspace) ClearSrc() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	err = os.RemoveAll(ws.ensureMarker())
	if err != nil {
		return errors.WithStack(err)
	}

	for pkg := range settings.LocalInstalls {
		err := ws.Uninstall(pkg, os.Stdout)
		if err != nil {
			return err
		}
	}

	return errors.WithStack(os.RemoveAll(ws.Src()))

}

func (ws *Workspace) UnpersistLocalInstall(pkg string) error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	if install, ok := settings.LocalInstalls[pkg]; ok {
		fmt.Printf("Removing %q from persistent local installs\n", pkg)
		install.Persistent = false

		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Remove test code<commit_after>package workspace

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"github.com/GetStream/vg/internal/utils"
	"github.com/pkg/errors"
)

func (ws *Workspace) InstallLocalPackagePersistently(pkg string, localPath string) error {
	err := ws.InstallLocalPackage(pkg, localPath)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Persistent = true

	fmt.Fprintf(os.Stderr, "Persisting the local install for %q\n", pkg)
	return ws.SaveSettings(settings)
}

func commandExists(command string) (bool, error) {
	_, err := exec.LookPath("bindfs")
	if err != nil {
		execErr, ok := err.(*exec.Error)
		if !ok || execErr.Err != exec.ErrNotFound {
			return false, errors.WithStack(err)
		}
		// Command doesn't exist
		return false, nil
	}
	return true, nil
}

func (ws *Workspace) InstallLocalPackage(pkg string, localPath string) error {
	pkgDir := filepath.Join(path.Split(pkg))
	target := filepath.Join(ws.Src(), pkgDir)

	err := ws.Uninstall(pkg, os.Stderr)
	if err != nil {
		return err
	}

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg] = &localInstall{
		Path: localPath,
	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	hasBindfs, err := commandExists("bindfs")
	if err != nil {
		return err
	}

	if hasBindfs {
		err = ws.installLocalPackageWithBindfs(pkg, localPath, target)
	} else {
		err = ws.installLocalPackageWithSymlink(pkg, localPath, target)
	}

	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Successful = true
	return ws.SaveSettings(settings)

}

func (ws *Workspace) installLocalPackageWithBindfs(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using bindfs\n", src, pkg)

	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	settings.LocalInstalls[pkg].Bindfs = true

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	err = os.MkdirAll(target, 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	cmd := exec.Command("bindfs", "--no-allow-other", src, target)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stderr
	return errors.WithStack(cmd.Run())
}

func (ws *Workspace) installLocalPackageWithSymlink(pkg, src, target string) error {
	_, _ = fmt.Fprintf(os.Stderr, "Installing local sources at %q in workspace as %q using symbolic links\n", src, pkg)

	err := os.MkdirAll(filepath.Dir(target), 0755)
	if err != nil {
		return errors.WithStack(err)
	}

	err = os.RemoveAll(target)
	if err != nil {
		return errors.WithStack(err)
	}

	return errors.WithStack(os.Symlink(src, target))
}

func (ws *Workspace) InstallSavedLocalPackages() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	for pkg, install := range settings.LocalInstalls {
		if install.Persistent {
			err = ws.InstallLocalPackagePersistently(pkg, install.Path)
		} else {
			err = ws.InstallLocalPackage(pkg, install.Path)
		}

		if err != nil {
			return err
		}
	}
	return nil
}

func (ws *Workspace) Uninstall(pkg string, logWriter io.Writer) error {
	pkgDir := utils.PkgToDir(pkg)
	pkgSrc := filepath.Join(ws.Src(), pkgDir)
	_, err := os.Stat(pkgSrc)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.WithStack(err)
	}

	fmt.Fprintf(logWriter, "Uninstalling %q from workspace\n", pkg)
	// Check if locally installed
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	install, localInstalled := settings.LocalInstalls[pkg]

	if localInstalled && install.Bindfs {
		fmt.Fprintf(logWriter, "  Unmounting bindfs mount at %q\n", pkgSrc)
		stderrBuff := &bytes.Buffer{}
		outputBuff := &bytes.Buffer{}

		var cmd *exec.Cmd
		var notMountedOutput string

		hasFusermount, err := commandExists("fusermount")
		if err != nil {
			return err
		}

		if hasFusermount {
			// Use fusermount if that exists
			cmd = exec.Command("fusermount", "-u", pkgSrc)
			notMountedOutput = fmt.Sprintf("fusermount: entry for %s not found", pkgSrc)
		} else {
			// Otherwise fallback to umount
			cmd = exec.Command("umount", pkgSrc)
			notMountedOutput = fmt.Sprintf("umount: %s: not mounted", pkgSrc)
		}

		cmd.Stderr = io.MultiWriter(stderrBuff, outputBuff)
		cmd.Stdout = outputBuff

		err = cmd.Run()
		if err != nil {
			if !strings.HasPrefix(stderrBuff.String(), notMountedOutput) {
				// We don't care if the write to stderr failed
				_, _ = io.Copy(os.Stderr, outputBuff)

				return errors.WithStack(err)
			}
		}

	}

	err = ws.SaveSettings(settings)
	if err != nil {
		return err
	}

	fmt.Fprintf(logWriter, "  Removing sources at %q\n", pkgSrc)
	err = os.RemoveAll(pkgSrc)
	if err != nil {
		return errors.Wrapf(err, "Couldn't remove package src %q", ws.Name())
	}

	if localInstalled && !install.Persistent {
		fmt.Fprintf(logWriter, "  Removing %q from locally installed packages\n", pkg)
		delete(settings.LocalInstalls, pkg)
		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	pkgInstalledDirs, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledDirs {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}

	pkgInstalledFiles, err := filepath.Glob(filepath.Join(ws.Pkg(), "*", pkgDir+".a"))
	if err != nil {
		return errors.Wrapf(err, "Something went wrong when globbing files for %q", pkg)
	}

	for _, path := range pkgInstalledFiles {
		fmt.Fprintf(logWriter, "  Removing %q\n", path)

		err = os.RemoveAll(path)
		if err != nil {
			return errors.Wrapf(err, "Couldn't remove installed package files for %q", pkg)
		}
	}
	return nil
}

func (ws *Workspace) ClearSrc() error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	err = os.RemoveAll(ws.ensureMarker())
	if err != nil {
		return errors.WithStack(err)
	}

	for pkg := range settings.LocalInstalls {
		err := ws.Uninstall(pkg, os.Stdout)
		if err != nil {
			return err
		}
	}

	return errors.WithStack(os.RemoveAll(ws.Src()))

}

func (ws *Workspace) UnpersistLocalInstall(pkg string) error {
	settings, err := ws.Settings()
	if err != nil {
		return err
	}

	if install, ok := settings.LocalInstalls[pkg]; ok {
		fmt.Printf("Removing %q from persistent local installs\n", pkg)
		install.Persistent = false

		err = ws.SaveSettings(settings)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package ioStatTransform
import(
   "fmt"
    "../diskStat"
   )
var LastRawStat = make(map[string]diskStat.DiskStat)

func TransformStat(channel <-chan diskStat.DiskStat) (err error) {
for {
		stat := <- channel
		prevStat,in := LastRawStat[stat.Device]

		if in {
			//do calculations
			// fmt.Println(stat)
			// fmt.Println(prevStat)
			if(stat.ReadsCompleted == 0 && stat.WritesCompleted == 0) {
				continue
			}
			timeDiff:= stat.RecordTime - prevStat.RecordTime
			rrq := stat.ReadsCompleted - prevStat.ReadsCompleted
			fmt.Println( stat.Device, " rrqm/s ", rrq, " timeDiff ", timeDiff / 1000000)

		}
			LastRawStat[stat.Device] = stat
	}
} <commit_msg>format<commit_after>package ioStatTransform
import(
   "fmt"
    "../diskStat"
   )
var LastRawStat = make(map[string]diskStat.DiskStat)

func TransformStat(channel <-chan diskStat.DiskStat) (err error) {
for {
		stat := <- channel
		prevStat,in := LastRawStat[stat.Device]

		if in {
			//do calculations
			// fmt.Println(stat)
			// fmt.Println(prevStat)
			if(stat.ReadsCompleted == 0 && stat.WritesCompleted == 0) {
				continue
			}
			timeDiff:= stat.RecordTime - prevStat.RecordTime
			rrq := stat.ReadsCompleted - prevStat.ReadsCompleted
			fmt.Println( stat.Device, " rrqm/s ", rrq, " timeDiff ", timeDiff / 1000000)

		}
		LastRawStat[stat.Device] = stat
	}
} <|endoftext|>


<commit_before>package router

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	e2e "k8s.io/kubernetes/test/e2e/framework"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	configv1 "github.com/openshift/api/config/v1"
	unidlingapi "github.com/openshift/api/unidling/v1alpha1"
	exutil "github.com/openshift/origin/test/extended/util"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var _ = g.Describe("[sig-network-edge][Conformance][Area:Networking][Feature:Router]", func() {
	defer g.GinkgoRecover()

	var (
		configPath = exutil.FixturePath("testdata", "router", "router-idle.yaml")
		oc         = exutil.NewCLI("router-idling")
	)

	// this hook must be registered before the framework namespace teardown
	// hook
	g.AfterEach(func() {
		if g.CurrentGinkgoTestDescription().Failed {
			exutil.DumpPodLogsStartingWithInNamespace("router", "openshift-ingress", oc.AsAdmin())
		}
	})

	g.Describe("The HAProxy router", func() {
		g.It("should be able to connect to a service that is idled because a GET on the route will unidle it", func() {
			infra, err := oc.AdminConfigClient().ConfigV1().Infrastructures().Get(context.Background(), "cluster", metav1.GetOptions{})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to get cluster-wide infrastructure")
			switch infra.Status.PlatformStatus.Type {
			case configv1.OvirtPlatformType, configv1.KubevirtPlatformType, configv1.LibvirtPlatformType, configv1.VSpherePlatformType:
				// Skip on platforms where the default
				// router is not exposed by a load
				// balancer service.
				g.Skip("https://bugzilla.redhat.com/show_bug.cgi?id=1933114")
			}

			timeout := 15 * time.Minute

			g.By(fmt.Sprintf("creating test fixture from a config file %q", configPath))
			err = oc.Run("new-app").Args("-f", configPath).Execute()
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to create test fixture")

			g.By("Waiting for pods to be running")
			err = waitForRunningPods(oc, 1, exutil.ParseLabelsOrDie("app=idle-test"), timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "pods not running")

			g.By("Getting a 200 status code when accessing the route")
			hostname, err := getHostnameForRoute(oc, "idle-test")
			o.Expect(err).NotTo(o.HaveOccurred())
			err = waitHTTPGetStatus(hostname, http.StatusOK, timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "expected status 200 from the GET request")

			g.By("Idling the service")
			_, err = oc.Run("idle").Args("idle-test").Output()
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to idle the service")

			var annotations map[string]string

			g.By("Fetching the endpoints and checking the idle annotations are present")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				endpoints, err := oc.KubeClient().CoreV1().Endpoints(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting endpoints: %v", err)
					return false, nil
				}
				annotations = endpoints.Annotations
				_, idledAt := annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := annotations[unidlingapi.UnidleTargetAnnotation]
				return idledAt && unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to fetch the endpoints")
			mustVerifyIdleAnnotationValues(annotations)

			g.By("Fetching the service and checking the idle annotations are present")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				service, err := oc.KubeClient().CoreV1().Services(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting service: %v", err)
					return false, nil
				}
				annotations = service.Annotations
				_, idledAt := annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := annotations[unidlingapi.UnidleTargetAnnotation]
				return idledAt && unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to fetch the service")
			mustVerifyIdleAnnotationValues(annotations)

			// wait for target deployment to actually scale down
			err = waitForRunningPods(oc, 0, exutil.ParseLabelsOrDie("app=idle-test"), timeout)
			o.Expect(err).NotTo(o.HaveOccurred())

			g.By("Unidling the service by making a GET request on the route")
			err = waitHTTPGetStatus(hostname, http.StatusOK, timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "expected status 200 from the GET request")

			g.By("Validating that the idle annotations have been removed from the endpoints")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				endpoints, err := oc.KubeClient().CoreV1().Endpoints(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting endpoints: %v", err)
					return false, nil
				}
				_, idledAt := endpoints.Annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := endpoints.Annotations[unidlingapi.UnidleTargetAnnotation]
				return !idledAt && !unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "idle annotations not removed from endpoints")

			g.By("Validating that the idle annotations have been removed from the service")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				service, err := oc.KubeClient().CoreV1().Services(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting service: %v", err)
					return false, nil
				}
				_, idledAt := service.Annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := service.Annotations[unidlingapi.UnidleTargetAnnotation]
				return !idledAt && !unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "idle annotations not removed from service")
		})
	})
})

func mustVerifyIdleAnnotationValues(annotations map[string]string) {
	o.Expect(annotations).To(o.HaveKey(unidlingapi.IdledAtAnnotation))
	o.Expect(annotations).To(o.HaveKey(unidlingapi.UnidleTargetAnnotation))

	idledAtAnnotation := annotations[unidlingapi.IdledAtAnnotation]
	idledAtTime, err := time.Parse(time.RFC3339, idledAtAnnotation)
	o.Expect(err).ToNot(o.HaveOccurred())
	o.Expect(idledAtTime).To(o.BeTemporally("~", time.Now(), 5*time.Minute))

	g.By("Checking the idle targets")
	unidleTargetAnnotation := annotations[unidlingapi.UnidleTargetAnnotation]
	var unidleTargets []unidlingapi.RecordedScaleReference
	err = json.Unmarshal([]byte(unidleTargetAnnotation), &unidleTargets)
	o.Expect(err).ToNot(o.HaveOccurred())
	o.Expect(unidleTargets).To(o.Equal([]unidlingapi.RecordedScaleReference{
		{
			Replicas: 1,
			CrossGroupObjectReference: unidlingapi.CrossGroupObjectReference{
				Kind:  "Deployment",
				Group: "apps",
				Name:  "idle-test",
			},
		},
	}))
}

// waitForRunningPods waits for podCount pods matching podSelector are
// in the running state. It retries the request every second and will
// return an error if the conditions are not met after the specified
// timeout.
func waitForRunningPods(oc *exutil.CLI, podCount int, podLabels labels.Selector, timeout time.Duration) error {
	ns := oc.KubeFramework().Namespace.Name

	if err := wait.PollImmediate(time.Second, timeout, func() (bool, error) {
		podList, err := oc.AdminKubeClient().CoreV1().Pods(ns).List(context.Background(), metav1.ListOptions{LabelSelector: podLabels.String()})
		if err != nil {
			e2e.Logf("Error listing pods: %v", err)
			return false, nil
		}
		return len(podList.Items) == podCount, nil
	}); err != nil {
		return fmt.Errorf("failed to list pods: %v", err)
	}

	e2e.Logf("Waiting for %d pods in namespace %s", podCount, ns)
	c := oc.AdminKubeClient()
	pods, err := exutil.WaitForPods(c.CoreV1().Pods(ns), podLabels, exutil.CheckPodIsRunning, podCount, timeout)
	if err != nil {
		return fmt.Errorf("error in pod wait: %v", err)
	} else if len(pods) < podCount {
		return fmt.Errorf("only got %v out of %v pods in %s (timeout)", len(pods), podCount, timeout)
	}

	e2e.Logf("All expected pods in namespace %s are running", ns)
	return nil
}

// waitHTTPGetStatus repeatedly makes a HTTP GET request to hostname
// until the GET response equals statusCode. It retries every second
// and will return an error if the conditions are not met after the
// specified timeout.
func waitHTTPGetStatus(hostname string, statusCode int, timeout time.Duration) error {
	client := makeHTTPClient(false, 30*time.Second)
	var attempt int

	url := "http://" + hostname

	return wait.Poll(time.Second, timeout, func() (bool, error) {
		attempt += 1
		resp, err := client.Get(url)
		if err != nil {
			e2e.Logf("GET#%v %q error=%v", attempt, url, err)
			return false, nil // could be 503 if service not ready
		}
		defer resp.Body.Close()
		io.Copy(ioutil.Discard, resp.Body)
		e2e.Logf("GET#%v %q status=%v", attempt, url, resp.StatusCode)
		return resp.StatusCode == statusCode, nil
	})
}
<commit_msg>test/extended/router/idle: Only run on OVNKubernetes or OpenShiftSDN<commit_after>package router

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"time"

	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
	e2e "k8s.io/kubernetes/test/e2e/framework"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	configv1 "github.com/openshift/api/config/v1"
	unidlingapi "github.com/openshift/api/unidling/v1alpha1"
	exutil "github.com/openshift/origin/test/extended/util"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

var _ = g.Describe("[sig-network-edge][Conformance][Area:Networking][Feature:Router]", func() {
	defer g.GinkgoRecover()

	var (
		configPath = exutil.FixturePath("testdata", "router", "router-idle.yaml")
		oc         = exutil.NewCLI("router-idling")
	)

	// this hook must be registered before the framework namespace teardown
	// hook
	g.AfterEach(func() {
		if g.CurrentGinkgoTestDescription().Failed {
			exutil.DumpPodLogsStartingWithInNamespace("router", "openshift-ingress", oc.AsAdmin())
		}
	})

	g.Describe("The HAProxy router", func() {
		g.It("should be able to connect to a service that is idled because a GET on the route will unidle it", func() {
			network, err := oc.AdminConfigClient().ConfigV1().Networks().Get(context.Background(), "cluster", metav1.GetOptions{})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to get cluster network configuration")
			if !(network.Status.NetworkType == "OVNKubernetes" || network.Status.NetworkType == "OpenShiftSDN") {
				g.Skip("idle feature only supported on OVNKubernetes or OpenShiftSDN")
				return
			}

			infra, err := oc.AdminConfigClient().ConfigV1().Infrastructures().Get(context.Background(), "cluster", metav1.GetOptions{})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to get cluster-wide infrastructure")
			switch infra.Status.PlatformStatus.Type {
			case configv1.OvirtPlatformType, configv1.KubevirtPlatformType, configv1.LibvirtPlatformType, configv1.VSpherePlatformType:
				// Skip on platforms where the default
				// router is not exposed by a load
				// balancer service.
				g.Skip("https://bugzilla.redhat.com/show_bug.cgi?id=1933114")
			}

			timeout := 15 * time.Minute

			g.By(fmt.Sprintf("creating test fixture from a config file %q", configPath))
			err = oc.Run("new-app").Args("-f", configPath).Execute()
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to create test fixture")

			g.By("Waiting for pods to be running")
			err = waitForRunningPods(oc, 1, exutil.ParseLabelsOrDie("app=idle-test"), timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "pods not running")

			g.By("Getting a 200 status code when accessing the route")
			hostname, err := getHostnameForRoute(oc, "idle-test")
			o.Expect(err).NotTo(o.HaveOccurred())
			err = waitHTTPGetStatus(hostname, http.StatusOK, timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "expected status 200 from the GET request")

			g.By("Idling the service")
			_, err = oc.Run("idle").Args("idle-test").Output()
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to idle the service")

			var annotations map[string]string

			g.By("Fetching the endpoints and checking the idle annotations are present")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				endpoints, err := oc.KubeClient().CoreV1().Endpoints(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting endpoints: %v", err)
					return false, nil
				}
				annotations = endpoints.Annotations
				_, idledAt := annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := annotations[unidlingapi.UnidleTargetAnnotation]
				return idledAt && unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to fetch the endpoints")
			mustVerifyIdleAnnotationValues(annotations)

			g.By("Fetching the service and checking the idle annotations are present")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				service, err := oc.KubeClient().CoreV1().Services(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting service: %v", err)
					return false, nil
				}
				annotations = service.Annotations
				_, idledAt := annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := annotations[unidlingapi.UnidleTargetAnnotation]
				return idledAt && unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "failed to fetch the service")
			mustVerifyIdleAnnotationValues(annotations)

			// wait for target deployment to actually scale down
			err = waitForRunningPods(oc, 0, exutil.ParseLabelsOrDie("app=idle-test"), timeout)
			o.Expect(err).NotTo(o.HaveOccurred())

			g.By("Unidling the service by making a GET request on the route")
			err = waitHTTPGetStatus(hostname, http.StatusOK, timeout)
			o.Expect(err).NotTo(o.HaveOccurred(), "expected status 200 from the GET request")

			g.By("Validating that the idle annotations have been removed from the endpoints")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				endpoints, err := oc.KubeClient().CoreV1().Endpoints(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting endpoints: %v", err)
					return false, nil
				}
				_, idledAt := endpoints.Annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := endpoints.Annotations[unidlingapi.UnidleTargetAnnotation]
				return !idledAt && !unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "idle annotations not removed from endpoints")

			g.By("Validating that the idle annotations have been removed from the service")
			err = wait.PollImmediate(time.Second, timeout, func() (bool, error) {
				service, err := oc.KubeClient().CoreV1().Services(oc.Namespace()).Get(context.Background(), "idle-test", metav1.GetOptions{})
				if err != nil {
					e2e.Logf("Error getting service: %v", err)
					return false, nil
				}
				_, idledAt := service.Annotations[unidlingapi.IdledAtAnnotation]
				_, unidleTarget := service.Annotations[unidlingapi.UnidleTargetAnnotation]
				return !idledAt && !unidleTarget, nil
			})
			o.Expect(err).NotTo(o.HaveOccurred(), "idle annotations not removed from service")
		})
	})
})

func mustVerifyIdleAnnotationValues(annotations map[string]string) {
	o.Expect(annotations).To(o.HaveKey(unidlingapi.IdledAtAnnotation))
	o.Expect(annotations).To(o.HaveKey(unidlingapi.UnidleTargetAnnotation))

	idledAtAnnotation := annotations[unidlingapi.IdledAtAnnotation]
	idledAtTime, err := time.Parse(time.RFC3339, idledAtAnnotation)
	o.Expect(err).ToNot(o.HaveOccurred())
	o.Expect(idledAtTime).To(o.BeTemporally("~", time.Now(), 5*time.Minute))

	g.By("Checking the idle targets")
	unidleTargetAnnotation := annotations[unidlingapi.UnidleTargetAnnotation]
	var unidleTargets []unidlingapi.RecordedScaleReference
	err = json.Unmarshal([]byte(unidleTargetAnnotation), &unidleTargets)
	o.Expect(err).ToNot(o.HaveOccurred())
	o.Expect(unidleTargets).To(o.Equal([]unidlingapi.RecordedScaleReference{
		{
			Replicas: 1,
			CrossGroupObjectReference: unidlingapi.CrossGroupObjectReference{
				Kind:  "Deployment",
				Group: "apps",
				Name:  "idle-test",
			},
		},
	}))
}

// waitForRunningPods waits for podCount pods matching podSelector are
// in the running state. It retries the request every second and will
// return an error if the conditions are not met after the specified
// timeout.
func waitForRunningPods(oc *exutil.CLI, podCount int, podLabels labels.Selector, timeout time.Duration) error {
	ns := oc.KubeFramework().Namespace.Name

	if err := wait.PollImmediate(time.Second, timeout, func() (bool, error) {
		podList, err := oc.AdminKubeClient().CoreV1().Pods(ns).List(context.Background(), metav1.ListOptions{LabelSelector: podLabels.String()})
		if err != nil {
			e2e.Logf("Error listing pods: %v", err)
			return false, nil
		}
		return len(podList.Items) == podCount, nil
	}); err != nil {
		return fmt.Errorf("failed to list pods: %v", err)
	}

	e2e.Logf("Waiting for %d pods in namespace %s", podCount, ns)
	c := oc.AdminKubeClient()
	pods, err := exutil.WaitForPods(c.CoreV1().Pods(ns), podLabels, exutil.CheckPodIsRunning, podCount, timeout)
	if err != nil {
		return fmt.Errorf("error in pod wait: %v", err)
	} else if len(pods) < podCount {
		return fmt.Errorf("only got %v out of %v pods in %s (timeout)", len(pods), podCount, timeout)
	}

	e2e.Logf("All expected pods in namespace %s are running", ns)
	return nil
}

// waitHTTPGetStatus repeatedly makes a HTTP GET request to hostname
// until the GET response equals statusCode. It retries every second
// and will return an error if the conditions are not met after the
// specified timeout.
func waitHTTPGetStatus(hostname string, statusCode int, timeout time.Duration) error {
	client := makeHTTPClient(false, 30*time.Second)
	var attempt int

	url := "http://" + hostname

	return wait.Poll(time.Second, timeout, func() (bool, error) {
		attempt += 1
		resp, err := client.Get(url)
		if err != nil {
			e2e.Logf("GET#%v %q error=%v", attempt, url, err)
			return false, nil // could be 503 if service not ready
		}
		defer resp.Body.Close()
		io.Copy(ioutil.Discard, resp.Body)
		e2e.Logf("GET#%v %q status=%v", attempt, url, resp.StatusCode)
		return resp.StatusCode == statusCode, nil
	})
}
<|endoftext|>


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

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/youtube/vitess/go/testfiles"
)

func TestValidFiles(t *testing.T) {
	inputs := testfiles.Glob("bson_test/input*.go")
	for _, input := range inputs {
		b, err := ioutil.ReadFile(input)
		if err != nil {
			t.Fatal(err)
		}
		out, err := generateCode(string(b), "MyType")
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			return
		}
		want, err := ioutil.ReadFile(strings.Replace(input, "input", "output", 1))
		if err != nil {
			t.Fatal(err)
		}
		// goimports is flaky. So, let's not test that part.
		d, err := diff(skip_imports(want), skip_imports(out))
		if len(d) != 0 {
			t.Errorf("Unexpected output for %s:\n%s", input, string(d))
			if testing.Verbose() {
				t.Logf("%s:\n%s", input, out)
			}
		}
	}
}

// diff copied from gofmt.go
func diff(b1, b2 []byte) (data []byte, err error) {
	f1, err := ioutil.TempFile("", "bsongen")
	if err != nil {
		return
	}
	defer os.Remove(f1.Name())
	defer f1.Close()

	f2, err := ioutil.TempFile("", "bsongen")
	if err != nil {
		return
	}
	defer os.Remove(f2.Name())
	defer f2.Close()

	f1.Write(b1)
	f2.Write(b2)

	data, err = exec.Command("diff", "-u", f1.Name(), f2.Name()).CombinedOutput()
	if len(data) > 0 {
		// diff exits with a non-zero status when the files don't match.
		// Ignore that failure as long as we get output.
		err = nil
	}
	return
}

func skip_imports(b []byte) []byte {
	buf := bytes.NewBuffer(b)
	for {
		line, err := buf.ReadBytes('\n')
		if err != nil {
			return b[:0]
		}
		if len(line) == 0 || line[0] != ')' {
			continue
		}
		return b[len(b)-buf.Len():]
	}
}

var invalidInputs = []struct{ title, input, err string }{
	{
		"func type",
		`package a; func MyType(){};`,
		"MyType not found",
	}, {
		"non-struct non-simple top level type",
		`package a; type MyType Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		// Maybe support this in the future?
		"map type",
		`package a; type MyType map[string]Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		// Maybe support this in the future?
		"slice type",
		`package a; type MyType []Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		"anonymous embed",
		`package a; type MyType struct{Custom};`,
		"anonymous embeds not supported: Custom",
	}, {
		"interface with methods",
		`package a; type MyType struct{Val interface{Custom}};`,
		"is not a simple type",
	}, {
		// Maybe support this in the future?
		"array",
		`package a; type MyType struct{Val [5]int};`,
		"is not a simple type",
	},
}

func TestInvalidInputs(t *testing.T) {
	for _, tcase := range invalidInputs {
		out, err := generateCode(tcase.input, "MyType")
		if err == nil {
			t.Errorf("Expecting error for %s:\n%s", tcase.title, string(out))
		}
		if !strings.Contains(err.Error(), tcase.err) {
			t.Errorf("%s: got '%v', error should contain '%s'", tcase.title, err, tcase.err)
		}
	}
}
<commit_msg>Fix bsongen_test skipImports function.<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"errors"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/youtube/vitess/go/testfiles"
)

func TestValidFiles(t *testing.T) {
	inputs := testfiles.Glob("bson_test/input*.go")
	for _, input := range inputs {
		b, err := ioutil.ReadFile(input)
		if err != nil {
			t.Fatalf("ioutil.ReadFile error: %v", err)
		}
		want, err := ioutil.ReadFile(strings.Replace(input, "input", "output", 1))
		if err != nil {
			t.Fatalf("ioutil.ReadFile error: %v", err)
		}

		out, err := generateCode(string(b), "MyType")
		if err != nil {
			t.Fatalf("generateCode error: %v", err)
		}

		// goimports is flaky. So, let's not test that part.
		want, err = skipImports(want)
		if err != nil {
			t.Fatalf("skipImports error: %v", err)
		}
		out, err = skipImports(out)
		if err != nil {
			t.Fatalf("skipImports error: %v", err)
		}

		d, err := diff(want, out)
		if err != nil {
			t.Fatalf("diff error: %v", err)
		}
		if len(d) != 0 {
			t.Errorf("Unexpected output for %s:\n%s", input, string(d))
			if testing.Verbose() {
				t.Logf("%s:\n%s", input, out)
			}
		}
	}
}

// diff copied from gofmt.go
func diff(b1, b2 []byte) (data []byte, err error) {
	f1, err := ioutil.TempFile("", "bsongen")
	if err != nil {
		return
	}
	defer os.Remove(f1.Name())
	defer f1.Close()

	f2, err := ioutil.TempFile("", "bsongen")
	if err != nil {
		return
	}
	defer os.Remove(f2.Name())
	defer f2.Close()

	f1.Write(b1)
	f2.Write(b2)

	data, err = exec.Command("diff", "-u", f1.Name(), f2.Name()).CombinedOutput()
	if len(data) > 0 {
		// diff exits with a non-zero status when the files don't match.
		// Ignore that failure as long as we get output.
		err = nil
	}
	return
}

func skipImports(b []byte) ([]byte, error) {
	begin := bytes.Index(b, []byte("\nimport (\n"))
	if begin < 0 {
		return nil, errors.New("couldn't find beginning of import block")
	}
	end := bytes.Index(b, []byte("\n)\n"))
	if end < 0 {
		return nil, errors.New("couldn't find end of imports block")
	}
	return append(b[:begin], b[end+3:]...), nil
}

var invalidInputs = []struct{ title, input, err string }{
	{
		"func type",
		`package a; func MyType(){};`,
		"MyType not found",
	}, {
		"non-struct non-simple top level type",
		`package a; type MyType Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		// Maybe support this in the future?
		"map type",
		`package a; type MyType map[string]Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		// Maybe support this in the future?
		"slice type",
		`package a; type MyType []Custom;`,
		"MyType is not a struct or a simple type",
	}, {
		"anonymous embed",
		`package a; type MyType struct{Custom};`,
		"anonymous embeds not supported: Custom",
	}, {
		"interface with methods",
		`package a; type MyType struct{Val interface{Custom}};`,
		"is not a simple type",
	}, {
		// Maybe support this in the future?
		"array",
		`package a; type MyType struct{Val [5]int};`,
		"is not a simple type",
	},
}

func TestInvalidInputs(t *testing.T) {
	for _, tcase := range invalidInputs {
		out, err := generateCode(tcase.input, "MyType")
		if err == nil {
			t.Errorf("Expecting error for %s:\n%s", tcase.title, string(out))
		}
		if !strings.Contains(err.Error(), tcase.err) {
			t.Errorf("%s: got '%v', error should contain '%s'", tcase.title, err, tcase.err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"koding/broker/cache"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*cache.SubscriptionStorage)
	globalMapMutex         sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}

	// Closed when AMQP connection and channel are setup
	amqpReady chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
		amqpReady:         make(chan struct{}),
	}
}

func main() {
	NewBroker().Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	b.registerToKontrol()

	go b.startAMQP()
	<-b.amqpReady
	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	b.ConsumeConn = amqputil.CreateConnection("broker")
	defer b.ConsumeConn.Close()

	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// signal that we are ready now
	close(b.amqpReady)

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	client := NewClient(session, b)
	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	err := client.ControlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", client.SocketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<commit_msg>Broker: remove socketSubscriptionsMap from goBroker<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log             = logger.New("broker")
	STORAGE_BACKEND = "redis"
	routeMap        = make(map[string]([]*sockjs.Session))
	globalMapMutex  sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}

	// Closed when AMQP connection and channel are setup
	amqpReady chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
		amqpReady:         make(chan struct{}),
	}
}

func main() {
	NewBroker().Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	b.registerToKontrol()

	go b.startAMQP()
	<-b.amqpReady
	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	b.ConsumeConn = amqputil.CreateConnection("broker")
	defer b.ConsumeConn.Close()

	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// signal that we are ready now
	close(b.amqpReady)

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	client := NewClient(session, b)
	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	err := client.ControlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", client.SocketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

// sessionGaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// sessionGaugeStart and calls the returned function in a defer statement.
func sessionGaugeStart(session *sockjs.Session) (sessionGaugeEnd func()) {
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge := logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", session.Tag)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	sessionGaugeEnd := sessionGaugeStart(session)
	defer sessionGaugeEnd()

	socketId := randomString()
	session.Tag = socketId
	log.Debug("Client connected: %v", socketId)

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<commit_msg>broker: put connected and disconnected logs to gauge function<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/streadway/amqp"
)

var (
	log                    = logger.New("broker")
	routeMap               = make(map[string]([]*sockjs.Session))
	socketSubscriptionsMap = make(map[string]*map[string]bool)
	globalMapMutex         sync.Mutex
)

type Broker struct {
	Hostname          string
	ServiceUniqueName string
	PublishConn       *amqp.Connection
}

func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(config.BrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := "broker" + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
	}
}

func main() {
	lifecycle.Startup("broker", false)
	logger.RunGaugesLoop(log)

	broker := NewBroker()
	broker.registerToKontrol()

	go broker.startSockJS()
	broker.startAMQP() // blocking

	time.Sleep(5 * time.Second) // give amqputil time to log connection error
}

func (b *Broker) registerToKontrol() {
	if err := kontrolhelper.RegisterToKontrol(
		"broker", // servicename
		"broker",
		b.ServiceUniqueName,
		config.Uuid,
		b.Hostname,
		config.Current.Broker.Port,
	); err != nil {
		panic(err)
	}
}

func (b *Broker) startAMQP() {
	b.PublishConn = amqputil.CreateConnection("broker")
	defer b.PublishConn.Close()

	consumeConn := amqputil.CreateConnection("broker")
	defer consumeConn.Close()

	consumeChannel := amqputil.CreateChannel(consumeConn)
	defer consumeChannel.Close()

	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		"broker",            // serviceType
		"broker",            // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", "broker", "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		panic(err)
	}

	if err := consumeChannel.ExchangeBind("broker", "", "updateInstances", false, nil); err != nil {
		panic(err)
	}

	// start to listen from "broker" topic exchange
	for amqpMessage := range stream {
		routingKey := amqpMessage.RoutingKey
		payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

		pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
		for pos != -1 && pos < len(routingKey) {
			index := strings.IndexRune(routingKey[pos+1:], '.')
			pos += index + 1
			if index == -1 {
				pos = len(routingKey)
			}
			prefix := routingKey[:pos]
			globalMapMutex.Lock()
			for _, routeSession := range routeMap[prefix] {
				sendToClient(routeSession, routingKey, &payload)
			}
			globalMapMutex.Unlock()
		}
	}
}

func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		config.Current.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	server := &http.Server{
		Handler: &sockjs.Mux{
			Handlers: map[string]http.Handler{
				"/subscribe": service,
				"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					w.Header().Set("Content-Type", "text/plain")
					w.Write([]byte(strconv.Itoa(config.Current.BuildNumber)))
				}),
			},
		},
	}

	var listener net.Listener
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(config.Current.Broker.IP), Port: config.Current.Broker.Port})
	if err != nil {
		log.Fatal(err)
	}

	if config.Current.Broker.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(config.Current.Broker.CertFile, config.Current.Broker.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		listener = tls.NewListener(listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	lastErrorTime := time.Now()
	for {
		err := server.Serve(listener)
		if err != nil {
			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}

func randomString() string {
	r := make([]byte, 128/8)
	rand.Read(r)
	return base64.StdEncoding.EncodeToString(r)
}

// sessionGaugeStart starts the gauge for a given session. It returns a new
// function which ends the gauge for the given session. Usually one invokes
// sessionGaugeStart and calls the returned function in a defer statement.
func sessionGaugeStart(session *sockjs.Session) (sessionGaugeEnd func()) {
	log.Debug("Client connected: %v", session.Tag)
	changeClientsGauge := lifecycle.CreateClientsGauge()
	changeNewClientsGauge := logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge := logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	changeClientsGauge(1)
	changeNewClientsGauge(1)
	if session.IsWebsocket {
		changeWebsocketClientsGauge(1)
	}

	return func() {
		log.Debug("Client disconnected: %v", session.Tag)
		changeClientsGauge(-1)
		if session.IsWebsocket {
			changeWebsocketClientsGauge(-1)
		}
	}
}

func (b *Broker) sockjsSession(session *sockjs.Session) {
	defer log.RecoverAndLog()

	socketId := randomString()
	session.Tag = socketId

	sessionGaugeEnd := sessionGaugeStart(session)
	defer sessionGaugeEnd()

	var controlChannel *amqp.Channel
	var lastPayload string
	resetControlChannel := func() {
		if controlChannel != nil {
			controlChannel.Close()
		}
		var err error
		controlChannel, err = b.PublishConn.Channel()
		if err != nil {
			panic(err)
		}
		go func() {
			defer log.RecoverAndLog()

			for amqpErr := range controlChannel.NotifyClose(make(chan *amqp.Error)) {
				if !(strings.Contains(amqpErr.Error(), "NOT_FOUND") && (strings.Contains(amqpErr.Error(), "koding-social-") || strings.Contains(amqpErr.Error(), "auth-"))) {
					log.Warning("AMQP channel: %v Last publish payload: %v", amqpErr.Error(), lastPayload)
				}

				sendToClient(session, "broker.error", map[string]interface{}{"code": amqpErr.Code, "reason": amqpErr.Reason, "server": amqpErr.Server, "recover": amqpErr.Recover})
			}
		}()
	}
	resetControlChannel()
	defer func() { controlChannel.Close() }()

	subscriptions := make(map[string]bool)
	globalMapMutex.Lock()
	socketSubscriptionsMap[socketId] = &subscriptions
	globalMapMutex.Unlock()

	removeFromRouteMap := func(routingKeyPrefix string) {
		routeSessions := routeMap[routingKeyPrefix]
		for i, routeSession := range routeSessions {
			if routeSession == session {
				routeSessions[i] = routeSessions[len(routeSessions)-1]
				routeSessions = routeSessions[:len(routeSessions)-1]
				break
			}
		}
		if len(routeSessions) == 0 {
			delete(routeMap, routingKeyPrefix)
			return
		}
		routeMap[routingKeyPrefix] = routeSessions
	}

	subscribe := func(routingKeyPrefix string) {
		if subscriptions[routingKeyPrefix] {
			log.Warning("Duplicate subscription to same routing key. %v %v", session.Tag, routingKeyPrefix)
			return
		}
		if len(subscriptions) > 0 && len(subscriptions)%2000 == 0 {
			log.Warning("Client with more than %v subscriptions %v", strconv.Itoa(len(subscriptions)), session.Tag)
		}
		routeMap[routingKeyPrefix] = append(routeMap[routingKeyPrefix], session)
		subscriptions[routingKeyPrefix] = true
	}

	unsubscribe := func(routingKeyPrefix string) {
		removeFromRouteMap(routingKeyPrefix)
		delete(subscriptions, routingKeyPrefix)
	}

	defer func() {
		globalMapMutex.Lock()
		for routingKeyPrefix := range subscriptions {
			removeFromRouteMap(routingKeyPrefix)
		}
		globalMapMutex.Unlock()

		time.AfterFunc(5*time.Minute, func() {
			globalMapMutex.Lock()
			delete(socketSubscriptionsMap, socketId)
			globalMapMutex.Unlock()
		})

		for {
			err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientDisconnected", false, false, amqp.Publishing{Body: []byte(socketId)})
			if err == nil {
				break
			}
			if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
				panic(err)
			}
			resetControlChannel()
		}
	}()

	err := controlChannel.Publish(config.Current.Broker.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(socketId)})
	if err != nil {
		panic(err)
	}

	sendToClient(session, "broker.connected", socketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		defer log.RecoverAndLog()

		message := data.(map[string]interface{})
		log.Debug("Received message: %v", message)

		action := message["action"]
		switch action {
		case "subscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				subscribe(routingKeyPrefix)
			}
			sendToClient(session, "broker.subscribed", message["routingKeyPrefix"])

		case "resubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			oldSubscriptions, found := socketSubscriptionsMap[message["socketId"].(string)]
			if found {
				for routingKeyPrefix := range *oldSubscriptions {
					subscribe(routingKeyPrefix)
				}
			}
			sendToClient(session, "broker.resubscribed", found)

		case "unsubscribe":
			globalMapMutex.Lock()
			defer globalMapMutex.Unlock()
			for _, routingKeyPrefix := range strings.Split(message["routingKeyPrefix"].(string), " ") {
				unsubscribe(routingKeyPrefix)
			}

		case "publish":
			exchange := message["exchange"].(string)
			routingKey := message["routingKey"].(string)
			if !strings.HasPrefix(routingKey, "client.") {
				log.Warning("Invalid routing key: message: %v socketId: %v", message, socketId)
				return
			}
			for {
				lastPayload = ""
				err := controlChannel.Publish(exchange, routingKey, false, false, amqp.Publishing{CorrelationId: socketId, Body: []byte(message["payload"].(string))})
				if err == nil {
					lastPayload = message["payload"].(string)
					break
				}
				if amqpError, isAmqpError := err.(*amqp.Error); !isAmqpError || amqpError.Code != 504 {
					log.Warning("payload: %v routing key: %v exchange: %v err: %v", message["payload"], message["routingKey"], message["exchange"], err)
				}
				time.Sleep(time.Second / 4) // penalty for crashing the AMQP channel
				resetControlChannel()
			}

		case "ping":
			sendToClient(session, "broker.pong", nil)

		default:
			log.Warning("Invalid action. message: %v socketId: %v", message, socketId)

		}
	}
}
<|endoftext|>


<commit_before>package config

import "encoding/json"

//go:generate ../../../bin/go-bindata -modtime 1446555960 -pkg config -o config.json.go config.json

// Builtin represents an embedded configuration.
var Builtin *Config

// Config contains default values for services.
type Config struct {
	Endpoints Endpoints `json:"endpoints"`
}

// Endpoint represents a single worker's endpoint
// mapped to environments
type Endpoint struct {
	Environment []string `json:"environment"`
	URL         string   `json:"url"`
}

// Endpoints describes all workers and
// their default configuration.
type Endpoints map[string][]*Endpoint

// URL gives a url address for the given worker and environment.
//
// If specified worker does not exist
func (e Endpoints) URL(worker, environment string) string {
	endpoints, ok := e[worker]
	if !ok {
		return ""
	}

	var url string

	for _, endpoint := range endpoints {
		// Default URL is an endpoint with no explicit
		// environments set.
		if len(endpoint.Environment) == 0 {
			url = endpoint.URL
			continue
		}

		for _, env := range endpoint.Environment {
			switch env {
			case environment:
				return endpoint.URL
			case "development", "devmanaged":
				// If there's no default endpoint we fallback
				// to a development one.
				if url == "" {
					url = endpoint.URL
				}
			}
		}
	}

	return url
}

func init() {
	p, err := Asset("config.json")
	if err != nil {
		panic(err)
	}

	Builtin = &Config{}

	if err := json.Unmarshal(p, Builtin); err != nil {
		panic(err)
	}
}
<commit_msg>wercker: fix running go generate koding/config<commit_after>package config

import "encoding/json"

//go:generate go-bindata -mode 420 -modtime 1446555960 -pkg config -o config.json.go config.json

// Builtin represents an embedded configuration.
var Builtin *Config

// Config contains default values for services.
type Config struct {
	Endpoints Endpoints `json:"endpoints"`
}

// Endpoint represents a single worker's endpoint
// mapped to environments
type Endpoint struct {
	Environment []string `json:"environment"`
	URL         string   `json:"url"`
}

// Endpoints describes all workers and
// their default configuration.
type Endpoints map[string][]*Endpoint

// URL gives a url address for the given worker and environment.
//
// If specified worker does not exist
func (e Endpoints) URL(worker, environment string) string {
	endpoints, ok := e[worker]
	if !ok {
		return ""
	}

	var url string

	for _, endpoint := range endpoints {
		// Default URL is an endpoint with no explicit
		// environments set.
		if len(endpoint.Environment) == 0 {
			url = endpoint.URL
			continue
		}

		for _, env := range endpoint.Environment {
			switch env {
			case environment:
				return endpoint.URL
			case "development", "devmanaged":
				// If there's no default endpoint we fallback
				// to a development one.
				if url == "" {
					url = endpoint.URL
				}
			}
		}
	}

	return url
}

func init() {
	p, err := Asset("config.json")
	if err != nil {
		panic(err)
	}

	Builtin = &Config{}

	if err := json.Unmarshal(p, Builtin); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package log

import (
	"bytes"
	"encoding/json"
	"fmt"
	"koding/tools/config"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"time"
)

type Gauge struct {
	Name   string  `json:"name"`
	Value  float64 `json:"value"`
	Source string  `json:"source"`
	input  func() float64
}

var loggrSource string
var libratoSource string
var tags string

var gauges = make([]Gauge, 0)
var GaugeChanges = make(chan func())

func Init(service string) {
	hostname, _ := os.Hostname()
	loggrSource = fmt.Sprintf("%s %d on %s", service, os.Getpid(), strings.Split(hostname, ".")[0])
	libratoSource = fmt.Sprintf("%s.%d:%s", service, os.Getpid(), hostname)
	tags = service + " " + config.Profile

	CreateGauge("goroutines", func() float64 {
		return float64(runtime.NumGoroutine())
	})
	CreateGauge("memory", func() float64 {
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		return float64(m.Alloc)
	})
}

func NewEvent(level int, text string, data ...interface{}) url.Values {
	event := url.Values{
		"source": {loggrSource},
		"tags":   {LEVEL_TAGS[level] + " " + tags},
		"text":   {text},
	}
	if len(data) != 0 {
		dataStrings := make([]string, len(data))
		for i, part := range data {
			if bytes, ok := part.([]byte); ok {
				dataStrings[i] = string(bytes)
				continue
			}
			dataStrings[i] = fmt.Sprint(part)
		}
		event.Add("data", strings.Join(dataStrings, "\n"))
	}
	return event
}

func Send(event url.Values) {
	if !config.Current.Loggr.Push {
		fmt.Printf("%-30s %s\n", "["+event.Get("tags")+"]", event.Get("text"))
		if event.Get("data") != "" {
			for _, line := range strings.Split(event.Get("data"), "\n") {
				fmt.Printf("%-30s %s\n", "", line)
			}
		}
		return
	}

	event.Add("apikey", config.Current.Loggr.ApiKey)
	resp, err := http.PostForm(config.Current.Loggr.Url, event)
	if err != nil || resp.StatusCode != http.StatusCreated {
		fmt.Println("logger error: http.PostForm failed.\n%v\n%v\n%v\n", event, resp, err)
	}
	if resp.Body != nil {
		resp.Body.Close()
	}
}

func Log(level int, text string, data ...interface{}) {
	if level == DEBUG && !config.LogDebug {
		return
	}
	Send(NewEvent(level, text, data...))
}

const (
	ERR = iota
	WARN
	INFO
	DEBUG
)

var LEVEL_TAGS = []string{"error", "warning", "info", "debug"}

func Err(text string, data ...interface{}) {
	Log(ERR, text, data...)
}

func Warn(text string, data ...interface{}) {
	Log(WARN, text, data...)
}

func Info(text string, data ...interface{}) {
	Log(INFO, text, data...)
}

func Debug(text string, data ...interface{}) {
	Log(DEBUG, text, data...)
}

func LogError(err interface{}, stackOffset int) {
	data := make([]interface{}, 0)
	for i := 1 + stackOffset; ; i++ {
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		name := "<unknown>"
		if fn := runtime.FuncForPC(pc); fn != nil {
			name = fn.Name()
		}
		data = append(data, fmt.Sprintf("at %s (%s:%d)", name, file, line))
	}
	Log(ERR, fmt.Sprint(err), data...)
}

func RecoverAndLog() {
	err := recover()
	if err != nil {
		LogError(err, 2)
	}
}

func CreateGauge(name string, input func() float64) {
	gauges = append(gauges, Gauge{name, 0, libratoSource, input})
}

func CreateCounterGauge(name string, resetOnReport bool) func(int) {
	value := new(int)
	CreateGauge(name, func() float64 {
		v := *value
		if resetOnReport {
			*value = 0
		}
		return float64(v)
	})
	return func(diff int) {
		GaugeChanges <- func() {
			*value += diff
		}
	}
}

func RunGaugesLoop() {
	go func() {
		reportInterval := time.Tick(time.Duration(config.Current.Librato.Interval) * time.Millisecond)
		for {
			select {
			case <-reportInterval:
				LogGauges()

			case change := <-GaugeChanges:
				change()
			}
		}
	}()
}

func LogGauges() {
	if !config.Current.Librato.Push {
		tagPrefix := "[gauges " + tags + "]"
		for _, gauge := range gauges {
			fmt.Printf("%-30s %s: %v\n", tagPrefix, gauge.Name, gauge.input())
			tagPrefix = ""
		}
		return
	}

	for _, gauge := range gauges {
		gauge.Value = gauge.input()
	}
	var event struct {
		Gauges []Gauge `json:"gauges"`
	}
	event.Gauges = gauges
	b, err := json.Marshal(event)
	if err != nil {
		fmt.Println("logger error: json.Marshal failed.", err)
		return
	}

	request, err := http.NewRequest("POST", "https://metrics-api.librato.com/v1/metrics", bytes.NewReader(b))
	if err != nil {
		fmt.Println("logger error: http.NewRequest failed.", err)
		return
	}
	request.SetBasicAuth(config.Current.Librato.Email, config.Current.Librato.Token)
	request.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(request)
	if err != nil || resp.StatusCode != http.StatusOK {
		fmt.Printf("logger error: http.Post failed.\n%v\n%v\n%v\n", string(b), resp, err)
	}
	if resp.Body != nil {
		resp.Body.Close()
	}
}
<commit_msg>Fixed librato.<commit_after>package log

import (
	"bytes"
	"encoding/json"
	"fmt"
	"koding/tools/config"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"time"
)

type Gauge struct {
	Name   string  `json:"name"`
	Value  float64 `json:"value"`
	Source string  `json:"source"`
	input  func() float64
}

var loggrSource string
var libratoSource string
var tags string

var gauges = make([]*Gauge, 0)
var GaugeChanges = make(chan func())

func Init(service string) {
	hostname, _ := os.Hostname()
	loggrSource = fmt.Sprintf("%s %d on %s", service, os.Getpid(), strings.Split(hostname, ".")[0])
	libratoSource = fmt.Sprintf("%s.%d:%s", service, os.Getpid(), hostname)
	tags = service + " " + config.Profile

	CreateGauge("goroutines", func() float64 {
		return float64(runtime.NumGoroutine())
	})
	CreateGauge("memory", func() float64 {
		var m runtime.MemStats
		runtime.ReadMemStats(&m)
		return float64(m.Alloc)
	})
}

func NewEvent(level int, text string, data ...interface{}) url.Values {
	event := url.Values{
		"source": {loggrSource},
		"tags":   {LEVEL_TAGS[level] + " " + tags},
		"text":   {text},
	}
	if len(data) != 0 {
		dataStrings := make([]string, len(data))
		for i, part := range data {
			if bytes, ok := part.([]byte); ok {
				dataStrings[i] = string(bytes)
				continue
			}
			dataStrings[i] = fmt.Sprint(part)
		}
		event.Add("data", strings.Join(dataStrings, "\n"))
	}
	return event
}

func Send(event url.Values) {
	if !config.Current.Loggr.Push {
		fmt.Printf("%-30s %s\n", "["+event.Get("tags")+"]", event.Get("text"))
		if event.Get("data") != "" {
			for _, line := range strings.Split(event.Get("data"), "\n") {
				fmt.Printf("%-30s %s\n", "", line)
			}
		}
		return
	}

	event.Add("apikey", config.Current.Loggr.ApiKey)
	resp, err := http.PostForm(config.Current.Loggr.Url, event)
	if err != nil || resp.StatusCode != http.StatusCreated {
		fmt.Println("logger error: http.PostForm failed.\n%v\n%v\n%v\n", event, resp, err)
	}
	if resp.Body != nil {
		resp.Body.Close()
	}
}

func Log(level int, text string, data ...interface{}) {
	if level == DEBUG && !config.LogDebug {
		return
	}
	Send(NewEvent(level, text, data...))
}

const (
	ERR = iota
	WARN
	INFO
	DEBUG
)

var LEVEL_TAGS = []string{"error", "warning", "info", "debug"}

func Err(text string, data ...interface{}) {
	Log(ERR, text, data...)
}

func Warn(text string, data ...interface{}) {
	Log(WARN, text, data...)
}

func Info(text string, data ...interface{}) {
	Log(INFO, text, data...)
}

func Debug(text string, data ...interface{}) {
	Log(DEBUG, text, data...)
}

func LogError(err interface{}, stackOffset int) {
	data := make([]interface{}, 0)
	for i := 1 + stackOffset; ; i++ {
		pc, file, line, ok := runtime.Caller(i)
		if !ok {
			break
		}
		name := "<unknown>"
		if fn := runtime.FuncForPC(pc); fn != nil {
			name = fn.Name()
		}
		data = append(data, fmt.Sprintf("at %s (%s:%d)", name, file, line))
	}
	Log(ERR, fmt.Sprint(err), data...)
}

func RecoverAndLog() {
	err := recover()
	if err != nil {
		LogError(err, 2)
	}
}

func CreateGauge(name string, input func() float64) {
	gauges = append(gauges, &Gauge{name, 0, libratoSource, input})
}

func CreateCounterGauge(name string, resetOnReport bool) func(int) {
	value := new(int)
	CreateGauge(name, func() float64 {
		v := *value
		if resetOnReport {
			*value = 0
		}
		return float64(v)
	})
	return func(diff int) {
		GaugeChanges <- func() {
			*value += diff
		}
	}
}

func RunGaugesLoop() {
	go func() {
		reportInterval := time.Tick(time.Duration(config.Current.Librato.Interval) * time.Millisecond)
		for {
			select {
			case <-reportInterval:
				LogGauges()

			case change := <-GaugeChanges:
				change()
			}
		}
	}()
}

func LogGauges() {
	if !config.Current.Librato.Push {
		tagPrefix := "[gauges " + tags + "]"
		for _, gauge := range gauges {
			fmt.Printf("%-30s %s: %v\n", tagPrefix, gauge.Name, gauge.input())
			tagPrefix = ""
		}
		return
	}

	for _, gauge := range gauges {
		gauge.Value = gauge.input()
	}
	var event struct {
		Gauges []*Gauge `json:"gauges"`
	}
	event.Gauges = gauges
	b, err := json.Marshal(event)
	if err != nil {
		fmt.Println("logger error: json.Marshal failed.", err)
		return
	}

	request, err := http.NewRequest("POST", "https://metrics-api.librato.com/v1/metrics", bytes.NewReader(b))
	if err != nil {
		fmt.Println("logger error: http.NewRequest failed.", err)
		return
	}
	request.SetBasicAuth(config.Current.Librato.Email, config.Current.Librato.Token)
	request.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(request)
	if err != nil || resp.StatusCode != http.StatusOK {
		fmt.Printf("logger error: http.Post failed.\n%v\n%v\n%v\n", string(b), resp, err)
	}
	if resp.Body != nil {
		resp.Body.Close()
	}
}
<|endoftext|>