<commit_before>// Command deploy-agent tracks single configuration against registry, updates
// local state as required and runs deploy script.
package main

import (
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/artyom/autoflags"
	"github.com/artyom/logger"
	"github.com/artyom/untar"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

func main() {
	args := &mainArgs{
		State:  "state.json",
		Key:    "id_ecdsa",
		Addr:   "localhost:2022",
		Dir:    ".",
		Script: "./deploy.sh",
	}
	autoflags.Define(args)
	flag.Parse()
	log := log.New(os.Stderr, "", log.LstdFlags)
	if err := run(args, log); err != nil {
		log.Fatal(err)
	}
}

type mainArgs struct {
	State    string `flag:"state,file to save state to"`
	Name     string `flag:"name,configuration to track"`
	Key      string `flag:"key,ssh private key to use"`
	Addr     string `flag:"addr,registry address (host:port)"`
	Fp       string `flag:"fp,registry server key fingerprint"`
	Dir      string `flag:"dir,directory to store downloaded and unpacked files"`
	Script   string `flag:"script,script to run on deploys"`
	CleanOld bool   `flag:"cleanold,try to remove previous state unpacked files after switching state"`
	Verbose  bool   `flag:"v,be more chatty about what's happening"`
}

func run(args *mainArgs, log logger.Interface) error {
	if args.Name == "" {
		return errors.New("empty configuration name")
	}
	if args.State == "" {
		return errors.New("no state file set")
	}
	cfg, err := newClientConfig(args.Key, args.Fp)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go callOnSignal(func(s os.Signal) { log.Println(s); cancel() }, syscall.SIGINT, syscall.SIGTERM)
	for {
		if err := cycle(ctx, args, cfg, log); err != nil {
			log.Println(err)
		}
		select {
		case <-time.After(15 * time.Second):
		case <-ctx.Done():
			return nil
		}
	}
}

func cycle(ctx context.Context, args *mainArgs, cfg *ssh.ClientConfig, log logger.Interface) error {
	state, err := readLocalConfiguration(args.State)
	switch {
	case err == nil:
	case os.IsNotExist(err):
		state = &configuration{}
	default:
		return err
	}
	client, err := ssh.Dial("tcp", args.Addr, cfg)
	if err != nil {
		return err
	}
	defer client.Close()
	sconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sconn.Close()
	newState, err := readRemoteConfiguration(path.Join("configs", args.Name+".json"), sconn)
	if err != nil {
		return err
	}
	if len(newState.Hash) != 64 {
		return errors.Errorf("invalid state hash value: %q, expecting sha256 sum", newState.Hash)
	}
	if newState.Hash == state.Hash {
		return nil
	}
	if args.Verbose {
		log.Println("configuration update available:", newState.Hash)
		for i, l := range newState.Layers {
			log.Printf("\tcomponent %d: %q, version: %q", i+1, l.Name, l.Version)
		}
	}
	if err := downloadMissing(ctx, sconn, filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	sconn.Close()
	if err := unpackLayers(ctx, filepath.Join(args.Dir, unpackedDir, filepath.Base(newState.Hash)),
		filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	if args.Script != "" {
		if err := runDeploy(ctx, args.Script, filepath.Join(args.Dir, unpackedDir), state, newState); err != nil {
			return err
		}
	}
	if err := saveState(args.State, 0644, newState); err != nil {
		return err
	}
	if args.Verbose {
		log.Println("switched to new configuration:", newState.Hash)
	}
	if err := cleanCache(filepath.Join(args.Dir, cacheDir), newState.Layers, time.Hour); err != nil {
		log.Println("error cleaning cache directory:", err)
	}
	if args.CleanOld && state.Hash != "" {
		oldDir := filepath.Join(args.Dir, unpackedDir, filepath.Base(state.Hash))
		if err := os.RemoveAll(oldDir); err != nil {
			log.Println("failed to remove old state directory:", err)
		}
		_ = os.Remove(oldDir + ".ok")
	}
	return nil
}

func cleanCache(dir string, layers []layer, age time.Duration) error {
	keep := make(map[string]struct{}, len(layers))
	for _, l := range layers {
		keep[l.Hash] = struct{}{}
	}
	fn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.WithStack(err)
		}
		if _, ok := keep[filepath.Base(path)]; ok || !info.Mode().IsRegular() {
			return nil
		}
		if strings.HasPrefix(filepath.Base(path), tempDlPrefix) {
			return errors.WithStack(os.Remove(path))
		}
		if time.Now().Sub(info.ModTime()) < age {
			return nil
		}
		return errors.WithStack(os.Remove(path))
	}
	return filepath.Walk(dir, fn)
}

func runDeploy(ctx context.Context, script, dir string, oldState, newState *configuration) error {
	stateFile, err := saveStateTemp("", 0644, newState)
	if err != nil {
		return err
	}
	defer os.Remove(stateFile)
	cmd := exec.CommandContext(ctx, script)
	cmd.Env = []string{
		"OLDID=" + oldState.Hash,
		"OLDROOT=" + filepath.Join(dir, oldState.Hash),
		"NEWID=" + newState.Hash,
		"NEWROOT=" + filepath.Join(dir, newState.Hash),
		"STATEFILE=" + stateFile,
	}
	for _, name := range []string{"PATH", "HOME", "TMPDIR", "USER", "LOGNAME", "LANG", "LC_ALL", "SHELL"} {
		if val, ok := os.LookupEnv(name); ok {
			cmd.Env = append(cmd.Env, name+"="+val)
		}
	}
	for i, l := range newState.Layers {
		cmd.Env = append(cmd.Env,
			fmt.Sprintf("LAYER_%d_NAME=%s", i, l.Name),
			fmt.Sprintf("LAYER_%d_VERSION=%s", i, l.Version),
		)
	}
	return errors.Wrapf(cmd.Run(), "%q call failure", script)
}

func saveStateTemp(dir string, mode os.FileMode, state *configuration) (string, error) {
	f, err := ioutil.TempFile(dir, "json-state-")
	if err != nil {
		return "", err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	if err := f.Chmod(mode); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	return f.Name(), f.Close()
}

func saveState(dst string, mode os.FileMode, state *configuration) error {
	f, err := ioutil.TempFile(filepath.Dir(dst), "temp-state-")
	if err != nil {
		return err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		return err
	}
	if err := f.Chmod(mode); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), dst)
}

func unpackLayers(ctx context.Context, dst, filesDir string, layers []layer) error {
	touchfile := dst + ".ok"
	if _, err := os.Stat(touchfile); err == nil {
		if _, err := os.Stat(dst); err == nil {
			return nil
		}
	}
	if err := os.MkdirAll(dst, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		srcName := filepath.Join(filesDir, filepath.Base(l.Hash))
		if err := untarTo(srcName, dst); err != nil {
			return err
		}
	}
	_ = touch(touchfile)
	return nil
}

func untarTo(name, dir string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	gr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gr.Close()
	defer syscall.Umask(syscall.Umask(0))
	return untar.Untar(gr, dir)
}

func downloadMissing(ctx context.Context, sconn *sftp.Client, dir string, layers []layer) error {
	if err := os.MkdirAll(dir, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		dstName := filepath.Join(dir, filepath.Base(l.Hash))
		if _, err := os.Stat(dstName); err == nil {
			continue
		}
		if err := downloadFile(sconn, dstName, l.File, l.Hash); err != nil {
			return err
		}
	}
	return nil
}

func downloadFile(sconn *sftp.Client, dst, src, contentHash string) error {
	dstFile, err := ioutil.TempFile(filepath.Dir(dst), tempDlPrefix)
	if err != nil {
		return err
	}
	defer dstFile.Close()
	defer os.Remove(dstFile.Name())
	srcFile, err := sconn.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	h := sha256.New()
	tr := io.TeeReader(srcFile, h)
	if _, err := io.Copy(dstFile, tr); err != nil {
		return err
	}
	if contentHash != fmt.Sprintf("%x", h.Sum(nil)) {
		return errors.Errorf("downloaded file hash mismatch: %v", src)
	}
	if err := dstFile.Close(); err != nil {
		return err
	}
	return os.Rename(dstFile.Name(), dst)
}

func readRemoteConfiguration(name string, sconn *sftp.Client) (*configuration, error) {
	f, err := sconn.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func readLocalConfiguration(name string) (*configuration, error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func decodeConfiguration(r io.Reader) (*configuration, error) {
	var cfg configuration
	if err := json.NewDecoder(r).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func newClientConfig(keyfile, fingerprint string) (*ssh.ClientConfig, error) {
	if !strings.HasPrefix(fingerprint, "SHA256:") {
		return nil, errors.New("invalid server key fingerprint, expecting SHA256 fingerprint")
	}
	key, err := readKey(keyfile)
	if err != nil {
		return nil, err
	}
	config := &ssh.ClientConfig{
		User: "deploy-agent",
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(key)},
		Timeout: 30 * time.Second,
	}
	return config, nil
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}

func touch(name string) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	return f.Close()
}

func callOnSignal(fn func(os.Signal), sig ...os.Signal) {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, sig...)
	defer signal.Stop(sigCh)
	fn(<-sigCh)
}

type configuration struct {
	Name   string
	Mtime  time.Time
	Hash   string
	Layers []layer
}

type layer struct {
	Name    string
	Version string
	File    string
	Hash    string
	Ctime   time.Time
}

const (
	cacheDir     = "cache"
	unpackedDir  = "unpacked"
	tempDlPrefix = "temp-"
)
<commit_msg>deploy-agent: clean cache at the start of iteration cycle<commit_after>// Command deploy-agent tracks single configuration against registry, updates
// local state as required and runs deploy script.
package main

import (
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/artyom/autoflags"
	"github.com/artyom/logger"
	"github.com/artyom/untar"
	"github.com/pkg/errors"
	"github.com/pkg/sftp"
)

func main() {
	args := &mainArgs{
		State:  "state.json",
		Key:    "id_ecdsa",
		Addr:   "localhost:2022",
		Dir:    ".",
		Script: "./deploy.sh",
	}
	autoflags.Define(args)
	flag.Parse()
	log := log.New(os.Stderr, "", log.LstdFlags)
	if err := run(args, log); err != nil {
		log.Fatal(err)
	}
}

type mainArgs struct {
	State    string `flag:"state,file to save state to"`
	Name     string `flag:"name,configuration to track"`
	Key      string `flag:"key,ssh private key to use"`
	Addr     string `flag:"addr,registry address (host:port)"`
	Fp       string `flag:"fp,registry server key fingerprint"`
	Dir      string `flag:"dir,directory to store downloaded and unpacked files"`
	Script   string `flag:"script,script to run on deploys"`
	CleanOld bool   `flag:"cleanold,try to remove previous state unpacked files after switching state"`
	Verbose  bool   `flag:"v,be more chatty about what's happening"`
}

func run(args *mainArgs, log logger.Interface) error {
	if args.Name == "" {
		return errors.New("empty configuration name")
	}
	if args.State == "" {
		return errors.New("no state file set")
	}
	cfg, err := newClientConfig(args.Key, args.Fp)
	if err != nil {
		return err
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go callOnSignal(func(s os.Signal) { log.Println(s); cancel() }, syscall.SIGINT, syscall.SIGTERM)
	for {
		if err := cycle(ctx, args, cfg, log); err != nil {
			log.Println(err)
		}
		select {
		case <-time.After(15 * time.Second):
		case <-ctx.Done():
			return nil
		}
	}
}

func cycle(ctx context.Context, args *mainArgs, cfg *ssh.ClientConfig, log logger.Interface) error {
	state, err := readLocalConfiguration(args.State)
	switch {
	case err == nil:
	case os.IsNotExist(err):
		state = &configuration{}
	default:
		return err
	}
	if err := cleanCache(filepath.Join(args.Dir, cacheDir), state.Layers, time.Hour); err != nil {
		log.Println("error cleaning cache directory:", err)
	}
	client, err := ssh.Dial("tcp", args.Addr, cfg)
	if err != nil {
		return err
	}
	defer client.Close()
	sconn, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sconn.Close()
	newState, err := readRemoteConfiguration(path.Join("configs", args.Name+".json"), sconn)
	if err != nil {
		return err
	}
	if len(newState.Hash) != 64 {
		return errors.Errorf("invalid state hash value: %q, expecting sha256 sum", newState.Hash)
	}
	if newState.Hash == state.Hash {
		return nil
	}
	if args.Verbose {
		log.Println("configuration update available:", newState.Hash)
		for i, l := range newState.Layers {
			log.Printf("\tcomponent %d: %q, version: %q", i+1, l.Name, l.Version)
		}
	}
	if err := downloadMissing(ctx, sconn, filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	sconn.Close()
	if err := unpackLayers(ctx, filepath.Join(args.Dir, unpackedDir, filepath.Base(newState.Hash)),
		filepath.Join(args.Dir, cacheDir), newState.Layers); err != nil {
		return err
	}
	if args.Script != "" {
		if err := runDeploy(ctx, args.Script, filepath.Join(args.Dir, unpackedDir), state, newState); err != nil {
			return err
		}
	}
	if err := saveState(args.State, 0644, newState); err != nil {
		return err
	}
	if args.Verbose {
		log.Println("switched to new configuration:", newState.Hash)
	}
	if args.CleanOld && state.Hash != "" {
		oldDir := filepath.Join(args.Dir, unpackedDir, filepath.Base(state.Hash))
		if err := os.RemoveAll(oldDir); err != nil {
			log.Println("failed to remove old state directory:", err)
		}
		_ = os.Remove(oldDir + ".ok")
	}
	return nil
}

func cleanCache(dir string, layers []layer, age time.Duration) error {
	keep := make(map[string]struct{}, len(layers))
	for _, l := range layers {
		keep[l.Hash] = struct{}{}
	}
	fn := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return errors.WithStack(err)
		}
		if _, ok := keep[filepath.Base(path)]; ok || !info.Mode().IsRegular() {
			return nil
		}
		if strings.HasPrefix(filepath.Base(path), tempDlPrefix) {
			return errors.WithStack(os.Remove(path))
		}
		if time.Now().Sub(info.ModTime()) < age {
			return nil
		}
		return errors.WithStack(os.Remove(path))
	}
	return filepath.Walk(dir, fn)
}

func runDeploy(ctx context.Context, script, dir string, oldState, newState *configuration) error {
	stateFile, err := saveStateTemp("", 0644, newState)
	if err != nil {
		return err
	}
	defer os.Remove(stateFile)
	cmd := exec.CommandContext(ctx, script)
	cmd.Env = []string{
		"OLDID=" + oldState.Hash,
		"OLDROOT=" + filepath.Join(dir, oldState.Hash),
		"NEWID=" + newState.Hash,
		"NEWROOT=" + filepath.Join(dir, newState.Hash),
		"STATEFILE=" + stateFile,
	}
	for _, name := range []string{"PATH", "HOME", "TMPDIR", "USER", "LOGNAME", "LANG", "LC_ALL", "SHELL"} {
		if val, ok := os.LookupEnv(name); ok {
			cmd.Env = append(cmd.Env, name+"="+val)
		}
	}
	for i, l := range newState.Layers {
		cmd.Env = append(cmd.Env,
			fmt.Sprintf("LAYER_%d_NAME=%s", i, l.Name),
			fmt.Sprintf("LAYER_%d_VERSION=%s", i, l.Version),
		)
	}
	return errors.Wrapf(cmd.Run(), "%q call failure", script)
}

func saveStateTemp(dir string, mode os.FileMode, state *configuration) (string, error) {
	f, err := ioutil.TempFile(dir, "json-state-")
	if err != nil {
		return "", err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	if err := f.Chmod(mode); err != nil {
		os.Remove(f.Name())
		return "", err
	}
	return f.Name(), f.Close()
}

func saveState(dst string, mode os.FileMode, state *configuration) error {
	f, err := ioutil.TempFile(filepath.Dir(dst), "temp-state-")
	if err != nil {
		return err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	enc := json.NewEncoder(f)
	enc.SetIndent("", "\t")
	if err := enc.Encode(state); err != nil {
		return err
	}
	if err := f.Chmod(mode); err != nil {
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(f.Name(), dst)
}

func unpackLayers(ctx context.Context, dst, filesDir string, layers []layer) error {
	touchfile := dst + ".ok"
	if _, err := os.Stat(touchfile); err == nil {
		if _, err := os.Stat(dst); err == nil {
			return nil
		}
	}
	if err := os.MkdirAll(dst, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		srcName := filepath.Join(filesDir, filepath.Base(l.Hash))
		if err := untarTo(srcName, dst); err != nil {
			return err
		}
	}
	_ = touch(touchfile)
	return nil
}

func untarTo(name, dir string) error {
	f, err := os.Open(name)
	if err != nil {
		return err
	}
	defer f.Close()
	gr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gr.Close()
	defer syscall.Umask(syscall.Umask(0))
	return untar.Untar(gr, dir)
}

func downloadMissing(ctx context.Context, sconn *sftp.Client, dir string, layers []layer) error {
	if err := os.MkdirAll(dir, 0750); err != nil {
		return err
	}
	for _, l := range layers {
		if err := ctx.Err(); err != nil {
			return err
		}
		if len(l.Hash) != 64 {
			return errors.Errorf("invalid length of layer hash: %v", l)
		}
		dstName := filepath.Join(dir, filepath.Base(l.Hash))
		if _, err := os.Stat(dstName); err == nil {
			continue
		}
		if err := downloadFile(sconn, dstName, l.File, l.Hash); err != nil {
			return err
		}
	}
	return nil
}

func downloadFile(sconn *sftp.Client, dst, src, contentHash string) error {
	dstFile, err := ioutil.TempFile(filepath.Dir(dst), tempDlPrefix)
	if err != nil {
		return err
	}
	defer dstFile.Close()
	defer os.Remove(dstFile.Name())
	srcFile, err := sconn.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	h := sha256.New()
	tr := io.TeeReader(srcFile, h)
	if _, err := io.Copy(dstFile, tr); err != nil {
		return err
	}
	if contentHash != fmt.Sprintf("%x", h.Sum(nil)) {
		return errors.Errorf("downloaded file hash mismatch: %v", src)
	}
	if err := dstFile.Close(); err != nil {
		return err
	}
	return os.Rename(dstFile.Name(), dst)
}

func readRemoteConfiguration(name string, sconn *sftp.Client) (*configuration, error) {
	f, err := sconn.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func readLocalConfiguration(name string) (*configuration, error) {
	f, err := os.Open(name)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return decodeConfiguration(f)
}

func decodeConfiguration(r io.Reader) (*configuration, error) {
	var cfg configuration
	if err := json.NewDecoder(r).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

func newClientConfig(keyfile, fingerprint string) (*ssh.ClientConfig, error) {
	if !strings.HasPrefix(fingerprint, "SHA256:") {
		return nil, errors.New("invalid server key fingerprint, expecting SHA256 fingerprint")
	}
	key, err := readKey(keyfile)
	if err != nil {
		return nil, err
	}
	config := &ssh.ClientConfig{
		User: "deploy-agent",
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			if hostFp := ssh.FingerprintSHA256(key); hostFp != fingerprint {
				return errors.Errorf("host key fingerprint mismatch: %v", hostFp)
			}
			return nil
		},
		Auth:    []ssh.AuthMethod{ssh.PublicKeys(key)},
		Timeout: 30 * time.Second,
	}
	return config, nil
}

func readKey(name string) (ssh.Signer, error) {
	privateBytes, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(privateBytes)
}

func touch(name string) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE, 0640)
	if err != nil {
		return err
	}
	return f.Close()
}

func callOnSignal(fn func(os.Signal), sig ...os.Signal) {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, sig...)
	defer signal.Stop(sigCh)
	fn(<-sigCh)
}

type configuration struct {
	Name   string
	Mtime  time.Time
	Hash   string
	Layers []layer
}

type layer struct {
	Name    string
	Version string
	File    string
	Hash    string
	Ctime   time.Time
}

const (
	cacheDir     = "cache"
	unpackedDir  = "unpacked"
	tempDlPrefix = "temp-"
)
<|endoftext|>


<commit_before>package mirror

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/kafka-cg/consumergroup"
)

func (this *Mirror) makePub(c2 *zk.ZkCluster) (sarama.AsyncProducer, error) {
	cf := sarama.NewConfig()
	cf.Metadata.RefreshFrequency = time.Minute * 10
	cf.Metadata.Retry.Max = 3
	cf.Metadata.Retry.Backoff = time.Second * 3

	cf.ChannelBufferSize = 1000

	cf.Producer.Flush.Frequency = time.Second // TODO
	cf.Producer.Flush.Messages = 2000
	cf.Producer.Flush.MaxMessages = 0 // unlimited
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Retry.Backoff = time.Second * 3
	cf.Producer.Retry.Max = 3

	switch this.Compress {
	case "gzip":
		cf.Producer.Compression = sarama.CompressionGZIP

	case "snappy":
		cf.Producer.Compression = sarama.CompressionSnappy
	}
	return sarama.NewAsyncProducer(c2.BrokerList(), cf)
}

func (this *Mirror) makeSub(c1 *zk.ZkCluster, group string, topics []string) (*consumergroup.ConsumerGroup, error) {
	cf := consumergroup.NewConfig()
	cf.Zookeeper.Chroot = c1.Chroot()
	cf.Offsets.CommitInterval = time.Second * 10
	cf.Offsets.ProcessingTimeout = time.Second
	cf.ChannelBufferSize = 100
	cf.Consumer.Return.Errors = true
	cf.OneToOne = false

	sub, err := consumergroup.JoinConsumerGroup(group, topics, c1.ZkZone().ZkAddrList(), cf)
	return sub, err
}
<commit_msg>explicitly config<commit_after>package mirror

import (
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/kafka-cg/consumergroup"
)

func (this *Mirror) makePub(c2 *zk.ZkCluster) (sarama.AsyncProducer, error) {
	cf := sarama.NewConfig()
	cf.Metadata.RefreshFrequency = time.Minute * 10
	cf.Metadata.Retry.Max = 3
	cf.Metadata.Retry.Backoff = time.Second * 3

	cf.ChannelBufferSize = 1000

	cf.Producer.Return.Errors = true
	cf.Producer.Flush.Messages = 2000         // 2000 message in batch
	cf.Producer.Flush.Frequency = time.Second // flush interval
	cf.Producer.Flush.MaxMessages = 0         // unlimited
	cf.Producer.RequiredAcks = sarama.WaitForLocal
	cf.Producer.Retry.Backoff = time.Second * 4
	cf.Producer.Retry.Max = 3
	cf.Net.DialTimeout = time.Second * 30
	cf.Net.WriteTimeout = time.Second * 30
	cf.Net.ReadTimeout = time.Second * 30

	switch this.Compress {
	case "gzip":
		cf.Producer.Compression = sarama.CompressionGZIP

	case "snappy":
		cf.Producer.Compression = sarama.CompressionSnappy
	}
	return sarama.NewAsyncProducer(c2.BrokerList(), cf)
}

func (this *Mirror) makeSub(c1 *zk.ZkCluster, group string, topics []string) (*consumergroup.ConsumerGroup, error) {
	cf := consumergroup.NewConfig()
	cf.Zookeeper.Chroot = c1.Chroot()
	cf.Offsets.CommitInterval = time.Second * 10
	cf.Offsets.ProcessingTimeout = time.Second
	cf.Consumer.Offsets.Initial = sarama.OffsetOldest
	cf.ChannelBufferSize = 256
	cf.Consumer.Return.Errors = true
	cf.OneToOne = false

	sub, err := consumergroup.JoinConsumerGroup(group, topics, c1.ZkZone().ZkAddrList(), cf)
	return sub, err
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// Package app implements a Server object for running the scheduler.
package app

import (
	"context"
	"fmt"
	"net/http"
	"os"
	goruntime "runtime"

	"github.com/spf13/cobra"

	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	genericapifilters "k8s.io/apiserver/pkg/endpoints/filters"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server"
	genericfilters "k8s.io/apiserver/pkg/server/filters"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/apiserver/pkg/server/mux"
	"k8s.io/apiserver/pkg/server/routes"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/events"
	"k8s.io/client-go/tools/leaderelection"
	cliflag "k8s.io/component-base/cli/flag"
	"k8s.io/component-base/cli/globalflag"
	"k8s.io/component-base/configz"
	"k8s.io/component-base/logs"
	"k8s.io/component-base/metrics/legacyregistry"
	"k8s.io/component-base/term"
	"k8s.io/component-base/version"
	"k8s.io/component-base/version/verflag"
	"k8s.io/klog/v2"
	schedulerserverconfig "k8s.io/kubernetes/cmd/kube-scheduler/app/config"
	"k8s.io/kubernetes/cmd/kube-scheduler/app/options"
	"k8s.io/kubernetes/pkg/scheduler"
	kubeschedulerconfig "k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/framework/runtime"
	"k8s.io/kubernetes/pkg/scheduler/metrics/resources"
	"k8s.io/kubernetes/pkg/scheduler/profile"
)

// Option configures a framework.Registry.
type Option func(runtime.Registry) error

// NewSchedulerCommand creates a *cobra.Command object with default parameters and registryOptions
func NewSchedulerCommand(registryOptions ...Option) *cobra.Command {
	opts, err := options.NewOptions()
	if err != nil {
		klog.Fatalf("unable to initialize command options: %v", err)
	}

	namedFlagSets := opts.Flags()
	cmd := &cobra.Command{
		Use: "kube-scheduler",
		Long: `The Kubernetes scheduler is a control plane process which assigns
Pods to Nodes. The scheduler determines which Nodes are valid placements for
each Pod in the scheduling queue according to constraints and available
resources. The scheduler then ranks each valid Node and binds the Pod to a
suitable Node. Multiple different schedulers may be used within a cluster;
kube-scheduler is the reference implementation.
See [scheduling](https://kubernetes.io/docs/concepts/scheduling-eviction/)
for more information about scheduling and the kube-scheduler component.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := opts.Complete(&namedFlagSets); err != nil {
				return err
			}
			if err := runCommand(cmd, opts, registryOptions...); err != nil {
				return err
			}
			return nil
		},
		Args: func(cmd *cobra.Command, args []string) error {
			for _, arg := range args {
				if len(arg) > 0 {
					return fmt.Errorf("%q does not take any arguments, got %q", cmd.CommandPath(), args)
				}
			}
			return nil
		},
	}
	fs := cmd.Flags()
	verflag.AddFlags(namedFlagSets.FlagSet("global"))
	globalflag.AddGlobalFlags(namedFlagSets.FlagSet("global"), cmd.Name())
	for _, f := range namedFlagSets.FlagSets {
		fs.AddFlagSet(f)
	}

	cols, _, _ := term.TerminalSize(cmd.OutOrStdout())
	cliflag.SetUsageAndHelpFunc(cmd, namedFlagSets, cols)

	cmd.MarkFlagFilename("config", "yaml", "yml", "json")

	return cmd
}

// runCommand runs the scheduler.
func runCommand(cmd *cobra.Command, opts *options.Options, registryOptions ...Option) error {
	verflag.PrintAndExitIfRequested()
	cliflag.PrintFlags(cmd.Flags())

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		stopCh := server.SetupSignalHandler()
		<-stopCh
		cancel()
	}()

	cc, sched, err := Setup(ctx, opts, registryOptions...)
	if err != nil {
		return err
	}

	return Run(ctx, cc, sched)
}

// Run executes the scheduler based on the given configuration. It only returns on error or when context is done.
func Run(ctx context.Context, cc *schedulerserverconfig.CompletedConfig, sched *scheduler.Scheduler) error {
	// To help debugging, immediately log version
	klog.V(1).InfoS("Starting Kubernetes Scheduler version", "version", version.Get())

	// Configz registration.
	if cz, err := configz.New("componentconfig"); err == nil {
		cz.Set(cc.ComponentConfig)
	} else {
		return fmt.Errorf("unable to register configz: %s", err)
	}

	// Prepare the event broadcaster.
	cc.EventBroadcaster.StartRecordingToSink(ctx.Done())

	// Setup healthz checks.
	var checks []healthz.HealthChecker
	if cc.ComponentConfig.LeaderElection.LeaderElect {
		checks = append(checks, cc.LeaderElection.WatchDog)
	}

	waitingForLeader := make(chan struct{})
	isLeader := func() bool {
		select {
		case _, ok := <-waitingForLeader:
			// if channel is closed, we are leading
			return !ok
		default:
			// channel is open, we are waiting for a leader
			return false
		}
	}

	// Start up the healthz server.
	if cc.InsecureServing != nil {
		separateMetrics := cc.InsecureMetricsServing != nil
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader, separateMetrics, checks...), nil, nil)
		if err := cc.InsecureServing.Serve(handler, 0, ctx.Done()); err != nil {
			return fmt.Errorf("failed to start healthz server: %v", err)
		}
	}
	if cc.InsecureMetricsServing != nil {
		handler := buildHandlerChain(newMetricsHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader), nil, nil)
		if err := cc.InsecureMetricsServing.Serve(handler, 0, ctx.Done()); err != nil {
			return fmt.Errorf("failed to start metrics server: %v", err)
		}
	}
	if cc.SecureServing != nil {
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader, false, checks...), cc.Authentication.Authenticator, cc.Authorization.Authorizer)
		// TODO: handle stoppedCh returned by c.SecureServing.Serve
		if _, err := cc.SecureServing.Serve(handler, 0, ctx.Done()); err != nil {
			// fail early for secure handlers, removing the old error loop from above
			return fmt.Errorf("failed to start secure server: %v", err)
		}
	}

	// Start all informers.
	cc.InformerFactory.Start(ctx.Done())

	// Wait for all caches to sync before scheduling.
	cc.InformerFactory.WaitForCacheSync(ctx.Done())

	// If leader election is enabled, runCommand via LeaderElector until done and exit.
	if cc.LeaderElection != nil {
		cc.LeaderElection.Callbacks = leaderelection.LeaderCallbacks{
			OnStartedLeading: func(ctx context.Context) {
				close(waitingForLeader)
				sched.Run(ctx)
			},
			OnStoppedLeading: func() {
				select {
				case <-ctx.Done():
					// We were asked to terminate. Exit 0.
					klog.Info("Requested to terminate. Exiting.")
					os.Exit(0)
				default:
					// We lost the lock.
					klog.Exitf("leaderelection lost")
				}
			},
		}
		leaderElector, err := leaderelection.NewLeaderElector(*cc.LeaderElection)
		if err != nil {
			return fmt.Errorf("couldn't create leader elector: %v", err)
		}

		leaderElector.Run(ctx)

		return fmt.Errorf("lost lease")
	}

	// Leader election is disabled, so runCommand inline until done.
	close(waitingForLeader)
	sched.Run(ctx)
	return fmt.Errorf("finished without leader elect")
}

// buildHandlerChain wraps the given handler with the standard filters.
func buildHandlerChain(handler http.Handler, authn authenticator.Request, authz authorizer.Authorizer) http.Handler {
	requestInfoResolver := &apirequest.RequestInfoFactory{}
	failedHandler := genericapifilters.Unauthorized(scheme.Codecs)

	handler = genericapifilters.WithAuthorization(handler, authz, scheme.Codecs)
	handler = genericapifilters.WithAuthentication(handler, authn, failedHandler, nil)
	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericapifilters.WithCacheControl(handler)
	handler = genericfilters.WithHTTPLogging(handler)
	handler = genericfilters.WithPanicRecovery(handler, requestInfoResolver)

	return handler
}

func installMetricHandler(pathRecorderMux *mux.PathRecorderMux, informers informers.SharedInformerFactory, isLeader func() bool) {
	configz.InstallHandler(pathRecorderMux)
	pathRecorderMux.Handle("/metrics", legacyregistry.HandlerWithReset())

	resourceMetricsHandler := resources.Handler(informers.Core().V1().Pods().Lister())
	pathRecorderMux.HandleFunc("/metrics/resources", func(w http.ResponseWriter, req *http.Request) {
		if !isLeader() {
			return
		}
		resourceMetricsHandler.ServeHTTP(w, req)
	})
}

// newMetricsHandler builds a metrics server from the config.
func newMetricsHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, informers informers.SharedInformerFactory, isLeader func() bool) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	installMetricHandler(pathRecorderMux, informers, isLeader)
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
		routes.DebugFlags{}.Install(pathRecorderMux, "v", routes.StringFlagPutHandler(logs.GlogSetter))
	}
	return pathRecorderMux
}

// newHealthzHandler creates a healthz server from the config, and will also
// embed the metrics handler if the healthz and metrics address configurations
// are the same.
func newHealthzHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, informers informers.SharedInformerFactory, isLeader func() bool, separateMetrics bool, checks ...healthz.HealthChecker) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	healthz.InstallHandler(pathRecorderMux, checks...)
	if !separateMetrics {
		installMetricHandler(pathRecorderMux, informers, isLeader)
	}
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
		routes.DebugFlags{}.Install(pathRecorderMux, "v", routes.StringFlagPutHandler(logs.GlogSetter))
	}
	return pathRecorderMux
}

func getRecorderFactory(cc *schedulerserverconfig.CompletedConfig) profile.RecorderFactory {
	return func(name string) events.EventRecorder {
		return cc.EventBroadcaster.NewRecorder(name)
	}
}

// WithPlugin creates an Option based on plugin name and factory. Please don't remove this function: it is used to register out-of-tree plugins,
// hence there are no references to it from the kubernetes scheduler code base.
func WithPlugin(name string, factory runtime.PluginFactory) Option {
	return func(registry runtime.Registry) error {
		return registry.Register(name, factory)
	}
}

// Setup creates a completed config and a scheduler based on the command args and options
func Setup(ctx context.Context, opts *options.Options, outOfTreeRegistryOptions ...Option) (*schedulerserverconfig.CompletedConfig, *scheduler.Scheduler, error) {
	if errs := opts.Validate(); len(errs) > 0 {
		return nil, nil, utilerrors.NewAggregate(errs)
	}

	c, err := opts.Config()
	if err != nil {
		return nil, nil, err
	}

	// Get the completed config
	cc := c.Complete()

	outOfTreeRegistry := make(runtime.Registry)
	for _, option := range outOfTreeRegistryOptions {
		if err := option(outOfTreeRegistry); err != nil {
			return nil, nil, err
		}
	}

	recorderFactory := getRecorderFactory(&cc)
	completedProfiles := make([]kubeschedulerconfig.KubeSchedulerProfile, 0)
	// Create the scheduler.
	sched, err := scheduler.New(cc.Client,
		cc.InformerFactory,
		recorderFactory,
		ctx.Done(),
		scheduler.WithComponentConfigVersion(cc.ComponentConfig.TypeMeta.APIVersion),
		scheduler.WithKubeConfig(cc.KubeConfig),
		scheduler.WithProfiles(cc.ComponentConfig.Profiles...),
		scheduler.WithLegacyPolicySource(cc.LegacyPolicySource),
		scheduler.WithPercentageOfNodesToScore(cc.ComponentConfig.PercentageOfNodesToScore),
		scheduler.WithFrameworkOutOfTreeRegistry(outOfTreeRegistry),
		scheduler.WithPodMaxBackoffSeconds(cc.ComponentConfig.PodMaxBackoffSeconds),
		scheduler.WithPodInitialBackoffSeconds(cc.ComponentConfig.PodInitialBackoffSeconds),
		scheduler.WithExtenders(cc.ComponentConfig.Extenders...),
		scheduler.WithParallelism(cc.ComponentConfig.Parallelism),
		scheduler.WithBuildFrameworkCapturer(func(profile kubeschedulerconfig.KubeSchedulerProfile) {
			// Profiles are processed during Framework instantiation to set default plugins and configurations. Capturing them for logging
			completedProfiles = append(completedProfiles, profile)
		}),
	)
	if err != nil {
		return nil, nil, err
	}
	if err := options.LogOrWriteConfig(opts.WriteConfigTo, &cc.ComponentConfig, completedProfiles); err != nil {
		return nil, nil, err
	}

	return &cc, sched, nil
}
<commit_msg>Fix: set SilenceErrors and SilenceUsage not to change error message<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// Package app implements a Server object for running the scheduler.
package app

import (
	"context"
	"fmt"
	"net/http"
	"os"
	goruntime "runtime"

	"github.com/spf13/cobra"

	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	genericapifilters "k8s.io/apiserver/pkg/endpoints/filters"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server"
	genericfilters "k8s.io/apiserver/pkg/server/filters"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/apiserver/pkg/server/mux"
	"k8s.io/apiserver/pkg/server/routes"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/tools/events"
	"k8s.io/client-go/tools/leaderelection"
	cliflag "k8s.io/component-base/cli/flag"
	"k8s.io/component-base/cli/globalflag"
	"k8s.io/component-base/configz"
	"k8s.io/component-base/logs"
	"k8s.io/component-base/metrics/legacyregistry"
	"k8s.io/component-base/term"
	"k8s.io/component-base/version"
	"k8s.io/component-base/version/verflag"
	"k8s.io/klog/v2"
	schedulerserverconfig "k8s.io/kubernetes/cmd/kube-scheduler/app/config"
	"k8s.io/kubernetes/cmd/kube-scheduler/app/options"
	"k8s.io/kubernetes/pkg/scheduler"
	kubeschedulerconfig "k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/framework/runtime"
	"k8s.io/kubernetes/pkg/scheduler/metrics/resources"
	"k8s.io/kubernetes/pkg/scheduler/profile"
)

// Option configures a framework.Registry.
type Option func(runtime.Registry) error

// NewSchedulerCommand creates a *cobra.Command object with default parameters and registryOptions
func NewSchedulerCommand(registryOptions ...Option) *cobra.Command {
	opts, err := options.NewOptions()
	if err != nil {
		klog.Fatalf("unable to initialize command options: %v", err)
	}

	namedFlagSets := opts.Flags()
	cmd := &cobra.Command{
		Use: "kube-scheduler",
		Long: `The Kubernetes scheduler is a control plane process which assigns
Pods to Nodes. The scheduler determines which Nodes are valid placements for
each Pod in the scheduling queue according to constraints and available
resources. The scheduler then ranks each valid Node and binds the Pod to a
suitable Node. Multiple different schedulers may be used within a cluster;
kube-scheduler is the reference implementation.
See [scheduling](https://kubernetes.io/docs/concepts/scheduling-eviction/)
for more information about scheduling and the kube-scheduler component.`,
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := opts.Complete(&namedFlagSets); err != nil {
				return err
			}
			if err := runCommand(cmd, opts, registryOptions...); err != nil {
				return err
			}
			return nil
		},
		SilenceErrors: true,
		SilenceUsage:  true,
		Args: func(cmd *cobra.Command, args []string) error {
			for _, arg := range args {
				if len(arg) > 0 {
					return fmt.Errorf("%q does not take any arguments, got %q", cmd.CommandPath(), args)
				}
			}
			return nil
		},
	}
	fs := cmd.Flags()
	verflag.AddFlags(namedFlagSets.FlagSet("global"))
	globalflag.AddGlobalFlags(namedFlagSets.FlagSet("global"), cmd.Name())
	for _, f := range namedFlagSets.FlagSets {
		fs.AddFlagSet(f)
	}

	cols, _, _ := term.TerminalSize(cmd.OutOrStdout())
	cliflag.SetUsageAndHelpFunc(cmd, namedFlagSets, cols)

	cmd.MarkFlagFilename("config", "yaml", "yml", "json")

	return cmd
}

// runCommand runs the scheduler.
func runCommand(cmd *cobra.Command, opts *options.Options, registryOptions ...Option) error {
	verflag.PrintAndExitIfRequested()
	cliflag.PrintFlags(cmd.Flags())

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		stopCh := server.SetupSignalHandler()
		<-stopCh
		cancel()
	}()

	cc, sched, err := Setup(ctx, opts, registryOptions...)
	if err != nil {
		return err
	}

	return Run(ctx, cc, sched)
}

// Run executes the scheduler based on the given configuration. It only returns on error or when context is done.
func Run(ctx context.Context, cc *schedulerserverconfig.CompletedConfig, sched *scheduler.Scheduler) error {
	// To help debugging, immediately log version
	klog.V(1).InfoS("Starting Kubernetes Scheduler version", "version", version.Get())

	// Configz registration.
	if cz, err := configz.New("componentconfig"); err == nil {
		cz.Set(cc.ComponentConfig)
	} else {
		return fmt.Errorf("unable to register configz: %s", err)
	}

	// Prepare the event broadcaster.
	cc.EventBroadcaster.StartRecordingToSink(ctx.Done())

	// Setup healthz checks.
	var checks []healthz.HealthChecker
	if cc.ComponentConfig.LeaderElection.LeaderElect {
		checks = append(checks, cc.LeaderElection.WatchDog)
	}

	waitingForLeader := make(chan struct{})
	isLeader := func() bool {
		select {
		case _, ok := <-waitingForLeader:
			// if channel is closed, we are leading
			return !ok
		default:
			// channel is open, we are waiting for a leader
			return false
		}
	}

	// Start up the healthz server.
	if cc.InsecureServing != nil {
		separateMetrics := cc.InsecureMetricsServing != nil
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader, separateMetrics, checks...), nil, nil)
		if err := cc.InsecureServing.Serve(handler, 0, ctx.Done()); err != nil {
			return fmt.Errorf("failed to start healthz server: %v", err)
		}
	}
	if cc.InsecureMetricsServing != nil {
		handler := buildHandlerChain(newMetricsHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader), nil, nil)
		if err := cc.InsecureMetricsServing.Serve(handler, 0, ctx.Done()); err != nil {
			return fmt.Errorf("failed to start metrics server: %v", err)
		}
	}
	if cc.SecureServing != nil {
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, cc.InformerFactory, isLeader, false, checks...), cc.Authentication.Authenticator, cc.Authorization.Authorizer)
		// TODO: handle stoppedCh returned by c.SecureServing.Serve
		if _, err := cc.SecureServing.Serve(handler, 0, ctx.Done()); err != nil {
			// fail early for secure handlers, removing the old error loop from above
			return fmt.Errorf("failed to start secure server: %v", err)
		}
	}

	// Start all informers.
	cc.InformerFactory.Start(ctx.Done())

	// Wait for all caches to sync before scheduling.
	cc.InformerFactory.WaitForCacheSync(ctx.Done())

	// If leader election is enabled, runCommand via LeaderElector until done and exit.
	if cc.LeaderElection != nil {
		cc.LeaderElection.Callbacks = leaderelection.LeaderCallbacks{
			OnStartedLeading: func(ctx context.Context) {
				close(waitingForLeader)
				sched.Run(ctx)
			},
			OnStoppedLeading: func() {
				select {
				case <-ctx.Done():
					// We were asked to terminate. Exit 0.
					klog.Info("Requested to terminate. Exiting.")
					os.Exit(0)
				default:
					// We lost the lock.
					klog.Exitf("leaderelection lost")
				}
			},
		}
		leaderElector, err := leaderelection.NewLeaderElector(*cc.LeaderElection)
		if err != nil {
			return fmt.Errorf("couldn't create leader elector: %v", err)
		}

		leaderElector.Run(ctx)

		return fmt.Errorf("lost lease")
	}

	// Leader election is disabled, so runCommand inline until done.
	close(waitingForLeader)
	sched.Run(ctx)
	return fmt.Errorf("finished without leader elect")
}

// buildHandlerChain wraps the given handler with the standard filters.
func buildHandlerChain(handler http.Handler, authn authenticator.Request, authz authorizer.Authorizer) http.Handler {
	requestInfoResolver := &apirequest.RequestInfoFactory{}
	failedHandler := genericapifilters.Unauthorized(scheme.Codecs)

	handler = genericapifilters.WithAuthorization(handler, authz, scheme.Codecs)
	handler = genericapifilters.WithAuthentication(handler, authn, failedHandler, nil)
	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericapifilters.WithCacheControl(handler)
	handler = genericfilters.WithHTTPLogging(handler)
	handler = genericfilters.WithPanicRecovery(handler, requestInfoResolver)

	return handler
}

func installMetricHandler(pathRecorderMux *mux.PathRecorderMux, informers informers.SharedInformerFactory, isLeader func() bool) {
	configz.InstallHandler(pathRecorderMux)
	pathRecorderMux.Handle("/metrics", legacyregistry.HandlerWithReset())

	resourceMetricsHandler := resources.Handler(informers.Core().V1().Pods().Lister())
	pathRecorderMux.HandleFunc("/metrics/resources", func(w http.ResponseWriter, req *http.Request) {
		if !isLeader() {
			return
		}
		resourceMetricsHandler.ServeHTTP(w, req)
	})
}

// newMetricsHandler builds a metrics server from the config.
func newMetricsHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, informers informers.SharedInformerFactory, isLeader func() bool) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	installMetricHandler(pathRecorderMux, informers, isLeader)
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
		routes.DebugFlags{}.Install(pathRecorderMux, "v", routes.StringFlagPutHandler(logs.GlogSetter))
	}
	return pathRecorderMux
}

// newHealthzHandler creates a healthz server from the config, and will also
// embed the metrics handler if the healthz and metrics address configurations
// are the same.
func newHealthzHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, informers informers.SharedInformerFactory, isLeader func() bool, separateMetrics bool, checks ...healthz.HealthChecker) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	healthz.InstallHandler(pathRecorderMux, checks...)
	if !separateMetrics {
		installMetricHandler(pathRecorderMux, informers, isLeader)
	}
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
		routes.DebugFlags{}.Install(pathRecorderMux, "v", routes.StringFlagPutHandler(logs.GlogSetter))
	}
	return pathRecorderMux
}

func getRecorderFactory(cc *schedulerserverconfig.CompletedConfig) profile.RecorderFactory {
	return func(name string) events.EventRecorder {
		return cc.EventBroadcaster.NewRecorder(name)
	}
}

// WithPlugin creates an Option based on plugin name and factory. Please don't remove this function: it is used to register out-of-tree plugins,
// hence there are no references to it from the kubernetes scheduler code base.
func WithPlugin(name string, factory runtime.PluginFactory) Option {
	return func(registry runtime.Registry) error {
		return registry.Register(name, factory)
	}
}

// Setup creates a completed config and a scheduler based on the command args and options
func Setup(ctx context.Context, opts *options.Options, outOfTreeRegistryOptions ...Option) (*schedulerserverconfig.CompletedConfig, *scheduler.Scheduler, error) {
	if errs := opts.Validate(); len(errs) > 0 {
		return nil, nil, utilerrors.NewAggregate(errs)
	}

	c, err := opts.Config()
	if err != nil {
		return nil, nil, err
	}

	// Get the completed config
	cc := c.Complete()

	outOfTreeRegistry := make(runtime.Registry)
	for _, option := range outOfTreeRegistryOptions {
		if err := option(outOfTreeRegistry); err != nil {
			return nil, nil, err
		}
	}

	recorderFactory := getRecorderFactory(&cc)
	completedProfiles := make([]kubeschedulerconfig.KubeSchedulerProfile, 0)
	// Create the scheduler.
	sched, err := scheduler.New(cc.Client,
		cc.InformerFactory,
		recorderFactory,
		ctx.Done(),
		scheduler.WithComponentConfigVersion(cc.ComponentConfig.TypeMeta.APIVersion),
		scheduler.WithKubeConfig(cc.KubeConfig),
		scheduler.WithProfiles(cc.ComponentConfig.Profiles...),
		scheduler.WithLegacyPolicySource(cc.LegacyPolicySource),
		scheduler.WithPercentageOfNodesToScore(cc.ComponentConfig.PercentageOfNodesToScore),
		scheduler.WithFrameworkOutOfTreeRegistry(outOfTreeRegistry),
		scheduler.WithPodMaxBackoffSeconds(cc.ComponentConfig.PodMaxBackoffSeconds),
		scheduler.WithPodInitialBackoffSeconds(cc.ComponentConfig.PodInitialBackoffSeconds),
		scheduler.WithExtenders(cc.ComponentConfig.Extenders...),
		scheduler.WithParallelism(cc.ComponentConfig.Parallelism),
		scheduler.WithBuildFrameworkCapturer(func(profile kubeschedulerconfig.KubeSchedulerProfile) {
			// Profiles are processed during Framework instantiation to set default plugins and configurations. Capturing them for logging
			completedProfiles = append(completedProfiles, profile)
		}),
	)
	if err != nil {
		return nil, nil, err
	}
	if err := options.LogOrWriteConfig(opts.WriteConfigTo, &cc.ComponentConfig, completedProfiles); err != nil {
		return nil, nil, err
	}

	return &cc, sched, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// Package app implements a Server object for running the scheduler.
package app

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	goruntime "runtime"

	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	genericapifilters "k8s.io/apiserver/pkg/endpoints/filters"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	genericfilters "k8s.io/apiserver/pkg/server/filters"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/apiserver/pkg/server/mux"
	"k8s.io/apiserver/pkg/server/routes"
	apiserverflag "k8s.io/apiserver/pkg/util/flag"
	"k8s.io/apiserver/pkg/util/globalflag"
	v1core "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/leaderelection"
	schedulerserverconfig "k8s.io/kubernetes/cmd/kube-scheduler/app/config"
	"k8s.io/kubernetes/cmd/kube-scheduler/app/options"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controller"
	"k8s.io/kubernetes/pkg/scheduler"
	"k8s.io/kubernetes/pkg/scheduler/algorithmprovider"
	kubeschedulerconfig "k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/metrics"
	"k8s.io/kubernetes/pkg/util/configz"
	utilflag "k8s.io/kubernetes/pkg/util/flag"
	"k8s.io/kubernetes/pkg/version"
	"k8s.io/kubernetes/pkg/version/verflag"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/spf13/cobra"
	"k8s.io/klog"
)

// NewSchedulerCommand creates a *cobra.Command object with default parameters
func NewSchedulerCommand() *cobra.Command {
	opts, err := options.NewOptions()
	if err != nil {
		klog.Fatalf("unable to initialize command options: %v", err)
	}

	cmd := &cobra.Command{
		Use: "kube-scheduler",
		Long: `The Kubernetes scheduler is a policy-rich, topology-aware,
workload-specific function that significantly impacts availability, performance,
and capacity. The scheduler needs to take into account individual and collective
resource requirements, quality of service requirements, hardware/software/policy
constraints, affinity and anti-affinity specifications, data locality, inter-workload
interference, deadlines, and so on. Workload-specific requirements will be exposed
through the API as necessary.`,
		Run: func(cmd *cobra.Command, args []string) {
			if err := runCommand(cmd, args, opts); err != nil {
				fmt.Fprintf(os.Stderr, "%v\n", err)
				os.Exit(1)
			}
		},
	}
	fs := cmd.Flags()
	namedFlagSets := opts.Flags()
	verflag.AddFlags(namedFlagSets.FlagSet("global"))
	globalflag.AddGlobalFlags(namedFlagSets.FlagSet("global"), cmd.Name())
	for _, f := range namedFlagSets.FlagSets {
		fs.AddFlagSet(f)
	}

	usageFmt := "Usage:\n  %s\n"
	cols, _, _ := apiserverflag.TerminalSize(cmd.OutOrStdout())
	cmd.SetUsageFunc(func(cmd *cobra.Command) error {
		fmt.Fprintf(cmd.OutOrStderr(), usageFmt, cmd.UseLine())
		apiserverflag.PrintSections(cmd.OutOrStderr(), namedFlagSets, cols)
		return nil
	})
	cmd.SetHelpFunc(func(cmd *cobra.Command, args []string) {
		fmt.Fprintf(cmd.OutOrStdout(), "%s\n\n"+usageFmt, cmd.Long, cmd.UseLine())
		apiserverflag.PrintSections(cmd.OutOrStdout(), namedFlagSets, cols)
	})
	cmd.MarkFlagFilename("config", "yaml", "yml", "json")

	return cmd
}

// runCommand runs the scheduler.
func runCommand(cmd *cobra.Command, args []string, opts *options.Options) error {
	verflag.PrintAndExitIfRequested()
	utilflag.PrintFlags(cmd.Flags())

	if len(args) != 0 {
		fmt.Fprint(os.Stderr, "arguments are not supported\n")
	}

	if errs := opts.Validate(); len(errs) > 0 {
		fmt.Fprintf(os.Stderr, "%v\n", utilerrors.NewAggregate(errs))
		os.Exit(1)
	}

	if len(opts.WriteConfigTo) > 0 {
		if err := options.WriteConfigFile(opts.WriteConfigTo, &opts.ComponentConfig); err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			os.Exit(1)
		}
		klog.Infof("Wrote configuration to: %s\n", opts.WriteConfigTo)
	}

	c, err := opts.Config()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	stopCh := make(chan struct{})

	// Get the completed config
	cc := c.Complete()

	// To help debugging, immediately log version
	klog.Infof("Version: %+v", version.Get())

	// Apply algorithms based on feature gates.
	// TODO: make configurable?
	algorithmprovider.ApplyFeatureGates()

	// Configz registration.
	if cz, err := configz.New("componentconfig"); err == nil {
		cz.Set(cc.ComponentConfig)
	} else {
		return fmt.Errorf("unable to register configz: %s", err)
	}

	return Run(cc, stopCh)
}

// Run executes the scheduler based on the given configuration. It only return on error or when stopCh is closed.
func Run(cc schedulerserverconfig.CompletedConfig, stopCh <-chan struct{}) error {
	// To help debugging, immediately log version
	klog.V(1).Infof("Starting Kubernetes Scheduler version %+v", version.Get())

	// Create the scheduler.
	sched, err := scheduler.New(cc.Client,
		cc.InformerFactory.Core().V1().Nodes(),
		cc.PodInformer,
		cc.InformerFactory.Core().V1().PersistentVolumes(),
		cc.InformerFactory.Core().V1().PersistentVolumeClaims(),
		cc.InformerFactory.Core().V1().ReplicationControllers(),
		cc.InformerFactory.Apps().V1().ReplicaSets(),
		cc.InformerFactory.Apps().V1().StatefulSets(),
		cc.InformerFactory.Core().V1().Services(),
		cc.InformerFactory.Policy().V1beta1().PodDisruptionBudgets(),
		cc.InformerFactory.Storage().V1().StorageClasses(),
		cc.Recorder,
		cc.ComponentConfig.AlgorithmSource,
		stopCh,
		scheduler.WithName(cc.ComponentConfig.SchedulerName),
		scheduler.WithHardPodAffinitySymmetricWeight(cc.ComponentConfig.HardPodAffinitySymmetricWeight),
		scheduler.WithPreemptionDisabled(cc.ComponentConfig.DisablePreemption),
		scheduler.WithPercentageOfNodesToScore(cc.ComponentConfig.PercentageOfNodesToScore),
		scheduler.WithBindTimeoutSeconds(*cc.ComponentConfig.BindTimeoutSeconds))
	if err != nil {
		return err
	}

	// Prepare the event broadcaster.
	if cc.Broadcaster != nil && cc.EventClient != nil {
		cc.Broadcaster.StartLogging(klog.V(6).Infof)
		cc.Broadcaster.StartRecordingToSink(&v1core.EventSinkImpl{Interface: cc.EventClient.Events("")})
	}

	// Setup healthz checks.
	var checks []healthz.HealthzChecker
	if cc.ComponentConfig.LeaderElection.LeaderElect {
		checks = append(checks, cc.LeaderElection.WatchDog)
	}

	// Start up the healthz server.
	if cc.InsecureServing != nil {
		separateMetrics := cc.InsecureMetricsServing != nil
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, separateMetrics, checks...), nil, nil)
		if err := cc.InsecureServing.Serve(handler, 0, stopCh); err != nil {
			return fmt.Errorf("failed to start healthz server: %v", err)
		}
	}
	if cc.InsecureMetricsServing != nil {
		handler := buildHandlerChain(newMetricsHandler(&cc.ComponentConfig), nil, nil)
		if err := cc.InsecureMetricsServing.Serve(handler, 0, stopCh); err != nil {
			return fmt.Errorf("failed to start metrics server: %v", err)
		}
	}
	if cc.SecureServing != nil {
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, false, checks...), cc.Authentication.Authenticator, cc.Authorization.Authorizer)
		// TODO: handle stoppedCh returned by c.SecureServing.Serve
		if _, err := cc.SecureServing.Serve(handler, 0, stopCh); err != nil {
			// fail early for secure handlers, removing the old error loop from above
			return fmt.Errorf("failed to start healthz server: %v", err)
		}
	}

	// Start all informers.
	go cc.PodInformer.Informer().Run(stopCh)
	cc.InformerFactory.Start(stopCh)

	// Wait for all caches to sync before scheduling.
	cc.InformerFactory.WaitForCacheSync(stopCh)
	controller.WaitForCacheSync("scheduler", stopCh, cc.PodInformer.Informer().HasSynced)

	// Prepare a reusable runCommand function.
	run := func(ctx context.Context) {
		sched.Run()
		<-ctx.Done()
	}

	ctx, cancel := context.WithCancel(context.TODO()) // TODO once Run() accepts a context, it should be used here
	defer cancel()

	go func() {
		select {
		case <-stopCh:
			cancel()
		case <-ctx.Done():
		}
	}()

	// If leader election is enabled, runCommand via LeaderElector until done and exit.
	if cc.LeaderElection != nil {
		cc.LeaderElection.Callbacks = leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				utilruntime.HandleError(fmt.Errorf("lost master"))
			},
		}
		leaderElector, err := leaderelection.NewLeaderElector(*cc.LeaderElection)
		if err != nil {
			return fmt.Errorf("couldn't create leader elector: %v", err)
		}

		leaderElector.Run(ctx)

		return fmt.Errorf("lost lease")
	}

	// Leader election is disabled, so runCommand inline until done.
	run(ctx)
	return fmt.Errorf("finished without leader elect")
}

// buildHandlerChain wraps the given handler with the standard filters.
func buildHandlerChain(handler http.Handler, authn authenticator.Request, authz authorizer.Authorizer) http.Handler {
	requestInfoResolver := &apirequest.RequestInfoFactory{}
	failedHandler := genericapifilters.Unauthorized(legacyscheme.Codecs, false)

	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericapifilters.WithAuthorization(handler, authz, legacyscheme.Codecs)
	handler = genericapifilters.WithAuthentication(handler, authn, failedHandler, nil)
	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericfilters.WithPanicRecovery(handler)

	return handler
}

func installMetricHandler(pathRecorderMux *mux.PathRecorderMux) {
	configz.InstallHandler(pathRecorderMux)
	defaultMetricsHandler := prometheus.Handler().ServeHTTP
	pathRecorderMux.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
		if req.Method == "DELETE" {
			metrics.Reset()
			io.WriteString(w, "metrics reset\n")
			return
		}
		defaultMetricsHandler(w, req)
	})
}

// newMetricsHandler builds a metrics server from the config.
func newMetricsHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	installMetricHandler(pathRecorderMux)
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
	}
	return pathRecorderMux
}

// newHealthzServer creates a healthz server from the config, and will also
// embed the metrics handler if the healthz and metrics address configurations
// are the same.
func newHealthzHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, separateMetrics bool, checks ...healthz.HealthzChecker) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	healthz.InstallHandler(pathRecorderMux, checks...)
	if !separateMetrics {
		installMetricHandler(pathRecorderMux)
	}
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
	}
	return pathRecorderMux
}
<commit_msg>Fix function comment to consistent with its name<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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

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

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

// Package app implements a Server object for running the scheduler.
package app

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	goruntime "runtime"

	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	genericapifilters "k8s.io/apiserver/pkg/endpoints/filters"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	genericfilters "k8s.io/apiserver/pkg/server/filters"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/apiserver/pkg/server/mux"
	"k8s.io/apiserver/pkg/server/routes"
	apiserverflag "k8s.io/apiserver/pkg/util/flag"
	"k8s.io/apiserver/pkg/util/globalflag"
	v1core "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/leaderelection"
	schedulerserverconfig "k8s.io/kubernetes/cmd/kube-scheduler/app/config"
	"k8s.io/kubernetes/cmd/kube-scheduler/app/options"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controller"
	"k8s.io/kubernetes/pkg/scheduler"
	"k8s.io/kubernetes/pkg/scheduler/algorithmprovider"
	kubeschedulerconfig "k8s.io/kubernetes/pkg/scheduler/apis/config"
	"k8s.io/kubernetes/pkg/scheduler/metrics"
	"k8s.io/kubernetes/pkg/util/configz"
	utilflag "k8s.io/kubernetes/pkg/util/flag"
	"k8s.io/kubernetes/pkg/version"
	"k8s.io/kubernetes/pkg/version/verflag"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/spf13/cobra"
	"k8s.io/klog"
)

// NewSchedulerCommand creates a *cobra.Command object with default parameters
func NewSchedulerCommand() *cobra.Command {
	opts, err := options.NewOptions()
	if err != nil {
		klog.Fatalf("unable to initialize command options: %v", err)
	}

	cmd := &cobra.Command{
		Use: "kube-scheduler",
		Long: `The Kubernetes scheduler is a policy-rich, topology-aware,
workload-specific function that significantly impacts availability, performance,
and capacity. The scheduler needs to take into account individual and collective
resource requirements, quality of service requirements, hardware/software/policy
constraints, affinity and anti-affinity specifications, data locality, inter-workload
interference, deadlines, and so on. Workload-specific requirements will be exposed
through the API as necessary.`,
		Run: func(cmd *cobra.Command, args []string) {
			if err := runCommand(cmd, args, opts); err != nil {
				fmt.Fprintf(os.Stderr, "%v\n", err)
				os.Exit(1)
			}
		},
	}
	fs := cmd.Flags()
	namedFlagSets := opts.Flags()
	verflag.AddFlags(namedFlagSets.FlagSet("global"))
	globalflag.AddGlobalFlags(namedFlagSets.FlagSet("global"), cmd.Name())
	for _, f := range namedFlagSets.FlagSets {
		fs.AddFlagSet(f)
	}

	usageFmt := "Usage:\n  %s\n"
	cols, _, _ := apiserverflag.TerminalSize(cmd.OutOrStdout())
	cmd.SetUsageFunc(func(cmd *cobra.Command) error {
		fmt.Fprintf(cmd.OutOrStderr(), usageFmt, cmd.UseLine())
		apiserverflag.PrintSections(cmd.OutOrStderr(), namedFlagSets, cols)
		return nil
	})
	cmd.SetHelpFunc(func(cmd *cobra.Command, args []string) {
		fmt.Fprintf(cmd.OutOrStdout(), "%s\n\n"+usageFmt, cmd.Long, cmd.UseLine())
		apiserverflag.PrintSections(cmd.OutOrStdout(), namedFlagSets, cols)
	})
	cmd.MarkFlagFilename("config", "yaml", "yml", "json")

	return cmd
}

// runCommand runs the scheduler.
func runCommand(cmd *cobra.Command, args []string, opts *options.Options) error {
	verflag.PrintAndExitIfRequested()
	utilflag.PrintFlags(cmd.Flags())

	if len(args) != 0 {
		fmt.Fprint(os.Stderr, "arguments are not supported\n")
	}

	if errs := opts.Validate(); len(errs) > 0 {
		fmt.Fprintf(os.Stderr, "%v\n", utilerrors.NewAggregate(errs))
		os.Exit(1)
	}

	if len(opts.WriteConfigTo) > 0 {
		if err := options.WriteConfigFile(opts.WriteConfigTo, &opts.ComponentConfig); err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			os.Exit(1)
		}
		klog.Infof("Wrote configuration to: %s\n", opts.WriteConfigTo)
	}

	c, err := opts.Config()
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(1)
	}

	stopCh := make(chan struct{})

	// Get the completed config
	cc := c.Complete()

	// To help debugging, immediately log version
	klog.Infof("Version: %+v", version.Get())

	// Apply algorithms based on feature gates.
	// TODO: make configurable?
	algorithmprovider.ApplyFeatureGates()

	// Configz registration.
	if cz, err := configz.New("componentconfig"); err == nil {
		cz.Set(cc.ComponentConfig)
	} else {
		return fmt.Errorf("unable to register configz: %s", err)
	}

	return Run(cc, stopCh)
}

// Run executes the scheduler based on the given configuration. It only return on error or when stopCh is closed.
func Run(cc schedulerserverconfig.CompletedConfig, stopCh <-chan struct{}) error {
	// To help debugging, immediately log version
	klog.V(1).Infof("Starting Kubernetes Scheduler version %+v", version.Get())

	// Create the scheduler.
	sched, err := scheduler.New(cc.Client,
		cc.InformerFactory.Core().V1().Nodes(),
		cc.PodInformer,
		cc.InformerFactory.Core().V1().PersistentVolumes(),
		cc.InformerFactory.Core().V1().PersistentVolumeClaims(),
		cc.InformerFactory.Core().V1().ReplicationControllers(),
		cc.InformerFactory.Apps().V1().ReplicaSets(),
		cc.InformerFactory.Apps().V1().StatefulSets(),
		cc.InformerFactory.Core().V1().Services(),
		cc.InformerFactory.Policy().V1beta1().PodDisruptionBudgets(),
		cc.InformerFactory.Storage().V1().StorageClasses(),
		cc.Recorder,
		cc.ComponentConfig.AlgorithmSource,
		stopCh,
		scheduler.WithName(cc.ComponentConfig.SchedulerName),
		scheduler.WithHardPodAffinitySymmetricWeight(cc.ComponentConfig.HardPodAffinitySymmetricWeight),
		scheduler.WithPreemptionDisabled(cc.ComponentConfig.DisablePreemption),
		scheduler.WithPercentageOfNodesToScore(cc.ComponentConfig.PercentageOfNodesToScore),
		scheduler.WithBindTimeoutSeconds(*cc.ComponentConfig.BindTimeoutSeconds))
	if err != nil {
		return err
	}

	// Prepare the event broadcaster.
	if cc.Broadcaster != nil && cc.EventClient != nil {
		cc.Broadcaster.StartLogging(klog.V(6).Infof)
		cc.Broadcaster.StartRecordingToSink(&v1core.EventSinkImpl{Interface: cc.EventClient.Events("")})
	}

	// Setup healthz checks.
	var checks []healthz.HealthzChecker
	if cc.ComponentConfig.LeaderElection.LeaderElect {
		checks = append(checks, cc.LeaderElection.WatchDog)
	}

	// Start up the healthz server.
	if cc.InsecureServing != nil {
		separateMetrics := cc.InsecureMetricsServing != nil
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, separateMetrics, checks...), nil, nil)
		if err := cc.InsecureServing.Serve(handler, 0, stopCh); err != nil {
			return fmt.Errorf("failed to start healthz server: %v", err)
		}
	}
	if cc.InsecureMetricsServing != nil {
		handler := buildHandlerChain(newMetricsHandler(&cc.ComponentConfig), nil, nil)
		if err := cc.InsecureMetricsServing.Serve(handler, 0, stopCh); err != nil {
			return fmt.Errorf("failed to start metrics server: %v", err)
		}
	}
	if cc.SecureServing != nil {
		handler := buildHandlerChain(newHealthzHandler(&cc.ComponentConfig, false, checks...), cc.Authentication.Authenticator, cc.Authorization.Authorizer)
		// TODO: handle stoppedCh returned by c.SecureServing.Serve
		if _, err := cc.SecureServing.Serve(handler, 0, stopCh); err != nil {
			// fail early for secure handlers, removing the old error loop from above
			return fmt.Errorf("failed to start secure server: %v", err)
		}
	}

	// Start all informers.
	go cc.PodInformer.Informer().Run(stopCh)
	cc.InformerFactory.Start(stopCh)

	// Wait for all caches to sync before scheduling.
	cc.InformerFactory.WaitForCacheSync(stopCh)
	controller.WaitForCacheSync("scheduler", stopCh, cc.PodInformer.Informer().HasSynced)

	// Prepare a reusable runCommand function.
	run := func(ctx context.Context) {
		sched.Run()
		<-ctx.Done()
	}

	ctx, cancel := context.WithCancel(context.TODO()) // TODO once Run() accepts a context, it should be used here
	defer cancel()

	go func() {
		select {
		case <-stopCh:
			cancel()
		case <-ctx.Done():
		}
	}()

	// If leader election is enabled, runCommand via LeaderElector until done and exit.
	if cc.LeaderElection != nil {
		cc.LeaderElection.Callbacks = leaderelection.LeaderCallbacks{
			OnStartedLeading: run,
			OnStoppedLeading: func() {
				utilruntime.HandleError(fmt.Errorf("lost master"))
			},
		}
		leaderElector, err := leaderelection.NewLeaderElector(*cc.LeaderElection)
		if err != nil {
			return fmt.Errorf("couldn't create leader elector: %v", err)
		}

		leaderElector.Run(ctx)

		return fmt.Errorf("lost lease")
	}

	// Leader election is disabled, so runCommand inline until done.
	run(ctx)
	return fmt.Errorf("finished without leader elect")
}

// buildHandlerChain wraps the given handler with the standard filters.
func buildHandlerChain(handler http.Handler, authn authenticator.Request, authz authorizer.Authorizer) http.Handler {
	requestInfoResolver := &apirequest.RequestInfoFactory{}
	failedHandler := genericapifilters.Unauthorized(legacyscheme.Codecs, false)

	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericapifilters.WithAuthorization(handler, authz, legacyscheme.Codecs)
	handler = genericapifilters.WithAuthentication(handler, authn, failedHandler, nil)
	handler = genericapifilters.WithRequestInfo(handler, requestInfoResolver)
	handler = genericfilters.WithPanicRecovery(handler)

	return handler
}

func installMetricHandler(pathRecorderMux *mux.PathRecorderMux) {
	configz.InstallHandler(pathRecorderMux)
	defaultMetricsHandler := prometheus.Handler().ServeHTTP
	pathRecorderMux.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
		if req.Method == "DELETE" {
			metrics.Reset()
			io.WriteString(w, "metrics reset\n")
			return
		}
		defaultMetricsHandler(w, req)
	})
}

// newMetricsHandler builds a metrics server from the config.
func newMetricsHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	installMetricHandler(pathRecorderMux)
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
	}
	return pathRecorderMux
}

// newHealthzHandler creates a healthz server from the config, and will also
// embed the metrics handler if the healthz and metrics address configurations
// are the same.
func newHealthzHandler(config *kubeschedulerconfig.KubeSchedulerConfiguration, separateMetrics bool, checks ...healthz.HealthzChecker) http.Handler {
	pathRecorderMux := mux.NewPathRecorderMux("kube-scheduler")
	healthz.InstallHandler(pathRecorderMux, checks...)
	if !separateMetrics {
		installMetricHandler(pathRecorderMux)
	}
	if config.EnableProfiling {
		routes.Profiling{}.Install(pathRecorderMux)
		if config.EnableContentionProfiling {
			goruntime.SetBlockProfileRate(1)
		}
	}
	return pathRecorderMux
}
<|endoftext|>


<commit_before>// query_benchmarker speed tests ElasticSearch using requests from stdin.
//
// It reads encoded Query objects from stdin, and makes concurrent requests
// to the provided HTTP endpoint. This program currently has no knowledge of
// the internals of the endpoint.
//
// TODO(rw): On my machine, this only decodes 700k/sec messages from stdin.
package main

import (
	"bufio"
	"encoding/gob"
	"flag"
	"fmt"
	"github.com/influxdata/influxdb-comparisons/bulk_query"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/influxdb-comparisons/util/report"
)

// Program option vars:
var (
	csvDaemonUrls        string
	daemonUrls           []string
	workers              int
	debug                int
	prettyPrintResponses bool
	limit                int64
	burnIn               uint64
	printInterval        uint64
	memProfile           string
	telemetryHost        string
	telemetryStderr      bool
	telemetryBatchSize   uint64
	telemetryTagsCSV     string
	telemetryBasicAuth   string
	reportDatabase       string
	reportHost           string
	reportUser           string
	reportPassword       string
	reportTagsCSV        string
)

// Global vars:
var (
	queryPool           sync.Pool
	queryChan           chan *Query
	statPool            sync.Pool
	statChan            chan *bulk_query.Stat
	workersGroup        sync.WaitGroup
	statGroup           sync.WaitGroup
	telemetryChanPoints chan *report.Point
	telemetryChanDone   chan struct{}
	telemetrySrcAddr    string
	telemetryTags       [][2]string
	statMapping         bulk_query.StatsMap
	reportTags          [][2]string
	reportHostname      string
)

const allQueriesLabel = "all queries"

// Parse args:
func init() {
	flag.StringVar(&csvDaemonUrls, "urls", "http://localhost:9200", "Daemon URLs, comma-separated. Will be used in a round-robin fashion.")
	flag.IntVar(&workers, "workers", 1, "Number of concurrent requests to make.")
	flag.IntVar(&debug, "debug", 0, "Whether to print debug messages.")
	flag.Int64Var(&limit, "limit", -1, "Limit the number of queries to send.")
	flag.Uint64Var(&burnIn, "burn-in", 0, "Number of queries to ignore before collecting statistics.")
	flag.Uint64Var(&printInterval, "print-interval", 100, "Print timing stats to stderr after this many queries (0 to disable)")
	flag.BoolVar(&prettyPrintResponses, "print-responses", false, "Pretty print JSON response bodies (for correctness checking) (default false).")
	flag.StringVar(&memProfile, "memprofile", "", "Write a memory profile to this file.")
	flag.StringVar(&telemetryHost, "telemetry-host", "", "InfluxDB host to write telegraf telemetry to (optional).")
	flag.StringVar(&telemetryTagsCSV, "telemetry-tags", "", "Tag(s) for telemetry. Format: key0:val0,key1:val1,...")
	flag.BoolVar(&telemetryStderr, "telemetry-stderr", false, "Whether to write telemetry also to stderr.")
	flag.Uint64Var(&telemetryBatchSize, "telemetry-batch-size", 1000, "Telemetry batch size (lines).")
	flag.StringVar(&telemetryBasicAuth, "telemetry-basic-auth", "", "basic auth (username:password) for telemetry.")
	flag.StringVar(&reportDatabase, "report-database", "database_benchmarks", "Database name where to store result metrics.")
	flag.StringVar(&reportHost, "report-host", "", "Host to send result metrics.")
	flag.StringVar(&reportUser, "report-user", "", "User for Host to send result metrics.")
	flag.StringVar(&reportPassword, "report-password", "", "User password for Host to send result metrics.")
	flag.StringVar(&reportTagsCSV, "report-tags", "", "Comma separated k:v tags to send  alongside result metrics.")

	flag.Parse()

	daemonUrls = strings.Split(csvDaemonUrls, ",")
	if len(daemonUrls) == 0 {
		log.Fatal("missing 'urls' flag")
	}
	fmt.Printf("daemon URLs: %v\n", daemonUrls)

	if telemetryHost != "" {
		fmt.Printf("telemetry destination: %v\n", telemetryHost)
		if telemetryBatchSize == 0 {
			panic("invalid telemetryBatchSize")
		}

		var err error
		telemetrySrcAddr, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("src addr for telemetry: %v\n", telemetrySrcAddr)

		if telemetryTagsCSV != "" {
			pairs := strings.Split(telemetryTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				telemetryTags = append(telemetryTags, tagpair)
			}
		}
		fmt.Printf("telemetry tags: %v\n", telemetryTags)
	}

	if reportHost != "" {
		fmt.Printf("results report destination: %v\n", reportHost)
		fmt.Printf("results report database: %v\n", reportDatabase)

		var err error
		reportHostname, err = os.Hostname()
		if err != nil {
			log.Fatalf("os.Hostname() error: %s", err.Error())
		}
		fmt.Printf("hostname for results report: %v\n", reportHostname)

		if reportTagsCSV != "" {
			pairs := strings.Split(reportTagsCSV, ",")
			for _, pair := range pairs {
				fields := strings.SplitN(pair, ":", 2)
				tagpair := [2]string{fields[0], fields[1]}
				reportTags = append(reportTags, tagpair)
			}
		}
		fmt.Printf("results report tags: %v\n", reportTags)
	}
}

func main() {
	// Make pools to minimize heap usage:
	queryPool = sync.Pool{
		New: func() interface{} {
			return &Query{
				HumanLabel:       make([]byte, 0, 1024),
				HumanDescription: make([]byte, 0, 1024),
				Method:           make([]byte, 0, 1024),
				Path:             make([]byte, 0, 1024),
				Body:             make([]byte, 0, 1024),
			}
		},
	}

	statPool = sync.Pool{
		New: func() interface{} {
			return &bulk_query.Stat{
				Label: make([]byte, 0, 1024),
				Value: 0.0,
			}
		},
	}

	// Make data and control channels:
	queryChan = make(chan *Query, workers)
	statChan = make(chan *bulk_query.Stat, workers)

	// Launch the stats processor:
	statGroup.Add(1)
	go processStats()

	if telemetryHost != "" {
		telemetryCollector := report.NewCollector(telemetryHost, "telegraf", reportUser, reportPassword)
		telemetryChanPoints, telemetryChanDone = report.TelemetryRunAsync(telemetryCollector, telemetryBatchSize, telemetryStderr, burnIn)
	}

	// Launch the query processors:
	for i := 0; i < workers; i++ {
		daemonUrl := daemonUrls[i%len(daemonUrls)]
		workersGroup.Add(1)
		w := NewHTTPClient(daemonUrl, debug)
		go processQueries(w, telemetryChanPoints, fmt.Sprintf("%d", i))
	}

	// Read in jobs, closing the job channel when done:
	input := bufio.NewReaderSize(os.Stdin, 1<<20)
	wallStart := time.Now()
	scan(input)
	close(queryChan)

	// Block for workers to finish sending requests, closing the stats
	// channel when done:
	workersGroup.Wait()
	close(statChan)

	// Wait on the stat collector to finish (and print its results):
	statGroup.Wait()

	wallEnd := time.Now()
	wallTook := wallEnd.Sub(wallStart)
	_, err := fmt.Printf("wall clock time: %fsec\n", float64(wallTook.Nanoseconds())/1e9)
	if err != nil {
		log.Fatal(err)
	}

	if telemetryHost != "" {
		fmt.Println("shutting down telemetry...")
		close(telemetryChanPoints)
		<-telemetryChanDone
		fmt.Println("done shutting down telemetry.")
	}

	// (Optional) create a memory profile:
	if memProfile != "" {
		f, err := os.Create(memProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}

	if reportHost != "" {
		reportParams := &report.QueryReportParams{
			ReportParams: report.ReportParams{
				DBType:             "ElasticSearch",
				ReportDatabaseName: reportDatabase,
				ReportHost:         reportHost,
				ReportUser:         reportUser,
				ReportPassword:     reportPassword,
				ReportTags:         reportTags,
				Hostname:           reportHostname,
				DestinationUrl:     csvDaemonUrls,
				Workers:            workers,
				ItemLimit:          int(limit),
			},
			BurnIn: int64(burnIn),
		}
		stat := statMapping[allQueriesLabel]
		err = report.ReportQueryResult(reportParams, allQueriesLabel, stat.Min, stat.Mean, stat.Max, stat.Count, wallTook)

		if err != nil {
			log.Fatal(err)
		}
	}
}

// scan reads encoded Queries and places them onto the workqueue.
func scan(r io.Reader) {
	dec := gob.NewDecoder(r)

	n := int64(0)
	for {
		if limit >= 0 && n >= limit {
			break
		}

		q := queryPool.Get().(*Query)
		err := dec.Decode(q)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		q.ID = n

		queryChan <- q

		n++

	}
}

// processQueries reads byte buffers from queryChan and writes them to the
// target server, while tracking latency.
func processQueries(w *HTTPClient, telemetrySink chan *report.Point, telemetryWorkerLabel string) {
	opts := &HTTPClientDoOptions{
		Debug:                debug,
		PrettyPrintResponses: prettyPrintResponses,
	}
	var queriesSeen int64
	for q := range queryChan {
		ts := time.Now().UnixNano()
		lagMillis, err := w.Do(q, opts)

		stat := statPool.Get().(*bulk_query.Stat)
		stat.Init(q.HumanLabel, lagMillis)
		statChan <- stat

		queryPool.Put(q)
		if err != nil {
			log.Fatalf("Error during request: %s\n", err.Error())
		}

		// Report telemetry, if applicable:
		if telemetrySink != nil {
			p := report.GetPointFromGlobalPool()
			p.Init("benchmark_query", ts)
			p.AddTag("src_addr", telemetrySrcAddr)
			p.AddTag("dst_addr", w.HostString)
			p.AddTag("worker_id", telemetryWorkerLabel)
			p.AddFloat64Field("rtt_ms", lagMillis)
			p.AddInt64Field("worker_req_num", queriesSeen)
			telemetrySink <- p
		}
		queriesSeen++
	}
	workersGroup.Done()
}

// processStats collects latency results, aggregating them into summary
// statistics. Optionally, they are printed to stderr at regular intervals.
func processStats() {
	statMapping = bulk_query.StatsMap{
		allQueriesLabel: &bulk_query.StatGroup{},
	}

	i := uint64(0)
	for stat := range statChan {
		if i < burnIn {
			i++
			statPool.Put(stat)
			continue
		} else if i == burnIn && burnIn > 0 {
			_, err := fmt.Fprintf(os.Stderr, "burn-in complete after %d queries with %d workers\n", burnIn, workers)
			if err != nil {
				log.Fatal(err)
			}
		}

		if _, ok := statMapping[string(stat.Label)]; !ok {
			statMapping[string(stat.Label)] = &bulk_query.StatGroup{}
		}

		statMapping[allQueriesLabel].Push(stat.Value)
		statMapping[string(stat.Label)].Push(stat.Value)

		statPool.Put(stat)

		i++

		// print stats to stderr (if printInterval is greater than zero):
		if printInterval > 0 && i > 0 && i%printInterval == 0 && (int64(i) < limit || limit < 0) {
			_, err := fmt.Fprintf(os.Stderr, "after %d queries with %d workers:\n", i-burnIn, workers)
			if err != nil {
				log.Fatal(err)
			}
			fprintStats(os.Stderr, statMapping)
			_, err = fmt.Fprintf(os.Stderr, "\n")
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	// the final stats output goes to stdout:
	_, err := fmt.Printf("run complete after %d queries with %d workers:\n", i-burnIn, workers)
	if err != nil {
		log.Fatal(err)
	}
	fprintStats(os.Stdout, statMapping)
	statGroup.Done()
}

// fprintStats pretty-prints stats to the given writer.
func fprintStats(w io.Writer, statGroups bulk_query.StatsMap) {
	maxKeyLength := 0
	keys := make([]string, 0, len(statGroups))
	for k := range statGroups {
		if len(k) > maxKeyLength {
			maxKeyLength = len(k)
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		v := statGroups[k]
		minRate := 1e3 / v.Min
		meanRate := 1e3 / v.Mean
		maxRate := 1e3 / v.Max
		paddedKey := fmt.Sprintf("%s", k)
		for len(paddedKey) < maxKeyLength {
			paddedKey += " "
		}
		_, err := fmt.Fprintf(w, "%s : min: %8.2fms (%7.2f/sec), mean: %8.2fms (%7.2f/sec), max: %7.2fms (%6.2f/sec), count: %8d, sum: %5.1fsec \n", paddedKey, v.Min, minRate, v.Mean, meanRate, v.Max, maxRate, v.Count, v.Sum/1e3)
		if err != nil {
			log.Fatal(err)
		}
	}

}
<commit_msg>ElasticSearch query bechmarker adapted to common query bechmarker interface<commit_after>// query_benchmarker speed tests ElasticSearch using requests from stdin.
//
// It reads encoded Query objects from stdin, and makes concurrent requests
// to the provided HTTP endpoint. This program currently has no knowledge of
// the internals of the endpoint.
//
// TODO(rw): On my machine, this only decodes 700k/sec messages from stdin.
package main

import (
	"encoding/gob"
	"flag"
	"fmt"
	"github.com/influxdata/influxdb-comparisons/bulk_query"
	"io"
	"log"
	"strings"
	"sync"

	"github.com/influxdata/influxdb-comparisons/util/report"
)

type ElasticQueryBenchmarker struct {
	// Program option vars:
	csvDaemonUrls string
	daemonUrls    []string
	// Global vars:
	queryPool    sync.Pool
	queryChan    chan *Query
	scanFinished bool
}

var querier = &ElasticQueryBenchmarker{}

// Parse args:
func init() {
	bulk_query.Benchmarker.Init()
	querier.Init()

	flag.Parse()

	bulk_query.Benchmarker.Validate()
	querier.Validate()
}

func main() {
	bulk_query.Benchmarker.RunBenchmark(querier)
}

func (b *ElasticQueryBenchmarker) Init() {
	flag.StringVar(&b.csvDaemonUrls, "urls", "http://localhost:9200", "Daemon URLs, comma-separated. Will be used in a round-robin fashion.")
}

func (b *ElasticQueryBenchmarker) Validate() {
	b.daemonUrls = strings.Split(b.csvDaemonUrls, ",")
	if len(b.daemonUrls) == 0 {
		log.Fatal("missing 'urls' flag")
	}
	fmt.Printf("daemon URLs: %v\n", b.daemonUrls)
}

func (b *ElasticQueryBenchmarker) Prepare() {
	// Make pools to minimize heap usage:
	b.queryPool = sync.Pool{
		New: func() interface{} {
			return &Query{
				HumanLabel:       make([]byte, 0, 1024),
				HumanDescription: make([]byte, 0, 1024),
				Method:           make([]byte, 0, 1024),
				Path:             make([]byte, 0, 1024),
				Body:             make([]byte, 0, 1024),
			}
		},
	}
	b.queryChan = make(chan *Query)
}

func (b *ElasticQueryBenchmarker) GetProcessor() bulk_query.Processor {
	return b
}

func (b *ElasticQueryBenchmarker) GetScanner() bulk_query.Scanner {
	return b
}

func (b *ElasticQueryBenchmarker) PrepareProcess(i int) {

}

func (b *ElasticQueryBenchmarker) RunProcess(i int, workersGroup *sync.WaitGroup, statPool sync.Pool, statChan chan *bulk_query.Stat) {
	daemonUrl := b.daemonUrls[i%len(b.daemonUrls)]
	w := NewHTTPClient(daemonUrl, bulk_query.Benchmarker.Debug())
	b.processQueries(w, workersGroup, statPool, statChan)
}

func (b *ElasticQueryBenchmarker) IsScanFinished() bool {
	return b.scanFinished
}

func (b *ElasticQueryBenchmarker) CleanUp() {
	close(b.queryChan)
}

func (b *ElasticQueryBenchmarker) UpdateReport(params *report.QueryReportParams, reportTags [][2]string, extraVals []report.ExtraVal) (updatedTags [][2]string, updatedExtraVals []report.ExtraVal) {
	params.DBType = "ElasticSearch"
	params.DestinationUrl = b.csvDaemonUrls
	updatedTags = reportTags
	updatedExtraVals = extraVals
	return
}

// scan reads encoded Queries and places them onto the workqueue.
func (b *ElasticQueryBenchmarker) RunScan(r io.Reader, closeChan chan int) {
	dec := gob.NewDecoder(r)

	n := int64(0)
loop:
	for {
		if bulk_query.Benchmarker.Limit() >= 0 && n >= bulk_query.Benchmarker.Limit() {
			break
		}

		q := b.queryPool.Get().(*Query)
		err := dec.Decode(q)
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal(err)
		}

		q.ID = n

		b.queryChan <- q

		n++
		select {
		case <-closeChan:
			fmt.Printf("Received finish request\n")
			break loop
		default:
		}
	}
	b.scanFinished = true

}

// processQueries reads byte buffers from queryChan and writes them to the
// target server, while tracking latency.
func (b *ElasticQueryBenchmarker) processQueries(w *HTTPClient, workersGroup *sync.WaitGroup, statPool sync.Pool, statChan chan *bulk_query.Stat) {
	opts := &HTTPClientDoOptions{
		Debug:                bulk_query.Benchmarker.Debug(),
		PrettyPrintResponses: bulk_query.Benchmarker.PrettyPrintResponses(),
	}
	var queriesSeen int64
	for q := range b.queryChan {
		lagMillis, err := w.Do(q, opts)

		stat := statPool.Get().(*bulk_query.Stat)
		stat.Init(q.HumanLabel, lagMillis)
		statChan <- stat

		b.queryPool.Put(q)
		if err != nil {
			log.Fatalf("Error during request: %s\n", err.Error())
		}

		queriesSeen++
	}
	workersGroup.Done()
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "net/http"
    "os"
    "reflect"
    "strings"

    "github.com/op/go-logging"
    "github.com/gorilla/websocket"
    "github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

// Get Consumer key/secret and Access token/secret from enviroment variables
var consumerKey = os.Getenv("TWITTER_CONSUMER_KEY")
var consumerSecret = os.Getenv("TWITTER_CONSUMER_SECRET")
var accessToken = os.Getenv("TWITTER_ACCESS_TOKEN")
var accessSecret = os.Getenv("TWITTER_ACCESS_SECRET")

var config = oauth1.NewConfig(consumerKey, consumerSecret)
var token = oauth1.NewToken(accessToken, accessSecret)

// Auto Oauth1
var httpClient = config.Client(oauth1.NoContext, token)

// Twitter Client
var client = twitter.NewClient(httpClient)

// Websocket Upgrade buffer
// Allow any domain can access stream API
var upgrader = websocket.Upgrader {
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool { return true },
}

// Store the tweet filter words
var filters = []string{};

// Store the WebSocket connections
var conns = make(map[string]*websocket.Conn)

var log = logging.MustGetLogger("Twitter Tweeds Log")

// Twitter stream
var globalTwitterStream *twitter.Stream

// Start the app then Listen and Serve
func main() {
    port := os.Getenv("PORT")

    if port == "" {
		log.Fatal("$PORT must be set")
	}

	// Serve static files
	fs := http.FileServer(http.Dir("static"))
    http.Handle("/", fs)

    http.HandleFunc("/twitter/stream", twitterStream)
    fmt.Println("TweetDeck started and ready to rock");
    http.ListenAndServe(":"+port,nil)
}

// Serve Websocket Twitter stream
func twitterStream(res http.ResponseWriter, req *http.Request) {
    log.Info("WS Twitter Request: ",req.URL.Query())
    filter := req.URL.Query()["filter"]
    if filter != nil {

        // Upgrades the http server connection to the websocket protocol
        conn, _ := upgrader.Upgrade(res, req, nil)

        filters = append(filters,filter[0]);
        conns[filter[0]] = conn;
        printConnsFiltes()

        // Stop the twitter stream in order to stream again with new filter
        stopGlobTwitterStream()

    	// Read only tweet from twitter stream
    	demux := twitter.NewSwitchDemux()
    	demux.Tweet = func(tweet *twitter.Tweet) {
    	    text := tweet.Text

    	    for _, v := range filters {
    	       if strings.Contains(text,v) && conns[v] != nil {
    	           wsWriter(conns[v],text)
    	       }
    	    }
    	}

    	fmt.Println("Starting Stream...")

    	// Filter
    	filterParams := &twitter.StreamFilterParams{
    		Track:         filters,
    		StallWarnings: twitter.Bool(true),
    	}

    	stream, err := client.Streams.Filter(filterParams)

    	globalTwitterStream = stream

    	fmt.Println(reflect.TypeOf(stream))
    	if err != nil {
    		log.Error(err)
    	}

        go wsReader(conn,filter[0],stream)      // Read messages from the websocket or stop websoket and twitter stream
    	go demux.HandleChan(stream.Messages)    // Receive twitter messages until stream quits

    } else {
        fmt.Fprintf(res, "Warning: Filter param is required to stream the data from Twitter")
    }
}

func stopGlobTwitterStream() {
    if globalTwitterStream != nil {
            globalTwitterStream.Stop()
    }
}

func printConnsFiltes() {
    fmt.Println("Filters:",filters)
    fmt.Println("Connections:",conns)
}

// Read messages from the websocket
func wsReader(conn *websocket.Conn,filter string,stream *twitter.Stream) {
    for {
        _, _, err := conn.ReadMessage()

        // Close the websocket connection
        if err != nil {
            // Close connection
            conn.Close()
             // Remove connections from connection list
            delete(conns, filter)

            for index, value := range filters {
                if value == filter {
                    // Remove filtered word from filter slice
                    filters = append(filters[:index], filters[index+1:]...)
                    break
                }
            }

            log.Info("wsReader/WS Connection Closed: "+filter)
            break
        } else {
            log.Error(err)
        }
    }
}

// Write messages to the websocket
func wsWriter(conn *websocket.Conn,filter string) {
    conn.WriteJSON(tweetStruct{
        Tweet: filter,
    })
}

type tweetStruct struct {
    Tweet string
}

func HelloTweetDeck() string {
    return "Hello World"
}<commit_msg>Fix HelloTweetDeck()'s return text<commit_after>package main

import (
    "fmt"
    "net/http"
    "os"
    "reflect"
    "strings"

    "github.com/op/go-logging"
    "github.com/gorilla/websocket"
    "github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

// Get Consumer key/secret and Access token/secret from enviroment variables
var consumerKey = os.Getenv("TWITTER_CONSUMER_KEY")
var consumerSecret = os.Getenv("TWITTER_CONSUMER_SECRET")
var accessToken = os.Getenv("TWITTER_ACCESS_TOKEN")
var accessSecret = os.Getenv("TWITTER_ACCESS_SECRET")

var config = oauth1.NewConfig(consumerKey, consumerSecret)
var token = oauth1.NewToken(accessToken, accessSecret)

// Auto Oauth1
var httpClient = config.Client(oauth1.NoContext, token)

// Twitter Client
var client = twitter.NewClient(httpClient)

// Websocket Upgrade buffer
// Allow any domain can access stream API
var upgrader = websocket.Upgrader {
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool { return true },
}

// Store the tweet filter words
var filters = []string{};

// Store the WebSocket connections
var conns = make(map[string]*websocket.Conn)

var log = logging.MustGetLogger("Twitter Tweeds Log")

// Twitter stream
var globalTwitterStream *twitter.Stream

// Start the app then Listen and Serve
func main() {
    port := os.Getenv("PORT")

    if port == "" {
		log.Fatal("$PORT must be set")
	}

	// Serve static files
	fs := http.FileServer(http.Dir("static"))
    http.Handle("/", fs)

    http.HandleFunc("/twitter/stream", twitterStream)
    fmt.Println("TweetDeck started and ready to rock");
    http.ListenAndServe(":"+port,nil)
}

// Serve Websocket Twitter stream
func twitterStream(res http.ResponseWriter, req *http.Request) {
    log.Info("WS Twitter Request: ",req.URL.Query())
    filter := req.URL.Query()["filter"]
    if filter != nil {

        // Upgrades the http server connection to the websocket protocol
        conn, _ := upgrader.Upgrade(res, req, nil)

        filters = append(filters,filter[0]);
        conns[filter[0]] = conn;
        printConnsFiltes()

        // Stop the twitter stream in order to stream again with new filter
        stopGlobTwitterStream()

    	// Read only tweet from twitter stream
    	demux := twitter.NewSwitchDemux()
    	demux.Tweet = func(tweet *twitter.Tweet) {
    	    text := tweet.Text

    	    for _, v := range filters {
    	       if strings.Contains(text,v) && conns[v] != nil {
    	           wsWriter(conns[v],text)
    	       }
    	    }
    	}

    	fmt.Println("Starting Stream...")

    	// Filter
    	filterParams := &twitter.StreamFilterParams{
    		Track:         filters,
    		StallWarnings: twitter.Bool(true),
    	}

    	stream, err := client.Streams.Filter(filterParams)

    	globalTwitterStream = stream

    	fmt.Println(reflect.TypeOf(stream))
    	if err != nil {
    		log.Error(err)
    	}

        go wsReader(conn,filter[0],stream)      // Read messages from the websocket or stop websoket and twitter stream
    	go demux.HandleChan(stream.Messages)    // Receive twitter messages until stream quits

    } else {
        fmt.Fprintf(res, "Warning: Filter param is required to stream the data from Twitter")
    }
}

func stopGlobTwitterStream() {
    if globalTwitterStream != nil {
            globalTwitterStream.Stop()
    }
}

func printConnsFiltes() {
    fmt.Println("Filters:",filters)
    fmt.Println("Connections:",conns)
}

// Read messages from the websocket
func wsReader(conn *websocket.Conn,filter string,stream *twitter.Stream) {
    for {
        _, _, err := conn.ReadMessage()

        // Close the websocket connection
        if err != nil {
            // Close connection
            conn.Close()
             // Remove connections from connection list
            delete(conns, filter)

            for index, value := range filters {
                if value == filter {
                    // Remove filtered word from filter slice
                    filters = append(filters[:index], filters[index+1:]...)
                    break
                }
            }

            log.Info("wsReader/WS Connection Closed: "+filter)
            break
        } else {
            log.Error(err)
        }
    }
}

// Write messages to the websocket
func wsWriter(conn *websocket.Conn,filter string) {
    conn.WriteJSON(tweetStruct{
        Tweet: filter,
    })
}

type tweetStruct struct {
    Tweet string
}

func HelloTweetDeck() string {
    return "Tweet Deck Test"
}<|endoftext|>