<commit_before>// Copyright 2015 The appc 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.

package lib

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"unsafe"

	"github.com/containers/build/registry"
	"github.com/containers/build/util"

	docker2aci "github.com/appc/docker2aci/lib"
	"github.com/appc/docker2aci/lib/common"
	"github.com/appc/spec/aci"
	"github.com/appc/spec/discovery"
	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/pkg/fileutil"
	"github.com/coreos/rkt/pkg/user"
)

var (
	placeholdername = "acbuild-unnamed"
)

// Begin will start a new build, storing the untarred ACI the build operates on
// at a.CurrentACIPath. If start is the empty string, the build will begin with
// an empty ACI, otherwise the ACI stored at start will be used at the starting
// point.
func (a *ACBuild) Begin(start string, insecure bool) (err error) {
	_, err = os.Stat(a.ContextPath)
	switch {
	case os.IsNotExist(err):
		break
	case err != nil:
		return err
	default:
		return fmt.Errorf("build already in progress in this working dir")
	}

	err = os.MkdirAll(a.ContextPath, 0755)
	if err != nil {
		return err
	}

	if err = a.lock(); err != nil {
		return err
	}
	defer func() {
		if err1 := a.unlock(); err == nil {
			err = err1
		}
	}()

	defer func() {
		// If there was an error while beginning, we don't want to produce an
		// unexpected build context
		if err != nil {
			os.RemoveAll(a.ContextPath)
		}
	}()

	if start != "" {
		err = os.MkdirAll(a.CurrentACIPath, 0755)
		if err != nil {
			return err
		}
		if start[0] == '.' || start[0] == '/' {
			finfo, err := os.Stat(start)
			switch {
			case os.IsNotExist(err):
				return fmt.Errorf("no such file or directory: %s", start)
			case err != nil:
				return err
			case finfo.IsDir():
				a.beginFromLocalDirectory(start)
			default:
				return a.beginFromLocalImage(start)
			}
		} else {
			dockerPrefix := "docker://"
			if strings.HasPrefix(start, dockerPrefix) {
				start = strings.TrimPrefix(start, dockerPrefix)
				return a.beginFromRemoteDockerImage(start, insecure)
			}
			return a.beginFromRemoteImage(start, insecure)
		}
	}
	return a.beginWithEmptyACI()
}

func (a *ACBuild) beginFromLocalImage(start string) error {
	finfo, err := os.Stat(start)
	if err != nil {
		return err
	}
	if finfo.IsDir() {
		return fmt.Errorf("provided starting ACI is a directory: %s", start)
	}
	err = util.ExtractImage(start, a.CurrentACIPath, nil)
	if err != nil {
		return err
	}

	for _, file := range []struct {
		FileName string
		FilePath string
	}{
		{"manifest file", path.Join(a.CurrentACIPath, aci.ManifestFile)},
		{"rootfs directory", path.Join(a.CurrentACIPath, aci.RootfsDir)},
	} {
		_, err = os.Stat(file.FilePath)
		switch {
		case os.IsNotExist(err):
			fmt.Fprintf(os.Stderr, "%s is missing, assuming build is beginning with a tar of a rootfs\n", file.FileName)
			return a.startedFromTar()
		case err != nil:
			return err
		}
	}
	return nil
}

func (a *ACBuild) startedFromTar() error {
	tmpPath := path.Join(a.ContextPath, "rootfs")
	err := os.Rename(a.CurrentACIPath, tmpPath)
	if err != nil {
		return err
	}
	err = a.beginWithEmptyACI()
	if err != nil {
		return err
	}
	err = os.Remove(path.Join(a.CurrentACIPath, aci.RootfsDir))
	if err != nil {
		return err
	}
	return os.Rename(tmpPath, path.Join(a.CurrentACIPath, aci.RootfsDir))
}

func (a *ACBuild) beginFromLocalDirectory(start string) error {
	err := os.MkdirAll(a.CurrentACIPath, 0755)
	if err != nil {
		return err
	}

	err = fileutil.CopyTree(start, path.Join(a.CurrentACIPath, aci.RootfsDir), user.NewBlankUidRange())
	if err != nil {
		return err
	}

	return a.writeEmptyManifest()
}

func (a *ACBuild) beginWithEmptyACI() error {
	err := os.MkdirAll(path.Join(a.CurrentACIPath, aci.RootfsDir), 0755)
	if err != nil {
		return err
	}

	return a.writeEmptyManifest()
}

func (a *ACBuild) writeEmptyManifest() error {
	acid, err := types.NewACIdentifier("acbuild-unnamed")
	if err != nil {
		return err
	}

	archlabel, err := types.NewACIdentifier("arch")
	if err != nil {
		return err
	}

	oslabel, err := types.NewACIdentifier("os")
	if err != nil {
		return err
	}

	archvalue := runtime.GOARCH
	if runtime.GOOS == "linux" && (archvalue == "arm" || archvalue == "arm64") {
		var x uint32 = 0x01020304
		test := *(*byte)(unsafe.Pointer(&x))
		switch {
		case test == 0x01 && archvalue == "arm":
			archvalue = "armv7b"
		case test == 0x04 && archvalue == "arm":
			archvalue = "armv7l"
		case test == 0x01 && archvalue == "arm64":
			archvalue = "aarch64_be"
		case test == 0x04 && archvalue == "arm64":
			archvalue = "aarch64"
		}
	}

	manifest := &schema.ImageManifest{
		ACKind:    schema.ImageManifestKind,
		ACVersion: schema.AppContainerVersion,
		Name:      *acid,
		Labels: types.Labels{
			types.Label{
				*archlabel,
				archvalue,
			},
			types.Label{
				*oslabel,
				runtime.GOOS,
			},
		},
	}

	manblob, err := manifest.MarshalJSON()
	if err != nil {
		return err
	}

	manfile, err := os.Create(path.Join(a.CurrentACIPath, aci.ManifestFile))
	if err != nil {
		return err
	}

	_, err = manfile.Write(manblob)
	if err != nil {
		return err
	}

	err = manfile.Close()
	if err != nil {
		return err
	}

	return nil
}

func (a *ACBuild) beginFromRemoteImage(start string, insecure bool) error {
	app, err := discovery.NewAppFromString(start)
	if err != nil {
		return err
	}
	labels, err := types.LabelsFromMap(app.Labels)
	if err != nil {
		return err
	}

	tmpDepStoreTarPath, err := ioutil.TempDir("", "acbuild-begin-tar")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDepStoreTarPath)

	tmpDepStoreExpandedPath, err := ioutil.TempDir("", "acbuild-begin-expanded")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDepStoreExpandedPath)

	reg := registry.Registry{
		DepStoreTarPath:      tmpDepStoreTarPath,
		DepStoreExpandedPath: tmpDepStoreExpandedPath,
		Insecure:             insecure,
		Debug:                a.Debug,
	}

	err = reg.Fetch(app.Name, labels, 0, false)
	if err != nil {
		if urlerr, ok := err.(*url.Error); ok {
			if operr, ok := urlerr.Err.(*net.OpError); ok {
				if dnserr, ok := operr.Err.(*net.DNSError); ok {
					if dnserr.Err == "no such host" {
						return fmt.Errorf("unknown host when fetching image, check your connection and local file paths must start with '/' or '.'")
					}
				}
			}
		}
		return err
	}

	files, err := ioutil.ReadDir(tmpDepStoreTarPath)
	if err != nil {
		return err
	}

	if len(files) != 1 {
		var filelist string
		for _, file := range files {
			if filelist == "" {
				filelist = file.Name()
			} else {
				filelist = filelist + ", " + file.Name()
			}
		}
		panic("unexpected number of files in store after download: " + filelist)
	}

	return util.ExtractImage(path.Join(tmpDepStoreTarPath, files[0].Name()), a.CurrentACIPath, nil)
}

func (a *ACBuild) beginFromRemoteDockerImage(start string, insecure bool) (err error) {
	outputDir, err := ioutil.TempDir("", "acbuild")
	if err != nil {
		return err
	}
	defer os.RemoveAll(outputDir)

	tempDir, err := ioutil.TempDir("", "acbuild")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	insecureConf := common.InsecureConfig{
		SkipVerify: insecure,
		AllowHTTP:  insecure,
	}

	config := docker2aci.RemoteConfig{
		CommonConfig: docker2aci.CommonConfig{
			Squash:      true,
			OutputDir:   outputDir,
			TmpDir:      tempDir,
			Compression: common.GzipCompression,
		},
		Username: "",
		Password: "",
		Insecure: insecureConf,
	}
	renderedACIs, err := docker2aci.ConvertRemoteRepo(start, config)
	if err != nil {
		return err
	}
	if len(renderedACIs) > 1 {
		return fmt.Errorf("internal error: docker2aci didn't squash the image")
	}
	if len(renderedACIs) == 0 {
		return fmt.Errorf("internal error: docker2aci didn't produce any images")
	}
	absRenderedACI, err := filepath.Abs(renderedACIs[0])
	if err != nil {
		return err
	}

	return util.ExtractImage(absRenderedACI, a.CurrentACIPath, nil)
}
<commit_msg>Fix Begin to return error generated when starting from local dir<commit_after>// Copyright 2015 The appc 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.

package lib

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"unsafe"

	"github.com/containers/build/registry"
	"github.com/containers/build/util"

	docker2aci "github.com/appc/docker2aci/lib"
	"github.com/appc/docker2aci/lib/common"
	"github.com/appc/spec/aci"
	"github.com/appc/spec/discovery"
	"github.com/appc/spec/schema"
	"github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/pkg/fileutil"
	"github.com/coreos/rkt/pkg/user"
)

var (
	placeholdername = "acbuild-unnamed"
)

// Begin will start a new build, storing the untarred ACI the build operates on
// at a.CurrentACIPath. If start is the empty string, the build will begin with
// an empty ACI, otherwise the ACI stored at start will be used at the starting
// point.
func (a *ACBuild) Begin(start string, insecure bool) (err error) {
	_, err = os.Stat(a.ContextPath)
	switch {
	case os.IsNotExist(err):
		break
	case err != nil:
		return err
	default:
		return fmt.Errorf("build already in progress in this working dir")
	}

	err = os.MkdirAll(a.ContextPath, 0755)
	if err != nil {
		return err
	}

	if err = a.lock(); err != nil {
		return err
	}
	defer func() {
		if err1 := a.unlock(); err == nil {
			err = err1
		}
	}()

	defer func() {
		// If there was an error while beginning, we don't want to produce an
		// unexpected build context
		if err != nil {
			os.RemoveAll(a.ContextPath)
		}
	}()

	if start != "" {
		err = os.MkdirAll(a.CurrentACIPath, 0755)
		if err != nil {
			return err
		}
		if start[0] == '.' || start[0] == '/' {
			finfo, err := os.Stat(start)
			switch {
			case os.IsNotExist(err):
				return fmt.Errorf("no such file or directory: %s", start)
			case err != nil:
				return err
			case finfo.IsDir():
				return a.beginFromLocalDirectory(start)
			default:
				return a.beginFromLocalImage(start)
			}
		} else {
			dockerPrefix := "docker://"
			if strings.HasPrefix(start, dockerPrefix) {
				start = strings.TrimPrefix(start, dockerPrefix)
				return a.beginFromRemoteDockerImage(start, insecure)
			}
			return a.beginFromRemoteImage(start, insecure)
		}
	}
	return a.beginWithEmptyACI()
}

func (a *ACBuild) beginFromLocalImage(start string) error {
	finfo, err := os.Stat(start)
	if err != nil {
		return err
	}
	if finfo.IsDir() {
		return fmt.Errorf("provided starting ACI is a directory: %s", start)
	}
	err = util.ExtractImage(start, a.CurrentACIPath, nil)
	if err != nil {
		return err
	}

	for _, file := range []struct {
		FileName string
		FilePath string
	}{
		{"manifest file", path.Join(a.CurrentACIPath, aci.ManifestFile)},
		{"rootfs directory", path.Join(a.CurrentACIPath, aci.RootfsDir)},
	} {
		_, err = os.Stat(file.FilePath)
		switch {
		case os.IsNotExist(err):
			fmt.Fprintf(os.Stderr, "%s is missing, assuming build is beginning with a tar of a rootfs\n", file.FileName)
			return a.startedFromTar()
		case err != nil:
			return err
		}
	}
	return nil
}

func (a *ACBuild) startedFromTar() error {
	tmpPath := path.Join(a.ContextPath, "rootfs")
	err := os.Rename(a.CurrentACIPath, tmpPath)
	if err != nil {
		return err
	}
	err = a.beginWithEmptyACI()
	if err != nil {
		return err
	}
	err = os.Remove(path.Join(a.CurrentACIPath, aci.RootfsDir))
	if err != nil {
		return err
	}
	return os.Rename(tmpPath, path.Join(a.CurrentACIPath, aci.RootfsDir))
}

func (a *ACBuild) beginFromLocalDirectory(start string) error {
	err := os.MkdirAll(a.CurrentACIPath, 0755)
	if err != nil {
		return err
	}

	err = fileutil.CopyTree(start, path.Join(a.CurrentACIPath, aci.RootfsDir), user.NewBlankUidRange())
	if err != nil {
		return err
	}

	return a.writeEmptyManifest()
}

func (a *ACBuild) beginWithEmptyACI() error {
	err := os.MkdirAll(path.Join(a.CurrentACIPath, aci.RootfsDir), 0755)
	if err != nil {
		return err
	}

	return a.writeEmptyManifest()
}

func (a *ACBuild) writeEmptyManifest() error {
	acid, err := types.NewACIdentifier("acbuild-unnamed")
	if err != nil {
		return err
	}

	archlabel, err := types.NewACIdentifier("arch")
	if err != nil {
		return err
	}

	oslabel, err := types.NewACIdentifier("os")
	if err != nil {
		return err
	}

	archvalue := runtime.GOARCH
	if runtime.GOOS == "linux" && (archvalue == "arm" || archvalue == "arm64") {
		var x uint32 = 0x01020304
		test := *(*byte)(unsafe.Pointer(&x))
		switch {
		case test == 0x01 && archvalue == "arm":
			archvalue = "armv7b"
		case test == 0x04 && archvalue == "arm":
			archvalue = "armv7l"
		case test == 0x01 && archvalue == "arm64":
			archvalue = "aarch64_be"
		case test == 0x04 && archvalue == "arm64":
			archvalue = "aarch64"
		}
	}

	manifest := &schema.ImageManifest{
		ACKind:    schema.ImageManifestKind,
		ACVersion: schema.AppContainerVersion,
		Name:      *acid,
		Labels: types.Labels{
			types.Label{
				*archlabel,
				archvalue,
			},
			types.Label{
				*oslabel,
				runtime.GOOS,
			},
		},
	}

	manblob, err := manifest.MarshalJSON()
	if err != nil {
		return err
	}

	manfile, err := os.Create(path.Join(a.CurrentACIPath, aci.ManifestFile))
	if err != nil {
		return err
	}

	_, err = manfile.Write(manblob)
	if err != nil {
		return err
	}

	err = manfile.Close()
	if err != nil {
		return err
	}

	return nil
}

func (a *ACBuild) beginFromRemoteImage(start string, insecure bool) error {
	app, err := discovery.NewAppFromString(start)
	if err != nil {
		return err
	}
	labels, err := types.LabelsFromMap(app.Labels)
	if err != nil {
		return err
	}

	tmpDepStoreTarPath, err := ioutil.TempDir("", "acbuild-begin-tar")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDepStoreTarPath)

	tmpDepStoreExpandedPath, err := ioutil.TempDir("", "acbuild-begin-expanded")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tmpDepStoreExpandedPath)

	reg := registry.Registry{
		DepStoreTarPath:      tmpDepStoreTarPath,
		DepStoreExpandedPath: tmpDepStoreExpandedPath,
		Insecure:             insecure,
		Debug:                a.Debug,
	}

	err = reg.Fetch(app.Name, labels, 0, false)
	if err != nil {
		if urlerr, ok := err.(*url.Error); ok {
			if operr, ok := urlerr.Err.(*net.OpError); ok {
				if dnserr, ok := operr.Err.(*net.DNSError); ok {
					if dnserr.Err == "no such host" {
						return fmt.Errorf("unknown host when fetching image, check your connection and local file paths must start with '/' or '.'")
					}
				}
			}
		}
		return err
	}

	files, err := ioutil.ReadDir(tmpDepStoreTarPath)
	if err != nil {
		return err
	}

	if len(files) != 1 {
		var filelist string
		for _, file := range files {
			if filelist == "" {
				filelist = file.Name()
			} else {
				filelist = filelist + ", " + file.Name()
			}
		}
		panic("unexpected number of files in store after download: " + filelist)
	}

	return util.ExtractImage(path.Join(tmpDepStoreTarPath, files[0].Name()), a.CurrentACIPath, nil)
}

func (a *ACBuild) beginFromRemoteDockerImage(start string, insecure bool) (err error) {
	outputDir, err := ioutil.TempDir("", "acbuild")
	if err != nil {
		return err
	}
	defer os.RemoveAll(outputDir)

	tempDir, err := ioutil.TempDir("", "acbuild")
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	insecureConf := common.InsecureConfig{
		SkipVerify: insecure,
		AllowHTTP:  insecure,
	}

	config := docker2aci.RemoteConfig{
		CommonConfig: docker2aci.CommonConfig{
			Squash:      true,
			OutputDir:   outputDir,
			TmpDir:      tempDir,
			Compression: common.GzipCompression,
		},
		Username: "",
		Password: "",
		Insecure: insecureConf,
	}
	renderedACIs, err := docker2aci.ConvertRemoteRepo(start, config)
	if err != nil {
		return err
	}
	if len(renderedACIs) > 1 {
		return fmt.Errorf("internal error: docker2aci didn't squash the image")
	}
	if len(renderedACIs) == 0 {
		return fmt.Errorf("internal error: docker2aci didn't produce any images")
	}
	absRenderedACI, err := filepath.Abs(renderedACIs[0])
	if err != nil {
		return err
	}

	return util.ExtractImage(absRenderedACI, a.CurrentACIPath, nil)
}
<|endoftext|>


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

package app

import (
	"context"
	"embed"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"time"

	"github.com/influxdata/influxdb-client-go/v2/api"
	"github.com/influxdata/influxdb-client-go/v2/api/query"
	"golang.org/x/build/internal/influx"
	"golang.org/x/build/third_party/bandchart"
)

// /dashboard/ displays a dashboard of benchmark results over time for
// performance monitoring.

//go:embed dashboard/*
var dashboardFS embed.FS

// dashboardRegisterOnMux registers the dashboard URLs on mux.
func (a *App) dashboardRegisterOnMux(mux *http.ServeMux) {
	mux.Handle("/dashboard/", http.FileServer(http.FS(dashboardFS)))
	mux.Handle("/dashboard/third_party/bandchart/", http.StripPrefix("/dashboard/third_party/bandchart/", http.FileServer(http.FS(bandchart.FS))))
	mux.HandleFunc("/dashboard/data.json", a.dashboardData)
}

// BenchmarkJSON contains the timeseries values for a single benchmark name +
// unit.
//
// We could try to shoehorn this into benchfmt.Result, but that isn't really
// the best fit for a graph.
type BenchmarkJSON struct {
	Name string
	Unit string

	// These will be sorted by CommitDate.
	Values []ValueJSON
}

type ValueJSON struct {
	CommitHash string
	CommitDate time.Time

	// These are pre-formatted as percent change.
	Low    float64
	Center float64
	High   float64
}

func fluxRecordToValue(rec *query.FluxRecord) (ValueJSON, error) {
	low, ok := rec.ValueByKey("low").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s low value got type %T want float64", rec, rec.ValueByKey("low"))
	}

	center, ok := rec.ValueByKey("center").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s center value got type %T want float64", rec, rec.ValueByKey("center"))
	}

	high, ok := rec.ValueByKey("high").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s high value got type %T want float64", rec, rec.ValueByKey("high"))
	}

	commit, ok := rec.ValueByKey("experiment-commit").(string)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s experiment-commit value got type %T want float64", rec, rec.ValueByKey("experiment-commit"))
	}

	return ValueJSON{
		CommitDate: rec.Time(),
		CommitHash: commit,
		Low:        low - 1,
		Center:     center - 1,
		High:       high - 1,
	}, nil
}

// validateRe is an allowlist of characters for a Flux string literal for
// benchmark names. The string will be quoted, so we must not allow ending the
// quote sequence.
var validateRe = regexp.MustCompile(`[a-zA-Z/_:-]+`)

func validateFluxString(s string) error {
	if !validateRe.MatchString(s) {
		return fmt.Errorf("malformed value %q", s)
	}
	return nil
}

var errBenchmarkNotFound = errors.New("benchmark not found")

// fetchNamedUnitBenchmark queries Influx for a specific name + unit benchmark.
func fetchNamedUnitBenchmark(ctx context.Context, qc api.QueryAPI, name, unit string) (*BenchmarkJSON, error) {
	if err := validateFluxString(name); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}
	if err := validateFluxString(unit); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}

	query := fmt.Sprintf(`
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["name"] == "%s")
  |> filter(fn: (r) => r["unit"] == "%s")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`, name, unit)

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	b, err := groupBenchmarkResults(res)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errBenchmarkNotFound
	}
	if len(b) > 1 {
		return nil, fmt.Errorf("query returned too many benchmarks: %+v", b)
	}
	return b[0], nil
}

// fetchDefaultBenchmarks queries Influx for the default benchmark set.
func fetchDefaultBenchmarks(ctx context.Context, qc api.QueryAPI) ([]*BenchmarkJSON, error) {
	// Keep benchmarks with the same name grouped together, which is
	// assumed by the JS.
	benchmarks := []struct{ name, unit string }{
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "sec/op",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "p90-latency-sec",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "average-RSS-bytes",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "peak-RSS-bytes",
		},
		{
			name: "GoBuildKubelet",
			unit: "sec/op",
		},
		{
			name: "GoBuildKubeletLink",
			unit: "sec/op",
		},
	}

	ret := make([]*BenchmarkJSON, 0, len(benchmarks))
	for _, bench := range benchmarks {
		b, err := fetchNamedUnitBenchmark(ctx, qc, bench.name, bench.unit)
		if err != nil {
			return nil, fmt.Errorf("error fetching benchmark %s/%s: %w", bench.name, bench.unit, err)
		}
		ret = append(ret, b)
	}

	return ret, nil
}

// fetchNamedBenchmark queries Influx for all benchmark results with the passed
// name (for all units).
func fetchNamedBenchmark(ctx context.Context, qc api.QueryAPI, name string) ([]*BenchmarkJSON, error) {
	if err := validateFluxString(name); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}

	query := fmt.Sprintf(`
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["name"] == "%s")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`, name)

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	b, err := groupBenchmarkResults(res)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errBenchmarkNotFound
	}
	return b, nil
}

// fetchAllBenchmarks queries Influx for all benchmark results.
func fetchAllBenchmarks(ctx context.Context, qc api.QueryAPI) ([]*BenchmarkJSON, error) {
	const query = `
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	return groupBenchmarkResults(res)
}

// groupBenchmarkResults groups all benchmark results from the passed query.
func groupBenchmarkResults(res *api.QueryTableResult) ([]*BenchmarkJSON, error) {
	type key struct {
		name string
		unit string
	}
	m := make(map[key]*BenchmarkJSON)

	for res.Next() {
		rec := res.Record()

		name, ok := rec.ValueByKey("name").(string)
		if !ok {
			return nil, fmt.Errorf("record %s name value got type %T want string", rec, rec.ValueByKey("name"))
		}

		unit, ok := rec.ValueByKey("unit").(string)
		if !ok {
			return nil, fmt.Errorf("record %s unit value got type %T want string", rec, rec.ValueByKey("unit"))
		}

		k := key{name, unit}
		b, ok := m[k]
		if !ok {
			b = &BenchmarkJSON{
				Name: name,
				Unit: unit,
			}
			m[k] = b
		}

		v, err := fluxRecordToValue(res.Record())
		if err != nil {
			return nil, err
		}

		b.Values = append(b.Values, v)
	}

	s := make([]*BenchmarkJSON, 0, len(m))
	for _, b := range m {
		s = append(s, b)
	}
	// Keep benchmarks with the same name grouped together, which is
	// assumed by the JS.
	sort.Slice(s, func(i, j int) bool {
		if s[i].Name == s[j].Name {
			return s[i].Unit < s[j].Unit
		}
		return s[i].Name < s[j].Name
	})

	for _, b := range s {
		sort.Slice(b.Values, func(i, j int) bool {
			return b.Values[i].CommitDate.Before(b.Values[j].CommitDate)
		})
	}

	return s, nil
}

// search handles /dashboard/data.json.
//
// TODO(prattmic): Consider caching Influx results in-memory for a few mintures
// to reduce load on Influx.
func (a *App) dashboardData(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	start := time.Now()
	defer func() {
		log.Printf("Dashboard total query time: %s", time.Since(start))
	}()

	ifxc, err := a.influxClient(ctx)
	if err != nil {
		log.Printf("Error getting Influx client: %v", err)
		http.Error(w, "Error connecting to Influx", 500)
		return
	}
	defer ifxc.Close()

	qc := ifxc.QueryAPI(influx.Org)

	benchmark := r.FormValue("benchmark")
	var benchmarks []*BenchmarkJSON
	if benchmark == "" {
		benchmarks, err = fetchDefaultBenchmarks(ctx, qc)
	} else if benchmark == "all" {
		benchmarks, err = fetchAllBenchmarks(ctx, qc)
	} else {
		benchmarks, err = fetchNamedBenchmark(ctx, qc, benchmark)
	}
	if err == errBenchmarkNotFound {
		log.Printf("Benchmark not found: %q", benchmark)
		http.Error(w, "Benchmark not found", 404)
		return
	}
	if err != nil {
		log.Printf("Error fetching benchmarks: %v", err)
		http.Error(w, "Error fetching benchmarks", 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	e := json.NewEncoder(w)
	e.SetIndent("", "\t")
	e.Encode(benchmarks)
}
<commit_msg>perf: add a few additional homepage benchmarks<commit_after>// Copyright 2022 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package app

import (
	"context"
	"embed"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"sort"
	"time"

	"github.com/influxdata/influxdb-client-go/v2/api"
	"github.com/influxdata/influxdb-client-go/v2/api/query"
	"golang.org/x/build/internal/influx"
	"golang.org/x/build/third_party/bandchart"
)

// /dashboard/ displays a dashboard of benchmark results over time for
// performance monitoring.

//go:embed dashboard/*
var dashboardFS embed.FS

// dashboardRegisterOnMux registers the dashboard URLs on mux.
func (a *App) dashboardRegisterOnMux(mux *http.ServeMux) {
	mux.Handle("/dashboard/", http.FileServer(http.FS(dashboardFS)))
	mux.Handle("/dashboard/third_party/bandchart/", http.StripPrefix("/dashboard/third_party/bandchart/", http.FileServer(http.FS(bandchart.FS))))
	mux.HandleFunc("/dashboard/data.json", a.dashboardData)
}

// BenchmarkJSON contains the timeseries values for a single benchmark name +
// unit.
//
// We could try to shoehorn this into benchfmt.Result, but that isn't really
// the best fit for a graph.
type BenchmarkJSON struct {
	Name string
	Unit string

	// These will be sorted by CommitDate.
	Values []ValueJSON
}

type ValueJSON struct {
	CommitHash string
	CommitDate time.Time

	// These are pre-formatted as percent change.
	Low    float64
	Center float64
	High   float64
}

func fluxRecordToValue(rec *query.FluxRecord) (ValueJSON, error) {
	low, ok := rec.ValueByKey("low").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s low value got type %T want float64", rec, rec.ValueByKey("low"))
	}

	center, ok := rec.ValueByKey("center").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s center value got type %T want float64", rec, rec.ValueByKey("center"))
	}

	high, ok := rec.ValueByKey("high").(float64)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s high value got type %T want float64", rec, rec.ValueByKey("high"))
	}

	commit, ok := rec.ValueByKey("experiment-commit").(string)
	if !ok {
		return ValueJSON{}, fmt.Errorf("record %s experiment-commit value got type %T want float64", rec, rec.ValueByKey("experiment-commit"))
	}

	return ValueJSON{
		CommitDate: rec.Time(),
		CommitHash: commit,
		Low:        low - 1,
		Center:     center - 1,
		High:       high - 1,
	}, nil
}

// validateRe is an allowlist of characters for a Flux string literal for
// benchmark names. The string will be quoted, so we must not allow ending the
// quote sequence.
var validateRe = regexp.MustCompile(`[a-zA-Z/_:-]+`)

func validateFluxString(s string) error {
	if !validateRe.MatchString(s) {
		return fmt.Errorf("malformed value %q", s)
	}
	return nil
}

var errBenchmarkNotFound = errors.New("benchmark not found")

// fetchNamedUnitBenchmark queries Influx for a specific name + unit benchmark.
func fetchNamedUnitBenchmark(ctx context.Context, qc api.QueryAPI, name, unit string) (*BenchmarkJSON, error) {
	if err := validateFluxString(name); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}
	if err := validateFluxString(unit); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}

	query := fmt.Sprintf(`
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["name"] == "%s")
  |> filter(fn: (r) => r["unit"] == "%s")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`, name, unit)

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	b, err := groupBenchmarkResults(res)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errBenchmarkNotFound
	}
	if len(b) > 1 {
		return nil, fmt.Errorf("query returned too many benchmarks: %+v", b)
	}
	return b[0], nil
}

// fetchDefaultBenchmarks queries Influx for the default benchmark set.
func fetchDefaultBenchmarks(ctx context.Context, qc api.QueryAPI) ([]*BenchmarkJSON, error) {
	// Keep benchmarks with the same name grouped together, which is
	// assumed by the JS.
	benchmarks := []struct{ name, unit string }{
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "sec/op",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "p50-latency-sec",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "p90-latency-sec",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "average-RSS-bytes",
		},
		{
			name: "Tile38WithinCircle100kmRequest",
			unit: "peak-RSS-bytes",
		},
		{
			name: "BleveQuery",
			unit: "sec/op",
		},
		{
			name: "BleveQuery",
			unit: "average-RSS-bytes",
		},
		{
			name: "BleveQuery",
			unit: "peak-RSS-bytes",
		},
		{
			name: "GoBuildKubelet",
			unit: "sec/op",
		},
		{
			name: "GoBuildKubeletLink",
			unit: "sec/op",
		},
		{
			name: "RegexMatch-8",
			unit: "sec/op",
		},
		{
			name: "BuildJSON-8",
			unit: "sec/op",
		},
		{
			name: "ZapJSON-8",
			unit: "sec/op",
		},
	}

	ret := make([]*BenchmarkJSON, 0, len(benchmarks))
	for _, bench := range benchmarks {
		b, err := fetchNamedUnitBenchmark(ctx, qc, bench.name, bench.unit)
		if err != nil {
			return nil, fmt.Errorf("error fetching benchmark %s/%s: %w", bench.name, bench.unit, err)
		}
		ret = append(ret, b)
	}

	return ret, nil
}

// fetchNamedBenchmark queries Influx for all benchmark results with the passed
// name (for all units).
func fetchNamedBenchmark(ctx context.Context, qc api.QueryAPI, name string) ([]*BenchmarkJSON, error) {
	if err := validateFluxString(name); err != nil {
		return nil, fmt.Errorf("invalid benchmark name: %w", err)
	}

	query := fmt.Sprintf(`
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["name"] == "%s")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`, name)

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	b, err := groupBenchmarkResults(res)
	if err != nil {
		return nil, err
	}
	if len(b) == 0 {
		return nil, errBenchmarkNotFound
	}
	return b, nil
}

// fetchAllBenchmarks queries Influx for all benchmark results.
func fetchAllBenchmarks(ctx context.Context, qc api.QueryAPI) ([]*BenchmarkJSON, error) {
	const query = `
from(bucket: "perf")
  |> range(start: -30d)
  |> filter(fn: (r) => r["_measurement"] == "benchmark-result")
  |> filter(fn: (r) => r["branch"] == "master")
  |> filter(fn: (r) => r["goos"] == "linux")
  |> filter(fn: (r) => r["goarch"] == "amd64")
  |> pivot(columnKey: ["_field"], rowKey: ["_time"], valueColumn: "_value")
  |> yield(name: "last")
`

	res, err := qc.Query(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("error performing query: %W", err)
	}

	return groupBenchmarkResults(res)
}

// groupBenchmarkResults groups all benchmark results from the passed query.
func groupBenchmarkResults(res *api.QueryTableResult) ([]*BenchmarkJSON, error) {
	type key struct {
		name string
		unit string
	}
	m := make(map[key]*BenchmarkJSON)

	for res.Next() {
		rec := res.Record()

		name, ok := rec.ValueByKey("name").(string)
		if !ok {
			return nil, fmt.Errorf("record %s name value got type %T want string", rec, rec.ValueByKey("name"))
		}

		unit, ok := rec.ValueByKey("unit").(string)
		if !ok {
			return nil, fmt.Errorf("record %s unit value got type %T want string", rec, rec.ValueByKey("unit"))
		}

		k := key{name, unit}
		b, ok := m[k]
		if !ok {
			b = &BenchmarkJSON{
				Name: name,
				Unit: unit,
			}
			m[k] = b
		}

		v, err := fluxRecordToValue(res.Record())
		if err != nil {
			return nil, err
		}

		b.Values = append(b.Values, v)
	}

	s := make([]*BenchmarkJSON, 0, len(m))
	for _, b := range m {
		s = append(s, b)
	}
	// Keep benchmarks with the same name grouped together, which is
	// assumed by the JS.
	sort.Slice(s, func(i, j int) bool {
		if s[i].Name == s[j].Name {
			return s[i].Unit < s[j].Unit
		}
		return s[i].Name < s[j].Name
	})

	for _, b := range s {
		sort.Slice(b.Values, func(i, j int) bool {
			return b.Values[i].CommitDate.Before(b.Values[j].CommitDate)
		})
	}

	return s, nil
}

// search handles /dashboard/data.json.
//
// TODO(prattmic): Consider caching Influx results in-memory for a few mintures
// to reduce load on Influx.
func (a *App) dashboardData(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	start := time.Now()
	defer func() {
		log.Printf("Dashboard total query time: %s", time.Since(start))
	}()

	ifxc, err := a.influxClient(ctx)
	if err != nil {
		log.Printf("Error getting Influx client: %v", err)
		http.Error(w, "Error connecting to Influx", 500)
		return
	}
	defer ifxc.Close()

	qc := ifxc.QueryAPI(influx.Org)

	benchmark := r.FormValue("benchmark")
	var benchmarks []*BenchmarkJSON
	if benchmark == "" {
		benchmarks, err = fetchDefaultBenchmarks(ctx, qc)
	} else if benchmark == "all" {
		benchmarks, err = fetchAllBenchmarks(ctx, qc)
	} else {
		benchmarks, err = fetchNamedBenchmark(ctx, qc, benchmark)
	}
	if err == errBenchmarkNotFound {
		log.Printf("Benchmark not found: %q", benchmark)
		http.Error(w, "Benchmark not found", 404)
		return
	}
	if err != nil {
		log.Printf("Error fetching benchmarks: %v", err)
		http.Error(w, "Error fetching benchmarks", 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	e := json.NewEncoder(w)
	e.SetIndent("", "\t")
	e.Encode(benchmarks)
}
<|endoftext|>


<commit_before>package log

import (
	"fmt"
	"io"
	"os"
	"reflect"
	"time"
)

type ioLogger struct {
	c   chan Event
	out io.Writer
	buf []byte
}

var (
	sub *Subscription
)

func init() {
	l := &ioLogger{c: make(chan Event, 100), out: os.Stderr}
	sub = Subscribe(func(evt Event) {
		l.c <- evt
	})
	go l.listenEvent()
}

func (l *ioLogger) listenEvent() {
	for true {
		e := <-l.c
		l.writeEvent(e)
	}
}

// Cheap integer to fixed-width decimal ASCII.  Give a negative width to avoid zero-padding.
func itoa(buf *[]byte, i int, wid int) {
	// Assemble decimal in reverse order.
	var b [20]byte
	bp := len(b) - 1
	for i >= 10 || wid > 1 {
		wid--
		q := i / 10
		b[bp] = byte('0' + i - q*10)
		bp--
		i = q
	}
	// i < 10
	b[bp] = byte('0' + i)
	*buf = append(*buf, b[bp:]...)
}

func (l *ioLogger) formatHeader(buf *[]byte, prefix string, t time.Time) {
	t = t.UTC()
	// Y/M/D
	year, month, day := t.Date()
	itoa(buf, year, 4)
	*buf = append(*buf, '/')
	itoa(buf, int(month), 2)
	*buf = append(*buf, '/')
	itoa(buf, day, 2)
	*buf = append(*buf, ' ')

	// H/M/S
	hour, min, sec := t.Clock()
	itoa(buf, hour, 2)
	*buf = append(*buf, ':')
	itoa(buf, min, 2)
	*buf = append(*buf, ':')
	itoa(buf, sec, 2)

	// no microseconds
	// *buf = append(*buf, '.')
	// itoa(buf, t.Nanosecond()/1e3, 6)

	*buf = append(*buf, ' ')
	if len(prefix) > 0 {
		*buf = append(*buf, prefix...)
		*buf = append(*buf, ' ')
	}
}

func (l *ioLogger) writeEvent(e Event) {
	l.buf = l.buf[:0]
	l.formatHeader(&l.buf, e.Prefix, e.Time)
	l.out.Write(l.buf)
	if len(e.Message) > 0 {
		l.out.Write([]byte(e.Message))
		l.out.Write([]byte{' '})
	}

	wr := ioEncoder{l.out}
	for _, f := range e.Context {
		f.Encode(wr)
		l.out.Write([]byte{' '})
	}
	for _, f := range e.Fields {
		f.Encode(wr)
		l.out.Write([]byte{' '})
	}
	wr.Write([]byte{'\n'})
}

type ioEncoder struct {
	io.Writer
}

func (e ioEncoder) EncodeBool(key string, val bool) {
	fmt.Fprintf(e, "%s=%t", key, val)
}

func (e ioEncoder) EncodeFloat64(key string, val float64) {
	fmt.Fprintf(e, "%s=%f", key, val)
}

func (e ioEncoder) EncodeInt(key string, val int) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeInt64(key string, val int64) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeDuration(key string, val time.Duration) {
	fmt.Fprintf(e, "%s=%s", key, val)
}

func (e ioEncoder) EncodeUint(key string, val uint) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeUint64(key string, val uint64) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeString(key string, val string) {
	fmt.Fprintf(e, "%s=%q", key, val)
}

func (e ioEncoder) EncodeObject(key string, val interface{}) {
	fmt.Fprintf(e, "%s=%q", key, val)
}

func (e ioEncoder) EncodeType(key string, val reflect.Type) {
	fmt.Fprintf(e, "%s=%v", key, val)
}
<commit_msg>Print log messages in one piece<commit_after>package log

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"reflect"
	"time"
)

type ioLogger struct {
	c   chan Event
	out io.Writer
	buf []byte
}

var (
	sub *Subscription
)

func init() {
	l := &ioLogger{c: make(chan Event, 100), out: os.Stderr}
	sub = Subscribe(func(evt Event) {
		l.c <- evt
	})
	go l.listenEvent()
}

func (l *ioLogger) listenEvent() {
	for true {
		e := <-l.c
		l.writeEvent(e)
	}
}

// Cheap integer to fixed-width decimal ASCII.  Give a negative width to avoid zero-padding.
func itoa(buf bytes.Buffer, i int, wid int) {
	// Assemble decimal in reverse order.
	var b [20]byte
	bp := len(b) - 1
	for i >= 10 || wid > 1 {
		wid--
		q := i / 10
		b[bp] = byte('0' + i - q*10)
		bp--
		i = q
	}
	// i < 10
	b[bp] = byte('0' + i)
	buf.Write(b[bp:])
}

func (l *ioLogger) formatHeader(buf bytes.Buffer, prefix string, t time.Time) {
	t = t.UTC()
	// Y/M/D
	year, month, day := t.Date()
	itoa(buf, year, 4)
	buf.WriteByte('/')
	itoa(buf, int(month), 2)
	buf.WriteByte('/')
	itoa(buf, day, 2)
	buf.WriteByte(' ')

	// H/M/S
	hour, min, sec := t.Clock()
	itoa(buf, hour, 2)
	buf.WriteByte(':')
	itoa(buf, min, 2)
	buf.WriteByte(':')
	itoa(buf, sec, 2)

	// no microseconds
	// *buf = append(*buf, '.')
	// itoa(buf, t.Nanosecond()/1e3, 6)

	buf.WriteByte(' ')
	if len(prefix) > 0 {
		buf.WriteString(prefix)
		buf.WriteByte(' ')
	}
}

func (l *ioLogger) writeEvent(e Event) {
	var buf = bytes.Buffer{}
	l.formatHeader(buf, e.Prefix, e.Time)
	if len(e.Message) > 0 {
		buf.WriteString(e.Message)
		buf.WriteByte(' ')
	}

	wr := ioEncoder{&buf}
	for _, f := range e.Context {
		f.Encode(wr)
		buf.WriteByte(' ')
	}
	for _, f := range e.Fields {
		f.Encode(wr)
		buf.WriteByte(' ')
	}
	buf.WriteByte('\n')
	l.out.Write(buf.Bytes())
}

type ioEncoder struct {
	io.Writer
}

func (e ioEncoder) EncodeBool(key string, val bool) {
	fmt.Fprintf(e, "%s=%t", key, val)
}

func (e ioEncoder) EncodeFloat64(key string, val float64) {
	fmt.Fprintf(e, "%s=%f", key, val)
}

func (e ioEncoder) EncodeInt(key string, val int) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeInt64(key string, val int64) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeDuration(key string, val time.Duration) {
	fmt.Fprintf(e, "%s=%s", key, val)
}

func (e ioEncoder) EncodeUint(key string, val uint) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeUint64(key string, val uint64) {
	fmt.Fprintf(e, "%s=%d", key, val)
}

func (e ioEncoder) EncodeString(key string, val string) {
	fmt.Fprintf(e, "%s=%q", key, val)
}

func (e ioEncoder) EncodeObject(key string, val interface{}) {
	fmt.Fprintf(e, "%s=%q", key, val)
}

func (e ioEncoder) EncodeType(key string, val reflect.Type) {
	fmt.Fprintf(e, "%s=%v", key, val)
}
<|endoftext|>


<commit_before>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package repo_indexer

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/url"
	"strings"

	"github.com/golang/protobuf/ptypes/wrappers"
	"k8s.io/helm/pkg/chartutil"
	"k8s.io/helm/pkg/proto/hapi/chart"
	"k8s.io/helm/pkg/repo"

	appClient "openpitrix.io/openpitrix/pkg/client/app"
	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb"
	"openpitrix.io/openpitrix/pkg/utils"
	"openpitrix.io/openpitrix/pkg/utils/sender"
	"openpitrix.io/openpitrix/pkg/utils/yaml"
)

// Reference: https://sourcegraph.com/github.com/kubernetes/helm@fe9d365/-/blob/pkg/repo/chartrepo.go#L111:27
func GetIndexFile(repoUrl string) (indexFile *repo.IndexFile, err error) {
	var indexURL string
	indexFile = &repo.IndexFile{}
	parsedURL, err := url.Parse(repoUrl)
	if err != nil {
		return
	}
	parsedURL.Path = strings.TrimSuffix(parsedURL.Path, "/") + "/index.yaml"
	indexURL = parsedURL.String()
	resp, err := utils.HttpGet(indexURL)
	if err != nil {
		return
	}
	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	err = yaml.Unmarshal(content, indexFile)
	if err != nil {
		return
	}
	indexFile.SortEntries()
	return
}

// Reference: https://sourcegraph.com/github.com/kubernetes/helm@fe9d365/-/blob/pkg/downloader/chart_downloader.go#L225:35
func GetPackageFile(chartVersion *repo.ChartVersion, repoUrl string) (*chart.Chart, error) {
	if len(chartVersion.URLs) == 0 {
		return nil, fmt.Errorf("chart [%s] has no downloadable URLs", chartVersion.Name)
	}
	u, err := url.Parse(chartVersion.URLs[0])
	if err != nil {
		return nil, fmt.Errorf("invalid chart URL format: %v", chartVersion.URLs)
	}

	// If the URL is relative (no scheme), prepend the chart repo's base URL
	if !u.IsAbs() {
		repoURL, err := url.Parse(repoUrl)
		if err != nil {
			return nil, err
		}
		q := repoURL.Query()
		// We need a trailing slash for ResolveReference to work, but make sure there isn't already one
		repoURL.Path = strings.TrimSuffix(repoURL.Path, "/") + "/"
		u = repoURL.ResolveReference(u)
		u.RawQuery = q.Encode()
	}
	resp, err := utils.HttpGet(u.String())
	if err != nil {
		return nil, err
	}
	return chartutil.LoadArchive(resp.Body)
}

func SyncAppInfo(repoId, owner, chartName string, chartVersions *repo.ChartVersions) (string, error) {
	var appId string
	logger.Debugf("chart [%s] has [%d] versions", chartName, chartVersions.Len())
	ctx := sender.NewContext(context.Background(), sender.GetSystemUser())
	appManagerClient, err := appClient.NewAppManagerClient(ctx)
	if err != nil {
		return appId, err
	}
	req := pb.DescribeAppsRequest{}
	req.RepoId = []string{repoId}
	req.Owner = []string{owner}
	req.ChartName = []string{chartName}
	res, err := appManagerClient.DescribeApps(ctx, &req)
	if err != nil {
		return appId, err
	}
	var description, icon, home, sources *wrappers.StringValue
	if chartVersions.Len() > 0 {
		chartVersion := (*chartVersions)[0]
		description = utils.ToProtoString(chartVersion.GetDescription())
		icon = utils.ToProtoString(chartVersion.GetIcon())
		home = utils.ToProtoString(chartVersion.GetHome())
		sources = utils.ToProtoString(strings.Join(chartVersion.Sources, ","))
	}
	if res.TotalCount == 0 {
		createReq := pb.CreateAppRequest{}
		createReq.RepoId = utils.ToProtoString(repoId)
		createReq.ChartName = utils.ToProtoString(chartName)
		createReq.Name = utils.ToProtoString(chartName)
		createReq.Description = description
		createReq.Icon = icon
		createReq.Home = home
		createReq.Sources = sources

		createRes, err := appManagerClient.CreateApp(ctx, &createReq)
		if err != nil {
			return appId, err
		}
		appId = createRes.GetApp().GetAppId().GetValue()
		return appId, err

	} else {
		modifyReq := pb.ModifyAppRequest{}
		modifyReq.AppId = res.AppSet[0].AppId
		modifyReq.Name = utils.ToProtoString(chartName)
		modifyReq.ChartName = utils.ToProtoString(chartName)
		modifyReq.Description = description
		modifyReq.Icon = icon
		modifyReq.Home = home
		modifyReq.Sources = sources

		modifyRes, err := appManagerClient.ModifyApp(ctx, &modifyReq)
		if err != nil {
			return appId, err
		}
		appId = modifyRes.GetApp().GetAppId().GetValue()
		return appId, err
	}
}

func SyncAppVersionInfo(appId, owner string, chartVersion *repo.ChartVersion) (string, error) {
	var versionId string
	ctx := sender.NewContext(context.Background(), sender.GetSystemUser())
	appManagerClient, err := app.NewAppManagerClient(ctx)
	if err != nil {
		return versionId, err
	}
	appVersionName := chartVersion.GetVersion()
	if chartVersion.GetAppVersion() != "" {
		appVersionName += fmt.Sprintf(" [%s]", chartVersion.GetAppVersion())
	}
	packageName := chartVersion.URLs[0]
	description := chartVersion.GetDescription()
	req := pb.DescribeAppVersionsRequest{}
	req.AppId = []string{appId}
	req.Owner = []string{owner}
	req.Name = []string{appVersionName}
	res, err := appManagerClient.DescribeAppVersions(ctx, &req)
	if err != nil {
		return versionId, err
	}
	if res.TotalCount == 0 {
		createReq := pb.CreateAppVersionRequest{}
		createReq.AppId = utils.ToProtoString(appId)
		createReq.Owner = utils.ToProtoString(owner)
		createReq.Name = utils.ToProtoString(appVersionName)
		createReq.PackageName = utils.ToProtoString(packageName)
		createReq.Description = utils.ToProtoString(description)

		createRes, err := appManagerClient.CreateAppVersion(ctx, &createReq)
		if err != nil {
			return versionId, err
		}
		versionId = createRes.GetAppVersion().GetVersionId().GetValue()
		return versionId, err
	} else {
		modifyReq := pb.ModifyAppVersionRequest{}
		modifyReq.VersionId = res.AppVersionSet[0].VersionId
		modifyReq.PackageName = utils.ToProtoString(packageName)
		modifyReq.Description = utils.ToProtoString(description)

		modifyRes, err := appManagerClient.ModifyAppVersion(ctx, &modifyReq)
		if err != nil {
			return versionId, err
		}
		versionId = modifyRes.GetAppVersion().GetVersionId().GetValue()
		return versionId, err
	}
}
<commit_msg>Fix travis ci error<commit_after>// Copyright 2018 The OpenPitrix Authors. All rights reserved.
// Use of this source code is governed by a Apache license
// that can be found in the LICENSE file.

package repo_indexer

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/url"
	"strings"

	"github.com/golang/protobuf/ptypes/wrappers"
	"k8s.io/helm/pkg/chartutil"
	"k8s.io/helm/pkg/proto/hapi/chart"
	"k8s.io/helm/pkg/repo"

	"openpitrix.io/openpitrix/pkg/client/app"
	"openpitrix.io/openpitrix/pkg/logger"
	"openpitrix.io/openpitrix/pkg/pb"
	"openpitrix.io/openpitrix/pkg/utils"
	"openpitrix.io/openpitrix/pkg/utils/sender"
	"openpitrix.io/openpitrix/pkg/utils/yaml"
)

// Reference: https://sourcegraph.com/github.com/kubernetes/helm@fe9d365/-/blob/pkg/repo/chartrepo.go#L111:27
func GetIndexFile(repoUrl string) (indexFile *repo.IndexFile, err error) {
	var indexURL string
	indexFile = &repo.IndexFile{}
	parsedURL, err := url.Parse(repoUrl)
	if err != nil {
		return
	}
	parsedURL.Path = strings.TrimSuffix(parsedURL.Path, "/") + "/index.yaml"
	indexURL = parsedURL.String()
	resp, err := utils.HttpGet(indexURL)
	if err != nil {
		return
	}
	content, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	err = yaml.Unmarshal(content, indexFile)
	if err != nil {
		return
	}
	indexFile.SortEntries()
	return
}

// Reference: https://sourcegraph.com/github.com/kubernetes/helm@fe9d365/-/blob/pkg/downloader/chart_downloader.go#L225:35
func GetPackageFile(chartVersion *repo.ChartVersion, repoUrl string) (*chart.Chart, error) {
	if len(chartVersion.URLs) == 0 {
		return nil, fmt.Errorf("chart [%s] has no downloadable URLs", chartVersion.Name)
	}
	u, err := url.Parse(chartVersion.URLs[0])
	if err != nil {
		return nil, fmt.Errorf("invalid chart URL format: %v", chartVersion.URLs)
	}

	// If the URL is relative (no scheme), prepend the chart repo's base URL
	if !u.IsAbs() {
		repoURL, err := url.Parse(repoUrl)
		if err != nil {
			return nil, err
		}
		q := repoURL.Query()
		// We need a trailing slash for ResolveReference to work, but make sure there isn't already one
		repoURL.Path = strings.TrimSuffix(repoURL.Path, "/") + "/"
		u = repoURL.ResolveReference(u)
		u.RawQuery = q.Encode()
	}
	resp, err := utils.HttpGet(u.String())
	if err != nil {
		return nil, err
	}
	return chartutil.LoadArchive(resp.Body)
}

func SyncAppInfo(repoId, owner, chartName string, chartVersions *repo.ChartVersions) (string, error) {
	var appId string
	logger.Debugf("chart [%s] has [%d] versions", chartName, chartVersions.Len())
	ctx := sender.NewContext(context.Background(), sender.GetSystemUser())
	appManagerClient, err := app.NewAppManagerClient(ctx)
	if err != nil {
		return appId, err
	}
	req := pb.DescribeAppsRequest{}
	req.RepoId = []string{repoId}
	req.Owner = []string{owner}
	req.ChartName = []string{chartName}
	res, err := appManagerClient.DescribeApps(ctx, &req)
	if err != nil {
		return appId, err
	}
	var description, icon, home, sources *wrappers.StringValue
	if chartVersions.Len() > 0 {
		chartVersion := (*chartVersions)[0]
		description = utils.ToProtoString(chartVersion.GetDescription())
		icon = utils.ToProtoString(chartVersion.GetIcon())
		home = utils.ToProtoString(chartVersion.GetHome())
		sources = utils.ToProtoString(strings.Join(chartVersion.Sources, ","))
	}
	if res.TotalCount == 0 {
		createReq := pb.CreateAppRequest{}
		createReq.RepoId = utils.ToProtoString(repoId)
		createReq.ChartName = utils.ToProtoString(chartName)
		createReq.Name = utils.ToProtoString(chartName)
		createReq.Description = description
		createReq.Icon = icon
		createReq.Home = home
		createReq.Sources = sources

		createRes, err := appManagerClient.CreateApp(ctx, &createReq)
		if err != nil {
			return appId, err
		}
		appId = createRes.GetApp().GetAppId().GetValue()
		return appId, err

	} else {
		modifyReq := pb.ModifyAppRequest{}
		modifyReq.AppId = res.AppSet[0].AppId
		modifyReq.Name = utils.ToProtoString(chartName)
		modifyReq.ChartName = utils.ToProtoString(chartName)
		modifyReq.Description = description
		modifyReq.Icon = icon
		modifyReq.Home = home
		modifyReq.Sources = sources

		modifyRes, err := appManagerClient.ModifyApp(ctx, &modifyReq)
		if err != nil {
			return appId, err
		}
		appId = modifyRes.GetApp().GetAppId().GetValue()
		return appId, err
	}
}

func SyncAppVersionInfo(appId, owner string, chartVersion *repo.ChartVersion) (string, error) {
	var versionId string
	ctx := sender.NewContext(context.Background(), sender.GetSystemUser())
	appManagerClient, err := app.NewAppManagerClient(ctx)
	if err != nil {
		return versionId, err
	}
	appVersionName := chartVersion.GetVersion()
	if chartVersion.GetAppVersion() != "" {
		appVersionName += fmt.Sprintf(" [%s]", chartVersion.GetAppVersion())
	}
	packageName := chartVersion.URLs[0]
	description := chartVersion.GetDescription()
	req := pb.DescribeAppVersionsRequest{}
	req.AppId = []string{appId}
	req.Owner = []string{owner}
	req.Name = []string{appVersionName}
	res, err := appManagerClient.DescribeAppVersions(ctx, &req)
	if err != nil {
		return versionId, err
	}
	if res.TotalCount == 0 {
		createReq := pb.CreateAppVersionRequest{}
		createReq.AppId = utils.ToProtoString(appId)
		createReq.Owner = utils.ToProtoString(owner)
		createReq.Name = utils.ToProtoString(appVersionName)
		createReq.PackageName = utils.ToProtoString(packageName)
		createReq.Description = utils.ToProtoString(description)

		createRes, err := appManagerClient.CreateAppVersion(ctx, &createReq)
		if err != nil {
			return versionId, err
		}
		versionId = createRes.GetAppVersion().GetVersionId().GetValue()
		return versionId, err
	} else {
		modifyReq := pb.ModifyAppVersionRequest{}
		modifyReq.VersionId = res.AppVersionSet[0].VersionId
		modifyReq.PackageName = utils.ToProtoString(packageName)
		modifyReq.Description = utils.ToProtoString(description)

		modifyRes, err := appManagerClient.ModifyAppVersion(ctx, &modifyReq)
		if err != nil {
			return versionId, err
		}
		versionId = modifyRes.GetAppVersion().GetVersionId().GetValue()
		return versionId, err
	}
}
<|endoftext|>


<commit_before><commit_msg>Implement method for adding tracks to a playlist<commit_after><|endoftext|>


<commit_before>package outbound

import (
	"crypto/md5"
	"crypto/rand"
	"net"
	"sync"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/common/alloc"
	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/internal"
	"github.com/v2ray/v2ray-core/proxy/vmess/protocol"
	"github.com/v2ray/v2ray-core/proxy/vmess/protocol/user"
	"github.com/v2ray/v2ray-core/transport/ray"
)

type VMessOutboundHandler struct {
	receiverManager *ReceiverManager
	space           app.Space
}

func (this *VMessOutboundHandler) Dispatch(firstPacket v2net.Packet, ray ray.OutboundRay) error {
	vNextAddress, vNextUser := this.receiverManager.PickReceiver()

	command := protocol.CmdTCP
	if firstPacket.Destination().IsUDP() {
		command = protocol.CmdUDP
	}
	request := &protocol.VMessRequest{
		Version: protocol.Version,
		User:    vNextUser,
		Command: command,
		Address: firstPacket.Destination().Address(),
		Port:    firstPacket.Destination().Port(),
	}

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()                             // Buffer is released after communication finishes.
	v2net.ReadAllBytes(rand.Reader, buffer.Value[:36]) // 16 + 16 + 4
	request.RequestIV = buffer.Value[:16]
	request.RequestKey = buffer.Value[16:32]
	request.ResponseHeader = buffer.Value[32:36]

	return startCommunicate(request, vNextAddress, ray, firstPacket)
}

func startCommunicate(request *protocol.VMessRequest, dest v2net.Destination, ray ray.OutboundRay, firstPacket v2net.Packet) error {
	var destIp net.IP
	if dest.Address().IsIPv4() || dest.Address().IsIPv6() {
		destIp = dest.Address().IP()
	} else {
		ips, err := net.LookupIP(dest.Address().Domain())
		if err != nil {
			return err
		}
		destIp = ips[0]
	}
	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   destIp,
		Port: int(dest.Port()),
	})
	if err != nil {
		log.Error("Failed to open %s: %v", dest.String(), err)
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return err
	}
	log.Info("VMessOut: Tunneling request to %s via %s", request.Address.String(), dest.String())

	defer conn.Close()

	input := ray.OutboundInput()
	output := ray.OutboundOutput()

	var requestFinish, responseFinish sync.Mutex
	requestFinish.Lock()
	responseFinish.Lock()

	go handleRequest(conn, request, firstPacket, input, &requestFinish)
	go handleResponse(conn, request, output, &responseFinish, (request.Command == protocol.CmdUDP))

	requestFinish.Lock()
	conn.CloseWrite()
	responseFinish.Lock()
	return nil
}

func handleRequest(conn net.Conn, request *protocol.VMessRequest, firstPacket v2net.Packet, input <-chan *alloc.Buffer, finish *sync.Mutex) {
	defer finish.Unlock()
	aesStream, err := v2crypto.NewAesEncryptionStream(request.RequestKey[:], request.RequestIV[:])
	if err != nil {
		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
		return
	}
	encryptRequestWriter := v2crypto.NewCryptionWriter(aesStream, conn)

	buffer := alloc.NewBuffer().Clear()
	defer buffer.Release()
	buffer, err = request.ToBytes(user.NewTimeHash(user.HMACHash{}), user.GenerateRandomInt64InRange, buffer)
	if err != nil {
		log.Error("VMessOut: Failed to serialize VMess request: %v", err)
		return
	}

	// Send first packet of payload together with request, in favor of small requests.
	firstChunk := firstPacket.Chunk()
	moreChunks := firstPacket.MoreChunks()

	for firstChunk == nil && moreChunks {
		firstChunk, moreChunks = <-input
	}

	if firstChunk == nil && !moreChunks {
		log.Warning("VMessOut: Nothing to send. Existing...")
		return
	}

	aesStream.XORKeyStream(firstChunk.Value, firstChunk.Value)
	buffer.Append(firstChunk.Value)
	firstChunk.Release()

	_, err = conn.Write(buffer.Value)
	if err != nil {
		log.Error("VMessOut: Failed to write VMess request: %v", err)
		return
	}

	if moreChunks {
		v2net.ChanToWriter(encryptRequestWriter, input)
	}
	return
}

func headerMatch(request *protocol.VMessRequest, responseHeader []byte) bool {
	return ((request.ResponseHeader[0] ^ request.ResponseHeader[1]) == responseHeader[0]) &&
		((request.ResponseHeader[2] ^ request.ResponseHeader[3]) == responseHeader[1])
}

func handleResponse(conn net.Conn, request *protocol.VMessRequest, output chan<- *alloc.Buffer, finish *sync.Mutex, isUDP bool) {
	defer finish.Unlock()
	defer close(output)
	responseKey := md5.Sum(request.RequestKey[:])
	responseIV := md5.Sum(request.RequestIV[:])

	aesStream, err := v2crypto.NewAesDecryptionStream(responseKey[:], responseIV[:])
	if err != nil {
		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
		return
	}
	decryptResponseReader := v2crypto.NewCryptionReader(aesStream, conn)

	buffer, err := v2net.ReadFrom(decryptResponseReader, nil)
	if err != nil {
		log.Error("VMessOut: Failed to read VMess response (%d bytes): %v", buffer.Len(), err)
		buffer.Release()
		return
	}
	if buffer.Len() < 4 || !headerMatch(request, buffer.Value[:2]) {
		log.Warning("VMessOut: unexepcted response header. The connection is probably hijacked.")
		return
	}
	log.Info("VMessOut received %d bytes from %s", buffer.Len()-4, conn.RemoteAddr().String())

	responseBegin := 4
	if buffer.Value[2] != 0 {
		dataLen := int(buffer.Value[3])
		if buffer.Len() < dataLen+4 { // Rare case
			diffBuffer := make([]byte, dataLen+4-buffer.Len())
			v2net.ReadAllBytes(decryptResponseReader, diffBuffer)
			buffer.Append(diffBuffer)
		}
		command := buffer.Value[2]
		data := buffer.Value[4 : 4+dataLen]
		go handleCommand(command, data)
		responseBegin = 4 + dataLen
	}

	buffer.SliceFrom(responseBegin)
	output <- buffer

	if !isUDP {
		v2net.ReaderToChan(output, decryptResponseReader)
	}

	return
}

func init() {
	if err := internal.RegisterOutboundConnectionHandlerFactory("vmess", func(space app.Space, rawConfig interface{}) (proxy.OutboundConnectionHandler, error) {
		vOutConfig := rawConfig.(Config)
		return &VMessOutboundHandler{
			space:           space,
			receiverManager: NewReceiverManager(vOutConfig.Receivers()),
		}, nil
	}); err != nil {
		panic(err)
	}
}
<commit_msg>First step to reduce number of bytes for response validation<commit_after>package outbound

import (
	"crypto/md5"
	"crypto/rand"
	"net"
	"sync"

	"github.com/v2ray/v2ray-core/app"
	"github.com/v2ray/v2ray-core/common/alloc"
	v2crypto "github.com/v2ray/v2ray-core/common/crypto"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
	"github.com/v2ray/v2ray-core/proxy"
	"github.com/v2ray/v2ray-core/proxy/internal"
	"github.com/v2ray/v2ray-core/proxy/vmess/protocol"
	"github.com/v2ray/v2ray-core/proxy/vmess/protocol/user"
	"github.com/v2ray/v2ray-core/transport/ray"
)

type VMessOutboundHandler struct {
	receiverManager *ReceiverManager
	space           app.Space
}

func (this *VMessOutboundHandler) Dispatch(firstPacket v2net.Packet, ray ray.OutboundRay) error {
	vNextAddress, vNextUser := this.receiverManager.PickReceiver()

	command := protocol.CmdTCP
	if firstPacket.Destination().IsUDP() {
		command = protocol.CmdUDP
	}
	request := &protocol.VMessRequest{
		Version: protocol.Version,
		User:    vNextUser,
		Command: command,
		Address: firstPacket.Destination().Address(),
		Port:    firstPacket.Destination().Port(),
	}

	buffer := alloc.NewSmallBuffer()
	defer buffer.Release()                             // Buffer is released after communication finishes.
	v2net.ReadAllBytes(rand.Reader, buffer.Value[:33]) // 16 + 16 + 1
	buffer.Value[33] = 0
	buffer.Value[34] = 0
	buffer.Value[35] = 0
	request.RequestIV = buffer.Value[:16]
	request.RequestKey = buffer.Value[16:32]
	request.ResponseHeader = buffer.Value[32:36]

	return startCommunicate(request, vNextAddress, ray, firstPacket)
}

func startCommunicate(request *protocol.VMessRequest, dest v2net.Destination, ray ray.OutboundRay, firstPacket v2net.Packet) error {
	var destIp net.IP
	if dest.Address().IsIPv4() || dest.Address().IsIPv6() {
		destIp = dest.Address().IP()
	} else {
		ips, err := net.LookupIP(dest.Address().Domain())
		if err != nil {
			return err
		}
		destIp = ips[0]
	}
	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   destIp,
		Port: int(dest.Port()),
	})
	if err != nil {
		log.Error("Failed to open %s: %v", dest.String(), err)
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return err
	}
	log.Info("VMessOut: Tunneling request to %s via %s", request.Address.String(), dest.String())

	defer conn.Close()

	input := ray.OutboundInput()
	output := ray.OutboundOutput()

	var requestFinish, responseFinish sync.Mutex
	requestFinish.Lock()
	responseFinish.Lock()

	go handleRequest(conn, request, firstPacket, input, &requestFinish)
	go handleResponse(conn, request, output, &responseFinish, (request.Command == protocol.CmdUDP))

	requestFinish.Lock()
	conn.CloseWrite()
	responseFinish.Lock()
	return nil
}

func handleRequest(conn net.Conn, request *protocol.VMessRequest, firstPacket v2net.Packet, input <-chan *alloc.Buffer, finish *sync.Mutex) {
	defer finish.Unlock()
	aesStream, err := v2crypto.NewAesEncryptionStream(request.RequestKey[:], request.RequestIV[:])
	if err != nil {
		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
		return
	}
	encryptRequestWriter := v2crypto.NewCryptionWriter(aesStream, conn)

	buffer := alloc.NewBuffer().Clear()
	defer buffer.Release()
	buffer, err = request.ToBytes(user.NewTimeHash(user.HMACHash{}), user.GenerateRandomInt64InRange, buffer)
	if err != nil {
		log.Error("VMessOut: Failed to serialize VMess request: %v", err)
		return
	}

	// Send first packet of payload together with request, in favor of small requests.
	firstChunk := firstPacket.Chunk()
	moreChunks := firstPacket.MoreChunks()

	for firstChunk == nil && moreChunks {
		firstChunk, moreChunks = <-input
	}

	if firstChunk == nil && !moreChunks {
		log.Warning("VMessOut: Nothing to send. Existing...")
		return
	}

	aesStream.XORKeyStream(firstChunk.Value, firstChunk.Value)
	buffer.Append(firstChunk.Value)
	firstChunk.Release()

	_, err = conn.Write(buffer.Value)
	if err != nil {
		log.Error("VMessOut: Failed to write VMess request: %v", err)
		return
	}

	if moreChunks {
		v2net.ChanToWriter(encryptRequestWriter, input)
	}
	return
}

func headerMatch(request *protocol.VMessRequest, responseHeader []byte) bool {
	return (request.ResponseHeader[0] == responseHeader[0])
}

func handleResponse(conn net.Conn, request *protocol.VMessRequest, output chan<- *alloc.Buffer, finish *sync.Mutex, isUDP bool) {
	defer finish.Unlock()
	defer close(output)
	responseKey := md5.Sum(request.RequestKey[:])
	responseIV := md5.Sum(request.RequestIV[:])

	aesStream, err := v2crypto.NewAesDecryptionStream(responseKey[:], responseIV[:])
	if err != nil {
		log.Error("VMessOut: Failed to create AES encryption stream: %v", err)
		return
	}
	decryptResponseReader := v2crypto.NewCryptionReader(aesStream, conn)

	buffer, err := v2net.ReadFrom(decryptResponseReader, nil)
	if err != nil {
		log.Error("VMessOut: Failed to read VMess response (%d bytes): %v", buffer.Len(), err)
		buffer.Release()
		return
	}
	if buffer.Len() < 4 || !headerMatch(request, buffer.Value[:2]) {
		log.Warning("VMessOut: unexepcted response header. The connection is probably hijacked.")
		return
	}
	log.Info("VMessOut received %d bytes from %s", buffer.Len()-4, conn.RemoteAddr().String())

	responseBegin := 4
	if buffer.Value[2] != 0 {
		dataLen := int(buffer.Value[3])
		if buffer.Len() < dataLen+4 { // Rare case
			diffBuffer := make([]byte, dataLen+4-buffer.Len())
			v2net.ReadAllBytes(decryptResponseReader, diffBuffer)
			buffer.Append(diffBuffer)
		}
		command := buffer.Value[2]
		data := buffer.Value[4 : 4+dataLen]
		go handleCommand(command, data)
		responseBegin = 4 + dataLen
	}

	buffer.SliceFrom(responseBegin)
	output <- buffer

	if !isUDP {
		v2net.ReaderToChan(output, decryptResponseReader)
	}

	return
}

func init() {
	if err := internal.RegisterOutboundConnectionHandlerFactory("vmess", func(space app.Space, rawConfig interface{}) (proxy.OutboundConnectionHandler, error) {
		vOutConfig := rawConfig.(Config)
		return &VMessOutboundHandler{
			space:           space,
			receiverManager: NewReceiverManager(vOutConfig.Receivers()),
		}, nil
	}); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  *redis.Client
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				var d time.Duration
				if d, err = time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			jobRequest, err := et.Infos().JobRequestWithEvent(event)
			if err != nil {
				s.log.Warnf("[scheduler] Could not encode realtime event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			_, err = s.broker.PushJob(jobRequest)
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			s.client.ZRem(SchedKey, results[0])
			return err
		}
		switch t := t.(type) {
		case *EventTrigger: // Debounced
			job := t.Infos().JobRequest()
			job.Debounced = true
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		case *AtTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		return s.client.HSet(hKey, t.ID(), t.Infos().Arguments).Err()
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	err := couchdb.ForeachDocs(db, consts.Triggers, func(data []byte) error {
		var t *TriggerInfos
		if err := json.Unmarshal(data, &t); err != nil {
			return err
		}
		infos = append(infos, t)
		return nil
	})
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

// RebuildRedis puts all the triggers in redis (idempotent)
func (s *RedisScheduler) RebuildRedis(domain string) error {
	triggers, err := s.GetAll(domain)
	if err != nil {
		return err
	}
	for _, t := range triggers {
		if err = s.addToRedis(t, time.Now()); err != nil {
			return err
		}
	}
	return nil
}

var _ Scheduler = &RedisScheduler{}
<commit_msg>Do not delete in redis a trigger when couchdb is not available<commit_after>package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/permissions"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/go-redis/redis"
	"github.com/sirupsen/logrus"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// eventLoopSize is the number of goroutines handling @events and triggering
// jobs.
const eventLoopSize = 50

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], s[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	broker  jobs.Broker
	client  *redis.Client
	closed  chan struct{}
	stopped chan struct{}
	log     *logrus.Entry
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client:  client,
		log:     logger.WithNamespace("scheduler-redis"),
		stopped: make(chan struct{}),
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + "/" + infos.TID
}

func eventsKey(domain string) string {
	return "events-" + domain
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.broker = b
	s.closed = make(chan struct{})
	s.startEventDispatcher()
	go s.pollLoop()
	return nil
}

func (s *RedisScheduler) pollLoop() {
	ticker := time.NewTicker(pollInterval)
	for {
		select {
		case <-s.closed:
			ticker.Stop()
			s.stopped <- struct{}{}
			return
		case <-ticker.C:
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				s.log.Warnf("[scheduler] Failed to poll redis: %s", err)
			}
		}
	}
}

func (s *RedisScheduler) startEventDispatcher() {
	eventsCh := make(chan *realtime.Event, 100)
	go func() {
		c := realtime.GetHub().SubscribeLocalAll()
		defer func() {
			c.Close()
			close(eventsCh)
		}()
		for {
			select {
			case <-s.closed:
				return
			case event := <-c.Channel:
				eventsCh <- event
			}
		}
	}()
	for i := 0; i < eventLoopSize; i++ {
		go s.eventLoop(eventsCh)
	}
}

func (s *RedisScheduler) eventLoop(eventsCh <-chan *realtime.Event) {
	for event := range eventsCh {
		key := eventsKey(event.Domain)
		m, err := s.client.HGetAll(key).Result()
		if err != nil {
			s.log.Errorf("[scheduler] Could not fetch redis set %s: %s",
				key, err.Error())
			continue
		}
		for triggerID, args := range m {
			rule, err := permissions.UnmarshalRuleString(args)
			if err != nil {
				s.log.Warnf("[scheduler] Coud not unmarshal rule %s: %s",
					key, err.Error())
				continue
			}
			if !eventMatchPermission(event, &rule) {
				continue
			}
			t, err := s.Get(event.Domain, triggerID)
			if err != nil {
				s.log.Warnf("[scheduler] Could not fetch @event trigger %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			et := t.(*EventTrigger)
			if et.infos.Debounce != "" {
				var d time.Duration
				if d, err = time.ParseDuration(et.infos.Debounce); err == nil {
					timestamp := time.Now().Add(d)
					s.client.ZAddNX(TriggersKey, redis.Z{
						Score:  float64(timestamp.UTC().Unix()),
						Member: redisKey(t.Infos()),
					})
					continue
				}
				s.log.Warnf("[scheduler] Trigger %s %s has an invalid debounce: %s",
					et.infos.Domain, et.infos.TID, et.infos.Debounce)
			}
			jobRequest, err := et.Infos().JobRequestWithEvent(event)
			if err != nil {
				s.log.Warnf("[scheduler] Could not encode realtime event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
			_, err = s.broker.PushJob(jobRequest)
			if err != nil {
				s.log.Warnf("[scheduler] Could not push job trigger by event %s %s: %s",
					event.Domain, triggerID, err.Error())
				continue
			}
		}
	}
}

// Shutdown the scheduling of triggers
func (s *RedisScheduler) Shutdown(ctx context.Context) error {
	if s.closed == nil {
		return nil
	}
	fmt.Print("  shutting down redis scheduler...")
	close(s.closed)
	select {
	case <-ctx.Done():
		fmt.Println("failed: ", ctx.Err())
		return ctx.Err()
	case <-s.stopped:
		fmt.Println("ok.")
	}
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), "/", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			if err == ErrNotFoundTrigger {
				s.client.ZRem(SchedKey, results[0])
			}
			return err
		}
		switch t := t.(type) {
		case *EventTrigger: // Debounced
			job := t.Infos().JobRequest()
			job.Debounced = true
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
		case *AtTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			if err = s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Infos().JobRequest()
			if _, err = s.broker.PushJob(job); err != nil {
				return err
			}
			score, err := strconv.ParseInt(results[1].(string), 10, 64)
			var prev time.Time
			if err != nil {
				prev = time.Now()
			} else {
				prev = time.Unix(score, 0)
			}
			if err := s.addToRedis(t, prev); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t, time.Now())
}

func (s *RedisScheduler) addToRedis(t Trigger, prev time.Time) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *EventTrigger:
		hKey := eventsKey(t.Infos().Domain)
		return s.client.HSet(hKey, t.ID(), t.Infos().Arguments).Err()
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(prev)
		now := time.Now()
		if timestamp.Before(now) {
			timestamp = t.NextExecution(now)
		}
	default:
		return errors.New("Not implemented yet")
	}
	pipe := s.client.Pipeline()
	pipe.ZAdd(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
	pipe.ZRem(SchedKey, redisKey(t.Infos()))
	_, err := pipe.Exec()
	return err
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	switch t.(type) {
	case *EventTrigger:
		return s.client.HDel(eventsKey(t.Infos().Domain), t.ID()).Err()
	case *AtTrigger, *CronTrigger:
		pipe := s.client.Pipeline()
		pipe.ZRem(TriggersKey, redisKey(t.Infos()))
		pipe.ZRem(SchedKey, redisKey(t.Infos()))
		_, err := pipe.Exec()
		return err
	}
	return nil
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	err := couchdb.ForeachDocs(db, consts.Triggers, func(data []byte) error {
		var t *TriggerInfos
		if err := json.Unmarshal(data, &t); err != nil {
			return err
		}
		infos = append(infos, t)
		return nil
	})
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

// RebuildRedis puts all the triggers in redis (idempotent)
func (s *RedisScheduler) RebuildRedis(domain string) error {
	triggers, err := s.GetAll(domain)
	if err != nil {
		return err
	}
	for _, t := range triggers {
		if err = s.addToRedis(t, time.Now()); err != nil {
			return err
		}
	}
	return nil
}

var _ Scheduler = &RedisScheduler{}
<|endoftext|>