<commit_before>package settings

import (
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"

	"github.com/argoproj/argo-cd/common"
	tlsutil "github.com/argoproj/argo-cd/util/tls"
	"github.com/ghodss/yaml"
	log "github.com/sirupsen/logrus"
	apiv1 "k8s.io/api/core/v1"
	apierr "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

// ArgoCDSettings holds in-memory runtime configuration options.
type ArgoCDSettings struct {
	// URL is the externally facing URL users will visit to reach ArgoCD.
	// The value here is used when configuring SSO. Omitting this value will disable SSO.
	URL string
	// DexConfig is contains portions of a dex config yaml
	DexConfig string
	// LocalUsers holds users local to (stored on) the server.  This is to be distinguished from any potential alternative future login providers (LDAP, SAML, etc.) that might ever be added.
	LocalUsers map[string]string
	// ServerSignature holds the key used to generate JWT tokens.
	ServerSignature []byte
	// Certificate holds the certificate/private key for the ArgoCD API server.
	// If nil, will run insecure without TLS.
	Certificate *tls.Certificate
	// WebhookGitLabSecret holds the shared secret for authenticating GitHub webhook events
	WebhookGitHubSecret string
	// WebhookGitLabSecret holds the shared secret for authenticating GitLab webhook events
	WebhookGitLabSecret string
	// WebhookBitbucketUUID holds the UUID for authenticating Bitbucket webhook events
	WebhookBitbucketUUID string
}

const (
	// settingAdminPasswordKey designates the key for a root password inside a Kubernetes secret.
	settingAdminPasswordKey = "admin.password"
	// settingServerSignatureKey designates the key for a server secret key inside a Kubernetes secret.
	settingServerSignatureKey = "server.secretkey"
	// settingServerCertificate designates the key for the public cert used in TLS
	settingServerCertificate = "server.crt"
	// settingServerPrivateKey designates the key for the private key used in TLS
	settingServerPrivateKey = "server.key"
	// settingURLKey designates the key where ArgoCDs external URL is set
	settingURLKey = "url"
	// settingDexConfigKey designates the key for the dex config
	settingDexConfigKey = "dex.config"
	// settingsWebhookGitHubSecret is the key for the GitHub shared webhook secret
	settingsWebhookGitHubSecretKey = "webhook.github.secret"
	// settingsWebhookGitLabSecret is the key for the GitLab shared webhook secret
	settingsWebhookGitLabSecretKey = "webhook.gitlab.secret"
	// settingsWebhookBitbucketUUID is the key for Bitbucket webhook UUID
	settingsWebhookBitbucketUUIDKey = "webhook.bitbucket.uuid"
)

// SettingsManager holds config info for a new manager with which to access Kubernetes ConfigMaps.
type SettingsManager struct {
	clientset kubernetes.Interface
	namespace string
}

// GetSettings retrieves settings from the ConfigManager.
func (mgr *SettingsManager) GetSettings() (*ArgoCDSettings, error) {
	argoCDCM, err := mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Get(common.ArgoCDConfigMapName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	argoCDSecret, err := mgr.clientset.CoreV1().Secrets(mgr.namespace).Get(common.ArgoCDSecretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	var settings ArgoCDSettings
	settings.DexConfig = argoCDCM.Data[settingDexConfigKey]
	settings.URL = argoCDCM.Data[settingURLKey]

	adminPasswordHash, ok := argoCDSecret.Data[settingAdminPasswordKey]
	if !ok {
		return nil, fmt.Errorf("admin user not found")
	}
	settings.LocalUsers = map[string]string{
		common.ArgoCDAdminUsername: string(adminPasswordHash),
	}
	secretKey, ok := argoCDSecret.Data[settingServerSignatureKey]
	if !ok {
		return nil, fmt.Errorf("server secret key not found")
	}
	settings.ServerSignature = secretKey
	if githubWebhookSecret := argoCDSecret.Data[settingsWebhookGitHubSecretKey]; len(githubWebhookSecret) > 0 {
		settings.WebhookGitHubSecret = string(githubWebhookSecret)
	}
	if gitlabWebhookSecret := argoCDSecret.Data[settingsWebhookGitLabSecretKey]; len(gitlabWebhookSecret) > 0 {
		settings.WebhookGitLabSecret = string(gitlabWebhookSecret)
	}
	if bitbucketWebhookUUID := argoCDSecret.Data[settingsWebhookBitbucketUUIDKey]; len(bitbucketWebhookUUID) > 0 {
		settings.WebhookBitbucketUUID = string(bitbucketWebhookUUID)
	}

	serverCert, certOk := argoCDSecret.Data[settingServerCertificate]
	serverKey, keyOk := argoCDSecret.Data[settingServerPrivateKey]
	if certOk && keyOk {
		cert, err := tls.X509KeyPair(serverCert, serverKey)
		if err != nil {
			return nil, fmt.Errorf("invalid x509 key pair %s/%s in secret: %s", settingServerCertificate, settingServerPrivateKey, err)
		}
		settings.Certificate = &cert
	}
	return &settings, nil
}

// SaveSettings serializes ArgoCD settings and upserts it into K8s secret/configmap
func (mgr *SettingsManager) SaveSettings(settings *ArgoCDSettings) error {
	// Upsert the config data
	argoCDCM, err := mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Get(common.ArgoCDConfigMapName, metav1.GetOptions{})
	createCM := false
	if err != nil {
		if !apierr.IsNotFound(err) {
			return err
		}
		argoCDCM = &apiv1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name: common.ArgoCDConfigMapName,
			},
			Data: make(map[string]string),
		}
		createCM = true
	}
	argoCDCM.Data[settingURLKey] = settings.URL
	argoCDCM.Data[settingDexConfigKey] = settings.DexConfig
	if createCM {
		_, err = mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Create(argoCDCM)
	} else {
		_, err = mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Update(argoCDCM)
	}
	if err != nil {
		return err
	}

	// Upsert the secret data. Ensure we do not delete any extra keys which user may have added
	argoCDSecret, err := mgr.clientset.CoreV1().Secrets(mgr.namespace).Get(common.ArgoCDSecretName, metav1.GetOptions{})
	createSecret := false
	if err != nil {
		if !apierr.IsNotFound(err) {
			return err
		}
		argoCDSecret = &apiv1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name: common.ArgoCDSecretName,
			},
			Data: make(map[string][]byte),
		}
		createSecret = true
	}
	argoCDSecret.StringData = make(map[string]string)
	argoCDSecret.StringData[settingServerSignatureKey] = string(settings.ServerSignature)
	argoCDSecret.StringData[settingAdminPasswordKey] = settings.LocalUsers[common.ArgoCDAdminUsername]
	if settings.WebhookGitHubSecret != "" {
		argoCDSecret.StringData[settingsWebhookGitHubSecretKey] = settings.WebhookGitHubSecret
	}
	if settings.WebhookGitLabSecret != "" {
		argoCDSecret.StringData[settingsWebhookGitLabSecretKey] = settings.WebhookGitLabSecret
	}
	if settings.WebhookBitbucketUUID != "" {
		argoCDSecret.StringData[settingsWebhookBitbucketUUIDKey] = settings.WebhookBitbucketUUID
	}
	if settings.Certificate != nil {
		certBytes, keyBytes := tlsutil.EncodeX509KeyPair(*settings.Certificate)
		argoCDSecret.StringData[settingServerCertificate] = string(certBytes)
		argoCDSecret.StringData[settingServerPrivateKey] = string(keyBytes)
	} else {
		delete(argoCDSecret.Data, settingServerCertificate)
		delete(argoCDSecret.Data, settingServerPrivateKey)
	}
	if createSecret {
		_, err = mgr.clientset.CoreV1().Secrets(mgr.namespace).Create(argoCDSecret)
	} else {
		_, err = mgr.clientset.CoreV1().Secrets(mgr.namespace).Update(argoCDSecret)
	}
	if err != nil {
		return err
	}
	return nil
}

// NewSettingsManager generates a new SettingsManager pointer and returns it
func NewSettingsManager(clientset kubernetes.Interface, namespace string) *SettingsManager {
	return &SettingsManager{
		clientset: clientset,
		namespace: namespace,
	}
}

func (a *ArgoCDSettings) IsSSOConfigured() bool {
	if a.URL == "" {
		return false
	}
	var dexCfg map[string]interface{}
	err := yaml.Unmarshal([]byte(a.DexConfig), &dexCfg)
	if err != nil {
		log.Warn("invalid dex yaml config")
		return false
	}
	return len(dexCfg) > 0
}

// TLSConfig returns a tls.Config with the configured certificates
func (a *ArgoCDSettings) TLSConfig() *tls.Config {
	if a.Certificate == nil {
		return nil
	}
	certPool := x509.NewCertPool()
	pemCertBytes, _ := tlsutil.EncodeX509KeyPair(*a.Certificate)
	ok := certPool.AppendCertsFromPEM(pemCertBytes)
	if !ok {
		panic("bad certs")
	}
	return &tls.Config{
		RootCAs: certPool,
	}
}

func (a *ArgoCDSettings) IssuerURL() string {
	return a.URL + common.DexAPIEndpoint
}

func (a *ArgoCDSettings) RedirectURL() string {
	return a.URL + common.CallbackEndpoint
}

// OAuth2ClientSecret calculates an arbitrary, but predictable OAuth2 client secret string derived
// from the server secret. This is called by the dex startup wrapper (argocd-util rundex), as well
// as the API server, such that they both independently come to the same conclusion of what the
// OAuth2 shared client secret should be.
func (a *ArgoCDSettings) OAuth2ClientSecret() string {
	h := sha256.New()
	_, err := h.Write(a.ServerSignature)
	if err != nil {
		panic(err)
	}
	sha := h.Sum(nil)
	return base64.URLEncoding.EncodeToString(sha)[:40]
}
<commit_msg>Handle potential panic when `argo install settings` run against an empty configmap<commit_after>package settings

import (
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"

	"github.com/argoproj/argo-cd/common"
	tlsutil "github.com/argoproj/argo-cd/util/tls"
	"github.com/ghodss/yaml"
	log "github.com/sirupsen/logrus"
	apiv1 "k8s.io/api/core/v1"
	apierr "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

// ArgoCDSettings holds in-memory runtime configuration options.
type ArgoCDSettings struct {
	// URL is the externally facing URL users will visit to reach ArgoCD.
	// The value here is used when configuring SSO. Omitting this value will disable SSO.
	URL string
	// DexConfig is contains portions of a dex config yaml
	DexConfig string
	// LocalUsers holds users local to (stored on) the server.  This is to be distinguished from any potential alternative future login providers (LDAP, SAML, etc.) that might ever be added.
	LocalUsers map[string]string
	// ServerSignature holds the key used to generate JWT tokens.
	ServerSignature []byte
	// Certificate holds the certificate/private key for the ArgoCD API server.
	// If nil, will run insecure without TLS.
	Certificate *tls.Certificate
	// WebhookGitLabSecret holds the shared secret for authenticating GitHub webhook events
	WebhookGitHubSecret string
	// WebhookGitLabSecret holds the shared secret for authenticating GitLab webhook events
	WebhookGitLabSecret string
	// WebhookBitbucketUUID holds the UUID for authenticating Bitbucket webhook events
	WebhookBitbucketUUID string
}

const (
	// settingAdminPasswordKey designates the key for a root password inside a Kubernetes secret.
	settingAdminPasswordKey = "admin.password"
	// settingServerSignatureKey designates the key for a server secret key inside a Kubernetes secret.
	settingServerSignatureKey = "server.secretkey"
	// settingServerCertificate designates the key for the public cert used in TLS
	settingServerCertificate = "server.crt"
	// settingServerPrivateKey designates the key for the private key used in TLS
	settingServerPrivateKey = "server.key"
	// settingURLKey designates the key where ArgoCDs external URL is set
	settingURLKey = "url"
	// settingDexConfigKey designates the key for the dex config
	settingDexConfigKey = "dex.config"
	// settingsWebhookGitHubSecret is the key for the GitHub shared webhook secret
	settingsWebhookGitHubSecretKey = "webhook.github.secret"
	// settingsWebhookGitLabSecret is the key for the GitLab shared webhook secret
	settingsWebhookGitLabSecretKey = "webhook.gitlab.secret"
	// settingsWebhookBitbucketUUID is the key for Bitbucket webhook UUID
	settingsWebhookBitbucketUUIDKey = "webhook.bitbucket.uuid"
)

// SettingsManager holds config info for a new manager with which to access Kubernetes ConfigMaps.
type SettingsManager struct {
	clientset kubernetes.Interface
	namespace string
}

// GetSettings retrieves settings from the ConfigManager.
func (mgr *SettingsManager) GetSettings() (*ArgoCDSettings, error) {
	argoCDCM, err := mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Get(common.ArgoCDConfigMapName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	argoCDSecret, err := mgr.clientset.CoreV1().Secrets(mgr.namespace).Get(common.ArgoCDSecretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}

	var settings ArgoCDSettings
	settings.DexConfig = argoCDCM.Data[settingDexConfigKey]
	settings.URL = argoCDCM.Data[settingURLKey]

	adminPasswordHash, ok := argoCDSecret.Data[settingAdminPasswordKey]
	if !ok {
		return nil, fmt.Errorf("admin user not found")
	}
	settings.LocalUsers = map[string]string{
		common.ArgoCDAdminUsername: string(adminPasswordHash),
	}
	secretKey, ok := argoCDSecret.Data[settingServerSignatureKey]
	if !ok {
		return nil, fmt.Errorf("server secret key not found")
	}
	settings.ServerSignature = secretKey
	if githubWebhookSecret := argoCDSecret.Data[settingsWebhookGitHubSecretKey]; len(githubWebhookSecret) > 0 {
		settings.WebhookGitHubSecret = string(githubWebhookSecret)
	}
	if gitlabWebhookSecret := argoCDSecret.Data[settingsWebhookGitLabSecretKey]; len(gitlabWebhookSecret) > 0 {
		settings.WebhookGitLabSecret = string(gitlabWebhookSecret)
	}
	if bitbucketWebhookUUID := argoCDSecret.Data[settingsWebhookBitbucketUUIDKey]; len(bitbucketWebhookUUID) > 0 {
		settings.WebhookBitbucketUUID = string(bitbucketWebhookUUID)
	}

	serverCert, certOk := argoCDSecret.Data[settingServerCertificate]
	serverKey, keyOk := argoCDSecret.Data[settingServerPrivateKey]
	if certOk && keyOk {
		cert, err := tls.X509KeyPair(serverCert, serverKey)
		if err != nil {
			return nil, fmt.Errorf("invalid x509 key pair %s/%s in secret: %s", settingServerCertificate, settingServerPrivateKey, err)
		}
		settings.Certificate = &cert
	}
	return &settings, nil
}

// SaveSettings serializes ArgoCD settings and upserts it into K8s secret/configmap
func (mgr *SettingsManager) SaveSettings(settings *ArgoCDSettings) error {
	// Upsert the config data
	argoCDCM, err := mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Get(common.ArgoCDConfigMapName, metav1.GetOptions{})
	createCM := false
	if err != nil {
		if !apierr.IsNotFound(err) {
			return err
		}
		argoCDCM = &apiv1.ConfigMap{
			ObjectMeta: metav1.ObjectMeta{
				Name: common.ArgoCDConfigMapName,
			},
		}
		createCM = true
	}
	if argoCDCM.Data == nil {
		argoCDCM.Data = make(map[string]string)
	}
	argoCDCM.Data[settingURLKey] = settings.URL
	argoCDCM.Data[settingDexConfigKey] = settings.DexConfig
	if createCM {
		_, err = mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Create(argoCDCM)
	} else {
		_, err = mgr.clientset.CoreV1().ConfigMaps(mgr.namespace).Update(argoCDCM)
	}
	if err != nil {
		return err
	}

	// Upsert the secret data. Ensure we do not delete any extra keys which user may have added
	argoCDSecret, err := mgr.clientset.CoreV1().Secrets(mgr.namespace).Get(common.ArgoCDSecretName, metav1.GetOptions{})
	createSecret := false
	if err != nil {
		if !apierr.IsNotFound(err) {
			return err
		}
		argoCDSecret = &apiv1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name: common.ArgoCDSecretName,
			},
			Data: make(map[string][]byte),
		}
		createSecret = true
	}
	argoCDSecret.StringData = make(map[string]string)
	argoCDSecret.StringData[settingServerSignatureKey] = string(settings.ServerSignature)
	argoCDSecret.StringData[settingAdminPasswordKey] = settings.LocalUsers[common.ArgoCDAdminUsername]
	if settings.WebhookGitHubSecret != "" {
		argoCDSecret.StringData[settingsWebhookGitHubSecretKey] = settings.WebhookGitHubSecret
	}
	if settings.WebhookGitLabSecret != "" {
		argoCDSecret.StringData[settingsWebhookGitLabSecretKey] = settings.WebhookGitLabSecret
	}
	if settings.WebhookBitbucketUUID != "" {
		argoCDSecret.StringData[settingsWebhookBitbucketUUIDKey] = settings.WebhookBitbucketUUID
	}
	if settings.Certificate != nil {
		certBytes, keyBytes := tlsutil.EncodeX509KeyPair(*settings.Certificate)
		argoCDSecret.StringData[settingServerCertificate] = string(certBytes)
		argoCDSecret.StringData[settingServerPrivateKey] = string(keyBytes)
	} else {
		delete(argoCDSecret.Data, settingServerCertificate)
		delete(argoCDSecret.Data, settingServerPrivateKey)
	}
	if createSecret {
		_, err = mgr.clientset.CoreV1().Secrets(mgr.namespace).Create(argoCDSecret)
	} else {
		_, err = mgr.clientset.CoreV1().Secrets(mgr.namespace).Update(argoCDSecret)
	}
	if err != nil {
		return err
	}
	return nil
}

// NewSettingsManager generates a new SettingsManager pointer and returns it
func NewSettingsManager(clientset kubernetes.Interface, namespace string) *SettingsManager {
	return &SettingsManager{
		clientset: clientset,
		namespace: namespace,
	}
}

func (a *ArgoCDSettings) IsSSOConfigured() bool {
	if a.URL == "" {
		return false
	}
	var dexCfg map[string]interface{}
	err := yaml.Unmarshal([]byte(a.DexConfig), &dexCfg)
	if err != nil {
		log.Warn("invalid dex yaml config")
		return false
	}
	return len(dexCfg) > 0
}

// TLSConfig returns a tls.Config with the configured certificates
func (a *ArgoCDSettings) TLSConfig() *tls.Config {
	if a.Certificate == nil {
		return nil
	}
	certPool := x509.NewCertPool()
	pemCertBytes, _ := tlsutil.EncodeX509KeyPair(*a.Certificate)
	ok := certPool.AppendCertsFromPEM(pemCertBytes)
	if !ok {
		panic("bad certs")
	}
	return &tls.Config{
		RootCAs: certPool,
	}
}

func (a *ArgoCDSettings) IssuerURL() string {
	return a.URL + common.DexAPIEndpoint
}

func (a *ArgoCDSettings) RedirectURL() string {
	return a.URL + common.CallbackEndpoint
}

// OAuth2ClientSecret calculates an arbitrary, but predictable OAuth2 client secret string derived
// from the server secret. This is called by the dex startup wrapper (argocd-util rundex), as well
// as the API server, such that they both independently come to the same conclusion of what the
// OAuth2 shared client secret should be.
func (a *ArgoCDSettings) OAuth2ClientSecret() string {
	h := sha256.New()
	_, err := h.Write(a.ServerSignature)
	if err != nil {
		panic(err)
	}
	sha := h.Sum(nil)
	return base64.URLEncoding.EncodeToString(sha)[:40]
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"math"
	"os"
	"strconv"

	"github.com/vladvelici/graph-dataset-tools/util"
)

// Define flags.
var (
	flagAction  = flag.String("action", "", "Possibile actions: details, components, remove")
	flagN       = flag.Float64("n", 0, "\\% of edges to remove from each graph.")
	flagOutput  = flag.String("o", "component_", "Output file prefix. It will be followed by the component number and an underscore.")
	flagVerbose = flag.Bool("verbose", false, "Whether to print lots of debug information on stdout.")
	flagAll     = flag.Bool("all", false, "Whether to print all edges (to->from and from->to) when removing, in the removed list.")
	flagForce   = flag.Bool("force", false, "Do not make safety checks. Might be faster.")
	flagHelp    = flag.Bool("help", false, "Show this help message")
	flagH       = flag.Bool("h", false, "Show this help message")
)

var helpMessage = `conncomp is a tool that deals with connected components tasks.

Possible uses (can optionally add -o and -verbose flags before filelist):

-action details filelist			Outputs details about the graphs in the files. Number of edges, number of nodes, directed/undirected, connected and no. of components.
-action components filelist			Splits the given graph in connected components files (prefix can be changed with -o flag).
-action remove -n P [-all] filelist	Removes at most P% random edges from the given graphs. Prefix of file controlled with -o flag.
-action force-undirected			Force the graph into an undirected graph. Writes to a different file.

-h or -help to display this message and quit.
`

func help() {
	fmt.Println(helpMessage)
	flag.PrintDefaults()
}

type Action func()

func main() {
	flag.Usage = help
	flag.Parse()

	if *flagHelp || *flagH {
		help()
		return
	}

	if *flagAction == "" {
		fmt.Println("Need to specify an action! See -help.")
		return
	}

	controller := map[string]Action{
		"details":          actionDetails,
		"components":       actionComponents,
		"remove":           actionRemove,
		"force-undirected": actionForceUndirected,
	}

	action, ok := controller[*flagAction]
	if !ok {
		fmt.Print("Invalid action! See -help for info. Available actions are: ")
		for txt, _ := range controller {
			fmt.Print(txt)
			fmt.Print(" ")
		}
		fmt.Println()
		return
	}

	action()
}

// Return whether the given graphs are directed on undirected.
func actionDetails() {
	files := flag.Args()
	fmt.Println("filename\t\tType\t\tConnected?\t#components\t#nodes\t#edges")
	fmt.Println("========\t\t====\t\t==========\t===========\t======\t======")
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s \t\t Cannot read graph (%s).\n", f, err.Error())
			continue
		}

		fmt.Printf("%s\t\t", f)
		undir := graph.IsUndirected()
		if undir {
			fmt.Print("Undirected\t")
		} else {
			fmt.Print("Directed\t")
		}

		components := graph.ConnectedGraphs()
		if len(components) <= 1 {
			fmt.Print("connected\t")
		} else {
			fmt.Print("disconnected\t")
		}

		fmt.Printf("%d\t\t", len(components))
		fmt.Printf("%d\t", len(graph.Nodes))

		var noEdges int
		edges := graph.EdgeList()
		if undir {
			noEdges = len(edges) / 2
		} else {
			noEdges = len(edges)
		}

		fmt.Printf("%d\n", noEdges)
	}
}

// Splits graphs into their connected compoments, and writes those components as separate files.
func actionComponents() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		components := graph.ConnectedGraphs()

		for i, comp := range components {
			fname := *flagOutput + strconv.Itoa(i) + "_" + f
			wr, err := os.Create(fname)
			if err != nil {
				fmt.Printf("%s: Cannot write to %s, skipping connected component #%d. (%s)\n", f, fname, i, err.Error())
			}
			writer := util.NewWriter(wr)
			comp.EachEdge(func(from, to *Node) bool {
				err := writer.Write(from.Id, to.Id, nil)
				if err != nil {
					fmt.Printf("%s: Cannot write edge to %s, skipping remaining of connected component #%d. (%s)\n", f, fname, i, err.Error())
					return false
				}
				return true
			})
			err = writer.Flush()
			if err != nil {
				fmt.Printf("%s: (FLUSH) Output might be corrupted.")
			}
			err = wr.Close()
			if err != nil {
				fmt.Printf("%s: (CLOSE) Output might be corrupted.")
			}
		}
	}
}

// Get (perhaps directed) graphs and make them directed by forcing all edges to be bi-directional.
func actionForceUndirected() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		fname := *flagOutput + f
		wr, err := os.Create(fname)
		if err != nil {
			fmt.Printf("%s: Cannot write to %s, skipping file. (%s)\n", f, fname, err.Error())
			continue
		}
		writer := util.NewWriter(wr)
		edges := graph.EdgeList()
		for _, e := range edges {
			// it has from -> to if we got here.
			// adding to -> from.
			graph.AddDirectedEdge(e.To, e.From)
		}
		graph.EachEdge(func(from, to *Node) bool {
			err := writer.Write(from.Id, to.Id, nil)
			if err != nil {
				fmt.Printf("%s: Cannot write edge to %s, aborting this graph. (%s)\n", f, fname, err.Error())
				return false
			}
			return true
		})
		err = writer.Flush()
		if err != nil {
			fmt.Printf("%s: (FLUSH) Output might be corrupted.", f)
		}
		err = wr.Close()
		if err != nil {
			fmt.Printf("%s: (CLOSE) Output might be corrupted.", f)
		}
	}
}

// write a graph
func writeGraph(graph *Graph, fname string) error {
	wr, err := os.Create(fname)
	if err != nil {
		return fmt.Errorf("Cannot write to %s, skipping file. (%s)", fname, err.Error())
	}
	writer := util.NewWriter(wr)
	err = nil
	graph.EachEdge(func(from, to *Node) bool {
		err := writer.Write(from.Id, to.Id, nil)
		if err != nil {
			err = fmt.Errorf("Cannot write edge to %s, aborting this graph. (%s)", fname, err.Error())
			return false
		}
		return true
	})
	if err != nil {
		return err
	}
	err = writer.Flush()
	if err != nil {
		return fmt.Errorf("(FLUSH) Output might be corrupted. (%s)", err.Error())
	}
	err = wr.Close()
	if err != nil {
		return fmt.Errorf("(CLOSE) Output might be corrupted. (%s)", err.Error())
	}
	return nil
}

// Remove random edges from a graph using a spanning tree to assure connectness.
func actionRemove() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		if !*flagForce && !graph.IsConnected() {
			fmt.Printf("%s: Graph not connected. Skipping. Use the --force to do it anyway.\n", f)
			continue
		}

		if !*flagForce && !graph.IsUndirected() {
			fmt.Printf("%s: Graph is directed. Skipping. Use the --force to do it anyway.\n", f)
			continue
		}

		edges := len(graph.EdgeList())
		remove := int(math.Floor(*flagN*float64(edges)/2 + 0.5))
		mst := graph.Mst()
		removed := graph.RemoveRandomEdges(remove, mst)

		// write out the processed graph
		fname := *flagOutput + f
		err = writeGraph(graph, fname)
		if err != nil {
			fmt.Printf("%s: %s\n", err.Error())
			continue
		}

		fname = *flagOutput + "removed_" + f
		wr, err := os.Create(fname)
		if err != nil {
			fmt.Printf("%s: Cannot write to %s, skipping file. (%s)", f, fname, err.Error())
			continue
		}

		writer := util.NewWriter(wr)
		for _, edge := range removed {
			err = writer.Write(edge.From, edge.To, nil)
			if err != nil {
				fmt.Printf("%s: Cannot write edge to %s. Skipping remaining of graph. (%s)", f, fname, err.Error())
				break
			}
			if *flagAll {
				err = writer.Write(edge.To, edge.From, nil)
				if err != nil {
					fmt.Printf("%s: Cannot write (reverse) edge to %s. Skipping remaining of graph. (%s)", f, fname, err.Error())
					break
				}
			}
		}

		err = writer.Flush()
		if err != nil {
			fmt.Printf("%s: (FLUSH) Output might be corrupted.", f)
		}
		err = wr.Close()
		if err != nil {
			fmt.Printf("%s: (CLOSE) Output might be corrupted.", f)
		}

	}
}
<commit_msg>conncomp help message update<commit_after>package main

import (
	"flag"
	"fmt"
	"math"
	"os"
	"strconv"

	"github.com/vladvelici/graph-dataset-tools/util"
)

// Define flags.
var (
	flagAction  = flag.String("action", "", "Possibile actions: details, components, remove")
	flagN       = flag.Float64("n", 0, "\\% of edges to remove from each graph.")
	flagOutput  = flag.String("o", "component_", "Output file prefix. It will be followed by the component number and an underscore.")
	flagVerbose = flag.Bool("verbose", false, "Whether to print lots of debug information on stdout.")
	flagAll     = flag.Bool("all", false, "Whether to print all edges (to->from and from->to) when removing, in the removed list.")
	flagForce   = flag.Bool("force", false, "Do not make safety checks. Might be faster.")
	flagHelp    = flag.Bool("help", false, "Show this help message")
	flagH       = flag.Bool("h", false, "Show this help message")
)

var helpMessage = `conncomp is a tool that deals with connected components tasks.

Usage:

conncomp -action <action> [flags] <list of edge lists>

Possible actions:

  details           Output details about the graphs.
  components        Splits the graph(s) in connected components.
  remove -n P       Removes at most P% random edges from graph(s).
  force-undirected  Force the graph(s) into undirected graph(s).

All actions except 'details' output in new files whose names are controlled
with the flag -o <file-prefix>.

Use the -h or -help flag to display this message and quit.

Full list of flags:

`

func help() {
	fmt.Println(helpMessage)
	flag.PrintDefaults()
}

type Action func()

func main() {
	flag.Usage = help
	flag.Parse()

	if *flagHelp || *flagH {
		help()
		return
	}

	if *flagAction == "" {
		fmt.Println("Need to specify an action! See -help.")
		return
	}

	controller := map[string]Action{
		"details":          actionDetails,
		"components":       actionComponents,
		"remove":           actionRemove,
		"force-undirected": actionForceUndirected,
	}

	action, ok := controller[*flagAction]
	if !ok {
		fmt.Print("Invalid action! See -help for info. Available actions are: ")
		for txt, _ := range controller {
			fmt.Print(txt)
			fmt.Print(" ")
		}
		fmt.Println()
		return
	}

	action()
}

// Return whether the given graphs are directed on undirected.
func actionDetails() {
	files := flag.Args()
	fmt.Println("filename\t\tType\t\tConnected?\t#components\t#nodes\t#edges")
	fmt.Println("========\t\t====\t\t==========\t===========\t======\t======")
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s \t\t Cannot read graph (%s).\n", f, err.Error())
			continue
		}

		fmt.Printf("%s\t\t", f)
		undir := graph.IsUndirected()
		if undir {
			fmt.Print("Undirected\t")
		} else {
			fmt.Print("Directed\t")
		}

		components := graph.ConnectedGraphs()
		if len(components) <= 1 {
			fmt.Print("connected\t")
		} else {
			fmt.Print("disconnected\t")
		}

		fmt.Printf("%d\t\t", len(components))
		fmt.Printf("%d\t", len(graph.Nodes))

		var noEdges int
		edges := graph.EdgeList()
		if undir {
			noEdges = len(edges) / 2
		} else {
			noEdges = len(edges)
		}

		fmt.Printf("%d\n", noEdges)
	}
}

// Splits graphs into their connected compoments, and writes those components as separate files.
func actionComponents() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		components := graph.ConnectedGraphs()

		for i, comp := range components {
			fname := *flagOutput + strconv.Itoa(i) + "_" + f
			wr, err := os.Create(fname)
			if err != nil {
				fmt.Printf("%s: Cannot write to %s, skipping connected component #%d. (%s)\n", f, fname, i, err.Error())
			}
			writer := util.NewWriter(wr)
			comp.EachEdge(func(from, to *Node) bool {
				err := writer.Write(from.Id, to.Id, nil)
				if err != nil {
					fmt.Printf("%s: Cannot write edge to %s, skipping remaining of connected component #%d. (%s)\n", f, fname, i, err.Error())
					return false
				}
				return true
			})
			err = writer.Flush()
			if err != nil {
				fmt.Printf("%s: (FLUSH) Output might be corrupted.")
			}
			err = wr.Close()
			if err != nil {
				fmt.Printf("%s: (CLOSE) Output might be corrupted.")
			}
		}
	}
}

// Get (perhaps directed) graphs and make them directed by forcing all edges to be bi-directional.
func actionForceUndirected() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		fname := *flagOutput + f
		wr, err := os.Create(fname)
		if err != nil {
			fmt.Printf("%s: Cannot write to %s, skipping file. (%s)\n", f, fname, err.Error())
			continue
		}
		writer := util.NewWriter(wr)
		edges := graph.EdgeList()
		for _, e := range edges {
			// it has from -> to if we got here.
			// adding to -> from.
			graph.AddDirectedEdge(e.To, e.From)
		}
		graph.EachEdge(func(from, to *Node) bool {
			err := writer.Write(from.Id, to.Id, nil)
			if err != nil {
				fmt.Printf("%s: Cannot write edge to %s, aborting this graph. (%s)\n", f, fname, err.Error())
				return false
			}
			return true
		})
		err = writer.Flush()
		if err != nil {
			fmt.Printf("%s: (FLUSH) Output might be corrupted.", f)
		}
		err = wr.Close()
		if err != nil {
			fmt.Printf("%s: (CLOSE) Output might be corrupted.", f)
		}
	}
}

// write a graph
func writeGraph(graph *Graph, fname string) error {
	wr, err := os.Create(fname)
	if err != nil {
		return fmt.Errorf("Cannot write to %s, skipping file. (%s)", fname, err.Error())
	}
	writer := util.NewWriter(wr)
	err = nil
	graph.EachEdge(func(from, to *Node) bool {
		err := writer.Write(from.Id, to.Id, nil)
		if err != nil {
			err = fmt.Errorf("Cannot write edge to %s, aborting this graph. (%s)", fname, err.Error())
			return false
		}
		return true
	})
	if err != nil {
		return err
	}
	err = writer.Flush()
	if err != nil {
		return fmt.Errorf("(FLUSH) Output might be corrupted. (%s)", err.Error())
	}
	err = wr.Close()
	if err != nil {
		return fmt.Errorf("(CLOSE) Output might be corrupted. (%s)", err.Error())
	}
	return nil
}

// Remove random edges from a graph using a spanning tree to assure connectness.
func actionRemove() {
	files := flag.Args()
	for _, f := range files {
		graph, err := ReadGraph(f)
		if err != nil {
			fmt.Printf("%s: Cannot read graph. Skipping. (%s)\n", f, err.Error())
			continue
		}

		if !*flagForce && !graph.IsConnected() {
			fmt.Printf("%s: Graph not connected. Skipping. Use the --force to do it anyway.\n", f)
			continue
		}

		if !*flagForce && !graph.IsUndirected() {
			fmt.Printf("%s: Graph is directed. Skipping. Use the --force to do it anyway.\n", f)
			continue
		}

		edges := len(graph.EdgeList())
		remove := int(math.Floor(*flagN*float64(edges)/2 + 0.5))
		mst := graph.Mst()
		removed := graph.RemoveRandomEdges(remove, mst)

		// write out the processed graph
		fname := *flagOutput + f
		err = writeGraph(graph, fname)
		if err != nil {
			fmt.Printf("%s: %s\n", err.Error())
			continue
		}

		fname = *flagOutput + "removed_" + f
		wr, err := os.Create(fname)
		if err != nil {
			fmt.Printf("%s: Cannot write to %s, skipping file. (%s)", f, fname, err.Error())
			continue
		}

		writer := util.NewWriter(wr)
		for _, edge := range removed {
			err = writer.Write(edge.From, edge.To, nil)
			if err != nil {
				fmt.Printf("%s: Cannot write edge to %s. Skipping remaining of graph. (%s)", f, fname, err.Error())
				break
			}
			if *flagAll {
				err = writer.Write(edge.To, edge.From, nil)
				if err != nil {
					fmt.Printf("%s: Cannot write (reverse) edge to %s. Skipping remaining of graph. (%s)", f, fname, err.Error())
					break
				}
			}
		}

		err = writer.Flush()
		if err != nil {
			fmt.Printf("%s: (FLUSH) Output might be corrupted.", f)
		}
		err = wr.Close()
		if err != nil {
			fmt.Printf("%s: (CLOSE) Output might be corrupted.", f)
		}

	}
}
<|endoftext|>


<commit_before>package rest_public_trades
<commit_msg>examples: add usage of rest public trades<commit_after>package main

import (
	"fmt"
	"log"
	"github.com/bitfinexcom/bitfinex-api-go/v2/rest"
	bitfinex "github.com/bitfinexcom/bitfinex-api-go/v2"
	"time"
)

func main() {
	c := rest.NewClient()

	// calculate start and end
	now := time.Now()
	millis := now.UnixNano() / 1000000
	prior := now.Add(time.Duration(-24) * time.Hour)
	millisStart := prior.UnixNano() / 1000000
	start := bitfinex.Mts(millisStart)
	end := bitfinex.Mts(millis)
	// send request
	trades, err := c.Trades.PublicHistoryWithQuery("tBTCUSD", start, end, 10, bitfinex.OldestFirst)
	if err != nil {
		log.Fatalf("%v", err)
	}
	for _, trade := range trades.Snapshot {
		fmt.Println(trade)
	}
}
<|endoftext|>


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

package main

import (
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/blevesearch/bleve"
	_ "github.com/blevesearch/bleve/config"
)

var indexPath = flag.String("index", "", "index path")
var keepExt = flag.Bool("keepExt", false, "keep extension in doc id")
var keepDir = flag.Bool("keepDir", false, "keep dir in doc id")

func main() {

	flag.Parse()

	if *indexPath == "" {
		log.Fatal("must specify index path")
	}

	// open the index
	index, err := bleve.Open(*indexPath)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		cerr := index.Close()
		if cerr != nil {
			log.Fatalf("error closing index: %v", err)
		}
	}()

	if flag.NArg() < 1 {
		log.Fatal("must specify at least one path to index")
	}

	for file := range handleArgs(flag.Args()) {
		// index the files
		docID := file.filename
		if !*keepDir {
			_, docID = filepath.Split(docID)
		}
		if !*keepExt {
			ext := filepath.Ext(docID)
			docID = docID[0 : len(docID)-len(ext)]
		}
		log.Printf("Indexing: %s", docID)
		err = index.Index(docID, file.contents)
		if err != nil {
			log.Fatal(err)
		}
	}
}

type file struct {
	filename string
	contents []byte
}

func handleArgs(args []string) chan file {
	rv := make(chan file)

	go func() {
		for _, arg := range args {
			arg = filepath.Clean(arg)
			handleArgRecursive(arg, rv)
		}
		close(rv)
	}()

	return rv
}

func handleArgRecursive(arg string, results chan file) {
	stat, err := os.Stat(arg)
	if err != nil {
		log.Print(err)
		return
	}
	if stat.IsDir() {
		// open the directory
		dirEntries, err := ioutil.ReadDir(arg)
		if err != nil {
			log.Fatal(err)
		}

		// walk the directory entries
		for _, dirEntry := range dirEntries {
			handleArgRecursive(arg+string(os.PathSeparator)+dirEntry.Name(), results)
		}
	} else {
		bytes, err := ioutil.ReadFile(arg)
		if err != nil {
			log.Fatal(err)
		}

		results <- file{
			filename: arg,
			contents: bytes,
		}
	}
}
<commit_msg>Use filepath.Walk instead of rolling our own func<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
//  except in compliance with the License. You may obtain a copy of the License at
//    http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software distributed under the
//  License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
//  either express or implied. See the License for the specific language governing permissions
//  and limitations under the License.

package main

import (
	"flag"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"

	"github.com/blevesearch/bleve"
	_ "github.com/blevesearch/bleve/config"
)

var (
	indexPath = flag.String("index", "", "index path")
	keepExt   = flag.Bool("keepExt", false, "keep extension in doc id")
	keepDir   = flag.Bool("keepDir", false, "keep dir in doc id")
	rv        chan file
)

func main() {

	flag.Parse()

	if *indexPath == "" {
		log.Fatal("must specify index path")
	}

	// open the index
	index, err := bleve.Open(*indexPath)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		cerr := index.Close()
		if cerr != nil {
			log.Fatalf("error closing index: %v", err)
		}
	}()

	if flag.NArg() < 1 {
		log.Fatal("must specify at least one path to index")
	}

	for file := range handleArgs(flag.Args()) {
		// index the files
		docID := file.filename
		if !*keepDir {
			_, docID = filepath.Split(docID)
		}
		if !*keepExt {
			ext := filepath.Ext(docID)
			docID = docID[0 : len(docID)-len(ext)]
		}
		log.Printf("Indexing: %s", docID)
		err = index.Index(docID, file.contents)
		if err != nil {
			log.Fatal(err)
		}
	}
}

type file struct {
	filename string
	contents []byte
}

func handleArgs(args []string) chan file {
	rv = make(chan file)

	go func() {
		for _, arg := range args {
			arg = filepath.Clean(arg)
			filepath.Walk(arg, getallfiles)
		}
		close(rv)
	}()

	return rv
}

func getallfiles(path string, finfo os.FileInfo, err error) error {
	if err != nil {
		log.Print(err)
		return err
	}
	if finfo.IsDir() {
		return nil
	}

	bytes, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatal(err)
	}
	rv <- file{
		filename: filepath.Base(path),
		contents: bytes,
	}

	return nil
}
<|endoftext|>


<commit_before>package miniredis

// 'Fail' methods.

import (
	"fmt"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
)

// T is implemented by Testing.T
type T interface {
	Fail()
}

// CheckGet does not call Errorf() iff there is a string key with the
// expected value. Normal use case is `m.CheckGet(t, "username", "theking")`.
func (m *Miniredis) CheckGet(t T, key, expected string) {
	found, err := m.Get(key)
	if err != nil {
		lError(t, "GET error, key %#v: %v", key, err)
		return
	}
	if found != expected {
		lError(t, "GET error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}

// CheckList does not call Errorf() iff there is a list key with the
// expected values.
// Normal use case is `m.CheckGet(t, "favorite_colors", "red", "green", "infrared")`.
func (m *Miniredis) CheckList(t T, key string, expected ...string) {
	found, err := m.List(key)
	if err != nil {
		lError(t, "List error, key %#v: %v", key, err)
		return
	}
	if !reflect.DeepEqual(expected, found) {
		lError(t, "List error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}

// CheckSet does not call Errorf() iff there is a set key with the
// expected values.
// Normal use case is `m.CheckSet(t, "visited", "Rome", "Stockholm", "Dublin")`.
func (m *Miniredis) CheckSet(t T, key string, expected ...string) {
	found, err := m.Members(key)
	if err != nil {
		lError(t, "Set error, key %#v: %v", key, err)
		return
	}
	sort.Strings(expected)
	if !reflect.DeepEqual(expected, found) {
		lError(t, "Set error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}

func lError(t T, format string, args ...interface{}) {
	_, file, line, _ := runtime.Caller(2)
	prefix := fmt.Sprintf("%s:%d: ", filepath.Base(file), line)
	fmt.Printf(prefix+format+"\n", args...)
	t.Fail()
}
<commit_msg>update the Check*() methods to use a t.Helper()<commit_after>package miniredis

import (
	"reflect"
	"sort"
)

// T is implemented by Testing.T
type T interface {
	Helper()
	Errorf(string, ...interface{})
}

// CheckGet does not call Errorf() iff there is a string key with the
// expected value. Normal use case is `m.CheckGet(t, "username", "theking")`.
func (m *Miniredis) CheckGet(t T, key, expected string) {
	t.Helper()

	found, err := m.Get(key)
	if err != nil {
		t.Errorf("GET error, key %#v: %v", key, err)
		return
	}
	if found != expected {
		t.Errorf("GET error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}

// CheckList does not call Errorf() iff there is a list key with the
// expected values.
// Normal use case is `m.CheckGet(t, "favorite_colors", "red", "green", "infrared")`.
func (m *Miniredis) CheckList(t T, key string, expected ...string) {
	t.Helper()

	found, err := m.List(key)
	if err != nil {
		t.Errorf("List error, key %#v: %v", key, err)
		return
	}
	if !reflect.DeepEqual(expected, found) {
		t.Errorf("List error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}

// CheckSet does not call Errorf() iff there is a set key with the
// expected values.
// Normal use case is `m.CheckSet(t, "visited", "Rome", "Stockholm", "Dublin")`.
func (m *Miniredis) CheckSet(t T, key string, expected ...string) {
	t.Helper()

	found, err := m.Members(key)
	if err != nil {
		t.Errorf("Set error, key %#v: %v", key, err)
		return
	}
	sort.Strings(expected)
	if !reflect.DeepEqual(expected, found) {
		t.Errorf("Set error, key %#v: Expected %#v, got %#v", key, expected, found)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/lowstz/slackhookgo"
)

func main() {
	var (
		color    string
		sentUp   = 0
		sentDown = 0
		protocol = flag.String("protocol", "tcp", "protocol tcp/udp")
		host     = flag.String("host", "ya.ru", "destination host")
		port     = flag.String("port", "80", "destination port")
		interval = flag.Uint("interval", 5, "interval check seconds")
		url      = flag.String("url", "", "hook url")
	)

	flag.Parse()
	checkFlags()

	address := fmt.Sprintf("%s:%s", *host, *port)

	for {
		conn, err := net.DialTimeout(*protocol, address, 3*time.Second)
		checkIfError(err)
		if err == nil {
			conn.Close()
			if sentUp == 0 {
				color = "good"
				text := "Destination host " + *host + ":" + *port + " reachable"
				sendToSlack(&color, &text, url)
				sentUp = 1
				sentDown = 0
			}
		} else {
			if sentDown == 0 {
				color = "danger"
				text := "Destination host " + *host + ":" + *port + " unreachable"
				sendToSlack(&color, &text, url)
				sentUp = 0
				sentDown = 1
			}
		}
		time.Sleep(time.Duration(*interval) * time.Second)
	}
}

func checkIfError(err error) {
	if err == nil {
		return
	}

	//fmt.Printf("\x1b[31;1m%s\x1b[0m\n", fmt.Sprintf("error: %s", err))
	log.Printf("error: %s", err)
	//os.Exit(1)
}

func sendToSlack(color, text, url *string) {
	err := slackhookgo.Send(
		*url,
		slackhookgo.NewSlackMessage(
			"username",
			"backup",
		).AddAttachment(
			slackhookgo.MessageAttachment{
				Color: *color,
				Text:  *text,
				Title: "<!channel>",
			},
		),
	)
	checkIfError(err)
}

func checkFlags() {
	flag.Parse()
	/*	if *host == "" {
			flag.PrintDefaults()
			log.Fatal("host missing, exiting.")
		}

		if *port == "" {
			flag.PrintDefaults()
			log.Fatal("port missing, exiting.")
		}
	*/
}
<commit_msg>add type alertMessage<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/lowstz/slackhookgo"
)

type alertMessage struct {
	url   string
	text  string
	color string
}

func (am *alertMessage) sentSlack() {
	err := slackhookgo.Send(
		am.url,
		slackhookgo.NewSlackMessage(
			"username",
			"backup",
		).AddAttachment(
			slackhookgo.MessageAttachment{
				Color: am.color,
				Text:  am.text,
				Title: "<!channel>",
			},
		),
	)
	checkIfError(err)
}

func checkIfError(err error) {
	if err == nil {
		return
	}
	log.Printf("error: %s", err)
}

func main() {
	var (
		sentUp   = 0
		sentDown = 0
		protocol = flag.String("protocol", "tcp", "protocol tcp/udp")
		host     = flag.String("host", "ya.ru", "destination host")
		port     = flag.String("port", "80", "destination port")
		interval = flag.Uint("interval", 5, "interval check seconds")
		url      = flag.String("url", "", "hook url")
	)

	flag.Parse()

	address := fmt.Sprintf("%s:%s", *host, *port)

	for {
		conn, err := net.DialTimeout(*protocol, address, 3*time.Second)
		checkIfError(err)
		if err == nil {
			conn.Close()
			if sentUp == 0 {
				am := alertMessage{
					color: "good",
					text:  "Destination host " + *host + ":" + *port + " reachable",
					url:   *url,
				}
				am.sentSlack()
				sentUp = 1
				sentDown = 0
			}
		} else {
			if sentDown == 0 {
				am := alertMessage{
					color: "danger",
					text:  "Destination host " + *host + ":" + *port + " unreachable",
					url:   *url,
				}
				am.sentSlack()
				sentUp = 0
				sentDown = 1
			}
		}
		time.Sleep(time.Duration(*interval) * time.Second)
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
)

func TestAccAWSELB_importBasic(t *testing.T) {
	resourceName := "aws_subnet.bar"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSELBDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSELBConfig,
			},

			resource.TestStep{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}
<commit_msg>provider/aws: Fix typo in ELB import test (#6886)<commit_after>package aws

import (
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
)

func TestAccAWSELB_importBasic(t *testing.T) {
	resourceName := "aws_elb.bar"

	resource.Test(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSELBDestroy,
		Steps: []resource.TestStep{
			resource.TestStep{
				Config: testAccAWSELBConfig,
			},

			resource.TestStep{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}
<|endoftext|>


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

package docker

import (
	"archive/tar"
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/url"
	"strings"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/tsuru/tsuru/app/image"
	"github.com/tsuru/tsuru/builder"
	"github.com/tsuru/tsuru/log"
	"github.com/tsuru/tsuru/net"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/dockercommon"
	"github.com/tsuru/tsuru/safe"
)

var _ builder.Builder = &dockerBuilder{}

func (b *dockerBuilder) PlatformAdd(opts builder.PlatformOptions) error {
	return b.buildPlatform(opts.Name, opts.Args, opts.Output, opts.Input)
}

func (b *dockerBuilder) PlatformUpdate(opts builder.PlatformOptions) error {
	return b.buildPlatform(opts.Name, opts.Args, opts.Output, opts.Input)
}

func (b *dockerBuilder) buildPlatform(name string, args map[string]string, w io.Writer, r io.Reader) error {
	var inputStream io.Reader
	var dockerfileURL string
	if r != nil {
		data, err := ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		var buf bytes.Buffer
		writer := tar.NewWriter(&buf)
		writer.WriteHeader(&tar.Header{
			Name: "Dockerfile",
			Mode: 0644,
			Size: int64(len(data)),
		})
		writer.Write(data)
		writer.Close()
		inputStream = &buf
	} else {
		dockerfileURL = args["dockerfile"]
		if dockerfileURL == "" {
			return errors.New("Dockerfile is required")
		}
		if _, err := url.ParseRequestURI(dockerfileURL); err != nil {
			return errors.New("Dockerfile parameter must be a URL")
		}
	}
	imageName := image.PlatformImageName(name)
	provisioners, err := provision.Registry()
	if err != nil {
		return err
	}
	var client *docker.Client
	for _, p := range provisioners {
		if provisioner, ok := p.(provision.BuilderDeploy); ok {
			client, err = provisioner.GetDockerClient(nil)
			if client != nil && err == nil {
				break
			}
		}
	}
	if client == nil {
		if err != nil {
			return err
		}
		return errors.New("No Docker nodes available")
	}
	buildOptions := docker.BuildImageOptions{
		Name:              imageName,
		Pull:              true,
		NoCache:           true,
		RmTmpContainer:    true,
		Remote:            dockerfileURL,
		InputStream:       inputStream,
		OutputStream:      w,
		InactivityTimeout: net.StreamInactivityTimeout,
		RawJSONStream:     true,
	}
	err = client.BuildImage(buildOptions)
	if err != nil {
		return err
	}
	parts := strings.Split(imageName, ":")
	var tag string
	if len(parts) > 2 {
		imageName = strings.Join(parts[:len(parts)-1], ":")
		tag = parts[len(parts)-1]
	} else if len(parts) > 1 {
		imageName = parts[0]
		tag = parts[1]
	} else {
		imageName = parts[0]
		tag = "latest"
	}
	var buf safe.Buffer
	pushOpts := docker.PushImageOptions{
		Name:              imageName,
		Tag:               tag,
		OutputStream:      &buf,
		InactivityTimeout: net.StreamInactivityTimeout,
		RawJSONStream:     true,
	}
	err = client.PushImage(pushOpts, dockercommon.RegistryAuthConfig())
	if err != nil {
		log.Errorf("[docker] Failed to push image %q (%s): %s", name, err, buf.String())
		return err
	}
	return nil
}

func (b *dockerBuilder) PlatformRemove(name string) error {
	provisioners, err := provision.Registry()
	if err != nil {
		return err
	}
	var client *docker.Client
	for _, p := range provisioners {
		if provisioner, ok := p.(provision.BuilderDeploy); ok {
			client, err = provisioner.GetDockerClient(nil)
			if err != nil {
				return err
			}
		}
	}
	img, err := client.InspectImage(image.PlatformImageName(name))
	if err != nil {
		return err
	}
	err = client.RemoveImage(img.ID)
	if err != nil && err == docker.ErrNoSuchImage {
		log.Errorf("error removing image %s from Docker: no such image", name)
		return nil
	}
	return err
}
<commit_msg>builder/docker: improve errors<commit_after>// Copyright 2017 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"archive/tar"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"

	docker "github.com/fsouza/go-dockerclient"
	"github.com/tsuru/tsuru/app/image"
	"github.com/tsuru/tsuru/builder"
	tsuruErrors "github.com/tsuru/tsuru/errors"
	"github.com/tsuru/tsuru/log"
	"github.com/tsuru/tsuru/net"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/dockercommon"
	"github.com/tsuru/tsuru/safe"
)

var _ builder.Builder = &dockerBuilder{}

func (b *dockerBuilder) PlatformAdd(opts builder.PlatformOptions) error {
	return b.buildPlatform(opts.Name, opts.Args, opts.Output, opts.Input)
}

func (b *dockerBuilder) PlatformUpdate(opts builder.PlatformOptions) error {
	return b.buildPlatform(opts.Name, opts.Args, opts.Output, opts.Input)
}

func (b *dockerBuilder) buildPlatform(name string, args map[string]string, w io.Writer, r io.Reader) error {
	var inputStream io.Reader
	var dockerfileURL string
	if r != nil {
		data, err := ioutil.ReadAll(r)
		if err != nil {
			return err
		}
		var buf bytes.Buffer
		writer := tar.NewWriter(&buf)
		writer.WriteHeader(&tar.Header{
			Name: "Dockerfile",
			Mode: 0644,
			Size: int64(len(data)),
		})
		writer.Write(data)
		writer.Close()
		inputStream = &buf
	} else {
		dockerfileURL = args["dockerfile"]
		if dockerfileURL == "" {
			return errors.New("Dockerfile is required")
		}
		if _, err := url.ParseRequestURI(dockerfileURL); err != nil {
			return errors.New("Dockerfile parameter must be a URL")
		}
	}
	imageName := image.PlatformImageName(name)
	provisioners, err := provision.Registry()
	if err != nil {
		return err
	}
	var client *docker.Client
	multiErr := tsuruErrors.NewMultiError()
	for _, p := range provisioners {
		if provisioner, ok := p.(provision.BuilderDeploy); ok {
			client, err = provisioner.GetDockerClient(nil)
			if err != nil {
				multiErr.Add(err)
			} else if client != nil {
				break
			}
		}
	}
	if client == nil {
		if multiErr.Len() > 0 {
			return fmt.Errorf("%s", multiErr)
		}
		return errors.New("No Docker nodes available")
	}
	buildOptions := docker.BuildImageOptions{
		Name:              imageName,
		Pull:              true,
		NoCache:           true,
		RmTmpContainer:    true,
		Remote:            dockerfileURL,
		InputStream:       inputStream,
		OutputStream:      w,
		InactivityTimeout: net.StreamInactivityTimeout,
		RawJSONStream:     true,
	}
	err = client.BuildImage(buildOptions)
	if err != nil {
		return err
	}
	parts := strings.Split(imageName, ":")
	var tag string
	if len(parts) > 2 {
		imageName = strings.Join(parts[:len(parts)-1], ":")
		tag = parts[len(parts)-1]
	} else if len(parts) > 1 {
		imageName = parts[0]
		tag = parts[1]
	} else {
		imageName = parts[0]
		tag = "latest"
	}
	var buf safe.Buffer
	pushOpts := docker.PushImageOptions{
		Name:              imageName,
		Tag:               tag,
		OutputStream:      &buf,
		InactivityTimeout: net.StreamInactivityTimeout,
		RawJSONStream:     true,
	}
	err = client.PushImage(pushOpts, dockercommon.RegistryAuthConfig())
	if err != nil {
		log.Errorf("[docker] Failed to push image %q (%s): %s", name, err, buf.String())
		return err
	}
	return nil
}

func (b *dockerBuilder) PlatformRemove(name string) error {
	provisioners, err := provision.Registry()
	if err != nil {
		return err
	}
	var client *docker.Client
	for _, p := range provisioners {
		if provisioner, ok := p.(provision.BuilderDeploy); ok {
			client, err = provisioner.GetDockerClient(nil)
			if err != nil {
				return err
			}
		}
	}
	img, err := client.InspectImage(image.PlatformImageName(name))
	if err != nil {
		return err
	}
	err = client.RemoveImage(img.ID)
	if err != nil && err == docker.ErrNoSuchImage {
		log.Errorf("error removing image %s from Docker: no such image", name)
		return nil
	}
	return err
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"github.com/julienschmidt/httprouter"
	"net/http"
)

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Fprint(w, "Welcome!")
}

func Hello(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	fmt.Fprintf(w, "hello, %s!\n", ps.ByName("name"))
}

// Router returns the router http handler for the package.
func Router() http.Handler {
	router := httprouter.New()

	router.GET("/", Index)
	router.GET("/hello/:name", Hello)

	// Authentication Routes
	router.GET("/api/users/redirect", RedirectUser)
	router.GET("/api/users/logout", LogoutUser)
	router.GET("/api/users/current", ServeCurrentUser)

	// API Routes
	router.GET("/api/listings", ServeRecentListings)
	router.POST("/api/listings", ServeAddListing)
	router.GET("/api/listings/:id", ServeListingById)
	router.GET("/api/listings/:id/photos", ServePhotosByListingId)

	router.GET("/api/seeks", ServeRecentSeeks)
	router.POST("/api/seeks", ServeAddSeek)
	router.GET("/api/seeks/:id", ServeSeekById)

	router.GET("/api/savedsearches", ServeRecentSavedSearches)
	router.POST("/api/savedsearches", ServeAddSeek)
	router.GET("/api/savedsearches/:id", ServeSeekById)

	router.GET("/api/search/:query", ServeSearch)

	return router
}
<commit_msg>Route saved searches to saved searches, not seeks. D'oh!<commit_after>package server

import (
	"fmt"
	"github.com/julienschmidt/httprouter"
	"net/http"
)

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Fprint(w, "Welcome!")
}

func Hello(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	fmt.Fprintf(w, "hello, %s!\n", ps.ByName("name"))
}

// Router returns the router http handler for the package.
func Router() http.Handler {
	router := httprouter.New()

	router.GET("/", Index)
	router.GET("/hello/:name", Hello)

	// Authentication Routes
	router.GET("/api/users/redirect", RedirectUser)
	router.GET("/api/users/logout", LogoutUser)
	router.GET("/api/users/current", ServeCurrentUser)

	// API Routes
	router.GET("/api/listings", ServeRecentListings)
	router.POST("/api/listings", ServeAddListing)
	router.GET("/api/listings/:id", ServeListingById)
	router.GET("/api/listings/:id/photos", ServePhotosByListingId)

	router.GET("/api/seeks", ServeRecentSeeks)
	router.POST("/api/seeks", ServeAddSeek)
	router.GET("/api/seeks/:id", ServeSeekById)

	router.GET("/api/savedsearches", ServeRecentSavedSearches)
	router.POST("/api/savedsearches", ServeAddSavedSearch)
	router.GET("/api/savedsearches/:id", ServeSavedSearchById)

	router.GET("/api/search/:query", ServeSearch)

	return router
}
<|endoftext|>


<commit_before>// This keeps track of the files which need to be written back

package vfscache

import (
	"container/heap"
	"context"
	"sync"
	"time"

	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/fserrors"
	"github.com/rclone/rclone/vfs/vfscommon"
)

const (
	maxUploadDelay = 5 * time.Minute // max delay betwen upload attempts
)

// putFn is the interface that item provides to store the data
type putFn func(context.Context) error

// writeBack keeps track of the items which need to be written back to the disk at some point
type writeBack struct {
	mu      sync.Mutex
	items   writeBackItems           // priority queue of *writeBackItem - writeBackItems are in here while awaiting transfer only
	lookup  map[*Item]*writeBackItem // for getting a *writeBackItem from a *Item - writeBackItems are in here until cancelled
	opt     *vfscommon.Options       // VFS options
	timer   *time.Timer              // next scheduled time for the uploader
	kick    chan struct{}            // send on this channel to wake up the uploader
	uploads int                      // number of uploads in progress
}

// make a new writeBack
//
// cancel the context to stop the background goroutine
func newWriteBack(ctx context.Context, opt *vfscommon.Options) *writeBack {
	wb := &writeBack{
		items:  writeBackItems{},
		lookup: make(map[*Item]*writeBackItem),
		opt:    opt,
		timer:  time.NewTimer(time.Second),
		kick:   make(chan struct{}, 1),
	}
	wb.timer.Stop()
	heap.Init(&wb.items)
	go wb.uploader(ctx)
	return wb
}

// writeBackItem stores an Item awaiting writeback
//
// These are stored on the items heap when awaiting transfer but
// removed from the items heap when transferring. They remain in the
// lookup map until cancelled.
//
// writeBack.mu must be held to manipulate this
type writeBackItem struct {
	name      string             // name of the item so we don't have to read it from item
	index     int                // index into the priority queue for update
	item      *Item              // Item that needs writeback
	expiry    time.Time          // When this expires we will write it back
	uploading bool               // True if item is being processed by upload() method
	onHeap    bool               // true if this item is on the items heap
	cancel    context.CancelFunc // To cancel the upload with
	done      chan struct{}      // closed when the cancellation completes
	putFn     putFn              // To write the object data
	tries     int                // number of times we have tried to upload
	delay     time.Duration      // delay between upload attempts
}

// A writeBackItems implements a priority queue by implementing
// heap.Interface and holds writeBackItems.
type writeBackItems []*writeBackItem

func (ws writeBackItems) Len() int { return len(ws) }

func (ws writeBackItems) Less(i, j int) bool {
	return ws[i].expiry.Sub(ws[j].expiry) < 0
}

func (ws writeBackItems) Swap(i, j int) {
	ws[i], ws[j] = ws[j], ws[i]
	ws[i].index = i
	ws[j].index = j
}

func (ws *writeBackItems) Push(x interface{}) {
	n := len(*ws)
	item := x.(*writeBackItem)
	item.index = n
	*ws = append(*ws, item)
}

func (ws *writeBackItems) Pop() interface{} {
	old := *ws
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // avoid memory leak
	item.index = -1 // for safety
	*ws = old[0 : n-1]
	return item
}

// update modifies the expiry of an Item in the queue.
//
// call with lock held
func (ws *writeBackItems) _update(item *writeBackItem, expiry time.Time) {
	item.expiry = expiry
	heap.Fix(ws, item.index)
}

// return a new expiry time based from now until the WriteBack timeout
//
// call with lock held
func (wb *writeBack) _newExpiry() time.Time {
	expiry := time.Now()
	if wb.opt.WriteBack > 0 {
		expiry = expiry.Add(wb.opt.WriteBack)
	}
	return expiry
}

// make a new writeBackItem
//
// call with the lock held
func (wb *writeBack) _newItem(item *Item, name string) *writeBackItem {
	wbItem := &writeBackItem{
		name:   name,
		item:   item,
		expiry: wb._newExpiry(),
		delay:  wb.opt.WriteBack,
	}
	wb._addItem(wbItem)
	wb._pushItem(wbItem)
	return wbItem
}

// add a writeBackItem to the lookup map
//
// call with the lock held
func (wb *writeBack) _addItem(wbItem *writeBackItem) {
	wb.lookup[wbItem.item] = wbItem
}

// delete a writeBackItem from the lookup map
//
// call with the lock held
func (wb *writeBack) _delItem(wbItem *writeBackItem) {
	delete(wb.lookup, wbItem.item)
}

// pop a writeBackItem from the items heap
//
// call with the lock held
func (wb *writeBack) _popItem() (wbItem *writeBackItem) {
	wbItem = heap.Pop(&wb.items).(*writeBackItem)
	wbItem.onHeap = false
	return wbItem
}

// push a writeBackItem onto the items heap
//
// call with the lock held
func (wb *writeBack) _pushItem(wbItem *writeBackItem) {
	if !wbItem.onHeap {
		heap.Push(&wb.items, wbItem)
		wbItem.onHeap = true
	}
}

// remove a writeBackItem from the items heap
//
// call with the lock held
func (wb *writeBack) _removeItem(wbItem *writeBackItem) {
	if wbItem.onHeap {
		heap.Remove(&wb.items, wbItem.index)
		wbItem.onHeap = false
	}
}

// peek the oldest writeBackItem - may be nil
//
// call with the lock held
func (wb *writeBack) _peekItem() (wbItem *writeBackItem) {
	if len(wb.items) == 0 {
		return nil
	}
	return wb.items[0]
}

// reset the timer which runs the expiries
func (wb *writeBack) _resetTimer() {
	wbItem := wb._peekItem()
	if wbItem == nil {
		wb.timer.Stop()
	} else {
		dt := time.Until(wbItem.expiry)
		if dt < 0 {
			dt = 0
		}
		wb.timer.Reset(dt)
	}
}

// add adds an item to the writeback queue or resets its timer if it
// is already there.
//
// if modified is false then it it doesn't a pending upload
func (wb *writeBack) add(item *Item, name string, modified bool, putFn putFn) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	wbItem, ok := wb.lookup[item]
	if !ok {
		wbItem = wb._newItem(item, name)
	} else {
		if wbItem.uploading && modified {
			// We are uploading already so cancel the upload
			wb._cancelUpload(wbItem)
		}
		// Kick the timer on
		wb.items._update(wbItem, wb._newExpiry())
	}
	wbItem.putFn = putFn
	wb._resetTimer()
}

// Call when a file is removed. This cancels a writeback if there is
// one and doesn't return the item to the queue.
func (wb *writeBack) remove(item *Item) (found bool) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	wbItem, found := wb.lookup[item]
	if found {
		fs.Debugf(wbItem.name, "vfs cache: cancelling writeback (uploading %v) %p item %p", wbItem.uploading, wbItem, wbItem.item)
		if wbItem.uploading {
			// We are uploading already so cancel the upload
			wb._cancelUpload(wbItem)
		}
		// Remove the item from the heap
		wb._removeItem(wbItem)
		// Remove the item from the lookup map
		wb._delItem(wbItem)
	}
	wb._resetTimer()
	return found
}

// kick the upload checker
//
// This should be called at the end of uploads just in case we had to
// pause uploades because max items was exceeded
//
// call with the lock held
func (wb *writeBack) _kickUploader() {
	select {
	case wb.kick <- struct{}{}:
	default:
	}
}

// upload the item - called as a goroutine
//
// uploading will have been incremented here already
func (wb *writeBack) upload(ctx context.Context, wbItem *writeBackItem) {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	putFn := wbItem.putFn
	wbItem.tries++

	wb.mu.Unlock()
	err := putFn(ctx)
	wb.mu.Lock()

	wbItem.cancel() // cancel context to release resources since store done

	fs.Debugf(wbItem.name, "uploading = false %p item %p", wbItem, wbItem.item)
	wbItem.uploading = false
	wb.uploads--

	if err != nil {
		// FIXME should this have a max number of transfer attempts?
		wbItem.delay *= 2
		if wbItem.delay > maxUploadDelay {
			wbItem.delay = maxUploadDelay
		}
		if _, uerr := fserrors.Cause(err); uerr == context.Canceled {
			fs.Infof(wbItem.name, "vfs cache: upload canceled sucessfully")
			// Upload was cancelled so reset timer
			wbItem.delay = wb.opt.WriteBack
		} else {
			fs.Errorf(wbItem.name, "vfs cache: failed to upload try #%d, will retry in %v: %v", wbItem.tries, wbItem.delay, err)
		}
		// push the item back on the queue for retry
		wb._pushItem(wbItem)
		wb.items._update(wbItem, time.Now().Add(wbItem.delay))
	} else {
		fs.Infof(wbItem.name, "vfs cache: upload succeeded try #%d", wbItem.tries)
		// show that we are done with the item
		wb._delItem(wbItem)
	}
	wb._kickUploader()
	close(wbItem.done)
}

// cancel the upload - the item should be on the heap after this returns
//
// call with lock held
func (wb *writeBack) _cancelUpload(wbItem *writeBackItem) {
	if !wbItem.uploading {
		return
	}
	fs.Infof(wbItem.name, "vfs cache: cancelling upload")
	if wbItem.cancel != nil {
		// Cancel the upload - this may or may not be effective
		wbItem.cancel()
		// wait for the uploader to finish
		//
		// we need to wait without the lock otherwise the
		// background part will never run.
		wb.mu.Unlock()
		<-wbItem.done
		wb.mu.Lock()
	}
	// uploading items are not on the heap so add them back
	wb._pushItem(wbItem)
	fs.Infof(wbItem.name, "vfs cache: cancelled upload")
}

// cancelUpload cancels the upload of the item if there is one in progress
//
// it returns true if there was an upload in progress
func (wb *writeBack) cancelUpload(item *Item) bool {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	wbItem, ok := wb.lookup[item]
	if !ok || !wbItem.uploading {
		return false
	}
	wb._cancelUpload(wbItem)
	return true
}

// this uploads as many items as possible
func (wb *writeBack) processItems(ctx context.Context) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	resetTimer := false
	for wbItem := wb._peekItem(); wbItem != nil && time.Until(wbItem.expiry) <= 0; wbItem = wb._peekItem() {
		// If reached transfer limit don't restart the timer
		if wb.uploads >= fs.Config.Transfers {
			fs.Debugf(wbItem.name, "vfs cache: delaying writeback as --transfers exceeded")
			resetTimer = false
			break
		}
		resetTimer = true
		// Pop the item, mark as uploading and start the uploader
		wbItem = wb._popItem()
		fs.Debugf(wbItem.name, "uploading = true %p item %p", wbItem, wbItem.item)
		wbItem.uploading = true
		wb.uploads++
		newCtx, cancel := context.WithCancel(ctx)
		wbItem.cancel = cancel
		wbItem.done = make(chan struct{})
		go wb.upload(newCtx, wbItem)
	}

	if resetTimer {
		wb._resetTimer()
	}
}

// Looks for items which need writing back and write them back until
// the context is cancelled
func (wb *writeBack) uploader(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			wb.timer.Stop()
			return
		case <-wb.timer.C:
			wb.processItems(ctx)
		case <-wb.kick:
			wb.processItems(ctx)
		}
	}
}

// return the number of uploads in progress
func (wb *writeBack) getStats() (uploadsInProgress, uploadsQueued int) {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	return wb.uploads, len(wb.items)
}
<commit_msg>vfs: use call after functions in writeback to simplify code<commit_after>// This keeps track of the files which need to be written back

package vfscache

import (
	"container/heap"
	"context"
	"sync"
	"time"

	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/fs/fserrors"
	"github.com/rclone/rclone/vfs/vfscommon"
)

const (
	maxUploadDelay = 5 * time.Minute // max delay betwen upload attempts
)

// putFn is the interface that item provides to store the data
type putFn func(context.Context) error

// writeBack keeps track of the items which need to be written back to the disk at some point
type writeBack struct {
	ctx     context.Context
	mu      sync.Mutex
	items   writeBackItems           // priority queue of *writeBackItem - writeBackItems are in here while awaiting transfer only
	lookup  map[*Item]*writeBackItem // for getting a *writeBackItem from a *Item - writeBackItems are in here until cancelled
	opt     *vfscommon.Options       // VFS options
	timer   *time.Timer              // next scheduled time for the uploader
	expiry  time.Time                // time the next item exires or IsZero
	uploads int                      // number of uploads in progress
}

// make a new writeBack
//
// cancel the context to stop the background goroutine
func newWriteBack(ctx context.Context, opt *vfscommon.Options) *writeBack {
	wb := &writeBack{
		ctx:    ctx,
		items:  writeBackItems{},
		lookup: make(map[*Item]*writeBackItem),
		opt:    opt,
	}
	heap.Init(&wb.items)
	return wb
}

// writeBackItem stores an Item awaiting writeback
//
// These are stored on the items heap when awaiting transfer but
// removed from the items heap when transferring. They remain in the
// lookup map until cancelled.
//
// writeBack.mu must be held to manipulate this
type writeBackItem struct {
	name      string             // name of the item so we don't have to read it from item
	index     int                // index into the priority queue for update
	item      *Item              // Item that needs writeback
	expiry    time.Time          // When this expires we will write it back
	uploading bool               // True if item is being processed by upload() method
	onHeap    bool               // true if this item is on the items heap
	cancel    context.CancelFunc // To cancel the upload with
	done      chan struct{}      // closed when the cancellation completes
	putFn     putFn              // To write the object data
	tries     int                // number of times we have tried to upload
	delay     time.Duration      // delay between upload attempts
}

// A writeBackItems implements a priority queue by implementing
// heap.Interface and holds writeBackItems.
type writeBackItems []*writeBackItem

func (ws writeBackItems) Len() int { return len(ws) }

func (ws writeBackItems) Less(i, j int) bool {
	return ws[i].expiry.Sub(ws[j].expiry) < 0
}

func (ws writeBackItems) Swap(i, j int) {
	ws[i], ws[j] = ws[j], ws[i]
	ws[i].index = i
	ws[j].index = j
}

func (ws *writeBackItems) Push(x interface{}) {
	n := len(*ws)
	item := x.(*writeBackItem)
	item.index = n
	*ws = append(*ws, item)
}

func (ws *writeBackItems) Pop() interface{} {
	old := *ws
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // avoid memory leak
	item.index = -1 // for safety
	*ws = old[0 : n-1]
	return item
}

// update modifies the expiry of an Item in the queue.
//
// call with lock held
func (ws *writeBackItems) _update(item *writeBackItem, expiry time.Time) {
	item.expiry = expiry
	heap.Fix(ws, item.index)
}

// return a new expiry time based from now until the WriteBack timeout
//
// call with lock held
func (wb *writeBack) _newExpiry() time.Time {
	expiry := time.Now()
	if wb.opt.WriteBack > 0 {
		expiry = expiry.Add(wb.opt.WriteBack)
	}
	return expiry
}

// make a new writeBackItem
//
// call with the lock held
func (wb *writeBack) _newItem(item *Item, name string) *writeBackItem {
	wbItem := &writeBackItem{
		name:   name,
		item:   item,
		expiry: wb._newExpiry(),
		delay:  wb.opt.WriteBack,
	}
	wb._addItem(wbItem)
	wb._pushItem(wbItem)
	return wbItem
}

// add a writeBackItem to the lookup map
//
// call with the lock held
func (wb *writeBack) _addItem(wbItem *writeBackItem) {
	wb.lookup[wbItem.item] = wbItem
}

// delete a writeBackItem from the lookup map
//
// call with the lock held
func (wb *writeBack) _delItem(wbItem *writeBackItem) {
	delete(wb.lookup, wbItem.item)
}

// pop a writeBackItem from the items heap
//
// call with the lock held
func (wb *writeBack) _popItem() (wbItem *writeBackItem) {
	wbItem = heap.Pop(&wb.items).(*writeBackItem)
	wbItem.onHeap = false
	return wbItem
}

// push a writeBackItem onto the items heap
//
// call with the lock held
func (wb *writeBack) _pushItem(wbItem *writeBackItem) {
	if !wbItem.onHeap {
		heap.Push(&wb.items, wbItem)
		wbItem.onHeap = true
	}
}

// remove a writeBackItem from the items heap
//
// call with the lock held
func (wb *writeBack) _removeItem(wbItem *writeBackItem) {
	if wbItem.onHeap {
		heap.Remove(&wb.items, wbItem.index)
		wbItem.onHeap = false
	}
}

// peek the oldest writeBackItem - may be nil
//
// call with the lock held
func (wb *writeBack) _peekItem() (wbItem *writeBackItem) {
	if len(wb.items) == 0 {
		return nil
	}
	return wb.items[0]
}

// reset the timer which runs the expiries
func (wb *writeBack) _resetTimer() {
	wbItem := wb._peekItem()
	if wbItem == nil {
		if wb.expiry.IsZero() {
			return
		}
		wb.expiry = time.Time{}
		fs.Debugf(nil, "resetTimer STOP")
		if wb.timer != nil {
			wb.timer.Stop()
			wb.timer = nil
		}
	} else {
		if wb.expiry.Equal(wbItem.expiry) {
			return
		}
		wb.expiry = wbItem.expiry
		dt := time.Until(wbItem.expiry)
		if dt < 0 {
			dt = 0
		}
		fs.Debugf(nil, "resetTimer dt=%v", dt)
		if wb.timer != nil {
			wb.timer.Stop()
		}
		wb.timer = time.AfterFunc(dt, func() {
			wb.processItems(wb.ctx)
		})
	}
}

// add adds an item to the writeback queue or resets its timer if it
// is already there.
//
// if modified is false then it it doesn't a pending upload
func (wb *writeBack) add(item *Item, name string, modified bool, putFn putFn) *writeBackItem {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	wbItem, ok := wb.lookup[item]
	if !ok {
		wbItem = wb._newItem(item, name)
	} else {
		if wbItem.uploading && modified {
			// We are uploading already so cancel the upload
			wb._cancelUpload(wbItem)
		}
		// Kick the timer on
		wb.items._update(wbItem, wb._newExpiry())
	}
	wbItem.putFn = putFn
	wb._resetTimer()
	return wbItem
}

// Call when a file is removed. This cancels a writeback if there is
// one and doesn't return the item to the queue.
func (wb *writeBack) remove(item *Item) (found bool) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	wbItem, found := wb.lookup[item]
	if found {
		fs.Debugf(wbItem.name, "vfs cache: cancelling writeback (uploading %v) %p item %p", wbItem.uploading, wbItem, wbItem.item)
		if wbItem.uploading {
			// We are uploading already so cancel the upload
			wb._cancelUpload(wbItem)
		}
		// Remove the item from the heap
		wb._removeItem(wbItem)
		// Remove the item from the lookup map
		wb._delItem(wbItem)
	}
	wb._resetTimer()
	return found
}

// upload the item - called as a goroutine
//
// uploading will have been incremented here already
func (wb *writeBack) upload(ctx context.Context, wbItem *writeBackItem) {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	putFn := wbItem.putFn
	wbItem.tries++

	wb.mu.Unlock()
	err := putFn(ctx)
	wb.mu.Lock()

	wbItem.cancel() // cancel context to release resources since store done

	//fs.Debugf(wbItem.name, "uploading = false %p item %p", wbItem, wbItem.item)
	wbItem.uploading = false
	wb.uploads--

	if err != nil {
		// FIXME should this have a max number of transfer attempts?
		wbItem.delay *= 2
		if wbItem.delay > maxUploadDelay {
			wbItem.delay = maxUploadDelay
		}
		if _, uerr := fserrors.Cause(err); uerr == context.Canceled {
			fs.Infof(wbItem.name, "vfs cache: upload canceled sucessfully")
			// Upload was cancelled so reset timer
			wbItem.delay = wb.opt.WriteBack
		} else {
			fs.Errorf(wbItem.name, "vfs cache: failed to upload try #%d, will retry in %v: %v", wbItem.tries, wbItem.delay, err)
		}
		// push the item back on the queue for retry
		wb._pushItem(wbItem)
		wb.items._update(wbItem, time.Now().Add(wbItem.delay))
	} else {
		fs.Infof(wbItem.name, "vfs cache: upload succeeded try #%d", wbItem.tries)
		// show that we are done with the item
		wb._delItem(wbItem)
	}
	wb._resetTimer()
	close(wbItem.done)
}

// cancel the upload - the item should be on the heap after this returns
//
// call with lock held
func (wb *writeBack) _cancelUpload(wbItem *writeBackItem) {
	if !wbItem.uploading {
		return
	}
	fs.Infof(wbItem.name, "vfs cache: cancelling upload")
	if wbItem.cancel != nil {
		// Cancel the upload - this may or may not be effective
		wbItem.cancel()
		// wait for the uploader to finish
		//
		// we need to wait without the lock otherwise the
		// background part will never run.
		wb.mu.Unlock()
		<-wbItem.done
		wb.mu.Lock()
	}
	// uploading items are not on the heap so add them back
	wb._pushItem(wbItem)
	fs.Infof(wbItem.name, "vfs cache: cancelled upload")
}

// cancelUpload cancels the upload of the item if there is one in progress
//
// it returns true if there was an upload in progress
func (wb *writeBack) cancelUpload(item *Item) bool {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	wbItem, ok := wb.lookup[item]
	if !ok || !wbItem.uploading {
		return false
	}
	wb._cancelUpload(wbItem)
	return true
}

// this uploads as many items as possible
func (wb *writeBack) processItems(ctx context.Context) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	if wb.ctx.Err() != nil {
		return
	}

	resetTimer := true
	for wbItem := wb._peekItem(); wbItem != nil && time.Until(wbItem.expiry) <= 0; wbItem = wb._peekItem() {
		// If reached transfer limit don't restart the timer
		if wb.uploads >= fs.Config.Transfers {
			fs.Debugf(wbItem.name, "vfs cache: delaying writeback as --transfers exceeded")
			resetTimer = false
			break
		}
		// Pop the item, mark as uploading and start the uploader
		wbItem = wb._popItem()
		//fs.Debugf(wbItem.name, "uploading = true %p item %p", wbItem, wbItem.item)
		wbItem.uploading = true
		wb.uploads++
		newCtx, cancel := context.WithCancel(ctx)
		wbItem.cancel = cancel
		wbItem.done = make(chan struct{})
		go wb.upload(newCtx, wbItem)
	}

	if resetTimer {
		wb._resetTimer()
	}
}

// return the number of uploads in progress
func (wb *writeBack) getStats() (uploadsInProgress, uploadsQueued int) {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	return wb.uploads, len(wb.items)
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
)

// Routez represents detail information on current routes
type Routez struct {
	NumRoutes int          `json:"num_routes"`
	Routes    []*RouteInfo `json:"routes"`
}

// RouteInfo has detailed information on a per connection basis.
type RouteInfo struct {
	Cid       uint64 `json:"cid"`
	URL       string `json:"url"`
	IP        string `json:"ip"`
	Port      int    `json:"port"`
	Solicited bool   `json:"solicited"`
	Subs      uint32 `json:"subscriptions"`
	Pending   int    `json:"pending_size"`
	InMsgs    int64  `json:"in_msgs"`
	OutMsgs   int64  `json:"out_msgs"`
	InBytes   int64  `json:"in_bytes"`
	OutBytes  int64  `json:"out_bytes"`
}

// HandleConnz process HTTP requests for connection information.
func (s *Server) HandleRoutez(w http.ResponseWriter, req *http.Request) {

	if req.Method == "GET" {
		r := Routez{Routes: []*RouteInfo{}}

		// Walk the list
		s.mu.Lock()
		for _, route := range s.routes {
			ri := &RouteInfo{
				Cid:       route.cid,
				Subs:      route.subs.Count(),
				Solicited: route.route.didSolicit,
				InMsgs:    route.inMsgs,
				OutMsgs:   route.outMsgs,
				InBytes:   route.inBytes,
				OutBytes:  route.outBytes,
			}

			if route.route.url != nil {
				ri.URL = route.route.url.String()
			}

			if ip, ok := route.nc.(*net.TCPConn); ok {
				addr := ip.RemoteAddr().(*net.TCPAddr)
				ri.Port = addr.Port
				ri.IP = addr.IP.String()
			}
			r.Routes = append(r.Routes, ri)
		}
		s.mu.Unlock()

		r.NumRoutes = len(r.Routes)
		b, err := json.MarshalIndent(r, "", "  ")
		if err != nil {
			Logf("Error marshalling response to /routez request: %v", err)
		}
		w.Write(b)
	} else if req.Method == "PUT" {
		body := make([]byte, 1024)
		req.Body.Read(body)
		routeURL, err := url.Parse(strings.Trim(string(body), "\x00"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte(fmt.Sprintf(`{"error": "could not parse URL: %v"}`, err)))
			return
		}

		s.connectToRoute(routeURL)
		w.Write([]byte(`{"status": "ok"}`))
	} else if req.Method == "DELETE" {
		body := make([]byte, 1024)
		req.Body.Read(body)
		url := strings.Trim(string(body), "\x00")

		s.mu.Lock()
		for _, route := range s.routes {
			if route.route.url.String() == url {
				route.mu.Lock()
				route.route.didSolicit = false // don't reconnect
				route.mu.Unlock()
				route.closeConnection()
				w.WriteHeader(200)
				w.Write([]byte(`{"status": "ok"}`))
				return
			}
		}
		s.mu.Unlock()
		w.WriteHeader(404)
		w.Write([]byte(`{"error": "could not find matching route"}`))
	}
}
<commit_msg>Route URL can be nil<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
)

// Routez represents detail information on current routes
type Routez struct {
	NumRoutes int          `json:"num_routes"`
	Routes    []*RouteInfo `json:"routes"`
}

// RouteInfo has detailed information on a per connection basis.
type RouteInfo struct {
	Cid       uint64 `json:"cid"`
	URL       string `json:"url"`
	IP        string `json:"ip"`
	Port      int    `json:"port"`
	Solicited bool   `json:"solicited"`
	Subs      uint32 `json:"subscriptions"`
	Pending   int    `json:"pending_size"`
	InMsgs    int64  `json:"in_msgs"`
	OutMsgs   int64  `json:"out_msgs"`
	InBytes   int64  `json:"in_bytes"`
	OutBytes  int64  `json:"out_bytes"`
}

// HandleConnz process HTTP requests for connection information.
func (s *Server) HandleRoutez(w http.ResponseWriter, req *http.Request) {

	if req.Method == "GET" {
		r := Routez{Routes: []*RouteInfo{}}

		// Walk the list
		s.mu.Lock()
		for _, route := range s.routes {
			ri := &RouteInfo{
				Cid:       route.cid,
				Subs:      route.subs.Count(),
				Solicited: route.route.didSolicit,
				InMsgs:    route.inMsgs,
				OutMsgs:   route.outMsgs,
				InBytes:   route.inBytes,
				OutBytes:  route.outBytes,
			}

			if route.route.url != nil {
				ri.URL = route.route.url.String()
			}

			if ip, ok := route.nc.(*net.TCPConn); ok {
				addr := ip.RemoteAddr().(*net.TCPAddr)
				ri.Port = addr.Port
				ri.IP = addr.IP.String()
			}
			r.Routes = append(r.Routes, ri)
		}
		s.mu.Unlock()

		r.NumRoutes = len(r.Routes)
		b, err := json.MarshalIndent(r, "", "  ")
		if err != nil {
			Logf("Error marshalling response to /routez request: %v", err)
		}
		w.Write(b)
	} else if req.Method == "PUT" {
		body := make([]byte, 1024)
		req.Body.Read(body)
		routeURL, err := url.Parse(strings.Trim(string(body), "\x00"))
		if err != nil {
			w.WriteHeader(400)
			w.Write([]byte(fmt.Sprintf(`{"error": "could not parse URL: %v"}`, err)))
			return
		}

		s.connectToRoute(routeURL)
		w.Write([]byte(`{"status": "ok"}`))
	} else if req.Method == "DELETE" {
		body := make([]byte, 1024)
		req.Body.Read(body)
		url := strings.Trim(string(body), "\x00")

		s.mu.Lock()
		for _, route := range s.routes {
			if route.route.url != nil && route.route.url.String() == url {
				route.mu.Lock()
				route.route.didSolicit = false // don't reconnect
				route.mu.Unlock()
				route.closeConnection()
				w.WriteHeader(200)
				w.Write([]byte(`{"status": "ok"}`))
				return
			}
		}
		s.mu.Unlock()
		w.WriteHeader(404)
		w.Write([]byte(`{"error": "could not find matching route"}`))
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 graph_tool authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Author <Ye Yin<hustcat@gmail.com>

package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/daemon/graphdriver/devmapper"
	"github.com/docker/docker/graph"
	"github.com/docker/docker/opts"
	"github.com/docker/docker/pkg/archive"
	"io"
	"os"
	"path"
)

var (
	root         string
	graphOptions []string
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s <flags>   [diff id parent]|[apply containerId]\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func initGraph() (*graph.Graph, error) {
	graphdriver.Register("devicemapper", devmapper.Init)
	graphdriver.DefaultDriver = "devicemapper"

	// Load storage driver
	driver, err := graphdriver.New(root, graphOptions)
	if err != nil {
		log.Errorf("Load storage driver error: %v", err)
		return nil, err
	}
	log.Debugf("Using graph driver %s", driver)

	log.Debugf("Creating images graph")
	g, err := graph.NewGraph(path.Join(root, "graph"), driver)
	if err != nil {
		log.Errorf("Creating images graph error: %v", err)
		return nil, err
	}
	return g, nil
}

func checkIsParent(id, parent string, g *graph.Graph) (bool, error) {
	isParent := false
	img, err := g.Get(id)
	if err != nil {
		return false, err
	}

	for {
		if img.Parent == "" {
			break
		}

		if img.Parent == parent {
			isParent = true
			break
		}

		img, err = g.Get(img.Parent)
		if err != nil {
			break
		}
	}
	return isParent, err
}

func exportDiff(id, parent string) error {
	g, err := initGraph()
	if err != nil {
		return err
	}

	b, err := checkIsParent(id, parent, g)
	if err != nil {
		return err
	}
	if !b {
		return fmt.Errorf("%s is not parent of %s", parent, id)
	}

	driver := g.Driver()
	fs, err := driver.Diff(id, parent)
	if err != nil {
		return err
	}
	defer fs.Close()

	_, err = io.Copy(os.Stdout, fs)
	if err != nil {
		return err
	}
	return nil
}

func applyLayer(containerId string) error {
	fs := os.Stdin
	dest := path.Join(root, "devicemapper", "mnt", containerId, "rootfs")
	fi, err := os.Stat(dest)
	if err != nil && !os.IsExist(err) {
		return err
	}

	if !fi.IsDir() {
		return fmt.Errorf(" Dest %s is not dir", dest)
	}

	err = archive.ApplyLayer(dest, fs)
	return err
}

func main() {
	flag.StringVar(&root, "r", "/var/lib/docker", "Docker root dir")
	opts.ListVar(&graphOptions, []string{"-storage-opt"}, "Set storage driver options")
	flDebug := flag.Bool("D", false, "Debug mode")

	flag.Parse()

	if *flDebug {
		//os.Setenv("DEBUG", "1")
		log.SetLevel(log.DebugLevel)
	}

	if flag.NArg() < 1 {
		usage()
	}

	args := flag.Args()

	switch args[0] {
	case "diff":
		if flag.NArg() < 3 {
			usage()
		}
		err := exportDiff(args[1], args[2])
		if err != nil {
			log.Errorf("Export diff  error: %v", err)
			os.Exit(1)
		}
		break
	case "apply":
		if flag.NArg() < 2 {
			usage()
		}
		err := applyLayer(args[1])
		if err != nil {
			log.Errorf("Apply diff  error: %v", err)
			os.Exit(1)
		} else {
			log.Infof("Apply diff success")
		}
		break
	default:
		fmt.Printf("Unknown command %s\n", args[0])
		usage()

		os.Exit(1)
	}
}
<commit_msg>add diffapply method<commit_after>// Copyright 2015 graph_tool authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Author <Ye Yin<hustcat@gmail.com>

package main

import (
	"flag"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/docker/docker/daemon/graphdriver"
	"github.com/docker/docker/daemon/graphdriver/devmapper"
	"github.com/docker/docker/graph"
	"github.com/docker/docker/opts"
	"github.com/docker/docker/pkg/archive"
	"io"
	"os"
	"path"
)

var (
	root         string
	graphOptions []string
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s <flags>   [diff id parent]|[apply containerId]|[diffapply id parent containerId]\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func initGraph() (*graph.Graph, error) {
	graphdriver.Register("devicemapper", devmapper.Init)
	graphdriver.DefaultDriver = "devicemapper"

	// Load storage driver
	driver, err := graphdriver.New(root, graphOptions)
	if err != nil {
		log.Errorf("Load storage driver error: %v", err)
		return nil, err
	}
	log.Debugf("Using graph driver %s", driver)

	log.Debugf("Creating images graph")
	g, err := graph.NewGraph(path.Join(root, "graph"), driver)
	if err != nil {
		log.Errorf("Creating images graph error: %v", err)
		return nil, err
	}
	return g, nil
}

func checkIsParent(id, parent string, g *graph.Graph) (bool, error) {
	isParent := false
	img, err := g.Get(id)
	if err != nil {
		return false, err
	}

	for {
		if img.Parent == "" {
			break
		}

		if img.Parent == parent {
			isParent = true
			break
		}

		img, err = g.Get(img.Parent)
		if err != nil {
			break
		}
	}
	return isParent, err
}

func exportDiff(id, parent string) error {
	g, err := initGraph()
	if err != nil {
		return err
	}

	b, err := checkIsParent(id, parent, g)
	if err != nil {
		return err
	}
	if !b {
		return fmt.Errorf("%s is not parent of %s", parent, id)
	}

	driver := g.Driver()
	fs, err := driver.Diff(id, parent)
	if err != nil {
		return err
	}
	defer fs.Close()

	_, err = io.Copy(os.Stdout, fs)
	if err != nil {
		return err
	}
	return nil
}

func applyLayer(containerId string) error {
	fs := os.Stdin
	dest := path.Join(root, "devicemapper", "mnt", containerId, "rootfs")
	fi, err := os.Stat(dest)
	if err != nil && !os.IsExist(err) {
		return err
	}

	if !fi.IsDir() {
		return fmt.Errorf(" Dest %s is not dir", dest)
	}

	err = archive.ApplyLayer(dest, fs)
	return err
}

func diffAndApply(id, parent, container string) error{
    g, err := initGraph()
    if err != nil {
        return err
    }

    b, err := checkIsParent(id, parent, g)
    if err != nil {
        return err
    }
    if !b {
        return fmt.Errorf("%s is not parent of %s", parent, id)
    }

    dest := path.Join(root, "devicemapper", "mnt", container, "rootfs")
    fi, err := os.Stat(dest)
    if err != nil && !os.IsExist(err) {
        return err
    }

    if !fi.IsDir() {
        return fmt.Errorf(" Dest %s is not dir", dest)
    }

    driver := g.Driver()
    fs, err := driver.Diff(id, parent)
    if err != nil {
        return err
    }
    defer fs.Close()

	err = archive.ApplyLayer(dest, fs)
    if err != nil {
        return err
    }
    return nil
}

func main() {
	flag.StringVar(&root, "r", "/var/lib/docker", "Docker root dir")
	opts.ListVar(&graphOptions, []string{"-storage-opt"}, "Set storage driver options")
	flDebug := flag.Bool("D", false, "Debug mode")

	flag.Parse()

	if *flDebug {
		//os.Setenv("DEBUG", "1")
		log.SetLevel(log.DebugLevel)
	}

	if flag.NArg() < 1 {
		usage()
	}

	args := flag.Args()

	switch args[0] {
	case "diff":
		if flag.NArg() < 3 {
			usage()
		}
		err := exportDiff(args[1], args[2])
		if err != nil {
			log.Errorf("Export diff  error: %v", err)
			os.Exit(1)
		}
		break
    case "diffapply":
        if flag.NArg() < 4 {
            usage()
        }
        err := diffAndApply(args[1], args[2], args[3])
        if err != nil {
            log.Errorf("Apply diff  error: %v", err)
            os.Exit(1)
        }else{
			log.Infof("Apply diff success")
		}
        break
	case "apply":
		if flag.NArg() < 2 {
			usage()
		}
		err := applyLayer(args[1])
		if err != nil {
			log.Errorf("Apply diff  error: %v", err)
			os.Exit(1)
		} else {
			log.Infof("Apply diff success")
		}
		break
	default:
		fmt.Printf("Unknown command %s\n", args[0])
		usage()

		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"github.com/nelhage/livegrep/client"
	"log"
)

type searchConnection struct {
	srv      *server
	ws       *websocket.Conn
	errors   chan error
	incoming chan Op
	outgoing chan Op
	shutdown bool
}

func (s *searchConnection) recvLoop() {
	var op Op
	for {
		if err := OpCodec.Receive(s.ws, &op); err != nil {
			log.Printf("Error in receive: %s\n", err.Error())
			if _, ok := err.(*ProtocolError); ok {
				continue
			} else {
				s.errors <- err
				break
			}
		}
		log.Printf("Incoming: %v", op)
		s.incoming <- op
		if s.shutdown {
			break
		}
	}
	close(s.incoming)
}

func (s *searchConnection) sendLoop() {
	for op := range s.outgoing {
		OpCodec.Send(s.ws, op)
	}
}

func query(q *OpQuery) *client.Query {
	return &client.Query{
		Line: q.Line,
		File: q.File,
		Repo: q.Repo,
	}
}

func (s *searchConnection) handle() {
	s.incoming = make(chan Op, 1)
	s.outgoing = make(chan Op, 1)
	s.errors = make(chan error, 1)

	go s.recvLoop()
	go s.sendLoop()
	defer close(s.outgoing)

	cl := client.ClientWithRetry(func() (client.Client, error) { return client.Dial("tcp", s.srv.config.Backends[0].Addr) })

	var nextQuery *OpQuery
	var inFlight *OpQuery

	var search client.Search
	var results <-chan *client.Result
	var err error

	for {
		select {
		case op := <-s.incoming:
			switch t := op.(type) {
			case *OpQuery:
				nextQuery = t
			default:
				s.outgoing <- &OpError{fmt.Sprintf("Invalid opcode %s", op.Opcode())}
				break
			}

		case e := <-s.errors:
			log.Printf("error reading from client: %s", e.Error())
			break
		case res, ok := <-results:
			if ok {
				s.outgoing <- &OpResult{inFlight.Id, res}
			} else {
				st, err := search.Close()
				if err == nil {
					s.outgoing <- &OpSearchDone{inFlight.Id, st}
				} else {
					s.outgoing <- &OpQueryError{inFlight.Id, err.Error()}
				}
				results = nil
				search = nil
				inFlight = nil
			}
		}
		if nextQuery != nil && results == nil {
			inFlight = nextQuery
			search, err = cl.Query(query(nextQuery))
			if err != nil {
				s.outgoing <- &OpQueryError{nextQuery.Id, err.Error()}
			} else {
				inFlight = nextQuery
				results = search.Results()
			}
			nextQuery = nil
		}
	}

	s.shutdown = true
}

func (s *server) HandleWebsocket(ws *websocket.Conn) {
	c := &searchConnection{
		srv: s,
		ws:  ws,
	}
	c.handle()
}
<commit_msg>Handle EOF from server<commit_after>package server

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"github.com/nelhage/livegrep/client"
	"log"
)

type searchConnection struct {
	srv      *server
	ws       *websocket.Conn
	errors   chan error
	incoming chan Op
	outgoing chan Op
	shutdown bool
}

func (s *searchConnection) recvLoop() {
	var op Op
	for {
		if err := OpCodec.Receive(s.ws, &op); err != nil {
			log.Printf("Error in receive: %s\n", err.Error())
			if _, ok := err.(*ProtocolError); ok {
				continue
			} else {
				s.errors <- err
				break
			}
		}
		log.Printf("Incoming: %v", op)
		s.incoming <- op
		if s.shutdown {
			break
		}
	}
	close(s.incoming)
}

func (s *searchConnection) sendLoop() {
	for op := range s.outgoing {
		OpCodec.Send(s.ws, op)
	}
}

func query(q *OpQuery) *client.Query {
	return &client.Query{
		Line: q.Line,
		File: q.File,
		Repo: q.Repo,
	}
}

func (s *searchConnection) handle() {
	s.incoming = make(chan Op, 1)
	s.outgoing = make(chan Op, 1)
	s.errors = make(chan error, 1)

	go s.recvLoop()
	go s.sendLoop()
	defer close(s.outgoing)

	cl := client.ClientWithRetry(func() (client.Client, error) { return client.Dial("tcp", s.srv.config.Backends[0].Addr) })

	var nextQuery *OpQuery
	var inFlight *OpQuery

	var search client.Search
	var results <-chan *client.Result
	var err error

	for {
		select {
		case op, ok := <-s.incoming:
			if !ok {
				break
			}
			switch t := op.(type) {
			case *OpQuery:
				nextQuery = t
			default:
				s.outgoing <- &OpError{fmt.Sprintf("Invalid opcode %s", op.Opcode())}
				break
			}

		case e := <-s.errors:
			log.Printf("error reading from client: %s", e.Error())
			break
		case res, ok := <-results:
			if ok {
				s.outgoing <- &OpResult{inFlight.Id, res}
			} else {
				st, err := search.Close()
				if err == nil {
					s.outgoing <- &OpSearchDone{inFlight.Id, st}
				} else {
					s.outgoing <- &OpQueryError{inFlight.Id, err.Error()}
				}
				results = nil
				search = nil
				inFlight = nil
			}
		}
		if nextQuery != nil && results == nil {
			inFlight = nextQuery
			search, err = cl.Query(query(nextQuery))
			if err != nil {
				s.outgoing <- &OpQueryError{nextQuery.Id, err.Error()}
			} else {
				inFlight = nextQuery
				results = search.Results()
			}
			nextQuery = nil
		}
	}

	s.shutdown = true
}

func (s *server) HandleWebsocket(ws *websocket.Conn) {
	c := &searchConnection{
		srv: s,
		ws:  ws,
	}
	c.handle()
}
<|endoftext|>