<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package dataflow

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/user"
	"path"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx"
	// Importing to get the side effect of the remote execution hook. See init().
	_ "github.com/apache/beam/sdks/go/pkg/beam/core/runtime/harness/init"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/hooks"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/protox"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/apache/beam/sdks/go/pkg/beam/options/gcpopts"
	"github.com/apache/beam/sdks/go/pkg/beam/options/jobopts"
	"github.com/apache/beam/sdks/go/pkg/beam/runners/universal/runnerlib"
	"github.com/apache/beam/sdks/go/pkg/beam/util/gcsx"
	"github.com/apache/beam/sdks/go/pkg/beam/x/hooks/perf"
	"github.com/golang/protobuf/proto"
	"golang.org/x/oauth2/google"
	df "google.golang.org/api/dataflow/v1b3"
	"google.golang.org/api/storage/v1"
)

// TODO(herohde) 5/16/2017: the Dataflow flags should match the other SDKs.

var (
	endpoint        = flag.String("dataflow_endpoint", "", "Dataflow endpoint (optional).")
	stagingLocation = flag.String("staging_location", "", "GCS staging location (required).")
	image           = flag.String("worker_harness_container_image", "", "Worker harness container image (required).")
	numWorkers      = flag.Int64("num_workers", 0, "Number of workers (optional).")
	zone            = flag.String("zone", "", "GCP zone (optional)")
	network         = flag.String("network", "", "GCP network (optional)")
	tempLocation    = flag.String("temp_location", "", "Temp location (optional)")
	machineType     = flag.String("worker_machine_type", "", "GCE machine type (optional)")

	dryRun         = flag.Bool("dry_run", false, "Dry run. Just print the job, but don't submit it.")
	teardownPolicy = flag.String("teardown_policy", "", "Job teardown policy (internal only).")

	// SDK options
	cpuProfiling     = flag.String("cpu_profiling", "", "Job records CPU profiles to this GCS location (optional)")
	sessionRecording = flag.String("session_recording", "", "Job records session transcripts")
)

func init() {
	// Note that we also _ import harness/init to setup the remote execution hook.
	beam.RegisterRunner("dataflow", Execute)

	perf.RegisterProfCaptureHook("gcs_profile_writer", gcsRecorderHook)
}

type dataflowOptions struct {
	PipelineURL string `json:"pipelineUrl"`
}

// Execute runs the given pipeline on Google Cloud Dataflow. It uses the
// default application credentials to submit the job.
func Execute(ctx context.Context, p *beam.Pipeline) error {
	project := *gcpopts.Project
	if project == "" {
		return errors.New("no Google Cloud project specified. Use --project=<project>")
	}
	if *stagingLocation == "" {
		return errors.New("no GCS staging location specified. Use --staging_location=gs://<bucket>/<path>")
	}
	if *image == "" {
		*image = jobopts.GetContainerImage(ctx)
	}
	jobName := jobopts.GetJobName()

	edges, _, err := p.Build()
	if err != nil {
		return err
	}

	if *cpuProfiling != "" {
		perf.EnableProfCaptureHook("gcs_profile_writer", *cpuProfiling)
	}

	if *sessionRecording != "" {
		// TODO(wcn): BEAM-4017
		// It's a bit inconvenient for GCS because the whole object is written in
		// one pass, whereas the session logs are constantly appended. We wouldn't
		// want to hold all the logs in memory to flush at the end of the pipeline
		// as we'd blow out memory on the worker. The implementation of the
		// CaptureHook should create an internal buffer and write chunks out to GCS
		// once they get to an appropriate size (50M or so?)
	}

	hooks.SerializeHooksToOptions()
	options := beam.PipelineOptions.Export()

	// (1) Upload Go binary and model to GCS.

	if *jobopts.WorkerBinary == "" {
		worker, err := runnerlib.BuildTempWorkerBinary(ctx)
		if err != nil {
			return err
		}
		defer os.Remove(worker)

		*jobopts.WorkerBinary = worker
	} else {
		log.Infof(ctx, "Using specified worker binary: '%v'", *jobopts.WorkerBinary)
	}

	binary, err := stageWorker(ctx, project, *stagingLocation, *jobopts.WorkerBinary)
	if err != nil {
		return err
	}

	model, err := graphx.Marshal(edges, &graphx.Options{ContainerImageURL: *image})
	if err != nil {
		return fmt.Errorf("failed to generate model pipeline: %v", err)
	}
	modelURL, err := stageModel(ctx, project, *stagingLocation, protox.MustEncode(model))
	if err != nil {
		return err
	}
	log.Info(ctx, proto.MarshalTextString(model))

	// (2) Translate pipeline to v1b3 speak.

	steps, err := translate(edges)
	if err != nil {
		return err
	}

	job := &df.Job{
		ProjectId: project,
		Name:      jobName,
		Type:      "JOB_TYPE_BATCH",
		Environment: &df.Environment{
			UserAgent: newMsg(userAgent{
				Name:    "Apache Beam SDK for Go",
				Version: "0.3.0",
			}),
			Version: newMsg(version{
				JobType: "FNAPI_BATCH",
				Major:   "6",
			}),
			SdkPipelineOptions: newMsg(pipelineOptions{
				DisplayData: findPipelineFlags(),
				Options: dataflowOptions{
					PipelineURL: modelURL,
				},
				GoOptions: options,
			}),
			WorkerPools: []*df.WorkerPool{{
				Kind: "harness",
				Packages: []*df.Package{{
					Location: binary,
					Name:     "worker",
				}},
				WorkerHarnessContainerImage: *image,
				NumWorkers:                  1,
				MachineType:                 *machineType,
				Network:                     *network,
				Zone:                        *zone,
			}},
			TempStoragePrefix: *stagingLocation + "/tmp",
			Experiments:       jobopts.GetExperiments(),
		},
		Steps: steps,
	}

	if *numWorkers > 0 {
		job.Environment.WorkerPools[0].NumWorkers = *numWorkers
	}
	if *teardownPolicy != "" {
		job.Environment.WorkerPools[0].TeardownPolicy = *teardownPolicy
	}
	if *tempLocation != "" {
		job.Environment.TempStoragePrefix = *tempLocation
	}
	printJob(ctx, job)

	if *dryRun {
		log.Info(ctx, "Dry-run: not submitting job!")
		return nil
	}

	// (4) Submit job.

	client, err := newClient(ctx, *endpoint)
	if err != nil {
		return err
	}
	upd, err := client.Projects.Jobs.Create(project, job).Do()
	if err != nil {
		return err
	}

	log.Infof(ctx, "Submitted job: %v", upd.Id)
	printJob(ctx, upd)
	if *endpoint == "" {
		log.Infof(ctx, "Console: https://console.cloud.google.com/dataflow/job/%v?project=%v", upd.Id, project)
	}
	log.Infof(ctx, "Logs: https://console.cloud.google.com/logs/viewer?project=%v&resource=dataflow_step%%2Fjob_id%%2F%v", project, upd.Id)

	if *jobopts.Async {
		return nil
	}

	time.Sleep(1 * time.Minute)
	for {
		j, err := client.Projects.Jobs.Get(project, upd.Id).Do()
		if err != nil {
			return fmt.Errorf("failed to get job: %v", err)
		}

		switch j.CurrentState {
		case "JOB_STATE_DONE":
			log.Info(ctx, "Job succeeded!")
			return nil

		case "JOB_STATE_FAILED":
			return fmt.Errorf("job %s failed", upd.Id)

		case "JOB_STATE_RUNNING":
			log.Info(ctx, "Job still running ...")

		default:
			log.Infof(ctx, "Job state: %v ...", j.CurrentState)
		}

		time.Sleep(30 * time.Second)
	}
}

// stageModel uploads the pipeline model to GCS as a unique object.
func stageModel(ctx context.Context, project, location string, model []byte) (string, error) {
	bucket, prefix, err := gcsx.ParseObject(location)
	if err != nil {
		return "", fmt.Errorf("invalid staging location %v: %v", location, err)
	}
	obj := path.Join(prefix, fmt.Sprintf("pipeline-%v", time.Now().UnixNano()))
	if *dryRun {
		full := fmt.Sprintf("gs://%v/%v", bucket, obj)
		log.Infof(ctx, "Dry-run: not uploading model %v", full)
		return full, nil
	}

	client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
	if err != nil {
		return "", err
	}
	return gcsx.Upload(client, project, bucket, obj, bytes.NewReader(model))
}

// stageWorker uploads the worker binary to GCS as a unique object.
func stageWorker(ctx context.Context, project, location, worker string) (string, error) {
	bucket, prefix, err := gcsx.ParseObject(location)
	if err != nil {
		return "", fmt.Errorf("invalid staging location %v: %v", location, err)
	}
	obj := path.Join(prefix, fmt.Sprintf("worker-%v", time.Now().UnixNano()))
	if *dryRun {
		full := fmt.Sprintf("gs://%v/%v", bucket, obj)
		log.Infof(ctx, "Dry-run: not uploading binary %v", full)
		return full, nil
	}

	client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
	if err != nil {
		return "", err
	}
	fd, err := os.Open(worker)
	if err != nil {
		return "", fmt.Errorf("failed to open worker binary %s: %v", worker, err)
	}
	defer fd.Close()
	defer os.Remove(worker)

	return gcsx.Upload(client, project, bucket, obj, fd)
}

func username() string {
	if u, err := user.Current(); err == nil {
		return u.Username
	}
	return "anon"
}

func findPipelineFlags() []*displayData {
	var ret []*displayData

	// TODO(herohde) 2/15/2017: decide if we want all set flags.
	flag.Visit(func(f *flag.Flag) {
		ret = append(ret, newDisplayData(f.Name, "", "flag", f.Value.(flag.Getter).Get()))
	})

	return ret
}

// newClient creates a new dataflow client with default application credentials
// and CloudPlatformScope. The Dataflow endpoint is optionally overridden.
func newClient(ctx context.Context, endpoint string) (*df.Service, error) {
	cl, err := google.DefaultClient(ctx, df.CloudPlatformScope)
	if err != nil {
		return nil, err
	}
	client, err := df.New(cl)
	if err != nil {
		return nil, err
	}
	if endpoint != "" {
		log.Infof(ctx, "Dataflow endpoint override: %s", endpoint)
		client.BasePath = endpoint
	}
	return client, nil
}

func printJob(ctx context.Context, job *df.Job) {
	str, err := json.MarshalIndent(job, "", "  ")
	if err != nil {
		log.Infof(ctx, "Failed to print job %v: %v", job.Id, err)
	}
	log.Info(ctx, string(str))
}

func gcsRecorderHook(opts []string) perf.CaptureHook {
	bucket, prefix, err := gcsx.ParseObject(opts[0])
	if err != nil {
		panic(fmt.Sprintf("Invalid hook configuration for gcsRecorderHook: %s", opts))
	}

	return func(ctx context.Context, spec string, r io.Reader) error {
		client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
		if err != nil {
			return fmt.Errorf("couldn't establish GCS client: %v", err)
		}
		return gcsx.WriteObject(client, bucket, path.Join(prefix, spec), r)
	}
}
<commit_msg>Add region argument to dataflow.go<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package dataflow

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/user"
	"path"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx"
	// Importing to get the side effect of the remote execution hook. See init().
	_ "github.com/apache/beam/sdks/go/pkg/beam/core/runtime/harness/init"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/hooks"
	"github.com/apache/beam/sdks/go/pkg/beam/core/util/protox"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	"github.com/apache/beam/sdks/go/pkg/beam/options/gcpopts"
	"github.com/apache/beam/sdks/go/pkg/beam/options/jobopts"
	"github.com/apache/beam/sdks/go/pkg/beam/runners/universal/runnerlib"
	"github.com/apache/beam/sdks/go/pkg/beam/util/gcsx"
	"github.com/apache/beam/sdks/go/pkg/beam/x/hooks/perf"
	"github.com/golang/protobuf/proto"
	"golang.org/x/oauth2/google"
	df "google.golang.org/api/dataflow/v1b3"
	"google.golang.org/api/storage/v1"
)

// TODO(herohde) 5/16/2017: the Dataflow flags should match the other SDKs.

var (
	endpoint        = flag.String("dataflow_endpoint", "", "Dataflow endpoint (optional).")
	stagingLocation = flag.String("staging_location", "", "GCS staging location (required).")
	image           = flag.String("worker_harness_container_image", "", "Worker harness container image (required).")
	numWorkers      = flag.Int64("num_workers", 0, "Number of workers (optional).")
	zone            = flag.String("zone", "", "GCP zone (optional)")
	region          = flag.String("region", "", "GCP Region (optional)")
	network         = flag.String("network", "", "GCP network (optional)")
	tempLocation    = flag.String("temp_location", "", "Temp location (optional)")
	machineType     = flag.String("worker_machine_type", "", "GCE machine type (optional)")

	dryRun         = flag.Bool("dry_run", false, "Dry run. Just print the job, but don't submit it.")
	teardownPolicy = flag.String("teardown_policy", "", "Job teardown policy (internal only).")

	// SDK options
	cpuProfiling     = flag.String("cpu_profiling", "", "Job records CPU profiles to this GCS location (optional)")
	sessionRecording = flag.String("session_recording", "", "Job records session transcripts")
)

func init() {
	// Note that we also _ import harness/init to setup the remote execution hook.
	beam.RegisterRunner("dataflow", Execute)

	perf.RegisterProfCaptureHook("gcs_profile_writer", gcsRecorderHook)
}

type dataflowOptions struct {
	PipelineURL string `json:"pipelineUrl"`
}

// Execute runs the given pipeline on Google Cloud Dataflow. It uses the
// default application credentials to submit the job.
func Execute(ctx context.Context, p *beam.Pipeline) error {
	project := *gcpopts.Project
	if project == "" {
		return errors.New("no Google Cloud project specified. Use --project=<project>")
	}
	if *stagingLocation == "" {
		return errors.New("no GCS staging location specified. Use --staging_location=gs://<bucket>/<path>")
	}
	if *image == "" {
		*image = jobopts.GetContainerImage(ctx)
	}
	jobName := jobopts.GetJobName()

	edges, _, err := p.Build()
	if err != nil {
		return err
	}

	if *cpuProfiling != "" {
		perf.EnableProfCaptureHook("gcs_profile_writer", *cpuProfiling)
	}

	if *sessionRecording != "" {
		// TODO(wcn): BEAM-4017
		// It's a bit inconvenient for GCS because the whole object is written in
		// one pass, whereas the session logs are constantly appended. We wouldn't
		// want to hold all the logs in memory to flush at the end of the pipeline
		// as we'd blow out memory on the worker. The implementation of the
		// CaptureHook should create an internal buffer and write chunks out to GCS
		// once they get to an appropriate size (50M or so?)
	}

	hooks.SerializeHooksToOptions()
	options := beam.PipelineOptions.Export()

	// (1) Upload Go binary and model to GCS.

	if *jobopts.WorkerBinary == "" {
		worker, err := runnerlib.BuildTempWorkerBinary(ctx)
		if err != nil {
			return err
		}
		defer os.Remove(worker)

		*jobopts.WorkerBinary = worker
	} else {
		log.Infof(ctx, "Using specified worker binary: '%v'", *jobopts.WorkerBinary)
	}

	binary, err := stageWorker(ctx, project, *stagingLocation, *jobopts.WorkerBinary)
	if err != nil {
		return err
	}

	model, err := graphx.Marshal(edges, &graphx.Options{ContainerImageURL: *image})
	if err != nil {
		return fmt.Errorf("failed to generate model pipeline: %v", err)
	}
	modelURL, err := stageModel(ctx, project, *stagingLocation, protox.MustEncode(model))
	if err != nil {
		return err
	}
	log.Info(ctx, proto.MarshalTextString(model))

	// (2) Translate pipeline to v1b3 speak.

	steps, err := translate(edges)
	if err != nil {
		return err
	}

	job := &df.Job{
		ProjectId: project,
		Name:      jobName,
		Type:      "JOB_TYPE_BATCH",
		Environment: &df.Environment{
			UserAgent: newMsg(userAgent{
				Name:    "Apache Beam SDK for Go",
				Version: "0.3.0",
			}),
			Version: newMsg(version{
				JobType: "FNAPI_BATCH",
				Major:   "6",
			}),
			SdkPipelineOptions: newMsg(pipelineOptions{
				DisplayData: findPipelineFlags(),
				Options: dataflowOptions{
					PipelineURL: modelURL,
				},
				GoOptions: options,
			}),
			WorkerPools: []*df.WorkerPool{{
				Kind: "harness",
				Packages: []*df.Package{{
					Location: binary,
					Name:     "worker",
				}},
				WorkerHarnessContainerImage: *image,
				NumWorkers:                  1,
				MachineType:                 *machineType,
				Network:                     *network,
				Zone:                        *zone,
				Region:                      *region,
			}},
			TempStoragePrefix: *stagingLocation + "/tmp",
			Experiments:       jobopts.GetExperiments(),
		},
		Steps: steps,
	}

	if *numWorkers > 0 {
		job.Environment.WorkerPools[0].NumWorkers = *numWorkers
	}
	if *teardownPolicy != "" {
		job.Environment.WorkerPools[0].TeardownPolicy = *teardownPolicy
	}
	if *tempLocation != "" {
		job.Environment.TempStoragePrefix = *tempLocation
	}
	printJob(ctx, job)

	if *dryRun {
		log.Info(ctx, "Dry-run: not submitting job!")
		return nil
	}

	// (4) Submit job.

	client, err := newClient(ctx, *endpoint)
	if err != nil {
		return err
	}
	upd, err := client.Projects.Jobs.Create(project, job).Do()
	if err != nil {
		return err
	}

	log.Infof(ctx, "Submitted job: %v", upd.Id)
	printJob(ctx, upd)
	if *endpoint == "" {
		log.Infof(ctx, "Console: https://console.cloud.google.com/dataflow/job/%v?project=%v", upd.Id, project)
	}
	log.Infof(ctx, "Logs: https://console.cloud.google.com/logs/viewer?project=%v&resource=dataflow_step%%2Fjob_id%%2F%v", project, upd.Id)

	if *jobopts.Async {
		return nil
	}

	time.Sleep(1 * time.Minute)
	for {
		j, err := client.Projects.Jobs.Get(project, upd.Id).Do()
		if err != nil {
			return fmt.Errorf("failed to get job: %v", err)
		}

		switch j.CurrentState {
		case "JOB_STATE_DONE":
			log.Info(ctx, "Job succeeded!")
			return nil

		case "JOB_STATE_FAILED":
			return fmt.Errorf("job %s failed", upd.Id)

		case "JOB_STATE_RUNNING":
			log.Info(ctx, "Job still running ...")

		default:
			log.Infof(ctx, "Job state: %v ...", j.CurrentState)
		}

		time.Sleep(30 * time.Second)
	}
}

// stageModel uploads the pipeline model to GCS as a unique object.
func stageModel(ctx context.Context, project, location string, model []byte) (string, error) {
	bucket, prefix, err := gcsx.ParseObject(location)
	if err != nil {
		return "", fmt.Errorf("invalid staging location %v: %v", location, err)
	}
	obj := path.Join(prefix, fmt.Sprintf("pipeline-%v", time.Now().UnixNano()))
	if *dryRun {
		full := fmt.Sprintf("gs://%v/%v", bucket, obj)
		log.Infof(ctx, "Dry-run: not uploading model %v", full)
		return full, nil
	}

	client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
	if err != nil {
		return "", err
	}
	return gcsx.Upload(client, project, bucket, obj, bytes.NewReader(model))
}

// stageWorker uploads the worker binary to GCS as a unique object.
func stageWorker(ctx context.Context, project, location, worker string) (string, error) {
	bucket, prefix, err := gcsx.ParseObject(location)
	if err != nil {
		return "", fmt.Errorf("invalid staging location %v: %v", location, err)
	}
	obj := path.Join(prefix, fmt.Sprintf("worker-%v", time.Now().UnixNano()))
	if *dryRun {
		full := fmt.Sprintf("gs://%v/%v", bucket, obj)
		log.Infof(ctx, "Dry-run: not uploading binary %v", full)
		return full, nil
	}

	client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
	if err != nil {
		return "", err
	}
	fd, err := os.Open(worker)
	if err != nil {
		return "", fmt.Errorf("failed to open worker binary %s: %v", worker, err)
	}
	defer fd.Close()
	defer os.Remove(worker)

	return gcsx.Upload(client, project, bucket, obj, fd)
}

func username() string {
	if u, err := user.Current(); err == nil {
		return u.Username
	}
	return "anon"
}

func findPipelineFlags() []*displayData {
	var ret []*displayData

	// TODO(herohde) 2/15/2017: decide if we want all set flags.
	flag.Visit(func(f *flag.Flag) {
		ret = append(ret, newDisplayData(f.Name, "", "flag", f.Value.(flag.Getter).Get()))
	})

	return ret
}

// newClient creates a new dataflow client with default application credentials
// and CloudPlatformScope. The Dataflow endpoint is optionally overridden.
func newClient(ctx context.Context, endpoint string) (*df.Service, error) {
	cl, err := google.DefaultClient(ctx, df.CloudPlatformScope)
	if err != nil {
		return nil, err
	}
	client, err := df.New(cl)
	if err != nil {
		return nil, err
	}
	if endpoint != "" {
		log.Infof(ctx, "Dataflow endpoint override: %s", endpoint)
		client.BasePath = endpoint
	}
	return client, nil
}

func printJob(ctx context.Context, job *df.Job) {
	str, err := json.MarshalIndent(job, "", "  ")
	if err != nil {
		log.Infof(ctx, "Failed to print job %v: %v", job.Id, err)
	}
	log.Info(ctx, string(str))
}

func gcsRecorderHook(opts []string) perf.CaptureHook {
	bucket, prefix, err := gcsx.ParseObject(opts[0])
	if err != nil {
		panic(fmt.Sprintf("Invalid hook configuration for gcsRecorderHook: %s", opts))
	}

	return func(ctx context.Context, spec string, r io.Reader) error {
		client, err := gcsx.NewClient(ctx, storage.DevstorageReadWriteScope)
		if err != nil {
			return fmt.Errorf("couldn't establish GCS client: %v", err)
		}
		return gcsx.WriteObject(client, bucket, path.Join(prefix, spec), r)
	}
}
<|endoftext|>


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

package main

import (
	"database/sql"
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/lib/pq"

	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/log"
)

const (
	maxTransfer = 999
	numAccounts = 999
	concurrency = 5
	aggregate   = true
	usePostgres = false
)

var numTransfers uint64

func moveMoney(db *sql.DB) {
	for {
		from, to := rand.Intn(numAccounts), rand.Intn(numAccounts)
		if from == to {
			continue
		}
		amount := rand.Intn(maxTransfer)
		tx, err := db.Begin()
		if err != nil {
			log.Fatal(err)
		}
		startTime := time.Now()
		// Query is very slow and will be fixed by https://github.com/cockroachdb/cockroach/issues/2140
		query := fmt.Sprintf("SELECT id, balance FROM accounts WHERE id IN (%d, %d)", from, to)
		rows, err := tx.Query(query)
		elapsed := time.Now().Sub(startTime)
		if elapsed > 10*time.Millisecond {
			log.Infof("%s took %v", query, elapsed)
		}
		if err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		var fromBalance, toBalance int
		for rows.Next() {
			var id, balance int
			if err = rows.Scan(&id, &balance); err != nil {
				log.Fatal(err)
			}
			switch id {
			case from:
				fromBalance = balance
			case to:
				toBalance = balance
			default:
				panic(fmt.Sprintf("got unexpected account %d", id))
			}
		}
		startTime = time.Now()
		update := fmt.Sprintf("UPDATE accounts SET balance=%d WHERE id=%d", fromBalance-amount, from)
		if _, err = tx.Exec(update); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		elapsed = time.Now().Sub(startTime)
		if elapsed > 50*time.Millisecond {
			log.Infof("%s took %v", update, elapsed)
		}
		update = fmt.Sprintf("UPDATE accounts SET balance=%d WHERE id=%d", toBalance+amount, to)
		if _, err = tx.Exec(update); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		elapsed = time.Now().Sub(startTime)
		if elapsed > 50*time.Millisecond {
			log.Infof("%s took %v", update, elapsed)
		}

		if err = tx.Commit(); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
		} else {
			atomic.AddUint64(&numTransfers, 1)
		}
	}
}

func verifyBank(db *sql.DB) {
	var sum int64
	if aggregate {
		if err := db.QueryRow("SELECT SUM(balance) FROM accounts").Scan(&sum); err != nil {
			log.Fatal(err)
		}
	} else {
		tx, err := db.Begin()
		if err != nil {
			log.Fatal(err)
		}
		rows, err := tx.Query("SELECT balance FROM accounts")
		if err != nil {
			log.Fatal(err)
		}
		for rows.Next() {
			var balance int64
			if err = rows.Scan(&balance); err != nil {
				log.Fatal(err)
			}
			sum += balance
		}
		if err = tx.Commit(); err != nil {
			log.Fatal(err)
		}
	}

	if sum == 0 {
		log.Info("The bank is in good order.")
	} else {
		log.Fatalf("The bank is not in good order. Total value: %d", sum)
	}
}

func main() {
	var db *sql.DB
	var err error
	var url string
	if usePostgres {
		if db, err = sql.Open("postgres", "dbname=postgres sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		url = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
		if db, err = sql.Open("cockroach", url); err != nil {
			log.Fatal(err)
		}
	}
	if _, err := db.Exec("CREATE DATABASE bank"); err != nil {
		if pqErr, ok := err.(*pq.Error); usePostgres && (!ok || pqErr.Code.Name() != "duplicate_database") {
			log.Fatal(err)
		}
	}
	db.Close()

	// Open db client with database settings.
	if usePostgres {
		if db, err = sql.Open("postgres", "dbname=bank sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		if db, err = sql.Open("cockroach", url+"&database=bank"); err != nil {
			log.Fatal(err)
		}
	}

	// concurrency + 1, for this thread and the "concurrency" number of
	// goroutines that move money
	db.SetMaxOpenConns(concurrency + 1)

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS accounts (id BIGINT PRIMARY KEY, balance BIGINT NOT NULL)"); err != nil {
		log.Fatal(err)
	}

	if _, err = db.Exec("TRUNCATE TABLE accounts"); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < numAccounts; i++ {
		if _, err = db.Exec("INSERT INTO accounts (id, balance) VALUES ($1, $2)", i, 0); err != nil {
			log.Fatal(err)
		}
	}

	verifyBank(db)

	var lastNumTransfers uint64
	lastNow := time.Now()

	for i := 0; i < concurrency; i++ {
		go moveMoney(db)
	}

	for range time.NewTicker(time.Second).C {
		now := time.Now()
		elapsed := time.Since(lastNow)
		numTransfers := atomic.LoadUint64(&numTransfers)
		log.Infof("%d transfers were executed at %.1f/second.", (numTransfers - lastNumTransfers), float64(numTransfers-lastNumTransfers)/elapsed.Seconds())
		verifyBank(db)
		lastNumTransfers = numTransfers
		lastNow = now
	}
}
<commit_msg>Use command line flags in sql bank example.<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Tamir Duberstein

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/lib/pq"

	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/security/securitytest"
	"github.com/cockroachdb/cockroach/server"
	"github.com/cockroachdb/cockroach/util/log"
)

var maxTransfer = flag.Int("max-transfer", 999, "Maximum amount to transfer in one transaction.")
var numAccounts = flag.Int("num-accounts", 999, "Number of accounts.")
var concurrency = flag.Int("concurrency", 5, "Number of concurrent actors moving money.")
var aggregate = flag.Bool("aggregate", true, "Use aggregate function to verify conservation of money.")
var usePostgres = flag.Bool("use-postgres", false, "Use postgres instead of cockroach.")

var numTransfers uint64

func moveMoney(db *sql.DB) {
	for {
		from, to := rand.Intn(*numAccounts), rand.Intn(*numAccounts)
		if from == to {
			continue
		}
		amount := rand.Intn(*maxTransfer)
		tx, err := db.Begin()
		if err != nil {
			log.Fatal(err)
		}
		startTime := time.Now()
		// Query is very slow and will be fixed by https://github.com/cockroachdb/cockroach/issues/2140
		query := fmt.Sprintf("SELECT id, balance FROM accounts WHERE id IN (%d, %d)", from, to)
		rows, err := tx.Query(query)
		elapsed := time.Now().Sub(startTime)
		if elapsed > 10*time.Millisecond {
			log.Infof("%s took %v", query, elapsed)
		}
		if err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		var fromBalance, toBalance int
		for rows.Next() {
			var id, balance int
			if err = rows.Scan(&id, &balance); err != nil {
				log.Fatal(err)
			}
			switch id {
			case from:
				fromBalance = balance
			case to:
				toBalance = balance
			default:
				panic(fmt.Sprintf("got unexpected account %d", id))
			}
		}
		startTime = time.Now()
		update := fmt.Sprintf("UPDATE accounts SET balance=%d WHERE id=%d", fromBalance-amount, from)
		if _, err = tx.Exec(update); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		elapsed = time.Now().Sub(startTime)
		if elapsed > 50*time.Millisecond {
			log.Infof("%s took %v", update, elapsed)
		}
		update = fmt.Sprintf("UPDATE accounts SET balance=%d WHERE id=%d", toBalance+amount, to)
		if _, err = tx.Exec(update); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
			if err = tx.Rollback(); err != nil {
				log.Fatal(err)
			}
			continue
		}
		elapsed = time.Now().Sub(startTime)
		if elapsed > 50*time.Millisecond {
			log.Infof("%s took %v", update, elapsed)
		}

		if err = tx.Commit(); err != nil {
			if log.V(1) {
				log.Warning(err)
			}
		} else {
			atomic.AddUint64(&numTransfers, 1)
		}
	}
}

func verifyBank(db *sql.DB) {
	var sum int64
	if *aggregate {
		if err := db.QueryRow("SELECT SUM(balance) FROM accounts").Scan(&sum); err != nil {
			log.Fatal(err)
		}
	} else {
		tx, err := db.Begin()
		if err != nil {
			log.Fatal(err)
		}
		rows, err := tx.Query("SELECT balance FROM accounts")
		if err != nil {
			log.Fatal(err)
		}
		for rows.Next() {
			var balance int64
			if err = rows.Scan(&balance); err != nil {
				log.Fatal(err)
			}
			sum += balance
		}
		if err = tx.Commit(); err != nil {
			log.Fatal(err)
		}
	}

	if sum == 0 {
		log.Info("The bank is in good order.")
	} else {
		log.Fatalf("The bank is not in good order. Total value: %d", sum)
	}
}

func main() {
	flag.Parse()
	var db *sql.DB
	var err error
	var url string
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=postgres sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		security.SetReadFileFn(securitytest.Asset)
		serv := server.StartTestServer(nil)
		defer serv.Stop()
		url = "https://root@" + serv.ServingAddr() + "?certs=test_certs"
		if db, err = sql.Open("cockroach", url); err != nil {
			log.Fatal(err)
		}
	}
	if _, err := db.Exec("CREATE DATABASE bank"); err != nil {
		if pqErr, ok := err.(*pq.Error); *usePostgres && (!ok || pqErr.Code.Name() != "duplicate_database") {
			log.Fatal(err)
		}
	}
	db.Close()

	// Open db client with database settings.
	if *usePostgres {
		if db, err = sql.Open("postgres", "dbname=bank sslmode=disable"); err != nil {
			log.Fatal(err)
		}
	} else {
		if db, err = sql.Open("cockroach", url+"&database=bank"); err != nil {
			log.Fatal(err)
		}
	}

	// concurrency + 1, for this thread and the "concurrency" number of
	// goroutines that move money
	db.SetMaxOpenConns(*concurrency + 1)

	if _, err = db.Exec("CREATE TABLE IF NOT EXISTS accounts (id BIGINT PRIMARY KEY, balance BIGINT NOT NULL)"); err != nil {
		log.Fatal(err)
	}

	if _, err = db.Exec("TRUNCATE TABLE accounts"); err != nil {
		log.Fatal(err)
	}

	for i := 0; i < *numAccounts; i++ {
		if _, err = db.Exec("INSERT INTO accounts (id, balance) VALUES ($1, $2)", i, 0); err != nil {
			log.Fatal(err)
		}
	}

	verifyBank(db)

	var lastNumTransfers uint64
	lastNow := time.Now()

	for i := 0; i < *concurrency; i++ {
		go moveMoney(db)
	}

	for range time.NewTicker(time.Second).C {
		now := time.Now()
		elapsed := time.Since(lastNow)
		numTransfers := atomic.LoadUint64(&numTransfers)
		log.Infof("%d transfers were executed at %.1f/second.", (numTransfers - lastNumTransfers), float64(numTransfers-lastNumTransfers)/elapsed.Seconds())
		verifyBank(db)
		lastNumTransfers = numTransfers
		lastNow = now
	}
}
<|endoftext|>


<commit_before>package vendor

import (
	"bufio"
	//	"fmt"
	"github.com/c0ze/golang-utils"
	"log"
	"os"
	"regexp"
	"strings"
)

// readLines reads a whole file into memory
// and returns a slice of its lines.
func readLines(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	return lines, scanner.Err()
}

var randomMac = regexp.MustCompile(`^([A-F0-9][2,6,A,E][A-F0-9]*)$`)
var VendorMap map[string]string

func Init() {
	VendorMap = make(map[string]string)
	var lines []string
	var err error
	lines, err = readLines("/etc/fup/oui.txt")
	if err != nil {
		lines, err = readLines("oui.txt")
		if err != nil {
			log.Println("readLines: %s", err)
			return
		}
	}

	for _, line := range lines {
		if len(line) > 1 {
			if string(line[0]) != " " && string(line[0]) != "" && string(line[0]) != "\t" && string(line[2]) != "-" {
				words := strings.Fields(line)
				if len(words) > 3 {
					//					VendorMap[words[0]] = fmt.Sprintf("%v\n", strings.Join(words[3:], " "))
					VendorMap[words[0]] = strings.Join(words[3:], " ")
				}
			}
		}
	}
}

func Lookup(mac string) string {
	vendor := ""

	sanitizedMac := utils.StripColon(mac)

	if len(sanitizedMac) > 5 {
		vendor = VendorMap[sanitizedMac[0:6]]
		if vendor == "" {
			if randomMac.MatchString(utils.StripColon(mac)) {
				vendor = "Random"
			} else {
				vendor = "Unknown"
			}
		}
	} else {
		vendor = "Malformed"
	}

	return vendor
}
<commit_msg>add documentation<commit_after>package vendor

import (
	"bufio"
	//	"fmt"
	"github.com/c0ze/golang-utils"
	"log"
	"os"
	"regexp"
	"strings"
)

// readLines reads a whole file into memory
// and returns a slice of its lines.
func readLines(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	return lines, scanner.Err()
}

var randomMac = regexp.MustCompile(`^([A-F0-9][2,6,A,E][A-F0-9]*)$`)

// VendorMap holds the Mac address -> Vendor name mapping
var VendorMap map[string]string

// Inits Wifi vendors.
// You must call this function once before performing any look ups.
func Init() {
	VendorMap = make(map[string]string)
	var lines []string
	var err error
	lines, err = readLines("/etc/fup/oui.txt")
	if err != nil {
		lines, err = readLines("oui.txt")
		if err != nil {
			log.Println("readLines: %s", err)
			return
		}
	}

	for _, line := range lines {
		if len(line) > 1 {
			if string(line[0]) != " " && string(line[0]) != "" && string(line[0]) != "\t" && string(line[2]) != "-" {
				words := strings.Fields(line)
				if len(words) > 3 {
					//					VendorMap[words[0]] = fmt.Sprintf("%v\n", strings.Join(words[3:], " "))
					VendorMap[words[0]] = strings.Join(words[3:], " ")
				}
			}
		}
	}
}

// Lookup function returns the vendor string matched from oui.txt.
// If the mac address is in random space, it will return Random.
// If the vendor is not in oui, it returns Unknown (you may want to update oui in this case).
// If the mac address is invalid, it returns Malformed.
func Lookup(mac string) string {
	vendor := ""

	sanitizedMac := utils.StripColon(mac)

	if len(sanitizedMac) > 5 {
		vendor = VendorMap[sanitizedMac[0:6]]
		if vendor == "" {
			if randomMac.MatchString(utils.StripColon(mac)) {
				vendor = "Random"
			} else {
				vendor = "Unknown"
			}
		}
	} else {
		vendor = "Malformed"
	}

	return vendor
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2016 Laurent Moussault. All rights reserved.
// Licensed under a simplified BSD license (see LICENSE file).

package gfx

//------------------------------------------------------------------------------

/*
#include <stdlib.h>
#include "glad.h"

static inline void Enable(GLenum c) {
	glEnable(c);
}
*/
import "C"

//------------------------------------------------------------------------------

// Enable an OpenGL capability
func Enable(c capability) {
	C.Enable((C.GLenum)(c))
}

//------------------------------------------------------------------------------

type capability C.GLenum

// OpenGL Capabilities
const (
	Blend                      capability = C.GL_BLEND
	ColorLogicOp               capability = C.GL_COLOR_LOGIC_OP
	CullFace                   capability = C.GL_CULL_FACE
	DebugOutput                capability = C.GL_DEBUG_OUTPUT
	DebugOutputSynchronous     capability = C.GL_DEBUG_OUTPUT_SYNCHRONOUS
	DepthClamp                 capability = C.GL_DEPTH_CLAMP
	DepthTest                  capability = C.GL_DEPTH_TEST
	Dither                     capability = C.GL_DITHER
	FramebufferSRGB            capability = C.GL_FRAMEBUFFER_SRGB
	LineSmooth                 capability = C.GL_LINE_SMOOTH
	Multisample                capability = C.GL_MULTISAMPLE
	PrimitiveRestart           capability = C.GL_PRIMITIVE_RESTART
	PrimitiveRestartFixedIndex capability = C.GL_PRIMITIVE_RESTART_FIXED_INDEX
	RasterixerDiscard          capability = C.GL_RASTERIZER_DISCARD
	SampleAlphaToCoverage      capability = C.GL_SAMPLE_ALPHA_TO_COVERAGE
	SampleAlphaToOne           capability = C.GL_SAMPLE_ALPHA_TO_ONE
	SampleCoverage             capability = C.GL_SAMPLE_COVERAGE
	SampleShading              capability = C.GL_SAMPLE_SHADING
	SampleMask                 capability = C.GL_SAMPLE_MASK
	ScissorTest                capability = C.GL_SCISSOR_TEST
	StencilTest                capability = C.GL_STENCIL_TEST
	TextureCubeMapSeamless     capability = C.GL_TEXTURE_CUBE_MAP_SEAMLESS
)

//------------------------------------------------------------------------------
<commit_msg>Add glDisable<commit_after>// Copyright (c) 2013-2016 Laurent Moussault. All rights reserved.
// Licensed under a simplified BSD license (see LICENSE file).

package gfx

//------------------------------------------------------------------------------

/*
#include <stdlib.h>
#include "glad.h"

static inline void Enable(GLenum c) {
	glEnable(c);
}

static inline void Disable(GLenum c) {
	glDisable(c);
}
*/
import "C"

//------------------------------------------------------------------------------

// Enable an OpenGL capability
func Enable(c capability) {
	C.Enable((C.GLenum)(c))
}

// Disable an OpenGL capability
func Disable(c capability) {
	C.Disable((C.GLenum)(c))
}

//------------------------------------------------------------------------------

type capability C.GLenum

// OpenGL Capabilities
const (
	Blend                      capability = C.GL_BLEND
	ColorLogicOp               capability = C.GL_COLOR_LOGIC_OP
	CullFace                   capability = C.GL_CULL_FACE
	DebugOutput                capability = C.GL_DEBUG_OUTPUT
	DebugOutputSynchronous     capability = C.GL_DEBUG_OUTPUT_SYNCHRONOUS
	DepthClamp                 capability = C.GL_DEPTH_CLAMP
	DepthTest                  capability = C.GL_DEPTH_TEST
	Dither                     capability = C.GL_DITHER
	FramebufferSRGB            capability = C.GL_FRAMEBUFFER_SRGB
	LineSmooth                 capability = C.GL_LINE_SMOOTH
	Multisample                capability = C.GL_MULTISAMPLE
	PrimitiveRestart           capability = C.GL_PRIMITIVE_RESTART
	PrimitiveRestartFixedIndex capability = C.GL_PRIMITIVE_RESTART_FIXED_INDEX
	RasterixerDiscard          capability = C.GL_RASTERIZER_DISCARD
	SampleAlphaToCoverage      capability = C.GL_SAMPLE_ALPHA_TO_COVERAGE
	SampleAlphaToOne           capability = C.GL_SAMPLE_ALPHA_TO_ONE
	SampleCoverage             capability = C.GL_SAMPLE_COVERAGE
	SampleShading              capability = C.GL_SAMPLE_SHADING
	SampleMask                 capability = C.GL_SAMPLE_MASK
	ScissorTest                capability = C.GL_SCISSOR_TEST
	StencilTest                capability = C.GL_STENCIL_TEST
	TextureCubeMapSeamless     capability = C.GL_TEXTURE_CUBE_MAP_SEAMLESS
)

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>package main

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

	"github.com/ansel1/merry"
	"github.com/dghubble/sling"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "update",
		Description: "update heroku-cli",
		Commands: CommandSet{
			{
				Topic:            "update",
				Hidden:           true,
				Description:      "updates the Heroku CLI",
				DisableAnalytics: true,
				Args:             []Arg{{Name: "channel", Optional: true}},
				Run: func(ctx *Context) {
					channel := ctx.Args.(map[string]string)["channel"]
					if channel == "" {
						channel = Channel
					}
					Update(channel)
				},
			},
		},
	})
}

// Autoupdate is a flag to enable/disable CLI autoupdating
var Autoupdate = "no"

// UpdateLockPath is the path to the updating lock file
var UpdateLockPath = filepath.Join(CacheHome, "updating.lock")
var autoupdateFile = filepath.Join(CacheHome, "autoupdate")

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	SubmitAnalytics()
	updateCLI(channel)
	UserPlugins.Update()
	truncate(ErrLogPath, 1000)
	cleanTmpDirs()
}

func updateCLI(channel string) {
	if Autoupdate != "yes" {
		return
	}
	manifest := GetUpdateManifest(channel)
	binExists, _ := FileExists(expectedBinPath())
	if binExists && manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	DownloadCLI(channel, filepath.Join(DataHome, "cli"), manifest)
	loadNewCLI()
}

// DownloadCLI downloads a CLI update to a given path
func DownloadCLI(channel, path string, manifest *Manifest) {
	locked, err := golock.IsLocked(UpdateLockPath)
	LogIfError(err)
	if locked {
		must(merry.Errorf("Update in progress"))
	}
	LogIfError(golock.Lock(UpdateLockPath))
	unlock := func() {
		golock.Unlock(UpdateLockPath)
	}
	defer unlock()
	hideCursor()
	downloadingMessage := fmt.Sprintf("heroku-cli: Updating to %s...", manifest.Version)
	if manifest.Channel != "stable" {
		downloadingMessage = fmt.Sprintf("%s (%s)", downloadingMessage, manifest.Channel)
	}
	Logln(downloadingMessage)
	build := manifest.Builds[runtime.GOOS+"-"+runtime.GOARCH]
	if build == nil {
		must(merry.Errorf("no build for %s", manifest.Channel))
	}
	reader, getSha, err := downloadXZ(build.URL, downloadingMessage)
	must(err)
	tmp := tmpDir(DataHome)
	must(extractTar(reader, tmp))
	sha := getSha()
	if sha != build.Sha256 {
		must(merry.Errorf("SHA mismatch: expected %s to be %s", sha, build.Sha256))
	}
	exists, _ := FileExists(path)
	if exists {
		must(os.Rename(path, filepath.Join(tmpDir(DataHome), "heroku")))
	}
	Errln(path)
	must(os.Rename(filepath.Join(tmp, "heroku"), path))
	Debugf("updated to %s\n", manifest.Version)
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded() bool {
	if exists, _ := FileExists(expectedBinPath()); !exists {
		return true
	}
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		if os.IsNotExist(err) {
			return true
		}
		LogIfError(err)
		return true
	}
	return time.Since(f.ModTime()) > 4*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	must(err)
	_, err = out.WriteString(time.Now().String())
	must(err)
	err = out.Close()
	must(err)
}

// Manifest is the manifest.json for releases
type Manifest struct {
	ReleasedAt string            `json:"released_at"`
	Version    string            `json:"version"`
	Channel    string            `json:"channel"`
	Builds     map[string]*Build `json:"builds"`
}

// Build is a part of a Manifest
type Build struct {
	URL    string `json:"url"`
	Sha256 string `json:"sha256"`
	Bytes  int64  `json:"bytes"`
}

var updateManifestRetrying = false

// GetUpdateManifest loads the manifest.json for a channel
func GetUpdateManifest(channel string) *Manifest {
	var m Manifest
	url := "https://cli-assets.heroku.com/branches/" + channel + "/manifest.json"
	rsp, err := sling.New().Get(url).ReceiveSuccess(&m)
	if err != nil && !updateManifestRetrying {
		updateManifestRetrying = true
		return GetUpdateManifest(channel)
	}
	must(err)
	must(getHTTPError(rsp))
	return &m
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if IsUpdateNeeded() {
		Debugln("triggering background update")
		touchAutoupdateFile()
		exec.Command(BinPath, "update").Start()
	}
}

func cleanTmpDirs() {
	clean := func(base string) {
		dir := filepath.Join(base, "tmp")
		if exists, _ := FileExists(dir); !exists {
			return
		}
		files, err := ioutil.ReadDir(dir)
		LogIfError(err)
		for _, file := range files {
			if time.Since(file.ModTime()) > 24*time.Hour {
				path := filepath.Join(dir, file.Name())
				Debugf("removing old tmp dir %s\n", path)
				LogIfError(os.RemoveAll(path))
			}
		}
	}
	clean(DataHome)
	clean(CacheHome)
}

func expectedBinPath() string {
	bin := filepath.Join(DataHome, "cli", "bin", "heroku")
	if runtime.GOOS == WINDOWS {
		bin = bin + ".exe"
	}
	return bin
}

func loadNewCLI() {
	if Autoupdate == "no" {
		return
	}
	expected := expectedBinPath()
	if BinPath == expected {
		return
	}
	if exists, _ := FileExists(expected); !exists {
		if exists, _ = FileExists(npmBinPath()); !exists {
			// uh oh, npm isn't where it should be.
			// The CLI probably isn't installed right so force an update
			Update(Channel)
		}
		return
	}
	execBin(expected, Args...)
}
<commit_msg>take out debug output<commit_after>package main

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

	"github.com/ansel1/merry"
	"github.com/dghubble/sling"
	"github.com/dickeyxxx/golock"
)

func init() {
	Topics = append(Topics, &Topic{
		Name:        "update",
		Description: "update heroku-cli",
		Commands: CommandSet{
			{
				Topic:            "update",
				Hidden:           true,
				Description:      "updates the Heroku CLI",
				DisableAnalytics: true,
				Args:             []Arg{{Name: "channel", Optional: true}},
				Run: func(ctx *Context) {
					channel := ctx.Args.(map[string]string)["channel"]
					if channel == "" {
						channel = Channel
					}
					Update(channel)
				},
			},
		},
	})
}

// Autoupdate is a flag to enable/disable CLI autoupdating
var Autoupdate = "no"

// UpdateLockPath is the path to the updating lock file
var UpdateLockPath = filepath.Join(CacheHome, "updating.lock")
var autoupdateFile = filepath.Join(CacheHome, "autoupdate")

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	SubmitAnalytics()
	updateCLI(channel)
	UserPlugins.Update()
	truncate(ErrLogPath, 1000)
	cleanTmpDirs()
}

func updateCLI(channel string) {
	if Autoupdate != "yes" {
		return
	}
	manifest := GetUpdateManifest(channel)
	binExists, _ := FileExists(expectedBinPath())
	if binExists && manifest.Version == Version && manifest.Channel == Channel {
		return
	}
	DownloadCLI(channel, filepath.Join(DataHome, "cli"), manifest)
	loadNewCLI()
}

// DownloadCLI downloads a CLI update to a given path
func DownloadCLI(channel, path string, manifest *Manifest) {
	locked, err := golock.IsLocked(UpdateLockPath)
	LogIfError(err)
	if locked {
		must(merry.Errorf("Update in progress"))
	}
	LogIfError(golock.Lock(UpdateLockPath))
	unlock := func() {
		golock.Unlock(UpdateLockPath)
	}
	defer unlock()
	hideCursor()
	downloadingMessage := fmt.Sprintf("heroku-cli: Updating to %s...", manifest.Version)
	if manifest.Channel != "stable" {
		downloadingMessage = fmt.Sprintf("%s (%s)", downloadingMessage, manifest.Channel)
	}
	Logln(downloadingMessage)
	build := manifest.Builds[runtime.GOOS+"-"+runtime.GOARCH]
	if build == nil {
		must(merry.Errorf("no build for %s", manifest.Channel))
	}
	reader, getSha, err := downloadXZ(build.URL, downloadingMessage)
	must(err)
	tmp := tmpDir(DataHome)
	must(extractTar(reader, tmp))
	sha := getSha()
	if sha != build.Sha256 {
		must(merry.Errorf("SHA mismatch: expected %s to be %s", sha, build.Sha256))
	}
	exists, _ := FileExists(path)
	if exists {
		must(os.Rename(path, filepath.Join(tmpDir(DataHome), "heroku")))
	}
	must(os.Rename(filepath.Join(tmp, "heroku"), path))
	Debugf("updated to %s\n", manifest.Version)
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded() bool {
	if exists, _ := FileExists(expectedBinPath()); !exists {
		return true
	}
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		if os.IsNotExist(err) {
			return true
		}
		LogIfError(err)
		return true
	}
	return time.Since(f.ModTime()) > 4*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	must(err)
	_, err = out.WriteString(time.Now().String())
	must(err)
	err = out.Close()
	must(err)
}

// Manifest is the manifest.json for releases
type Manifest struct {
	ReleasedAt string            `json:"released_at"`
	Version    string            `json:"version"`
	Channel    string            `json:"channel"`
	Builds     map[string]*Build `json:"builds"`
}

// Build is a part of a Manifest
type Build struct {
	URL    string `json:"url"`
	Sha256 string `json:"sha256"`
	Bytes  int64  `json:"bytes"`
}

var updateManifestRetrying = false

// GetUpdateManifest loads the manifest.json for a channel
func GetUpdateManifest(channel string) *Manifest {
	var m Manifest
	url := "https://cli-assets.heroku.com/branches/" + channel + "/manifest.json"
	rsp, err := sling.New().Get(url).ReceiveSuccess(&m)
	if err != nil && !updateManifestRetrying {
		updateManifestRetrying = true
		return GetUpdateManifest(channel)
	}
	must(err)
	must(getHTTPError(rsp))
	return &m
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if IsUpdateNeeded() {
		Debugln("triggering background update")
		touchAutoupdateFile()
		exec.Command(BinPath, "update").Start()
	}
}

func cleanTmpDirs() {
	clean := func(base string) {
		dir := filepath.Join(base, "tmp")
		if exists, _ := FileExists(dir); !exists {
			return
		}
		files, err := ioutil.ReadDir(dir)
		LogIfError(err)
		for _, file := range files {
			if time.Since(file.ModTime()) > 24*time.Hour {
				path := filepath.Join(dir, file.Name())
				Debugf("removing old tmp dir %s\n", path)
				LogIfError(os.RemoveAll(path))
			}
		}
	}
	clean(DataHome)
	clean(CacheHome)
}

func expectedBinPath() string {
	bin := filepath.Join(DataHome, "cli", "bin", "heroku")
	if runtime.GOOS == WINDOWS {
		bin = bin + ".exe"
	}
	return bin
}

func loadNewCLI() {
	if Autoupdate == "no" {
		return
	}
	expected := expectedBinPath()
	if BinPath == expected {
		return
	}
	if exists, _ := FileExists(expected); !exists {
		if exists, _ = FileExists(npmBinPath()); !exists {
			// uh oh, npm isn't where it should be.
			// The CLI probably isn't installed right so force an update
			Update(Channel)
		}
		return
	}
	execBin(expected, Args...)
}
<|endoftext|>


<commit_before>package chevalier

import (
	zmq "github.com/pebbe/zmq4"
)

// GetContents list for origin from a Vaultaire
// readerd listening on endpoint, returning it as a DataSourceBurst.
func GetContents(endpoint, origin string) (*DataSourceBurst, error) {
	sock, err := zmq.NewSocket(zmq.REQ)
	if err != nil {
		return nil, err
	}
	err = sock.Connect(endpoint)
	if err != nil {
		return nil, err
	}
	_, err = sock.Send("", 0)
	if err != nil {
		return nil, err
	}
	b, err := sock.RecvBytes(0)
	if err != nil {
		return nil, err
	}
	burst, err := UnmarshalSourceBurst(b)
	if err != nil {
		return nil, err
	}
	return burst, nil
}
<commit_msg>Send contents list request<commit_after>package chevalier

import (
	zmq "github.com/pebbe/zmq4"
)

// GetContents list for origin from a Vaultaire
// readerd listening on endpoint, returning it as a DataSourceBurst.
func GetContents(endpoint, origin string) (*DataSourceBurst, error) {
	sock, err := zmq.NewSocket(zmq.REQ)
	if err != nil {
		return nil, err
	}
	err = sock.Connect(endpoint)
	if err != nil {
		return nil, err
	}
	request := make([]byte, 1)
	request[0] = byte(ContentsListRequest)
	_, err = sock.SendBytes(request, 0)
	if err != nil {
		return nil, err
	}
	b, err := sock.RecvBytes(0)
	if err != nil {
		return nil, err
	}
	burst, err := UnmarshalSourceBurst(b)
	if err != nil {
		return nil, err
	}
	return burst, nil
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/nightlyone/lockfile"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		Errf("updating plugins... ")
		if err := node.UpdatePackages(); err != nil {
			panic(err)
		}
		Errln("done")
		manifest := getUpdateManifest(channel)
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		Errf("updating to %s (%s)... ", manifest.Version, manifest.Channel)
		update(build.URL, build.Sha1)
		Errln("done")
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	node.UpdatePackages()
	manifest := getUpdateManifest(Channel)
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		return false
	}
	return f.ModTime().Add(60 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("http://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	lock, err := lockfile.New(updateLockPath)
	if err != nil {
		Errln("Cannot initialize update lockfile.")
		panic(err)
	}
	err = lock.TryLock()
	if err != nil {
		panic(err)
	}
	defer lock.Unlock()
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath) // on windows you can't rename on top of an existing file
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}
<commit_msg>show warning when cannot autoupdate<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	"github.com/nightlyone/lockfile"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args["channel"]
		if channel == "" {
			channel = "master"
		}
		Errf("updating plugins... ")
		if err := node.UpdatePackages(); err != nil {
			panic(err)
		}
		Errln("done")
		manifest := getUpdateManifest(channel)
		build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
		Errf("updating to %s (%s)... ", manifest.Version, manifest.Channel)
		update(build.URL, build.Sha1)
		Errln("done")
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// UpdateIfNeeded checks for and performs an autoupdate if there is a new version out.
func UpdateIfNeeded() {
	if !updateNeeded() {
		return
	}
	node.UpdatePackages()
	manifest := getUpdateManifest(Channel)
	if manifest.Version == Version {
		// Set timestamp of bin so we don't update again
		os.Chtimes(binPath, time.Now(), time.Now())
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out.\n", Version, manifest.Version)
		return
	}
	// Leave out updating text until heroku-cli is used in place of ruby cli
	// So it doesn't confuse users with 2 different version numbers
	//Errf("Updating to %s... ", manifest.Version)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	update(build.URL, build.Sha1)
	//Errln("done")
	execBin()
	os.Exit(0)
}

func updateNeeded() bool {
	if Version == "dev" {
		return false
	}
	f, err := os.Stat(binPath)
	if err != nil {
		Errln("WARNING: cannot autoupdate. Try running `heroku update` to manually trigger an update.")
		return false
	}
	return f.ModTime().Add(60 * time.Minute).Before(time.Now())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("http://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func update(url, sha1 string) {
	lock, err := lockfile.New(updateLockPath)
	if err != nil {
		Errln("Cannot initialize update lockfile.")
		panic(err)
	}
	err = lock.TryLock()
	if err != nil {
		panic(err)
	}
	defer lock.Unlock()
	tmp, err := downloadBin(url)
	if err != nil {
		panic(err)
	}
	if fileSha1(tmp) != sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath) // on windows you can't rename on top of an existing file
	if err := os.Rename(tmp, binPath); err != nil {
		panic(err)
	}
}

func downloadBin(url string) (string, error) {
	out, err := os.OpenFile(binPath+"~", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return "", err
	}
	defer out.Close()
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return "", err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return "", err
	}
	_, err = io.Copy(out, uncompressed)
	return out.Name(), err
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func execBin() {
	cmd := exec.Command(binPath, os.Args[1:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}
<|endoftext|>


<commit_before>package chevalier

import (
	zmq "github.com/pebbe/zmq4"
)

// GetContents list for origin from a Vaultaire
// readerd listening on endpoint, returning it as a DataSourceBurst.
func GetContents(endpoint, origin string) (*DataSourceBurst, error) {
	sock, err := zmq.NewSocket(zmq.REQ)
	if err != nil {
		return nil, err
	}
	err = sock.Connect(endpoint)
	if err != nil {
		return nil, err
	}
	request := make([]byte, 1)
	request[0] = byte(ContentsListRequest)
	_, err = sock.SendBytes(request, 0)
	if err != nil {
		return nil, err
	}
	b, err := sock.RecvBytes(0)
	if err != nil {
		return nil, err
	}
	burst, err := UnmarshalSourceBurst(b)
	if err != nil {
		return nil, err
	}
	return burst, nil
}
<commit_msg>More unpacking source responses<commit_after>package chevalier

import (
	zmq "github.com/pebbe/zmq4"
)

// GetContents list for origin from a Vaultaire
// readerd listening on endpoint, returning it as a DataSourceBurst.
func GetContents(endpoint, origin string) (*DataSourceBurst, error) {
	sources := make([]*DataSourceResponse, 0)
	sock, err := zmq.NewSocket(zmq.REQ)
	if err != nil {
		return nil, err
	}
	err = sock.Connect(endpoint)
	if err != nil {
		return nil, err
	}
	request := make([]byte, 1)
	request[0] = byte(ContentsListRequest)
	_, err = sock.SendBytes(request, 0)
	if err != nil {
		return nil, err
	}
	for res := new(ContentsResponse); notStopResponse(res); res, err = recvContentsMessage(sock) {
		if err != nil {
			return nil, err
		}
		res_ := unpackSourceResponse(res)
		sources = append(sources, res_)
	}
	return nil, nil
}

func notStopResponse(res *ContentsResponse) bool {
	if res == nil {
		return false // first iteration or error
	}
	if res.opCode == ContentsListEntry {
		return false // data response, continue
	}
	return true
}

func recvContentsMessage(sock *zmq.Socket) (*ContentsResponse, error) {
	bs, err := sock.RecvBytes(0)
	if err != nil {
		return nil, err
	}
	return unpackContentsResponse(bs)
}

func unpackSourceResponse(res *ContentsResponse) *DataSourceResponse {
	return nil
}
<|endoftext|>


<commit_before>package network

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/rancher/os/pkg/log"
	"github.com/rancher/os/pkg/util"
)

const (
	DefaultRoutesCheckTimeout = 20000 // 20 second

	ipv4RouteFile = "/proc/net/route"
	ipv6RouteFile = "/proc/net/ipv6_route"

	ipv4DefaultGWFlags = "0003"
	ipv6DefaultGWFlags = "00450003"
)

func checkIPv4GW() bool {
	if _, err := os.Stat(ipv4RouteFile); os.IsNotExist(err) {
		return false
	}
	file, err := os.Open(ipv4RouteFile)
	if err != nil {
		return false
	}
	defer file.Close()
	scanner := bufio.NewReader(file)
	for {
		line, err := scanner.ReadString('\n')
		if err == io.EOF {
			break
		}
		//ignore the headers in the route info
		if strings.HasPrefix(line, "Iface") {
			continue
		}
		fields := strings.Fields(line)
		// Interested in fields:
		//  3 - flags
		if fields[3] == ipv4DefaultGWFlags {
			return true
		}
	}

	return false
}

func checkIPv6GW() bool {
	if _, err := os.Stat(ipv6RouteFile); os.IsNotExist(err) {
		return false
	}
	file, err := os.Open(ipv6RouteFile)
	if err != nil {
		return false
	}
	defer file.Close()
	scanner := bufio.NewReader(file)
	for {
		line, err := scanner.ReadString('\n')
		if err == io.EOF {
			break
		}
		fields := strings.Fields(line)
		// Interested in fields:
		//  3 - flags
		if fields[3] == ipv6DefaultGWFlags {
			return true
		}
	}

	return false
}

func checkAllDefaultGW() bool {
	return checkIPv4GW() || checkIPv6GW()
}

func AllDefaultGWOK(timeout int) error {
	backoff := util.Backoff{
		MaxMillis: timeout,
	}
	defer backoff.Close()

	var err error
	for ok := range backoff.Start() {
		if !ok {
			err = fmt.Errorf("Timeout waiting for the default gateway ready")
			break
		}
		if checkAllDefaultGW() {
			break
		}
		log.Debug("Waiting for the default gateway ready")
	}

	if err != nil {
		return err
	}

	log.Debug("The default gateway is ready")

	return nil
}
<commit_msg>Retry max timeout with the 20s when waitting for default gateway<commit_after>package network

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/rancher/os/pkg/log"
	"github.com/rancher/os/pkg/util"
)

const (
	DefaultRoutesCheckTimeout = 20000 // 20 second

	ipv4RouteFile = "/proc/net/route"
	ipv6RouteFile = "/proc/net/ipv6_route"

	ipv4DefaultGWFlags = "0003"
	ipv6DefaultGWFlags = "00450003"
)

func checkIPv4GW() bool {
	if _, err := os.Stat(ipv4RouteFile); os.IsNotExist(err) {
		return false
	}
	file, err := os.Open(ipv4RouteFile)
	if err != nil {
		return false
	}
	defer file.Close()
	scanner := bufio.NewReader(file)
	for {
		line, err := scanner.ReadString('\n')
		if err == io.EOF {
			break
		}
		//ignore the headers in the route info
		if strings.HasPrefix(line, "Iface") {
			continue
		}
		fields := strings.Fields(line)
		// Interested in fields:
		//  3 - flags
		if fields[3] == ipv4DefaultGWFlags {
			return true
		}
	}

	return false
}

func checkIPv6GW() bool {
	if _, err := os.Stat(ipv6RouteFile); os.IsNotExist(err) {
		return false
	}
	file, err := os.Open(ipv6RouteFile)
	if err != nil {
		return false
	}
	defer file.Close()
	scanner := bufio.NewReader(file)
	for {
		line, err := scanner.ReadString('\n')
		if err == io.EOF {
			break
		}
		fields := strings.Fields(line)
		// Interested in fields:
		//  3 - flags
		if fields[3] == ipv6DefaultGWFlags {
			return true
		}
	}

	return false
}

func checkAllDefaultGW() bool {
	return checkIPv4GW() || checkIPv6GW()
}

func AllDefaultGWOK(timeout int) error {
	backoff := util.Backoff{
		StartMillis:       100,
		MaxIntervalMillis: 2000,
		MaxMillis:         timeout,
	}
	defer backoff.Close()

	var err error
	for ok := range backoff.Start() {
		if !ok {
			err = fmt.Errorf("Timeout waiting for the default gateway ready")
			break
		}
		if checkAllDefaultGW() {
			break
		}
		log.Debug("Waiting for the default gateway ready")
	}

	if err != nil {
		return err
	}

	log.Debug("The default gateway is ready")

	return nil
}
<|endoftext|>


<commit_before>// Package varint encodes and decodes SQLite4 variable length integers
// which is an encoding of 64-bit unsigned integers into 1-9 bytes
// (varuint).
// See: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
package varint

// Uint64 decodes the received varint encoded byte slice; returning
// the value and the amount of bytes used.
//
// From: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
//
// Decode:
//
//    If A0 is between 0 and 240 inclusive, then the result is the value of A0.
//    If A0 is between 241 and 248 inclusive, then the result is 240+256*(A0-241)+A1.
//    If A0 is 249 then the result is 2288+256*A1+A2.
//    If A0 is 250 then the result is A1..A3 as a 3-byte big-ending integer.
//    If A0 is 251 then the result is A1..A4 as a 4-byte big-ending integer.
//    If A0 is 252 then the result is A1..A5 as a 5-byte big-ending integer.
//    If A0 is 253 then the result is A1..A6 as a 6-byte big-ending integer.
//    If A0 is 254 then the result is A1..A7 as a 7-byte big-ending integer.
//    If A0 is 255 then the result is A1..A8 as a 8-byte big-ending integer.
func Uint64(buf []byte) (uint64, int) {
	// check the first byte
	b := buf[0]
	if b <= 0xF0 {
		return uint64(buf[0]), 1
	}
	if b <= 0xF8 {
		return 240 + 256*(uint64(buf[0])-241) + uint64(buf[1]), 2
	}
	if b == 0xF9 {
		return 2288 + 256*uint64(buf[1]) + uint64(buf[2]), 3
	}
	if b == 0xFA {
		return uint64(buf[1])<<16 | uint64(buf[2])<<8 | uint64(buf[3]), 4
	}
	if b == 0xFB {
		return uint64(buf[1])<<24 | uint64(buf[2])<<16 | uint64(buf[3])<<8 | uint64(buf[4]), 5
	}
	if b == 0xFC {
		return uint64(buf[1])<<32 | uint64(buf[2])<<24 | uint64(buf[3])<<16 | uint64(buf[4])<<8 | uint64(buf[5]), 6
	}
	if b == 0xFD {
		return uint64(buf[1])<<40 | uint64(buf[2])<<32 | uint64(buf[3])<<24 | uint64(buf[4])<<16 | uint64(buf[5])<<8 | uint64(buf[6]), 7
	}
	if b == 0xFE {
		return uint64(buf[1])<<48 | uint64(buf[2])<<40 | uint64(buf[3])<<32 | uint64(buf[4])<<24 | uint64(buf[5])<<16 | uint64(buf[6])<<8 | uint64(buf[7]), 8
	}
	return uint64(buf[1])<<56 | uint64(buf[2])<<48 | uint64(buf[3])<<40 | uint64(buf[4])<<32 | uint64(buf[5])<<24 | uint64(buf[6])<<16 | uint64(buf[7])<<8 | uint64(buf[8]), 9
}

// PutUint64 encodes the received uint64 into varint using the minimum
// necessary bytes.  The number of bytes written is returned.
//
// From: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
//
// Encode:
//
//    Let the input value be V.
//
//    If V<=240 then output a single by A0 equal to V.
//    If V<=2287 then output A0 as (V-240)/256 + 241 and A1 as (V-240)%256.
//    If V<=67823 then output A0 as 249, A1 as (V-2288)/256, and A2 as (V-2288)%256.
//    If V<=16777215 then output A0 as 250 and A1 through A3 as a big-endian 3-byte integer.
//    If V<=4294967295 then output A0 as 251 and A1..A4 as a big-ending 4-byte integer.
//    If V<=1099511627775 then output A0 as 252 and A1..A5 as a big-ending 5-byte integer.
//    If V<=281474976710655 then output A0 as 253 and A1..A6 as a big-ending 6-byte integer.
//    If V<=72057594037927935 then output A0 as 254 and A1..A7 as a big-ending 7-byte integer.
//    Otherwise then output A0 as 255 and A1..A8 as a big-ending 8-byte integer.
func PutUint64(buf []byte, x uint64) int {
	if x < 241 {
		buf[0] = byte(x)
		return 1
	}
	if x < 2288 {
		buf[0] = byte((x-240)/256 + 241)
		buf[1] = byte((x - 240) % 256)
		return 2
	}
	if x < 67824 {
		buf[0] = 0xF9
		buf[1] = byte((x - 2288) / 256)
		buf[2] = byte((x - 2288) % 256)
		return 3
	}
	if x < 1<<24 {
		buf[0] = 0xFA
		buf[1] = byte(x >> 16)
		buf[2] = byte(x >> 8)
		buf[3] = byte(x)
		return 4
	}
	if x < 1<<32 {
		buf[0] = 0xFB
		buf[1] = byte(x >> 24)
		buf[2] = byte(x >> 16)
		buf[3] = byte(x >> 8)
		buf[4] = byte(x)
		return 5
	}
	if x < 1<<40 {
		buf[0] = 0xFC
		buf[1] = byte(x >> 32)
		buf[2] = byte(x >> 24)
		buf[3] = byte(x >> 16)
		buf[4] = byte(x >> 8)
		buf[5] = byte(x)
		return 6
	}
	if x < 1<<48 {
		buf[0] = 0xFD
		buf[1] = byte(x >> 40)
		buf[2] = byte(x >> 32)
		buf[3] = byte(x >> 24)
		buf[4] = byte(x >> 16)
		buf[5] = byte(x >> 8)
		buf[6] = byte(x)
		return 7
	}
	if x < 1<<56 {
		buf[0] = 0xFE
		buf[1] = byte(x >> 48)
		buf[2] = byte(x >> 40)
		buf[3] = byte(x >> 32)
		buf[4] = byte(x >> 24)
		buf[5] = byte(x >> 16)
		buf[6] = byte(x >> 8)
		buf[7] = byte(x)
		return 8
	}
	buf[0] = 0xFF
	buf[1] = byte(x >> 56)
	buf[2] = byte(x >> 48)
	buf[3] = byte(x >> 40)
	buf[4] = byte(x >> 32)
	buf[5] = byte(x >> 24)
	buf[6] = byte(x >> 16)
	buf[7] = byte(x >> 8)
	buf[8] = byte(x)
	return 9
}
<commit_msg>big-ending -> big-endian; it's incorrect in the sqlite4 docs<commit_after>// Package varint encodes and decodes SQLite4 variable length integers
// which is an encoding of 64-bit unsigned integers into 1-9 bytes
// (varuint).
// See: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
package varint

// Uint64 decodes the received varint encoded byte slice; returning
// the value and the amount of bytes used.
//
// From: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
//
// Decode:
//
//    If A0 is between 0 and 240 inclusive, then the result is the value of A0.
//    If A0 is between 241 and 248 inclusive, then the result is 240+256*(A0-241)+A1.
//    If A0 is 249 then the result is 2288+256*A1+A2.
//    If A0 is 250 then the result is A1..A3 as a 3-byte big-endian integer.
//    If A0 is 251 then the result is A1..A4 as a 4-byte big-endian integer.
//    If A0 is 252 then the result is A1..A5 as a 5-byte big-endian integer.
//    If A0 is 253 then the result is A1..A6 as a 6-byte big-endian integer.
//    If A0 is 254 then the result is A1..A7 as a 7-byte big-endian integer.
//    If A0 is 255 then the result is A1..A8 as a 8-byte big-endian integer.
func Uint64(buf []byte) (uint64, int) {
	// check the first byte
	b := buf[0]
	if b <= 0xF0 {
		return uint64(buf[0]), 1
	}
	if b <= 0xF8 {
		return 240 + 256*(uint64(buf[0])-241) + uint64(buf[1]), 2
	}
	if b == 0xF9 {
		return 2288 + 256*uint64(buf[1]) + uint64(buf[2]), 3
	}
	if b == 0xFA {
		return uint64(buf[1])<<16 | uint64(buf[2])<<8 | uint64(buf[3]), 4
	}
	if b == 0xFB {
		return uint64(buf[1])<<24 | uint64(buf[2])<<16 | uint64(buf[3])<<8 | uint64(buf[4]), 5
	}
	if b == 0xFC {
		return uint64(buf[1])<<32 | uint64(buf[2])<<24 | uint64(buf[3])<<16 | uint64(buf[4])<<8 | uint64(buf[5]), 6
	}
	if b == 0xFD {
		return uint64(buf[1])<<40 | uint64(buf[2])<<32 | uint64(buf[3])<<24 | uint64(buf[4])<<16 | uint64(buf[5])<<8 | uint64(buf[6]), 7
	}
	if b == 0xFE {
		return uint64(buf[1])<<48 | uint64(buf[2])<<40 | uint64(buf[3])<<32 | uint64(buf[4])<<24 | uint64(buf[5])<<16 | uint64(buf[6])<<8 | uint64(buf[7]), 8
	}
	return uint64(buf[1])<<56 | uint64(buf[2])<<48 | uint64(buf[3])<<40 | uint64(buf[4])<<32 | uint64(buf[5])<<24 | uint64(buf[6])<<16 | uint64(buf[7])<<8 | uint64(buf[8]), 9
}

// PutUint64 encodes the received uint64 into varint using the minimum
// necessary bytes.  The number of bytes written is returned.
//
// From: http://www.sqlite.org/src4/doc/trunk/www/varint.wiki
//
// Encode:
//
//    Let the input value be V.
//
//    If V<=240 then output a single by A0 equal to V.
//    If V<=2287 then output A0 as (V-240)/256 + 241 and A1 as (V-240)%256.
//    If V<=67823 then output A0 as 249, A1 as (V-2288)/256, and A2 as (V-2288)%256.
//    If V<=16777215 then output A0 as 250 and A1 through A3 as a big-endian 3-byte integer.
//    If V<=4294967295 then output A0 as 251 and A1..A4 as a big-endian 4-byte integer.
//    If V<=1099511627775 then output A0 as 252 and A1..A5 as a big-endian 5-byte integer.
//    If V<=281474976710655 then output A0 as 253 and A1..A6 as a big-endian 6-byte integer.
//    If V<=72057594037927935 then output A0 as 254 and A1..A7 as a big-endian 7-byte integer.
//    Otherwise then output A0 as 255 and A1..A8 as a big-endian 8-byte integer.
func PutUint64(buf []byte, x uint64) int {
	if x < 241 {
		buf[0] = byte(x)
		return 1
	}
	if x < 2288 {
		buf[0] = byte((x-240)/256 + 241)
		buf[1] = byte((x - 240) % 256)
		return 2
	}
	if x < 67824 {
		buf[0] = 0xF9
		buf[1] = byte((x - 2288) / 256)
		buf[2] = byte((x - 2288) % 256)
		return 3
	}
	if x < 1<<24 {
		buf[0] = 0xFA
		buf[1] = byte(x >> 16)
		buf[2] = byte(x >> 8)
		buf[3] = byte(x)
		return 4
	}
	if x < 1<<32 {
		buf[0] = 0xFB
		buf[1] = byte(x >> 24)
		buf[2] = byte(x >> 16)
		buf[3] = byte(x >> 8)
		buf[4] = byte(x)
		return 5
	}
	if x < 1<<40 {
		buf[0] = 0xFC
		buf[1] = byte(x >> 32)
		buf[2] = byte(x >> 24)
		buf[3] = byte(x >> 16)
		buf[4] = byte(x >> 8)
		buf[5] = byte(x)
		return 6
	}
	if x < 1<<48 {
		buf[0] = 0xFD
		buf[1] = byte(x >> 40)
		buf[2] = byte(x >> 32)
		buf[3] = byte(x >> 24)
		buf[4] = byte(x >> 16)
		buf[5] = byte(x >> 8)
		buf[6] = byte(x)
		return 7
	}
	if x < 1<<56 {
		buf[0] = 0xFE
		buf[1] = byte(x >> 48)
		buf[2] = byte(x >> 40)
		buf[3] = byte(x >> 32)
		buf[4] = byte(x >> 24)
		buf[5] = byte(x >> 16)
		buf[6] = byte(x >> 8)
		buf[7] = byte(x)
		return 8
	}
	buf[0] = 0xFF
	buf[1] = byte(x >> 56)
	buf[2] = byte(x >> 48)
	buf[3] = byte(x >> 40)
	buf[4] = byte(x >> 32)
	buf[5] = byte(x >> 24)
	buf[6] = byte(x >> 16)
	buf[7] = byte(x >> 8)
	buf[8] = byte(x)
	return 9
}
<|endoftext|>


<commit_before>package volmex

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/docker/go-plugins-helpers/volume"
)

// Volume holds volmex related volume configurations
// - Volume from go-plugins-helpers/volume
// - Options from 'docker volume create' (volume.Request from Create)
type Volume struct {
	volume.Volume
	Options map[string]string
}

// Driver represents the driver's configuration
// - State holds the drivers's volume configurations
// - mountBase is the base folder for volume storage directories
type Driver struct {
	state     State
	mountBase string
}

// NewDriver creates a new Driver instance
func NewDriver(state State, mountBase string) *Driver {
	return &Driver{
		state:     state,
		mountBase: mountBase,
	}
}

// the following methods implement the docker volume plugin protocol (v1)
// https://docs.docker.com/engine/extend/plugins_volume/#volumedriver

// Create is called by docker upon 'docker volume create' and creates a new volmex.Volume
func (d *Driver) Create(req volume.Request) volume.Response {
	fmt.Printf("received 'Create' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// check if a command was given
	if req.Options["cmd"] == "" {
		return volume.Response{
			Err: "no volmex mount command. specify with -o cmd=\"command\"",
		}
	}

	// check if the volume storage folder is missing and create it if necessary
	if _, err := os.Stat(d.mountBase + "/" + req.Name); os.IsNotExist(err) {
		os.Mkdir(d.mountBase+"/"+req.Name, 0777)
	}

	// store the volume configuration in a volmex.Volume and add it to the driver's state
	v := &Volume{
		Volume: volume.Volume{
			Name:       req.Name,
			Mountpoint: d.mountBase + "/" + req.Name,
		},
		Options: req.Options,
	}
	d.state.Put(v.Name, v)

	// save driver state
	err = d.state.Save()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	fmt.Printf("\tmountpoint: %v\n", v.Volume.Mountpoint)
	fmt.Printf("\tcommand: %v\n", v.Options["cmd"])

	return volume.Response{}
}

// Get returns the volume configuration for a given volume name
func (d *Driver) Get(req volume.Request) volume.Response {
	fmt.Printf("received 'Get' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	return volume.Response{
		Volume: &v.Volume,
	}
}

// List returns all known volume configurations
func (d *Driver) List(req volume.Request) volume.Response {
	fmt.Printf("received 'List' request\n")

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// retrieve volume configurations from state
	var vs []*volume.Volume
	for _, v := range d.state.List() {
		vs = append(vs, &v.Volume)
	}

	fmt.Printf("\t%v known volumes\n", len(vs))
	for _, v := range vs {
		fmt.Printf("\t\t%v\n", v.Name)
	}

	return volume.Response{
		Volumes: vs,
	}
}

// Remove is issued by docker when a user requests to delete a volume
// however, we don't actually remove any files or the storage folder
func (d *Driver) Remove(req volume.Request) volume.Response {
	fmt.Printf("received 'Remove' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// remove volume from state
	d.state.Remove(req.Name)

	// save driver state
	err = d.state.Save()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	return volume.Response{}
}

// Path returns the mountpoint for a given volume name
func (d *Driver) Path(req volume.Request) volume.Response {
	fmt.Printf("received 'Path' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: "volume unknown",
		}
	}

	return volume.Response{
		Mountpoint: v.Mountpoint,
	}
}

// Mount is called by docker before a container using a volmex volume is started
// since we don't actually do anything related to storage, we only execute the specified volmex command and return the mountpoint
func (d *Driver) Mount(req volume.MountRequest) volume.Response {
	fmt.Printf("received 'Mount' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: "volume unknown",
		}
	}

	fmt.Println("\texecuting volmex command " + v.Options["cmd"])

	// prepare command
	cmdString := strings.TrimSpace(v.Options["cmd"])
	cmdParts := strings.Split(cmdString, " ")
	cmd := exec.Command(cmdParts[0], cmdParts[1:]...)

	// set volmex environment for the new command
	cmd.Env = []string{
		"VOLMEX_NAME=" + v.Name,
		"VOLMEX_MOUNTPOINT=" + v.Mountpoint,
		"VOLMEX_CMD=" + cmdString,
	}

	// execute command and print output
	out, err := cmd.CombinedOutput()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}
	fmt.Println(string(out))

	return volume.Response{
		Mountpoint: v.Mountpoint,
	}
}

// Unmount is called by docker after a container using a volmex volume was stopped
// Currently we don't want to execute a command after a container was stopped, so we don't need to do anything here
func (d *Driver) Unmount(req volume.UnmountRequest) volume.Response {
	fmt.Printf("received 'Unmount' request for volume: %v\n", req.Name)
	return volume.Response{}
}

// Capabilities is called by docker to get certain driver options (atm only the 'scope')
func (d *Driver) Capabilities(req volume.Request) volume.Response {
	fmt.Printf("received 'Capabilities' request for volume: %v\n", req.Name)
	localCap := volume.Capability{Scope: "local"}
	return volume.Response{Capabilities: localCap}
}
<commit_msg>minor changes<commit_after>package volmex

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"github.com/docker/go-plugins-helpers/volume"
)

// Volume holds volmex related volume configurations
// - Volume from go-plugins-helpers/volume
// - Options from 'docker volume create' (volume.Request from Create)
type Volume struct {
	volume.Volume
	Options map[string]string
}

// Driver represents the driver's configuration
// - State holds the drivers's volume configurations
// - mountBase is the base folder for volume storage directories
type Driver struct {
	state     State
	mountBase string
}

// NewDriver creates a new Driver instance
func NewDriver(state State, mountBase string) *Driver {
	return &Driver{
		state:     state,
		mountBase: mountBase,
	}
}

// the following methods implement the docker volume plugin protocol (v1)
// https://docs.docker.com/engine/extend/plugins_volume/#volumedriver

// Create is called by docker upon 'docker volume create' and creates a new volmex.Volume
func (d *Driver) Create(req volume.Request) volume.Response {
	fmt.Printf("received 'Create' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// check if a command was given
	if req.Options["cmd"] == "" {
		return volume.Response{
			Err: "no volmex mount command. specify with -o cmd=\"command\"",
		}
	}

	// check if the volume storage folder is missing and create it if necessary
	if _, err := os.Stat(d.mountBase + "/" + req.Name); os.IsNotExist(err) {
		os.Mkdir(d.mountBase+"/"+req.Name, 0777)
	}

	// store the volume configuration in a volmex.Volume and add it to the driver's state
	v := &Volume{
		Volume: volume.Volume{
			Name:       req.Name,
			Mountpoint: d.mountBase + "/" + req.Name,
		},
		Options: req.Options,
	}
	d.state.Put(v.Name, v)

	// save driver state
	err = d.state.Save()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	fmt.Printf("\tmountpoint: %v\n", v.Volume.Mountpoint)
	fmt.Printf("\tcommand: %v\n", v.Options["cmd"])

	return volume.Response{}
}

// Get returns the volume configuration for a given volume name
func (d *Driver) Get(req volume.Request) volume.Response {
	fmt.Printf("received 'Get' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	return volume.Response{
		Volume: &v.Volume,
	}
}

// List returns all known volume configurations
func (d *Driver) List(req volume.Request) volume.Response {
	fmt.Printf("received 'List' request\n")

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// retrieve volume configurations from state
	var vs []*volume.Volume
	for _, v := range d.state.List() {
		vs = append(vs, &v.Volume)
	}

	fmt.Printf("\t%v known volumes\n", len(vs))
	for _, v := range vs {
		fmt.Printf("\t\t%v\n", v.Name)
	}

	return volume.Response{
		Volumes: vs,
	}
}

// Remove is issued by docker when a user requests to delete a volume
// however, we don't actually remove any files or the storage folder
func (d *Driver) Remove(req volume.Request) volume.Response {
	fmt.Printf("received 'Remove' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// remove volume from state
	d.state.Remove(req.Name)

	// save driver state
	err = d.state.Save()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	return volume.Response{}
}

// Path returns the mountpoint for a given volume name
func (d *Driver) Path(req volume.Request) volume.Response {
	fmt.Printf("received 'Path' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: "volume unknown",
		}
	}

	return volume.Response{
		Mountpoint: v.Mountpoint,
	}
}

// Mount is called by docker before a container using a volmex volume is started
// since we don't actually do anything related to storage, we only execute the specified volmex command and return the mountpoint
func (d *Driver) Mount(req volume.MountRequest) volume.Response {
	fmt.Printf("received 'Mount' request for volume: %v\n", req.Name)

	// load driver state
	err := d.state.Load()
	if err != nil {
		return volume.Response{
			Err: err.Error(),
		}
	}

	// try to retrieve volume configuration from driver state
	v, err := d.state.Get(req.Name)
	if err != nil {
		return volume.Response{
			Err: "volume unknown",
		}
	}

	fmt.Println("\texecuting volmex command: " + v.Options["cmd"])

	// prepare command
	cmdString := strings.TrimSpace(v.Options["cmd"])
	cmdParts := strings.Split(cmdString, " ")
	cmd := exec.Command(cmdParts[0], cmdParts[1:]...)

	// set volmex environment for the new command
	cmd.Env = []string{
		"VOLMEX_NAME=" + v.Name,
		"VOLMEX_MOUNTPOINT=" + v.Mountpoint,
		"VOLMEX_CMD=" + cmdString,
	}

	// execute command and print output
	out, err := cmd.CombinedOutput()
	if err != nil {
		return volume.Response{
			Err: "error from volmex command: " + err.Error(),
		}
	}
	fmt.Println(string(out))

	return volume.Response{
		Mountpoint: v.Mountpoint,
	}
}

// Unmount is called by docker after a container using a volmex volume was stopped
// Currently we don't want to execute a command after a container was stopped, so we don't need to do anything here
func (d *Driver) Unmount(req volume.UnmountRequest) volume.Response {
	fmt.Printf("received 'Unmount' request for volume: %v\n", req.Name)
	return volume.Response{}
}

// Capabilities is called by docker to get certain driver options (atm only the 'scope')
func (d *Driver) Capabilities(req volume.Request) volume.Response {
	fmt.Printf("received 'Capabilities' request for volume: %v\n", req.Name)
	localCap := volume.Capability{Scope: "local"}
	return volume.Response{Capabilities: localCap}
}
<|endoftext|>


<commit_before>package controllers

import (
	"github.com/freeznet/tomato/cloud"
	"github.com/freeznet/tomato/errs"
	"github.com/freeznet/tomato/rest"
	"github.com/freeznet/tomato/types"
	"github.com/freeznet/tomato/utils"
	"time"
)

// CloudCodeController ...
type CloudCodeController struct {
	ClassesController
}

func formatJobSchedule(jobSchedule types.M) types.M  {
	if _, ok := jobSchedule["startAfter"]; !ok {
		jobSchedule["startAfter"] = utils.TimetoString(time.Now())
	}
	return jobSchedule
}

func (c *CloudCodeController) validateJobSchedule()  {
	jobs := cloud.GetJobs()
	if v, ok := c.JSONBody["jobName"]; ok {
		if _, ok := jobs[v.(string)]; !ok {
			c.HandleError(errs.E(errs.InternalServerError, "Cannot Schedule a job that is not deployed"), 0)
		}
	}
}

//Prepare ...
func (c *CloudCodeController) Prepare()  {
	c.ClassesController.Prepare()
	if c.Ctx.ResponseWriter.Started == false {
		c.EnforceMasterKeyAccess()
	}
}

// HandleFind ...
// @router /jobs [get]
func (c *CloudCodeController) GetJobs()  {
	response, err := rest.Find(c.Auth, "_JobSchedule", types.M{}, types.M{}, c.Info.ClientSDK)

	if err != nil {
		c.HandleError(err, 0)
		return
	}

	results := utils.A(response["results"])

	c.Data["json"] = results[0]
	c.ServeJSON()
}

// HandleGet ...
// @router /jobs/data [get]
func (c *CloudCodeController) GetJobsData() {
	jobs := cloud.GetJobs()
	jobNames := []string{}
	for n := range jobs {
		jobNames = append(jobNames, n)
	}

	response, err := rest.Find(c.Auth, "_JobSchedule", types.M{}, types.M{}, nil)

	if err != nil {
		c.HandleError(err, 0)
		return
	}

	results := utils.A(response["results"])
	resultsJobNames := []string{}
	for _, v := range results {
		result := utils.M(v)
		resultsJobNames = append(resultsJobNames, result["jobName"].(string))
	}

	c.Data["json"] = types.M{"in_use": resultsJobNames, "jobs": jobNames}
	c.ServeJSON()
}

// HandleCreate ...
// @router /jobs [create]
func (c *CloudCodeController) HandleCreate()  {
	c.validateJobSchedule()

	c.ClassName = "_JobSchedule"
	c.JSONBody = formatJobSchedule(c.JSONBody)
	c.ClassesController.HandleCreate()
}

// HandleUpdate ...
// @router /jobs/:objectId [put]
func (c *CloudCodeController) HandleUpdate()  {
	c.validateJobSchedule()

	c.ClassName = "_JobSchedule"
	c.ObjectID = c.Ctx.Input.Param(":objectId")
	c.JSONBody = formatJobSchedule(c.JSONBody)
	c.ClassesController.HandleUpdate()
}

// HandleDelete ...
// @router / [delete]
func (c *CloudCodeController) HandleDelete()  {
	c.ClassName = "_JobSchedule"
	c.ClassesController.HandleDelete()
}
<commit_msg>BKD-8-restores-the-job-scheduling-endpoints：函数名称与parse同步<commit_after>package controllers

import (
	"github.com/freeznet/tomato/cloud"
	"github.com/freeznet/tomato/errs"
	"github.com/freeznet/tomato/rest"
	"github.com/freeznet/tomato/types"
	"github.com/freeznet/tomato/utils"
	"time"
)

// CloudCodeController ...
type CloudCodeController struct {
	ClassesController
}

func formatJobSchedule(jobSchedule types.M) types.M  {
	if _, ok := jobSchedule["startAfter"]; !ok {
		jobSchedule["startAfter"] = utils.TimetoString(time.Now())
	}
	return jobSchedule
}

func (c *CloudCodeController) validateJobSchedule()  {
	jobs := cloud.GetJobs()
	if v, ok := c.JSONBody["jobName"]; ok {
		if _, ok := jobs[v.(string)]; !ok {
			c.HandleError(errs.E(errs.InternalServerError, "Cannot Schedule a job that is not deployed"), 0)
		}
	}
}

//Prepare ...
func (c *CloudCodeController) Prepare()  {
	c.ClassesController.Prepare()
	if c.Ctx.ResponseWriter.Started == false {
		c.EnforceMasterKeyAccess()
	}
}

// HandleFind ...
// @router /jobs [get]
func (c *CloudCodeController) GetJobs()  {
	response, err := rest.Find(c.Auth, "_JobSchedule", types.M{}, types.M{}, c.Info.ClientSDK)

	if err != nil {
		c.HandleError(err, 0)
		return
	}

	results := utils.A(response["results"])

	c.Data["json"] = results[0]
	c.ServeJSON()
}

// HandleGet ...
// @router /jobs/data [get]
func (c *CloudCodeController) GetJobsData() {
	jobs := cloud.GetJobs()
	jobNames := []string{}
	for n := range jobs {
		jobNames = append(jobNames, n)
	}

	response, err := rest.Find(c.Auth, "_JobSchedule", types.M{}, types.M{}, nil)

	if err != nil {
		c.HandleError(err, 0)
		return
	}

	results := utils.A(response["results"])
	resultsJobNames := []string{}
	for _, v := range results {
		result := utils.M(v)
		resultsJobNames = append(resultsJobNames, result["jobName"].(string))
	}

	c.Data["json"] = types.M{"in_use": resultsJobNames, "jobs": jobNames}
	c.ServeJSON()
}

// CreateJob ...
// @router /jobs [create]
func (c *CloudCodeController) CreateJob()  {
	c.validateJobSchedule()

	c.ClassName = "_JobSchedule"
	c.JSONBody = formatJobSchedule(c.JSONBody)
	c.ClassesController.HandleCreate()
}

// EditJob ...
// @router /jobs/:objectId [put]
func (c *CloudCodeController) EditJob()  {
	c.validateJobSchedule()

	c.ClassName = "_JobSchedule"
	c.ObjectID = c.Ctx.Input.Param(":objectId")
	c.JSONBody = formatJobSchedule(c.JSONBody)
	c.ClassesController.HandleUpdate()
}

// DeleteJob ...
// @router / [delete]
func (c *CloudCodeController) DeleteJob()  {
	c.ClassName = "_JobSchedule"
	c.ClassesController.HandleDelete()
}
<|endoftext|>