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

// Package menu displays a Terminal UI based text menu to choose boot options
// from.
package menu

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/sh"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/sys/unix"
)

var (
	initialTimeout    = 10 * time.Second
	subsequentTimeout = 60 * time.Second
)

// Entry is a menu entry.
type Entry interface {
	// Label is the string displayed to the user in the menu.
	Label() string

	// Edit the kernel command line if possible. Must be called prior to
	// Load.
	Edit(func(cmdline string) string)

	// Load is called when the entry is chosen, but does not transfer
	// execution to another process or kernel.
	Load() error

	// Exec transfers execution to another process or kernel.
	//
	// Exec either returns an error or does not return at all.
	Exec() error

	// IsDefault indicates that this action should be run by default if the
	// user didn't make an entry choice.
	IsDefault() bool
}

func parseBootNum(choice string, entries []Entry) (int, error) {
	num, err := strconv.Atoi(choice)
	if err != nil {
		return -1, fmt.Errorf("%s is not a valid entry number: %v", choice, err)
	}
	if num < 1 || num > len(entries) {
		return -1, fmt.Errorf("%s is not a valid entry number", choice)
	}
	return num, nil
}

// Choose presents the user a menu on input to choose an entry from and returns that entry.
func Choose(input *os.File, entries ...Entry) Entry {
	fmt.Println("")
	for i, e := range entries {
		fmt.Printf("%02d. %s\n\n", i+1, e.Label())
	}
	fmt.Println("\r")

	oldState, err := terminal.MakeRaw(int(input.Fd()))
	if err != nil {
		log.Printf("BUG: Please report: We cannot actually let you choose from menu (MakeRaw failed): %v", err)
		return nil
	}
	defer terminal.Restore(int(input.Fd()), oldState)

	// TODO(chrisko): reduce this timeout a la GRUB. 3 seconds, and hitting
	// any button resets the timeout. We could save 7 seconds here.
	t := time.NewTimer(initialTimeout)

	boot := make(chan Entry, 1)

	go func() {
		// Note that term is in raw mode. Write \r\n whenever you would
		// write a \n. When testing in qemu, it might look fine because
		// there might be another tty cooking the newlines. In for
		// example minicom, the behavior is different. And you would
		// see something like:
		//
		//     Select a boot option to edit:
		//                                  >
		//
		// Instead of:
		//
		//     Select a boot option to edit:
		//      >
		term := terminal.NewTerminal(input, "")

		term.AutoCompleteCallback = func(line string, pos int, key rune) (string, int, bool) {
			// We ain't gonna autocomplete, but we'll reset the countdown timer when you press a key.
			t.Reset(subsequentTimeout)
			return "", 0, false
		}

		for {
			term.SetPrompt("Enter an option ('01' is the default, 'e' to edit kernel cmdline):\r\n > ")
			choice, err := term.ReadLine()
			if err != nil {
				if err != io.EOF {
					fmt.Printf("BUG: Please report: Terminal read error: %v.\n", err)
				}
				boot <- nil
				return
			}

			if choice == "e" {
				// Edit command line.
				term.SetPrompt("Select a boot option to edit:\r\n > ")
				choice, err := term.ReadLine()
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				num, err := parseBootNum(choice, entries)
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				entries[num-1].Edit(func(cmdline string) string {
					fmt.Fprintf(term, "The current quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					fmt.Fprintln(term, ` * Note the cmdline is c-style quoted. Ex: \n => newline, \\ => \`)
					term.SetPrompt("Enter an option:\r\n * (a)ppend, (o)verwrite, (r)eturn to main menu\r\n > ")
					choice, err := term.ReadLine()
					if err != nil {
						fmt.Fprintln(term, err)
						return cmdline
					}
					switch choice {
					case "a":
						term.SetPrompt("Enter unquoted cmdline to append:\r\n > ")
						appendCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						if appendCmdline != "" {
							cmdline += " " + appendCmdline
						}
					case "o":
						term.SetPrompt("Enter new unquoted cmdline:\r\n > ")
						newCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						cmdline = newCmdline
					case "r":
					default:
						fmt.Fprintf(term, "Unrecognized choice %q", choice)
					}
					fmt.Fprintf(term, "The new quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					return cmdline
				})
				fmt.Fprintln(term, "Returning to main menu...")
				continue
			}
			if choice == "" {
				// nil will result in the default order.
				boot <- nil
				return
			}
			num, err := parseBootNum(choice, entries)
			if err != nil {
				fmt.Fprintln(term, err)
				continue
			}
			boot <- entries[num-1]
			return
		}
	}()

	select {
	case entry := <-boot:
		if entry != nil {
			fmt.Printf("Chosen option %s.\r\n\r\n", entry.Label())
		}
		return entry

	case <-t.C:
		return nil
	}
}

// ShowMenuAndLoad lets the user choose one of entries and loads it. If no
// entry is chosen by the user, an entry whose IsDefault() is true will be
// returned.
//
// The user is left to call Entry.Exec when this function returns.
func ShowMenuAndLoad(input *os.File, entries ...Entry) Entry {
	// Clear the screen (ANSI terminal escape code for screen clear).
	fmt.Printf("\033[1;1H\033[2J\n\n")
	fmt.Printf("Welcome to NERF's Boot Menu\n\n")
	fmt.Printf("Enter a number to boot a kernel:\n")

	for {
		// Allow the user to choose.
		entry := Choose(input, entries...)
		if entry == nil {
			// This only returns something if the user explicitly
			// entered something.
			//
			// If nothing was entered, fall back to default.
			break
		}
		if err := entry.Load(); err != nil {
			log.Printf("Failed to load %s: %v", entry.Label(), err)
			continue
		}

		// Entry was successfully loaded. Leave it to the caller to
		// exec, so the caller can clean up the OS before rebooting or
		// kexecing (e.g. unmount file systems).
		return entry
	}

	fmt.Println("")

	// We only get one shot at actually booting, so boot the first kernel
	// that can be loaded correctly.
	for _, e := range entries {
		// Only perform actions that are default actions. I.e. don't
		// drop to shell.
		if e.IsDefault() {
			fmt.Printf("Attempting to boot %s.\n\n", e)

			if err := e.Load(); err != nil {
				log.Printf("Failed to load %s: %v", e.Label(), err)
				continue
			}

			// Entry was successfully loaded. Leave it to the
			// caller to exec, so the caller can clean up the OS
			// before rebooting or kexecing (e.g. unmount file
			// systems).
			return e
		}
	}
	return nil
}

// OSImages returns menu entries for the given OSImages.
func OSImages(verbose bool, imgs ...boot.OSImage) []Entry {
	var menu []Entry
	for _, img := range imgs {
		menu = append(menu, &OSImageAction{
			OSImage: img,
			Verbose: verbose,
		})
	}
	return menu
}

// OSImageAction is a menu.Entry that boots an OSImage.
type OSImageAction struct {
	boot.OSImage
	Verbose bool
}

// Load implements Entry.Load by loading the OS image into memory.
func (oia OSImageAction) Load() error {
	if err := oia.OSImage.Load(oia.Verbose); err != nil {
		return fmt.Errorf("could not load image %s: %v", oia.OSImage, err)
	}
	return nil
}

// Exec executes the loaded image.
func (oia OSImageAction) Exec() error {
	return boot.Execute()
}

// IsDefault returns true -- this action should be performed in order by
// default if the user did not choose a boot entry.
func (OSImageAction) IsDefault() bool { return true }

// StartShell is a menu.Entry that starts a LinuxBoot shell.
type StartShell struct{}

// Label is the label to show to the user.
func (StartShell) Label() string {
	return "Enter a LinuxBoot shell"
}

// Edit does nothing.
func (StartShell) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (StartShell) Load() error {
	return nil
}

// Exec implements Entry.Exec by running /bin/defaultsh.
func (StartShell) Exec() error {
	// Reset signal handler for SIGINT to enable user interrupts again
	signal.Reset(syscall.SIGINT)
	return sh.RunWithLogs("/bin/defaultsh")
}

// IsDefault indicates that this should not be run as a default action.
func (StartShell) IsDefault() bool { return false }

// Reboot is a menu.Entry that reboots the machine.
type Reboot struct{}

// Label is the label to show to the user.
func (Reboot) Label() string {
	return "Reboot"
}

// Edit does nothing.
func (Reboot) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (Reboot) Load() error {
	return nil
}

// Exec reboots the machine using sys_reboot.
func (Reboot) Exec() error {
	unix.Sync()
	return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
}

// IsDefault indicates that this should not be run as a default action.
func (Reboot) IsDefault() bool { return false }
<commit_msg>menu: Change 'NERF Boot Menu' to 'LinuxBoot Menu'<commit_after>// Copyright 2020 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package menu displays a Terminal UI based text menu to choose boot options
// from.
package menu

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/sh"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/sys/unix"
)

var (
	initialTimeout    = 10 * time.Second
	subsequentTimeout = 60 * time.Second
)

// Entry is a menu entry.
type Entry interface {
	// Label is the string displayed to the user in the menu.
	Label() string

	// Edit the kernel command line if possible. Must be called prior to
	// Load.
	Edit(func(cmdline string) string)

	// Load is called when the entry is chosen, but does not transfer
	// execution to another process or kernel.
	Load() error

	// Exec transfers execution to another process or kernel.
	//
	// Exec either returns an error or does not return at all.
	Exec() error

	// IsDefault indicates that this action should be run by default if the
	// user didn't make an entry choice.
	IsDefault() bool
}

func parseBootNum(choice string, entries []Entry) (int, error) {
	num, err := strconv.Atoi(choice)
	if err != nil {
		return -1, fmt.Errorf("%s is not a valid entry number: %v", choice, err)
	}
	if num < 1 || num > len(entries) {
		return -1, fmt.Errorf("%s is not a valid entry number", choice)
	}
	return num, nil
}

// Choose presents the user a menu on input to choose an entry from and returns that entry.
func Choose(input *os.File, entries ...Entry) Entry {
	fmt.Println("")
	for i, e := range entries {
		fmt.Printf("%02d. %s\n\n", i+1, e.Label())
	}
	fmt.Println("\r")

	oldState, err := terminal.MakeRaw(int(input.Fd()))
	if err != nil {
		log.Printf("BUG: Please report: We cannot actually let you choose from menu (MakeRaw failed): %v", err)
		return nil
	}
	defer terminal.Restore(int(input.Fd()), oldState)

	// TODO(chrisko): reduce this timeout a la GRUB. 3 seconds, and hitting
	// any button resets the timeout. We could save 7 seconds here.
	t := time.NewTimer(initialTimeout)

	boot := make(chan Entry, 1)

	go func() {
		// Note that term is in raw mode. Write \r\n whenever you would
		// write a \n. When testing in qemu, it might look fine because
		// there might be another tty cooking the newlines. In for
		// example minicom, the behavior is different. And you would
		// see something like:
		//
		//     Select a boot option to edit:
		//                                  >
		//
		// Instead of:
		//
		//     Select a boot option to edit:
		//      >
		term := terminal.NewTerminal(input, "")

		term.AutoCompleteCallback = func(line string, pos int, key rune) (string, int, bool) {
			// We ain't gonna autocomplete, but we'll reset the countdown timer when you press a key.
			t.Reset(subsequentTimeout)
			return "", 0, false
		}

		for {
			term.SetPrompt("Enter an option ('01' is the default, 'e' to edit kernel cmdline):\r\n > ")
			choice, err := term.ReadLine()
			if err != nil {
				if err != io.EOF {
					fmt.Printf("BUG: Please report: Terminal read error: %v.\n", err)
				}
				boot <- nil
				return
			}

			if choice == "e" {
				// Edit command line.
				term.SetPrompt("Select a boot option to edit:\r\n > ")
				choice, err := term.ReadLine()
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				num, err := parseBootNum(choice, entries)
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				entries[num-1].Edit(func(cmdline string) string {
					fmt.Fprintf(term, "The current quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					fmt.Fprintln(term, ` * Note the cmdline is c-style quoted. Ex: \n => newline, \\ => \`)
					term.SetPrompt("Enter an option:\r\n * (a)ppend, (o)verwrite, (r)eturn to main menu\r\n > ")
					choice, err := term.ReadLine()
					if err != nil {
						fmt.Fprintln(term, err)
						return cmdline
					}
					switch choice {
					case "a":
						term.SetPrompt("Enter unquoted cmdline to append:\r\n > ")
						appendCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						if appendCmdline != "" {
							cmdline += " " + appendCmdline
						}
					case "o":
						term.SetPrompt("Enter new unquoted cmdline:\r\n > ")
						newCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						cmdline = newCmdline
					case "r":
					default:
						fmt.Fprintf(term, "Unrecognized choice %q", choice)
					}
					fmt.Fprintf(term, "The new quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					return cmdline
				})
				fmt.Fprintln(term, "Returning to main menu...")
				continue
			}
			if choice == "" {
				// nil will result in the default order.
				boot <- nil
				return
			}
			num, err := parseBootNum(choice, entries)
			if err != nil {
				fmt.Fprintln(term, err)
				continue
			}
			boot <- entries[num-1]
			return
		}
	}()

	select {
	case entry := <-boot:
		if entry != nil {
			fmt.Printf("Chosen option %s.\r\n\r\n", entry.Label())
		}
		return entry

	case <-t.C:
		return nil
	}
}

// ShowMenuAndLoad lets the user choose one of entries and loads it. If no
// entry is chosen by the user, an entry whose IsDefault() is true will be
// returned.
//
// The user is left to call Entry.Exec when this function returns.
func ShowMenuAndLoad(input *os.File, entries ...Entry) Entry {
	// Clear the screen (ANSI terminal escape code for screen clear).
	fmt.Printf("\033[1;1H\033[2J\n\n")
	fmt.Printf("Welcome to LinuxBoot's Menu\n\n")
	fmt.Printf("Enter a number to boot a kernel:\n")

	for {
		// Allow the user to choose.
		entry := Choose(input, entries...)
		if entry == nil {
			// This only returns something if the user explicitly
			// entered something.
			//
			// If nothing was entered, fall back to default.
			break
		}
		if err := entry.Load(); err != nil {
			log.Printf("Failed to load %s: %v", entry.Label(), err)
			continue
		}

		// Entry was successfully loaded. Leave it to the caller to
		// exec, so the caller can clean up the OS before rebooting or
		// kexecing (e.g. unmount file systems).
		return entry
	}

	fmt.Println("")

	// We only get one shot at actually booting, so boot the first kernel
	// that can be loaded correctly.
	for _, e := range entries {
		// Only perform actions that are default actions. I.e. don't
		// drop to shell.
		if e.IsDefault() {
			fmt.Printf("Attempting to boot %s.\n\n", e)

			if err := e.Load(); err != nil {
				log.Printf("Failed to load %s: %v", e.Label(), err)
				continue
			}

			// Entry was successfully loaded. Leave it to the
			// caller to exec, so the caller can clean up the OS
			// before rebooting or kexecing (e.g. unmount file
			// systems).
			return e
		}
	}
	return nil
}

// OSImages returns menu entries for the given OSImages.
func OSImages(verbose bool, imgs ...boot.OSImage) []Entry {
	var menu []Entry
	for _, img := range imgs {
		menu = append(menu, &OSImageAction{
			OSImage: img,
			Verbose: verbose,
		})
	}
	return menu
}

// OSImageAction is a menu.Entry that boots an OSImage.
type OSImageAction struct {
	boot.OSImage
	Verbose bool
}

// Load implements Entry.Load by loading the OS image into memory.
func (oia OSImageAction) Load() error {
	if err := oia.OSImage.Load(oia.Verbose); err != nil {
		return fmt.Errorf("could not load image %s: %v", oia.OSImage, err)
	}
	return nil
}

// Exec executes the loaded image.
func (oia OSImageAction) Exec() error {
	return boot.Execute()
}

// IsDefault returns true -- this action should be performed in order by
// default if the user did not choose a boot entry.
func (OSImageAction) IsDefault() bool { return true }

// StartShell is a menu.Entry that starts a LinuxBoot shell.
type StartShell struct{}

// Label is the label to show to the user.
func (StartShell) Label() string {
	return "Enter a LinuxBoot shell"
}

// Edit does nothing.
func (StartShell) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (StartShell) Load() error {
	return nil
}

// Exec implements Entry.Exec by running /bin/defaultsh.
func (StartShell) Exec() error {
	// Reset signal handler for SIGINT to enable user interrupts again
	signal.Reset(syscall.SIGINT)
	return sh.RunWithLogs("/bin/defaultsh")
}

// IsDefault indicates that this should not be run as a default action.
func (StartShell) IsDefault() bool { return false }

// Reboot is a menu.Entry that reboots the machine.
type Reboot struct{}

// Label is the label to show to the user.
func (Reboot) Label() string {
	return "Reboot"
}

// Edit does nothing.
func (Reboot) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (Reboot) Load() error {
	return nil
}

// Exec reboots the machine using sys_reboot.
func (Reboot) Exec() error {
	unix.Sync()
	return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
}

// IsDefault indicates that this should not be run as a default action.
func (Reboot) IsDefault() bool { return false }
<|endoftext|>


<commit_before>// Copyright (c) 2013 The Go Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd.

// Package lintutil provides helpers for writing linter command lines.
package lintutil // import "honnef.co/go/tools/lint/lintutil"

import (
	"crypto/sha256"
	"errors"
	"flag"
	"fmt"
	"go/build"
	"go/token"
	"io"
	"log"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync/atomic"

	"honnef.co/go/tools/config"
	"honnef.co/go/tools/internal/cache"
	"honnef.co/go/tools/lint"
	"honnef.co/go/tools/lint/lintutil/format"
	"honnef.co/go/tools/version"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/packages"
)

func NewVersionFlag() flag.Getter {
	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}
	return version
}

type VersionFlag int

func (v *VersionFlag) String() string {
	return fmt.Sprintf("1.%d", *v)

}

func (v *VersionFlag) Set(s string) error {
	if len(s) < 3 {
		return errors.New("invalid Go version")
	}
	if s[0] != '1' {
		return errors.New("invalid Go version")
	}
	if s[1] != '.' {
		return errors.New("invalid Go version")
	}
	i, err := strconv.Atoi(s[2:])
	*v = VersionFlag(i)
	return err
}

func (v *VersionFlag) Get() interface{} {
	return int(*v)
}

func usage(name string, flags *flag.FlagSet) func() {
	return func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] # runs on package in current directory\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] packages\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] directory\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] files... # must be a single package\n", name)
		fmt.Fprintf(os.Stderr, "Flags:\n")
		flags.PrintDefaults()
	}
}

type list []string

func (list *list) String() string {
	return `"` + strings.Join(*list, ",") + `"`
}

func (list *list) Set(s string) error {
	if s == "" {
		*list = nil
		return nil
	}

	*list = strings.Split(s, ",")
	return nil
}

func FlagSet(name string) *flag.FlagSet {
	flags := flag.NewFlagSet("", flag.ExitOnError)
	flags.Usage = usage(name, flags)
	flags.String("tags", "", "List of `build tags`")
	flags.String("ignore", "", "Deprecated: use linter directives instead")
	flags.Bool("tests", true, "Include tests")
	flags.Bool("version", false, "Print version and exit")
	flags.Bool("show-ignored", false, "Don't filter ignored problems")
	flags.String("f", "text", "Output `format` (valid choices are 'stylish', 'text' and 'json')")
	flags.String("explain", "", "Print description of `check`")

	flags.String("debug.cpuprofile", "", "Write CPU profile to `file`")
	flags.String("debug.memprofile", "", "Write memory profile to `file`")
	flags.Bool("debug.version", false, "Print detailed version information about this program")

	checks := list{"inherit"}
	fail := list{"all"}
	flags.Var(&checks, "checks", "Comma-separated list of `checks` to enable.")
	flags.Var(&fail, "fail", "Comma-separated list of `checks` that can cause a non-zero exit status.")

	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}

	flags.Var(version, "go", "Target Go `version` in the format '1.x'")
	return flags
}

func findCheck(cs []*analysis.Analyzer, check string) (*analysis.Analyzer, bool) {
	for _, c := range cs {
		if c.Name == check {
			return c, true
		}
	}
	return nil, false
}

func ProcessFlagSet(cs []*analysis.Analyzer, cums []lint.CumulativeChecker, fs *flag.FlagSet) {
	tags := fs.Lookup("tags").Value.(flag.Getter).Get().(string)
	ignore := fs.Lookup("ignore").Value.(flag.Getter).Get().(string)
	tests := fs.Lookup("tests").Value.(flag.Getter).Get().(bool)
	goVersion := fs.Lookup("go").Value.(flag.Getter).Get().(int)
	formatter := fs.Lookup("f").Value.(flag.Getter).Get().(string)
	printVersion := fs.Lookup("version").Value.(flag.Getter).Get().(bool)
	showIgnored := fs.Lookup("show-ignored").Value.(flag.Getter).Get().(bool)
	explain := fs.Lookup("explain").Value.(flag.Getter).Get().(string)

	cpuProfile := fs.Lookup("debug.cpuprofile").Value.(flag.Getter).Get().(string)
	memProfile := fs.Lookup("debug.memprofile").Value.(flag.Getter).Get().(string)
	debugVersion := fs.Lookup("debug.version").Value.(flag.Getter).Get().(bool)

	cfg := config.Config{}
	cfg.Checks = *fs.Lookup("checks").Value.(*list)

	exit := func(code int) {
		if cpuProfile != "" {
			pprof.StopCPUProfile()
		}
		if memProfile != "" {
			f, err := os.Create(memProfile)
			if err != nil {
				panic(err)
			}
			runtime.GC()
			pprof.WriteHeapProfile(f)
		}
		os.Exit(code)
	}
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}

	if debugVersion {
		version.Verbose()
		exit(0)
	}

	if printVersion {
		version.Print()
		exit(0)
	}

	if explain != "" {
		var haystack []*analysis.Analyzer
		haystack = append(haystack, cs...)
		for _, cum := range cums {
			haystack = append(haystack, cum.Analyzer())
		}
		check, ok := findCheck(haystack, explain)
		if !ok {
			fmt.Fprintln(os.Stderr, "Couldn't find check", explain)
			exit(1)
		}
		if check.Doc == "" {
			fmt.Fprintln(os.Stderr, explain, "has no documentation")
			exit(1)
		}
		fmt.Println(check.Doc)
		exit(0)
	}

	ps, err := Lint(cs, cums, fs.Args(), &Options{
		Tags:      strings.Fields(tags),
		LintTests: tests,
		Ignores:   ignore,
		GoVersion: goVersion,
		Config:    cfg,
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		exit(1)
	}

	var f format.Formatter
	switch formatter {
	case "text":
		f = format.Text{W: os.Stdout}
	case "stylish":
		f = &format.Stylish{W: os.Stdout}
	case "json":
		f = format.JSON{W: os.Stdout}
	default:
		fmt.Fprintf(os.Stderr, "unsupported output format %q\n", formatter)
		exit(2)
	}

	var (
		total    int
		errors   int
		warnings int
	)

	fail := *fs.Lookup("fail").Value.(*list)
	analyzers := make([]*analysis.Analyzer, len(cs), len(cs)+len(cums))
	copy(analyzers, cs)
	for _, cum := range cums {
		analyzers = append(analyzers, cum.Analyzer())
	}
	shouldExit := lint.FilterChecks(analyzers, fail)
	shouldExit["compile"] = true

	total = len(ps)
	for _, p := range ps {
		if p.Severity == lint.Ignored && !showIgnored {
			continue
		}
		if shouldExit[p.Check] {
			errors++
		} else {
			p.Severity = lint.Warning
			warnings++
		}
		f.Format(p)
	}
	if f, ok := f.(format.Statter); ok {
		f.Stats(total, errors, warnings)
	}
	if errors > 0 {
		exit(1)
	}
	exit(0)
}

type Options struct {
	Config config.Config

	Tags      []string
	LintTests bool
	Ignores   string
	GoVersion int
}

func computeSalt() ([]byte, error) {
	if version.Version != "devel" {
		return []byte(version.Version), nil
	}
	p, err := os.Executable()
	if err != nil {
		return nil, err
	}
	f, err := os.Open(p)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return nil, err
	}
	return h.Sum(nil), nil
}

func Lint(cs []*analysis.Analyzer, cums []lint.CumulativeChecker, paths []string, opt *Options) ([]lint.Problem, error) {
	salt, err := computeSalt()
	if err != nil {
		return nil, fmt.Errorf("could not compute salt for cache: %s", err)
	}
	cache.SetSalt(salt)

	if opt == nil {
		opt = &Options{}
	}

	l := &lint.Linter{
		Checkers:           cs,
		CumulativeCheckers: cums,
		GoVersion:          opt.GoVersion,
		Config:             opt.Config,
	}
	cfg := &packages.Config{}
	if opt.LintTests {
		cfg.Tests = true
	}

	printStats := func() {
		// Individual stats are read atomically, but overall there
		// is no synchronisation. For printing rough progress
		// information, this doesn't matter.
		switch atomic.LoadUint64(&l.Stats.State) {
		case lint.StateInitializing:
			fmt.Fprintln(os.Stderr, "Status: initializing")
		case lint.StateGraph:
			fmt.Fprintln(os.Stderr, "Status: loading package graph")
		case lint.StateProcessing:
			fmt.Fprintf(os.Stderr, "Packages: %d/%d initial, %d/%d total; Workers: %d/%d; Problems: %d\n",
				atomic.LoadUint64(&l.Stats.ProcessedInitialPackages),
				atomic.LoadUint64(&l.Stats.InitialPackages),
				atomic.LoadUint64(&l.Stats.ProcessedPackages),
				atomic.LoadUint64(&l.Stats.TotalPackages),
				atomic.LoadUint64(&l.Stats.ActiveWorkers),
				atomic.LoadUint64(&l.Stats.TotalWorkers),
				atomic.LoadUint64(&l.Stats.Problems),
			)
		case lint.StateCumulative:
			fmt.Fprintln(os.Stderr, "Status: processing cumulative checkers")
		}
	}
	if len(infoSignals) > 0 {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, infoSignals...)
		defer signal.Stop(ch)
		go func() {
			for range ch {
				printStats()
			}
		}()
	}

	return l.Lint(cfg, paths)
}

var posRe = regexp.MustCompile(`^(.+?):(\d+)(?::(\d+)?)?$`)

func parsePos(pos string) token.Position {
	if pos == "-" || pos == "" {
		return token.Position{}
	}
	parts := posRe.FindStringSubmatch(pos)
	if parts == nil {
		panic(fmt.Sprintf("internal error: malformed position %q", pos))
	}
	file := parts[1]
	line, _ := strconv.Atoi(parts[2])
	col, _ := strconv.Atoi(parts[3])
	return token.Position{
		Filename: file,
		Line:     line,
		Column:   col,
	}
}
<commit_msg>lint/lintutil: add debug flag for hiding compile errors<commit_after>// Copyright (c) 2013 The Go Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd.

// Package lintutil provides helpers for writing linter command lines.
package lintutil // import "honnef.co/go/tools/lint/lintutil"

import (
	"crypto/sha256"
	"errors"
	"flag"
	"fmt"
	"go/build"
	"go/token"
	"io"
	"log"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"sync/atomic"

	"honnef.co/go/tools/config"
	"honnef.co/go/tools/internal/cache"
	"honnef.co/go/tools/lint"
	"honnef.co/go/tools/lint/lintutil/format"
	"honnef.co/go/tools/version"

	"golang.org/x/tools/go/analysis"
	"golang.org/x/tools/go/packages"
)

func NewVersionFlag() flag.Getter {
	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}
	return version
}

type VersionFlag int

func (v *VersionFlag) String() string {
	return fmt.Sprintf("1.%d", *v)

}

func (v *VersionFlag) Set(s string) error {
	if len(s) < 3 {
		return errors.New("invalid Go version")
	}
	if s[0] != '1' {
		return errors.New("invalid Go version")
	}
	if s[1] != '.' {
		return errors.New("invalid Go version")
	}
	i, err := strconv.Atoi(s[2:])
	*v = VersionFlag(i)
	return err
}

func (v *VersionFlag) Get() interface{} {
	return int(*v)
}

func usage(name string, flags *flag.FlagSet) func() {
	return func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] # runs on package in current directory\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] packages\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] directory\n", name)
		fmt.Fprintf(os.Stderr, "\t%s [flags] files... # must be a single package\n", name)
		fmt.Fprintf(os.Stderr, "Flags:\n")
		flags.PrintDefaults()
	}
}

type list []string

func (list *list) String() string {
	return `"` + strings.Join(*list, ",") + `"`
}

func (list *list) Set(s string) error {
	if s == "" {
		*list = nil
		return nil
	}

	*list = strings.Split(s, ",")
	return nil
}

func FlagSet(name string) *flag.FlagSet {
	flags := flag.NewFlagSet("", flag.ExitOnError)
	flags.Usage = usage(name, flags)
	flags.String("tags", "", "List of `build tags`")
	flags.String("ignore", "", "Deprecated: use linter directives instead")
	flags.Bool("tests", true, "Include tests")
	flags.Bool("version", false, "Print version and exit")
	flags.Bool("show-ignored", false, "Don't filter ignored problems")
	flags.String("f", "text", "Output `format` (valid choices are 'stylish', 'text' and 'json')")
	flags.String("explain", "", "Print description of `check`")

	flags.String("debug.cpuprofile", "", "Write CPU profile to `file`")
	flags.String("debug.memprofile", "", "Write memory profile to `file`")
	flags.Bool("debug.version", false, "Print detailed version information about this program")
	flags.Bool("debug.no-compile-errors", false, "Don't print compile errors")

	checks := list{"inherit"}
	fail := list{"all"}
	flags.Var(&checks, "checks", "Comma-separated list of `checks` to enable.")
	flags.Var(&fail, "fail", "Comma-separated list of `checks` that can cause a non-zero exit status.")

	tags := build.Default.ReleaseTags
	v := tags[len(tags)-1][2:]
	version := new(VersionFlag)
	if err := version.Set(v); err != nil {
		panic(fmt.Sprintf("internal error: %s", err))
	}

	flags.Var(version, "go", "Target Go `version` in the format '1.x'")
	return flags
}

func findCheck(cs []*analysis.Analyzer, check string) (*analysis.Analyzer, bool) {
	for _, c := range cs {
		if c.Name == check {
			return c, true
		}
	}
	return nil, false
}

func ProcessFlagSet(cs []*analysis.Analyzer, cums []lint.CumulativeChecker, fs *flag.FlagSet) {
	tags := fs.Lookup("tags").Value.(flag.Getter).Get().(string)
	ignore := fs.Lookup("ignore").Value.(flag.Getter).Get().(string)
	tests := fs.Lookup("tests").Value.(flag.Getter).Get().(bool)
	goVersion := fs.Lookup("go").Value.(flag.Getter).Get().(int)
	formatter := fs.Lookup("f").Value.(flag.Getter).Get().(string)
	printVersion := fs.Lookup("version").Value.(flag.Getter).Get().(bool)
	showIgnored := fs.Lookup("show-ignored").Value.(flag.Getter).Get().(bool)
	explain := fs.Lookup("explain").Value.(flag.Getter).Get().(string)

	cpuProfile := fs.Lookup("debug.cpuprofile").Value.(flag.Getter).Get().(string)
	memProfile := fs.Lookup("debug.memprofile").Value.(flag.Getter).Get().(string)
	debugVersion := fs.Lookup("debug.version").Value.(flag.Getter).Get().(bool)
	debugNoCompile := fs.Lookup("debug.no-compile-errors").Value.(flag.Getter).Get().(bool)

	cfg := config.Config{}
	cfg.Checks = *fs.Lookup("checks").Value.(*list)

	exit := func(code int) {
		if cpuProfile != "" {
			pprof.StopCPUProfile()
		}
		if memProfile != "" {
			f, err := os.Create(memProfile)
			if err != nil {
				panic(err)
			}
			runtime.GC()
			pprof.WriteHeapProfile(f)
		}
		os.Exit(code)
	}
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}

	if debugVersion {
		version.Verbose()
		exit(0)
	}

	if printVersion {
		version.Print()
		exit(0)
	}

	if explain != "" {
		var haystack []*analysis.Analyzer
		haystack = append(haystack, cs...)
		for _, cum := range cums {
			haystack = append(haystack, cum.Analyzer())
		}
		check, ok := findCheck(haystack, explain)
		if !ok {
			fmt.Fprintln(os.Stderr, "Couldn't find check", explain)
			exit(1)
		}
		if check.Doc == "" {
			fmt.Fprintln(os.Stderr, explain, "has no documentation")
			exit(1)
		}
		fmt.Println(check.Doc)
		exit(0)
	}

	ps, err := Lint(cs, cums, fs.Args(), &Options{
		Tags:      strings.Fields(tags),
		LintTests: tests,
		Ignores:   ignore,
		GoVersion: goVersion,
		Config:    cfg,
	})
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		exit(1)
	}

	var f format.Formatter
	switch formatter {
	case "text":
		f = format.Text{W: os.Stdout}
	case "stylish":
		f = &format.Stylish{W: os.Stdout}
	case "json":
		f = format.JSON{W: os.Stdout}
	default:
		fmt.Fprintf(os.Stderr, "unsupported output format %q\n", formatter)
		exit(2)
	}

	var (
		total    int
		errors   int
		warnings int
	)

	fail := *fs.Lookup("fail").Value.(*list)
	analyzers := make([]*analysis.Analyzer, len(cs), len(cs)+len(cums))
	copy(analyzers, cs)
	for _, cum := range cums {
		analyzers = append(analyzers, cum.Analyzer())
	}
	shouldExit := lint.FilterChecks(analyzers, fail)
	shouldExit["compile"] = true

	total = len(ps)
	for _, p := range ps {
		if p.Check == "compile" && debugNoCompile {
			continue
		}
		if p.Severity == lint.Ignored && !showIgnored {
			continue
		}
		if shouldExit[p.Check] {
			errors++
		} else {
			p.Severity = lint.Warning
			warnings++
		}
		f.Format(p)
	}
	if f, ok := f.(format.Statter); ok {
		f.Stats(total, errors, warnings)
	}
	if errors > 0 {
		exit(1)
	}
	exit(0)
}

type Options struct {
	Config config.Config

	Tags      []string
	LintTests bool
	Ignores   string
	GoVersion int
}

func computeSalt() ([]byte, error) {
	if version.Version != "devel" {
		return []byte(version.Version), nil
	}
	p, err := os.Executable()
	if err != nil {
		return nil, err
	}
	f, err := os.Open(p)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return nil, err
	}
	return h.Sum(nil), nil
}

func Lint(cs []*analysis.Analyzer, cums []lint.CumulativeChecker, paths []string, opt *Options) ([]lint.Problem, error) {
	salt, err := computeSalt()
	if err != nil {
		return nil, fmt.Errorf("could not compute salt for cache: %s", err)
	}
	cache.SetSalt(salt)

	if opt == nil {
		opt = &Options{}
	}

	l := &lint.Linter{
		Checkers:           cs,
		CumulativeCheckers: cums,
		GoVersion:          opt.GoVersion,
		Config:             opt.Config,
	}
	cfg := &packages.Config{}
	if opt.LintTests {
		cfg.Tests = true
	}

	printStats := func() {
		// Individual stats are read atomically, but overall there
		// is no synchronisation. For printing rough progress
		// information, this doesn't matter.
		switch atomic.LoadUint64(&l.Stats.State) {
		case lint.StateInitializing:
			fmt.Fprintln(os.Stderr, "Status: initializing")
		case lint.StateGraph:
			fmt.Fprintln(os.Stderr, "Status: loading package graph")
		case lint.StateProcessing:
			fmt.Fprintf(os.Stderr, "Packages: %d/%d initial, %d/%d total; Workers: %d/%d; Problems: %d\n",
				atomic.LoadUint64(&l.Stats.ProcessedInitialPackages),
				atomic.LoadUint64(&l.Stats.InitialPackages),
				atomic.LoadUint64(&l.Stats.ProcessedPackages),
				atomic.LoadUint64(&l.Stats.TotalPackages),
				atomic.LoadUint64(&l.Stats.ActiveWorkers),
				atomic.LoadUint64(&l.Stats.TotalWorkers),
				atomic.LoadUint64(&l.Stats.Problems),
			)
		case lint.StateCumulative:
			fmt.Fprintln(os.Stderr, "Status: processing cumulative checkers")
		}
	}
	if len(infoSignals) > 0 {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, infoSignals...)
		defer signal.Stop(ch)
		go func() {
			for range ch {
				printStats()
			}
		}()
	}

	return l.Lint(cfg, paths)
}

var posRe = regexp.MustCompile(`^(.+?):(\d+)(?::(\d+)?)?$`)

func parsePos(pos string) token.Position {
	if pos == "-" || pos == "" {
		return token.Position{}
	}
	parts := posRe.FindStringSubmatch(pos)
	if parts == nil {
		panic(fmt.Sprintf("internal error: malformed position %q", pos))
	}
	file := parts[1]
	line, _ := strconv.Atoi(parts[2])
	col, _ := strconv.Atoi(parts[3])
	return token.Position{
		Filename: file,
		Line:     line,
		Column:   col,
	}
}
<|endoftext|>


<commit_before>package writesplitter

import (
	"errors"
	"io"
	"os"
	"time"
)

const (
	Kilobyte  = 1024                // const for specifying ByteLimit
	Megabyte  = Kilobyte * Kilobyte // const for specifying ByteLimit
	formatStr = "2006-01-02T15.04.05.999999999Z0700.log"
)

var (
	ErrNotAFile = errors.New("WriteSplitter: invalid memory address or nil pointer dereference") // a custom error to signal that no file was closed
)

// WriteSplitter represents a disk bound io.WriteCloser that splits the input
// across consecutively named files based on either the number of bytes or the
// number of lines. Splitting does not guarantee true byte/line split
// precision as it does not parse the incoming data. The decision to split is
// before the underlying write operation based on the previous invocation. In
// other words, if a []byte sent to `Write()` contains enough bytes or new
// lines ('\n') to exceed the given limit, a new file won't be generated until
// the *next* invocation of `Write()`. If both LineLimit and ByteLimit are set,
// preference is given to LineLimit. By default, no splitting occurs because
// both LineLimit and ByteLimit are zero (0).
type WriteSplitter struct {
	LineLimit int            // how many write ops (typically one per line) before splitting the file
	ByteLimit int            // how many bytes before splitting the file
	Prefix    string         // files are named: $prefix + $nano-precision-timestamp + '.log'
	numBytes  int            // internal byte count
	numLines  int            // internal line count
	handle    io.WriteCloser // embedded file
}

// LineSplitter returns a WriteSplitter set to split at the given number of lines
func LineSplitter(limit int, prefix string) io.WriteCloser {
	return &WriteSplitter{LineLimit: limit, Prefix: prefix}
}

// ByteSplitter returns a WriteSplitter set to split at the given number of bytes
func ByteSplitter(limit int, prefix string) io.WriteCloser {
	return &WriteSplitter{ByteLimit: limit, Prefix: prefix}
}

// Close is a passthru and satisfies io.Closer. Subsequent writes will return an
// error.
func (ws *WriteSplitter) Close() error {
	if ws.handle != nil { // do not try to close nil
		return ws.handle.Close()
	}
	return ErrNotAFile // do not hide errors, but signal it's a WriteSplit error as opposed to an underlying os.* error
}

// Write satisfies io.Writer and internally manages file io. Write also limits
// each WriteSplitter to only one open file at a time.
func (ws *WriteSplitter) Write(p []byte) (int, error) {

	var n int
	var e error

	if ws.handle == nil {
		ws.handle, e = newFile(ws.Prefix)
	}

	switch {
	case ws.LineLimit > 0 && ws.numLines >= ws.LineLimit:
		fallthrough
	case ws.ByteLimit > 0 && ws.numBytes >= ws.ByteLimit:
		ws.Close()
		ws.handle, e = newFile(ws.Prefix)
		ws.numLines, ws.numBytes = 0, 0
	}

	if e != nil {
		return 0, e
	}

	n, e = ws.handle.Write(p)
	ws.numLines += 1
	ws.numBytes += n
	return n, e
}

// newFile creates a new file with the given prefix
func newFile(prefix string) (io.WriteCloser, error) {
	fn := prefix + time.Now().Format(formatStr)
	// fs is an abstraction layer for os allowing us to mock the filesystem for testing
	return createFile(fn)
}

// TestFileIO creates and removes a file to ensure that the location is writable.
func TestFileIO(prefix string) error {
	fn := prefix + "test.log"
	// It doesn't use the fs layer because it should be used to test the
	// writability of the actual filesystem. This test is unnecessary for mock filesystems
	if _, err := createFile(fn); err != nil {
		return err
	}
	removeFile(fn)
	return nil
}

/// This is for mocking the file IO. Used exclusively for testing
///-----------------------------------------------------------------------------

// createFile is the file creating function that wraps os.Create
var createFile = func(name string) (io.WriteCloser, error) {
	return os.Create(name)
}

// removeFile is the file removing function that wraps os.Remove
var removeFile = func(name string) error {
	return os.Remove(name)
}
<commit_msg>fixes func name<commit_after>package writesplitter

import (
	"errors"
	"io"
	"os"
	"time"
)

const (
	Kilobyte  = 1024                // const for specifying ByteLimit
	Megabyte  = Kilobyte * Kilobyte // const for specifying ByteLimit
	formatStr = "2006-01-02T15.04.05.999999999Z0700.log"
)

var (
	ErrNotAFile = errors.New("WriteSplitter: invalid memory address or nil pointer dereference") // a custom error to signal that no file was closed
)

// WriteSplitter represents a disk bound io.WriteCloser that splits the input
// across consecutively named files based on either the number of bytes or the
// number of lines. Splitting does not guarantee true byte/line split
// precision as it does not parse the incoming data. The decision to split is
// before the underlying write operation based on the previous invocation. In
// other words, if a []byte sent to `Write()` contains enough bytes or new
// lines ('\n') to exceed the given limit, a new file won't be generated until
// the *next* invocation of `Write()`. If both LineLimit and ByteLimit are set,
// preference is given to LineLimit. By default, no splitting occurs because
// both LineLimit and ByteLimit are zero (0).
type WriteSplitter struct {
	LineLimit int            // how many write ops (typically one per line) before splitting the file
	ByteLimit int            // how many bytes before splitting the file
	Prefix    string         // files are named: $prefix + $nano-precision-timestamp + '.log'
	numBytes  int            // internal byte count
	numLines  int            // internal line count
	handle    io.WriteCloser // embedded file
}

// LineSplitter returns a WriteSplitter set to split at the given number of lines
func LineSplitter(limit int, prefix string) io.WriteCloser {
	return &WriteSplitter{LineLimit: limit, Prefix: prefix}
}

// ByteSplitter returns a WriteSplitter set to split at the given number of bytes
func ByteSplitter(limit int, prefix string) io.WriteCloser {
	return &WriteSplitter{ByteLimit: limit, Prefix: prefix}
}

// Close is a passthru and satisfies io.Closer. Subsequent writes will return an
// error.
func (ws *WriteSplitter) Close() error {
	if ws.handle != nil { // do not try to close nil
		return ws.handle.Close()
	}
	return ErrNotAFile // do not hide errors, but signal it's a WriteSplit error as opposed to an underlying os.* error
}

// Write satisfies io.Writer and internally manages file io. Write also limits
// each WriteSplitter to only one open file at a time.
func (ws *WriteSplitter) Write(p []byte) (int, error) {

	var n int
	var e error

	if ws.handle == nil {
		ws.handle, e = createFile(ws.Prefix)
	}

	switch {
	case ws.LineLimit > 0 && ws.numLines >= ws.LineLimit:
		fallthrough
	case ws.ByteLimit > 0 && ws.numBytes >= ws.ByteLimit:
		ws.Close()
		ws.handle, e = createFile(ws.Prefix)
		ws.numLines, ws.numBytes = 0, 0
	}

	if e != nil {
		return 0, e
	}

	n, e = ws.handle.Write(p)
	ws.numLines += 1
	ws.numBytes += n
	return n, e
}

// TestFileIO creates and removes a file to ensure that the location is writable.
func TestFileIO(prefix string) error {
	fn := prefix + "test.log"
	// It doesn't use the fs layer because it should be used to test the
	// writability of the actual filesystem. This test is unnecessary for mock filesystems
	if _, err := createFile(fn); err != nil {
		return err
	}
	removeFile(fn)
	return nil
}

/// This is for mocking the file IO. Used exclusively for testing
///-----------------------------------------------------------------------------

// createFile is the file creating function that wraps os.Create
var createFile = func(prefix string) (io.WriteCloser, error) {
	name := prefix + time.Now().Format(formatStr)
	return os.Create(name)
}

// removeFile is the file removing function that wraps os.Remove
var removeFile = func(name string) error {
	return os.Remove(name)
}
<|endoftext|>


<commit_before>// Copyright 2018 The Kubeflow 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 controller

import (
	"testing"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeclientset "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/controller"

	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
	tfjobclientset "github.com/kubeflow/tf-operator/pkg/client/clientset/versioned"
)

func TestAddTFJob(t *testing.T) {
	// Prepare the clientset and controller for the test.
	kubeClientSet := kubeclientset.NewForConfigOrDie(&rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &v1.SchemeGroupVersion,
		},
	},
	)
	config := &rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &tfv1alpha2.SchemeGroupVersion,
		},
	}
	tfJobClientSet := tfjobclientset.NewForConfigOrDie(config)
	ctr, _, _ := newTFJobController(config, kubeClientSet, tfJobClientSet, controller.NoResyncPeriodFunc)
	ctr.tfJobInformerSynced = alwaysReady
	ctr.podInformerSynced = alwaysReady
	ctr.serviceInformerSynced = alwaysReady
	tfJobIndexer := ctr.tfJobInformer.GetIndexer()

	stopCh := make(chan struct{})
	run := func(<-chan struct{}) {
		ctr.Run(threadCount, stopCh)
	}
	go run(stopCh)

	var key string
	syncChan := make(chan string)
	ctr.syncHandler = func(tfJobKey string) (bool, error) {
		key = tfJobKey
		<-syncChan
		return true, nil
	}
	ctr.updateStatusHandler = func(tfjob *tfv1alpha2.TFJob) error {
		return nil
	}

	tfJob := newTFJob(1, 0)
	unstructured, err := convertTFJobToUnstructured(tfJob)
	if err != nil {
		t.Errorf("Failed to convert the TFJob to Unstructured: %v", err)
	}
	if err := tfJobIndexer.Add(unstructured); err != nil {
		t.Errorf("Failed to add tfjob to tfJobIndexer: %v", err)
	}
	ctr.addTFJob(unstructured)

	syncChan <- "sync"
	if key != getKey(tfJob, t) {
		t.Errorf("Failed to enqueue the TFJob %s: expected %s, got %s", tfJob.Name, getKey(tfJob, t), key)
	}
	close(stopCh)
}

func newTFJobWithChief(worker, ps int) *tfv1alpha2.TFJob {
	tfJob := newTFJob(worker, ps)
	tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeChief] = &tfv1alpha2.TFReplicaSpec{
		Template: newTFReplicaSpecTemplate(),
	}
	return tfJob
}

func newTFJob(worker, ps int) *tfv1alpha2.TFJob {
	tfJob := &tfv1alpha2.TFJob{
		TypeMeta: metav1.TypeMeta{
			Kind: tfv1alpha2.Kind,
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      testTFJobName,
			Namespace: metav1.NamespaceDefault,
		},
		Spec: tfv1alpha2.TFJobSpec{
			TFReplicaSpecs: make(map[tfv1alpha2.TFReplicaType]*tfv1alpha2.TFReplicaSpec),
		},
	}

	if worker > 0 {
		worker := int32(worker)
		workerReplicaSpec := &tfv1alpha2.TFReplicaSpec{
			Replicas: &worker,
			Template: newTFReplicaSpecTemplate(),
		}
		tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker] = workerReplicaSpec
	}

	if ps > 0 {
		ps := int32(ps)
		psReplicaSpec := &tfv1alpha2.TFReplicaSpec{
			Replicas: &ps,
			Template: newTFReplicaSpecTemplate(),
		}
		tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS] = psReplicaSpec
	}
	return tfJob
}

func getKey(tfJob *tfv1alpha2.TFJob, t *testing.T) string {
	key, err := KeyFunc(tfJob)
	if err != nil {
		t.Errorf("Unexpected error getting key for job %v: %v", tfJob.Name, err)
		return ""
	}
	return key
}

func checkCondition(tfJob *tfv1alpha2.TFJob, condition tfv1alpha2.TFJobConditionType, reason string) bool {
	for _, v := range tfJob.Status.Conditions {
		if v.Type == condition && v.Status == v1.ConditionTrue && v.Reason == reason {
			return true
		}
	}
	return false
}

func newTFReplicaSpecTemplate() v1.PodTemplateSpec {
	return v1.PodTemplateSpec{
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				v1.Container{
					Name:  tfv1alpha2.DefaultContainerName,
					Image: testImageName,
					Args:  []string{"Fake", "Fake"},
					Ports: []v1.ContainerPort{
						v1.ContainerPort{
							Name:          tfv1alpha2.DefaultPortName,
							ContainerPort: tfv1alpha2.DefaultPort,
						},
					},
				},
			},
		},
	}
}
<commit_msg>tfjob: Add test for labels and annotations (#658)<commit_after>// Copyright 2018 The Kubeflow 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 controller

import (
	"testing"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kubeclientset "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/controller"

	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
	tfjobclientset "github.com/kubeflow/tf-operator/pkg/client/clientset/versioned"
)

func TestAddTFJob(t *testing.T) {
	// Prepare the clientset and controller for the test.
	kubeClientSet := kubeclientset.NewForConfigOrDie(&rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &v1.SchemeGroupVersion,
		},
	},
	)
	config := &rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &tfv1alpha2.SchemeGroupVersion,
		},
	}
	tfJobClientSet := tfjobclientset.NewForConfigOrDie(config)
	ctr, _, _ := newTFJobController(config, kubeClientSet, tfJobClientSet, controller.NoResyncPeriodFunc)
	ctr.tfJobInformerSynced = alwaysReady
	ctr.podInformerSynced = alwaysReady
	ctr.serviceInformerSynced = alwaysReady
	tfJobIndexer := ctr.tfJobInformer.GetIndexer()

	stopCh := make(chan struct{})
	run := func(<-chan struct{}) {
		ctr.Run(threadCount, stopCh)
	}
	go run(stopCh)

	var key string
	syncChan := make(chan string)
	ctr.syncHandler = func(tfJobKey string) (bool, error) {
		key = tfJobKey
		<-syncChan
		return true, nil
	}
	ctr.updateStatusHandler = func(tfjob *tfv1alpha2.TFJob) error {
		return nil
	}

	tfJob := newTFJob(1, 0)
	unstructured, err := convertTFJobToUnstructured(tfJob)
	if err != nil {
		t.Errorf("Failed to convert the TFJob to Unstructured: %v", err)
	}
	if err := tfJobIndexer.Add(unstructured); err != nil {
		t.Errorf("Failed to add tfjob to tfJobIndexer: %v", err)
	}
	ctr.addTFJob(unstructured)

	syncChan <- "sync"
	if key != getKey(tfJob, t) {
		t.Errorf("Failed to enqueue the TFJob %s: expected %s, got %s", tfJob.Name, getKey(tfJob, t), key)
	}
	close(stopCh)
}

func TestCopyLabelsAndAnnotation(t *testing.T) {
	// Prepare the clientset and controller for the test.
	kubeClientSet := kubeclientset.NewForConfigOrDie(&rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &v1.SchemeGroupVersion,
		},
	},
	)
	config := &rest.Config{
		Host: "",
		ContentConfig: rest.ContentConfig{
			GroupVersion: &tfv1alpha2.SchemeGroupVersion,
		},
	}
	tfJobClientSet := tfjobclientset.NewForConfigOrDie(config)
	ctr, _, _ := newTFJobController(config, kubeClientSet, tfJobClientSet, controller.NoResyncPeriodFunc)
	fakePodControl := &controller.FakePodControl{}
	ctr.podControl = fakePodControl
	ctr.tfJobInformerSynced = alwaysReady
	ctr.podInformerSynced = alwaysReady
	ctr.serviceInformerSynced = alwaysReady
	tfJobIndexer := ctr.tfJobInformer.GetIndexer()

	stopCh := make(chan struct{})
	run := func(<-chan struct{}) {
		ctr.Run(threadCount, stopCh)
	}
	go run(stopCh)

	ctr.updateStatusHandler = func(tfJob *tfv1alpha2.TFJob) error {
		return nil
	}

	tfJob := newTFJob(1, 0)
	annotations := map[string]string{
		"annotation1": "1",
	}
	labels := map[string]string{
		"label1": "1",
	}
	tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Template.Labels = labels
	tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].Template.Annotations = annotations
	unstructured, err := convertTFJobToUnstructured(tfJob)
	if err != nil {
		t.Errorf("Failed to convert the TFJob to Unstructured: %v", err)
	}

	if err := tfJobIndexer.Add(unstructured); err != nil {
		t.Errorf("Failed to add tfjob to tfJobIndexer: %v", err)
	}

	_, err = ctr.syncTFJob(getKey(tfJob, t))
	if err != nil {
		t.Errorf("%s: unexpected error when syncing jobs %v", tfJob.Name, err)
	}

	if len(fakePodControl.Templates) != 1 {
		t.Errorf("Expected to create 1 pod while got %d", len(fakePodControl.Templates))
	}
	actual := fakePodControl.Templates[0]
	v, exist := actual.Labels["label1"]
	if !exist {
		t.Errorf("Labels does not exist")
	}
	if v != "1" {
		t.Errorf("Labels value do not equal")
	}

	v, exist = actual.Annotations["annotation1"]
	if !exist {
		t.Errorf("Annotations does not exist")
	}
	if v != "1" {
		t.Errorf("Annotations value does not equal")
	}

	close(stopCh)
}

func newTFJobWithChief(worker, ps int) *tfv1alpha2.TFJob {
	tfJob := newTFJob(worker, ps)
	tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeChief] = &tfv1alpha2.TFReplicaSpec{
		Template: newTFReplicaSpecTemplate(),
	}
	return tfJob
}

func newTFJob(worker, ps int) *tfv1alpha2.TFJob {
	tfJob := &tfv1alpha2.TFJob{
		TypeMeta: metav1.TypeMeta{
			Kind: tfv1alpha2.Kind,
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      testTFJobName,
			Namespace: metav1.NamespaceDefault,
		},
		Spec: tfv1alpha2.TFJobSpec{
			TFReplicaSpecs: make(map[tfv1alpha2.TFReplicaType]*tfv1alpha2.TFReplicaSpec),
		},
	}

	if worker > 0 {
		worker := int32(worker)
		workerReplicaSpec := &tfv1alpha2.TFReplicaSpec{
			Replicas: &worker,
			Template: newTFReplicaSpecTemplate(),
		}
		tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker] = workerReplicaSpec
	}

	if ps > 0 {
		ps := int32(ps)
		psReplicaSpec := &tfv1alpha2.TFReplicaSpec{
			Replicas: &ps,
			Template: newTFReplicaSpecTemplate(),
		}
		tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypePS] = psReplicaSpec
	}
	return tfJob
}

func getKey(tfJob *tfv1alpha2.TFJob, t *testing.T) string {
	key, err := KeyFunc(tfJob)
	if err != nil {
		t.Errorf("Unexpected error getting key for job %v: %v", tfJob.Name, err)
		return ""
	}
	return key
}

func checkCondition(tfJob *tfv1alpha2.TFJob, condition tfv1alpha2.TFJobConditionType, reason string) bool {
	for _, v := range tfJob.Status.Conditions {
		if v.Type == condition && v.Status == v1.ConditionTrue && v.Reason == reason {
			return true
		}
	}
	return false
}

func newTFReplicaSpecTemplate() v1.PodTemplateSpec {
	return v1.PodTemplateSpec{
		Spec: v1.PodSpec{
			Containers: []v1.Container{
				v1.Container{
					Name:  tfv1alpha2.DefaultContainerName,
					Image: testImageName,
					Args:  []string{"Fake", "Fake"},
					Ports: []v1.ContainerPort{
						v1.ContainerPort{
							Name:          tfv1alpha2.DefaultPortName,
							ContainerPort: tfv1alpha2.DefaultPort,
						},
					},
				},
			},
		},
	}
}
<|endoftext|>


<commit_before>package dialer

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"os"
	"strings"
	"time"

	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"

	"github.com/rancher/norman/types/slice"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/tunnelserver"
	"github.com/rancher/rancher/pkg/types/config"
	"github.com/rancher/rancher/pkg/types/config/dialer"
	"github.com/rancher/remotedialer"
	"github.com/rancher/rke/k8s"
	"github.com/rancher/rke/services"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
)

func NewFactory(apiContext *config.ScaledContext) (*Factory, error) {
	authorizer := tunnelserver.NewAuthorizer(apiContext)
	tunneler := tunnelserver.NewTunnelServer(authorizer)

	return &Factory{
		clusterLister:    apiContext.Management.Clusters("").Controller().Lister(),
		nodeLister:       apiContext.Management.Nodes("").Controller().Lister(),
		TunnelServer:     tunneler,
		TunnelAuthorizer: authorizer,
	}, nil
}

type Factory struct {
	nodeLister       v3.NodeLister
	clusterLister    v3.ClusterLister
	TunnelServer     *remotedialer.Server
	TunnelAuthorizer *tunnelserver.Authorizer
}

func (f *Factory) ClusterDialer(clusterName string) (dialer.Dialer, error) {
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		d, err := f.clusterDialer(clusterName, address)
		if err != nil {
			return nil, err
		}
		return d(ctx, network, address)
	}, nil
}

func isCloudDriver(cluster *v3.Cluster) bool {
	return !cluster.Spec.Internal &&
		cluster.Status.Driver != v32.ClusterDriverImported &&
		cluster.Status.Driver != v32.ClusterDriverRKE &&
		cluster.Status.Driver != v32.ClusterDriverK3s &&
		cluster.Status.Driver != v32.ClusterDriverK3os &&
		cluster.Status.Driver != v32.ClusterDriverRke2
}

func (f *Factory) translateClusterAddress(cluster *v3.Cluster, clusterHostPort, address string) string {
	if clusterHostPort != address {
		logrus.Tracef("dialerFactory: apiEndpoint clusterHostPort [%s] is not equal to address [%s]", clusterHostPort, address)
		return address
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return address
	}

	// Make sure that control plane node we are connecting to is not bad, also use internal address
	nodes, err := f.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		logrus.Debugf("Error listing nodes while translating cluster address, returning address [%s], error: %v", address, err)
		return address
	}

	clusterGood := v32.ClusterConditionReady.IsTrue(cluster)
	logrus.Tracef("dialerFactory: ClusterConditionReady for cluster [%s] is [%t]", cluster.Spec.DisplayName, clusterGood)
	lastGoodHost := ""
	logrus.Trace("dialerFactory: finding a node to tunnel the cluster connection")
	for _, node := range nodes {
		var (
			publicIP  = node.Status.NodeAnnotations[k8s.ExternalAddressAnnotation]
			privateIP = node.Status.NodeAnnotations[k8s.InternalAddressAnnotation]
		)

		fakeNode := &v1.Node{
			Status: node.Status.InternalNodeStatus,
		}

		nodeGood := v32.NodeConditionRegistered.IsTrue(node) && v32.NodeConditionProvisioned.IsTrue(node) &&
			!v32.NodeConditionReady.IsUnknown(fakeNode) && node.DeletionTimestamp == nil

		if !nodeGood {
			logrus.Tracef("dialerFactory: Skipping node [%s] for tunneling the cluster connection because nodeConditions are not as expected", node.Spec.RequestedHostname)
			logrus.Tracef("dialerFactory: Node conditions for node [%s]: %+v", node.Status.NodeName, node.Status.Conditions)
			continue
		}
		if privateIP == "" {
			logrus.Tracef("dialerFactory: Skipping node [%s] for tunneling the cluster connection because privateIP is empty", node.Status.NodeName)
			continue
		}

		logrus.Tracef("dialerFactory: IP addresses for node [%s]: publicIP [%s], privateIP [%s]", node.Status.NodeName, publicIP, privateIP)

		if publicIP == host {
			logrus.Tracef("dialerFactory: publicIP [%s] for node [%s] matches apiEndpoint host [%s], checking if cluster condition Ready is True", publicIP, node.Status.NodeName, host)
			if clusterGood {
				logrus.Trace("dialerFactory: cluster condition Ready is True")
				host = privateIP
				logrus.Tracef("dialerFactory: Using privateIP [%s] of node [%s] as node to tunnel the cluster connection", privateIP, node.Status.NodeName)
				return fmt.Sprintf("%s:%s", host, port)
			}
			logrus.Debug("dialerFactory: cluster condition Ready is False")
		} else if node.Status.NodeConfig != nil && slice.ContainsString(node.Status.NodeConfig.Role, services.ControlRole) {
			logrus.Tracef("dialerFactory: setting node [%s] with privateIP [%s] as option for the connection as it is a controlplane node", node.Status.NodeName, privateIP)
			lastGoodHost = privateIP
		}
	}

	if lastGoodHost != "" {
		logrus.Tracef("dialerFactory: returning [%s:%s] as last good option to tunnel the cluster connection", lastGoodHost, port)
		return fmt.Sprintf("%s:%s", lastGoodHost, port)
	}

	logrus.Tracef("dialerFactory: returning [%s], as no good option was found (no match with apiEndpoint or a controlplane node with correct conditions", address)
	return address
}

func (f *Factory) clusterDialer(clusterName, address string) (dialer.Dialer, error) {
	cluster, err := f.clusterLister.Get("", clusterName)
	if err != nil {
		return nil, err
	}

	if cluster.Spec.Internal {
		// For local (embedded, or import) we just assume we can connect directly
		return native()
	}

	hostPort := hostPort(cluster)
	logrus.Tracef("dialerFactory: apiEndpoint hostPort for cluster [%s] is [%s]", clusterName, hostPort)
	if (address == hostPort || isProxyAddress(address)) && isCloudDriver(cluster) {
		// For cloud drivers we just connect directly to the k8s API, not through the tunnel.  All other go through tunnel
		return native()
	}

	if f.TunnelServer.HasSession(cluster.Name) {
		logrus.Tracef("dialerFactory: tunnel session found for cluster [%s]", cluster.Name)
		cd := f.TunnelServer.Dialer(cluster.Name)
		return func(ctx context.Context, network, address string) (net.Conn, error) {
			if cluster.Status.Driver == v32.ClusterDriverRKE {
				address = f.translateClusterAddress(cluster, hostPort, address)
			}
			logrus.Tracef("dialerFactory: returning network [%s] and address [%s] as clusterDialer", network, address)
			return cd(ctx, network, address)
		}, nil
	}
	logrus.Tracef("dialerFactory: no tunnel session found for cluster [%s], falling back to nodeDialer", cluster.Name)

	// Try to connect to a node for the cluster dialer
	nodes, err := f.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		return nil, err
	}

	var localAPIEndpoint bool
	if cluster.Status.Driver == v32.ClusterDriverRKE {
		localAPIEndpoint = true
	}

	for _, node := range nodes {
		if node.DeletionTimestamp == nil && v32.NodeConditionProvisioned.IsTrue(node) {
			logrus.Tracef("dialerFactory: using node [%s]/[%s] for nodeDialer",
				node.Labels["management.cattle.io/nodename"], node.Name)
			if nodeDialer, err := f.nodeDialer(clusterName, node.Name); err == nil {
				return func(ctx context.Context, network, address string) (net.Conn, error) {
					if address == hostPort && localAPIEndpoint {
						logrus.Trace("dialerFactory: rewriting address/port to 127.0.0.1:6443 as node may not" +
							" have direct kube-api access")
						// The node dialer may not have direct access to kube-api so we hit localhost:6443 instead
						address = "127.0.0.1:6443"
					}
					logrus.Tracef("dialerFactory: Returning network [%s] and address [%s] as nodeDialer", network, address)
					return nodeDialer(ctx, network, address)
				}, nil
			}
		}
	}

	logrus.Debugf("No active connection for cluster [%s], will wait for about 30 seconds", cluster.Name)
	for i := 0; i < 4; i++ {
		if f.TunnelServer.HasSession(cluster.Name) {
			logrus.Debugf("Cluster [%s] has reconnected, resuming", cluster.Name)
			cd := f.TunnelServer.Dialer(cluster.Name)
			return func(ctx context.Context, network, address string) (net.Conn, error) {
				if cluster.Status.Driver == v32.ClusterDriverRKE {
					address = f.translateClusterAddress(cluster, hostPort, address)
				}
				logrus.Tracef("dialerFactory: returning network [%s] and address [%s] as clusterDialer", network, address)
				return cd(ctx, network, address)
			}, nil
		}
		time.Sleep(wait.Jitter(5*time.Second, 1))
	}

	return nil, fmt.Errorf("waiting for cluster [%s] agent to connect", cluster.Name)
}

func hostPort(cluster *v3.Cluster) string {
	u, err := url.Parse(cluster.Status.APIEndpoint)
	if err != nil {
		return ""
	}

	if strings.Contains(u.Host, ":") {
		return u.Host
	}
	return u.Host + ":443"
}

func native() (dialer.Dialer, error) {
	netDialer := net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}
	return netDialer.DialContext, nil
}

func (f *Factory) DockerDialer(clusterName, machineName string) (dialer.Dialer, error) {
	machine, err := f.nodeLister.Get(clusterName, machineName)
	if err != nil {
		return nil, err
	}

	sessionKey := machineSessionKey(machine)
	if f.TunnelServer.HasSession(sessionKey) {
		network, address := "unix", "/var/run/docker.sock"
		if machine.Status.InternalNodeStatus.NodeInfo.OperatingSystem == "windows" {
			network, address = "npipe", "//./pipe/docker_engine"
		}
		d := f.TunnelServer.Dialer(sessionKey)
		return func(ctx context.Context, _ string, _ string) (net.Conn, error) {
			return d(ctx, network, address)
		}, nil
	}

	return nil, fmt.Errorf("can not build dialer to [%s:%s]", clusterName, machineName)
}

func (f *Factory) NodeDialer(clusterName, machineName string) (dialer.Dialer, error) {
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		d, err := f.nodeDialer(clusterName, machineName)
		if err != nil {
			return nil, err
		}
		return d(ctx, network, address)
	}, nil
}

func (f *Factory) nodeDialer(clusterName, machineName string) (dialer.Dialer, error) {
	machine, err := f.nodeLister.Get(clusterName, machineName)
	if err != nil {
		return nil, err
	}

	sessionKey := machineSessionKey(machine)
	if f.TunnelServer.HasSession(sessionKey) {
		d := f.TunnelServer.Dialer(sessionKey)
		return dialer.Dialer(d), nil
	}

	return nil, fmt.Errorf("can not build dialer to [%s:%s]", clusterName, machineName)
}

func machineSessionKey(machine *v3.Node) string {
	return fmt.Sprintf("%s:%s", machine.Namespace, machine.Name)
}

func isProxyAddress(address string) bool {
	proxy := getEnvAny("HTTP_PROXY", "http_proxy")
	if proxy == "" {
		proxy = getEnvAny("HTTPS_PROXY", "https_proxy")
	}

	if proxy == "" {
		return false
	}

	parsed, err := parseProxy(proxy)
	if err != nil {
		logrus.Warnf("Failed to parse http_proxy url %s: %v", proxy, err)
		return false
	}
	return parsed.Host == address
}

func getEnvAny(names ...string) string {
	for _, n := range names {
		if val := os.Getenv(n); val != "" {
			return val
		}
	}
	return ""
}

func parseProxy(proxy string) (*url.URL, error) {
	if proxy == "" {
		return nil, nil
	}

	proxyURL, err := url.Parse(proxy)
	if err != nil ||
		(proxyURL.Scheme != "http" &&
			proxyURL.Scheme != "https" &&
			proxyURL.Scheme != "socks5") {
		// proxy was bogus. Try pre-pending "http://" to it and
		// see if that parses correctly. If not, fall through
		if proxyURL, err := url.Parse("http://" + proxy); err == nil {
			return proxyURL, nil
		}
	}
	if err != nil {
		return nil, fmt.Errorf("invalid proxy address %q: %v", proxy, err)
	}
	return proxyURL, nil
}
<commit_msg>As a corner case don't treat no driver as a cloud driver<commit_after>package dialer

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"os"
	"strings"
	"time"

	v32 "github.com/rancher/rancher/pkg/apis/management.cattle.io/v3"

	"github.com/rancher/norman/types/slice"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/rancher/rancher/pkg/tunnelserver"
	"github.com/rancher/rancher/pkg/types/config"
	"github.com/rancher/rancher/pkg/types/config/dialer"
	"github.com/rancher/remotedialer"
	"github.com/rancher/rke/k8s"
	"github.com/rancher/rke/services"
	"github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/wait"
)

func NewFactory(apiContext *config.ScaledContext) (*Factory, error) {
	authorizer := tunnelserver.NewAuthorizer(apiContext)
	tunneler := tunnelserver.NewTunnelServer(authorizer)

	return &Factory{
		clusterLister:    apiContext.Management.Clusters("").Controller().Lister(),
		nodeLister:       apiContext.Management.Nodes("").Controller().Lister(),
		TunnelServer:     tunneler,
		TunnelAuthorizer: authorizer,
	}, nil
}

type Factory struct {
	nodeLister       v3.NodeLister
	clusterLister    v3.ClusterLister
	TunnelServer     *remotedialer.Server
	TunnelAuthorizer *tunnelserver.Authorizer
}

func (f *Factory) ClusterDialer(clusterName string) (dialer.Dialer, error) {
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		d, err := f.clusterDialer(clusterName, address)
		if err != nil {
			return nil, err
		}
		return d(ctx, network, address)
	}, nil
}

func isCloudDriver(cluster *v3.Cluster) bool {
	return !cluster.Spec.Internal &&
		cluster.Status.Driver != "" &&
		cluster.Status.Driver != v32.ClusterDriverImported &&
		cluster.Status.Driver != v32.ClusterDriverRKE &&
		cluster.Status.Driver != v32.ClusterDriverK3s &&
		cluster.Status.Driver != v32.ClusterDriverK3os &&
		cluster.Status.Driver != v32.ClusterDriverRke2
}

func (f *Factory) translateClusterAddress(cluster *v3.Cluster, clusterHostPort, address string) string {
	if clusterHostPort != address {
		logrus.Tracef("dialerFactory: apiEndpoint clusterHostPort [%s] is not equal to address [%s]", clusterHostPort, address)
		return address
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return address
	}

	// Make sure that control plane node we are connecting to is not bad, also use internal address
	nodes, err := f.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		logrus.Debugf("Error listing nodes while translating cluster address, returning address [%s], error: %v", address, err)
		return address
	}

	clusterGood := v32.ClusterConditionReady.IsTrue(cluster)
	logrus.Tracef("dialerFactory: ClusterConditionReady for cluster [%s] is [%t]", cluster.Spec.DisplayName, clusterGood)
	lastGoodHost := ""
	logrus.Trace("dialerFactory: finding a node to tunnel the cluster connection")
	for _, node := range nodes {
		var (
			publicIP  = node.Status.NodeAnnotations[k8s.ExternalAddressAnnotation]
			privateIP = node.Status.NodeAnnotations[k8s.InternalAddressAnnotation]
		)

		fakeNode := &v1.Node{
			Status: node.Status.InternalNodeStatus,
		}

		nodeGood := v32.NodeConditionRegistered.IsTrue(node) && v32.NodeConditionProvisioned.IsTrue(node) &&
			!v32.NodeConditionReady.IsUnknown(fakeNode) && node.DeletionTimestamp == nil

		if !nodeGood {
			logrus.Tracef("dialerFactory: Skipping node [%s] for tunneling the cluster connection because nodeConditions are not as expected", node.Spec.RequestedHostname)
			logrus.Tracef("dialerFactory: Node conditions for node [%s]: %+v", node.Status.NodeName, node.Status.Conditions)
			continue
		}
		if privateIP == "" {
			logrus.Tracef("dialerFactory: Skipping node [%s] for tunneling the cluster connection because privateIP is empty", node.Status.NodeName)
			continue
		}

		logrus.Tracef("dialerFactory: IP addresses for node [%s]: publicIP [%s], privateIP [%s]", node.Status.NodeName, publicIP, privateIP)

		if publicIP == host {
			logrus.Tracef("dialerFactory: publicIP [%s] for node [%s] matches apiEndpoint host [%s], checking if cluster condition Ready is True", publicIP, node.Status.NodeName, host)
			if clusterGood {
				logrus.Trace("dialerFactory: cluster condition Ready is True")
				host = privateIP
				logrus.Tracef("dialerFactory: Using privateIP [%s] of node [%s] as node to tunnel the cluster connection", privateIP, node.Status.NodeName)
				return fmt.Sprintf("%s:%s", host, port)
			}
			logrus.Debug("dialerFactory: cluster condition Ready is False")
		} else if node.Status.NodeConfig != nil && slice.ContainsString(node.Status.NodeConfig.Role, services.ControlRole) {
			logrus.Tracef("dialerFactory: setting node [%s] with privateIP [%s] as option for the connection as it is a controlplane node", node.Status.NodeName, privateIP)
			lastGoodHost = privateIP
		}
	}

	if lastGoodHost != "" {
		logrus.Tracef("dialerFactory: returning [%s:%s] as last good option to tunnel the cluster connection", lastGoodHost, port)
		return fmt.Sprintf("%s:%s", lastGoodHost, port)
	}

	logrus.Tracef("dialerFactory: returning [%s], as no good option was found (no match with apiEndpoint or a controlplane node with correct conditions", address)
	return address
}

func (f *Factory) clusterDialer(clusterName, address string) (dialer.Dialer, error) {
	cluster, err := f.clusterLister.Get("", clusterName)
	if err != nil {
		return nil, err
	}

	if cluster.Spec.Internal {
		// For local (embedded, or import) we just assume we can connect directly
		return native()
	}

	hostPort := hostPort(cluster)
	logrus.Tracef("dialerFactory: apiEndpoint hostPort for cluster [%s] is [%s]", clusterName, hostPort)
	if (address == hostPort || isProxyAddress(address)) && isCloudDriver(cluster) {
		// For cloud drivers we just connect directly to the k8s API, not through the tunnel.  All other go through tunnel
		return native()
	}

	if f.TunnelServer.HasSession(cluster.Name) {
		logrus.Tracef("dialerFactory: tunnel session found for cluster [%s]", cluster.Name)
		cd := f.TunnelServer.Dialer(cluster.Name)
		return func(ctx context.Context, network, address string) (net.Conn, error) {
			if cluster.Status.Driver == v32.ClusterDriverRKE {
				address = f.translateClusterAddress(cluster, hostPort, address)
			}
			logrus.Tracef("dialerFactory: returning network [%s] and address [%s] as clusterDialer", network, address)
			return cd(ctx, network, address)
		}, nil
	}
	logrus.Tracef("dialerFactory: no tunnel session found for cluster [%s], falling back to nodeDialer", cluster.Name)

	// Try to connect to a node for the cluster dialer
	nodes, err := f.nodeLister.List(cluster.Name, labels.Everything())
	if err != nil {
		return nil, err
	}

	var localAPIEndpoint bool
	if cluster.Status.Driver == v32.ClusterDriverRKE {
		localAPIEndpoint = true
	}

	for _, node := range nodes {
		if node.DeletionTimestamp == nil && v32.NodeConditionProvisioned.IsTrue(node) {
			logrus.Tracef("dialerFactory: using node [%s]/[%s] for nodeDialer",
				node.Labels["management.cattle.io/nodename"], node.Name)
			if nodeDialer, err := f.nodeDialer(clusterName, node.Name); err == nil {
				return func(ctx context.Context, network, address string) (net.Conn, error) {
					if address == hostPort && localAPIEndpoint {
						logrus.Trace("dialerFactory: rewriting address/port to 127.0.0.1:6443 as node may not" +
							" have direct kube-api access")
						// The node dialer may not have direct access to kube-api so we hit localhost:6443 instead
						address = "127.0.0.1:6443"
					}
					logrus.Tracef("dialerFactory: Returning network [%s] and address [%s] as nodeDialer", network, address)
					return nodeDialer(ctx, network, address)
				}, nil
			}
		}
	}

	logrus.Debugf("No active connection for cluster [%s], will wait for about 30 seconds", cluster.Name)
	for i := 0; i < 4; i++ {
		if f.TunnelServer.HasSession(cluster.Name) {
			logrus.Debugf("Cluster [%s] has reconnected, resuming", cluster.Name)
			cd := f.TunnelServer.Dialer(cluster.Name)
			return func(ctx context.Context, network, address string) (net.Conn, error) {
				if cluster.Status.Driver == v32.ClusterDriverRKE {
					address = f.translateClusterAddress(cluster, hostPort, address)
				}
				logrus.Tracef("dialerFactory: returning network [%s] and address [%s] as clusterDialer", network, address)
				return cd(ctx, network, address)
			}, nil
		}
		time.Sleep(wait.Jitter(5*time.Second, 1))
	}

	return nil, fmt.Errorf("waiting for cluster [%s] agent to connect", cluster.Name)
}

func hostPort(cluster *v3.Cluster) string {
	u, err := url.Parse(cluster.Status.APIEndpoint)
	if err != nil {
		return ""
	}

	if strings.Contains(u.Host, ":") {
		return u.Host
	}
	return u.Host + ":443"
}

func native() (dialer.Dialer, error) {
	netDialer := net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}
	return netDialer.DialContext, nil
}

func (f *Factory) DockerDialer(clusterName, machineName string) (dialer.Dialer, error) {
	machine, err := f.nodeLister.Get(clusterName, machineName)
	if err != nil {
		return nil, err
	}

	sessionKey := machineSessionKey(machine)
	if f.TunnelServer.HasSession(sessionKey) {
		network, address := "unix", "/var/run/docker.sock"
		if machine.Status.InternalNodeStatus.NodeInfo.OperatingSystem == "windows" {
			network, address = "npipe", "//./pipe/docker_engine"
		}
		d := f.TunnelServer.Dialer(sessionKey)
		return func(ctx context.Context, _ string, _ string) (net.Conn, error) {
			return d(ctx, network, address)
		}, nil
	}

	return nil, fmt.Errorf("can not build dialer to [%s:%s]", clusterName, machineName)
}

func (f *Factory) NodeDialer(clusterName, machineName string) (dialer.Dialer, error) {
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		d, err := f.nodeDialer(clusterName, machineName)
		if err != nil {
			return nil, err
		}
		return d(ctx, network, address)
	}, nil
}

func (f *Factory) nodeDialer(clusterName, machineName string) (dialer.Dialer, error) {
	machine, err := f.nodeLister.Get(clusterName, machineName)
	if err != nil {
		return nil, err
	}

	sessionKey := machineSessionKey(machine)
	if f.TunnelServer.HasSession(sessionKey) {
		d := f.TunnelServer.Dialer(sessionKey)
		return dialer.Dialer(d), nil
	}

	return nil, fmt.Errorf("can not build dialer to [%s:%s]", clusterName, machineName)
}

func machineSessionKey(machine *v3.Node) string {
	return fmt.Sprintf("%s:%s", machine.Namespace, machine.Name)
}

func isProxyAddress(address string) bool {
	proxy := getEnvAny("HTTP_PROXY", "http_proxy")
	if proxy == "" {
		proxy = getEnvAny("HTTPS_PROXY", "https_proxy")
	}

	if proxy == "" {
		return false
	}

	parsed, err := parseProxy(proxy)
	if err != nil {
		logrus.Warnf("Failed to parse http_proxy url %s: %v", proxy, err)
		return false
	}
	return parsed.Host == address
}

func getEnvAny(names ...string) string {
	for _, n := range names {
		if val := os.Getenv(n); val != "" {
			return val
		}
	}
	return ""
}

func parseProxy(proxy string) (*url.URL, error) {
	if proxy == "" {
		return nil, nil
	}

	proxyURL, err := url.Parse(proxy)
	if err != nil ||
		(proxyURL.Scheme != "http" &&
			proxyURL.Scheme != "https" &&
			proxyURL.Scheme != "socks5") {
		// proxy was bogus. Try pre-pending "http://" to it and
		// see if that parses correctly. If not, fall through
		if proxyURL, err := url.Parse("http://" + proxy); err == nil {
			return proxyURL, nil
		}
	}
	if err != nil {
		return nil, fmt.Errorf("invalid proxy address %q: %v", proxy, err)
	}
	return proxyURL, nil
}
<|endoftext|>


<commit_before>package ingester

import (
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"

	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/common/model"
	"github.com/weaveworks/common/user"

	"github.com/cortexproject/cortex/pkg/chunk"
	"github.com/cortexproject/cortex/pkg/ingester/client"
	"github.com/cortexproject/cortex/pkg/util"
	"github.com/grafana/loki/pkg/chunkenc"
)

const (
	// Backoff for retrying 'immediate' flushes. Only counts for queue
	// position, not wallclock time.
	flushBackoff = 1 * time.Second

	nameLabel = model.LabelName("__name__")
	logsValue = model.LabelValue("logs")
)

// Flush triggers a flush of all the chunks and closes the flush queues.
// Called from the Lifecycler as part of the ingester shutdown.
func (i *Ingester) Flush() {
	i.sweepUsers(true)

	// Close the flush queues, to unblock waiting workers.
	for _, flushQueue := range i.flushQueues {
		flushQueue.Close()
	}

	i.flushQueuesDone.Wait()
}

// FlushHandler triggers a flush of all in memory chunks.  Mainly used for
// local testing.
func (i *Ingester) FlushHandler(w http.ResponseWriter, _ *http.Request) {
	i.sweepUsers(true)
	w.WriteHeader(http.StatusNoContent)
}

type flushOp struct {
	from      model.Time
	userID    string
	fp        model.Fingerprint
	immediate bool
}

func (o *flushOp) Key() string {
	return fmt.Sprintf("%s-%s-%v", o.userID, o.fp, o.immediate)
}

func (o *flushOp) Priority() int64 {
	return -int64(o.from)
}

// sweepUsers periodically schedules series for flushing and garbage collects users with no series
func (i *Ingester) sweepUsers(immediate bool) {
	instances := i.getInstances()

	for _, instance := range instances {
		i.sweepInstance(instance, immediate)
	}
}

func (i *Ingester) sweepInstance(instance *instance, immediate bool) {
	instance.streamsMtx.Lock()
	defer instance.streamsMtx.Unlock()

	for _, stream := range instance.streams {
		i.sweepStream(instance, stream, immediate)
		i.removeFlushedChunks(instance, stream)
	}
}

func (i *Ingester) sweepStream(instance *instance, stream *stream, immediate bool) {
	if len(stream.chunks) == 0 {
		return
	}

	lastChunk := stream.chunks[len(stream.chunks)-1]
	if len(stream.chunks) == 1 && time.Since(lastChunk.lastUpdated) < i.cfg.MaxChunkIdle && !immediate {
		return
	}

	flushQueueIndex := int(uint64(stream.fp) % uint64(i.cfg.ConcurrentFlushes))
	firstTime, _ := stream.chunks[0].chunk.Bounds()
	i.flushQueues[flushQueueIndex].Enqueue(&flushOp{
		model.TimeFromUnixNano(firstTime.UnixNano()), instance.instanceID,
		stream.fp, immediate,
	})
}

func (i *Ingester) flushLoop(j int) {
	defer func() {
		level.Debug(util.Logger).Log("msg", "Ingester.flushLoop() exited")
		i.flushQueuesDone.Done()
	}()

	for {
		o := i.flushQueues[j].Dequeue()
		if o == nil {
			return
		}
		op := o.(*flushOp)

		level.Debug(util.Logger).Log("msg", "flushing stream", "userid", op.userID, "fp", op.fp, "immediate", op.immediate)

		err := i.flushUserSeries(op.userID, op.fp, op.immediate)
		if err != nil {
			level.Error(util.WithUserID(op.userID, util.Logger)).Log("msg", "failed to flush user", "err", err)
		}

		// If we're exiting & we failed to flush, put the failed operation
		// back in the queue at a later point.
		if op.immediate && err != nil {
			op.from = op.from.Add(flushBackoff)
			i.flushQueues[j].Enqueue(op)
		}
	}
}

func (i *Ingester) flushUserSeries(userID string, fp model.Fingerprint, immediate bool) error {
	instance, ok := i.getInstanceByID(userID)
	if !ok {
		return nil
	}

	chunks, labels := i.collectChunksToFlush(instance, fp, immediate)
	if len(chunks) < 1 {
		return nil
	}

	ctx := user.InjectOrgID(context.Background(), userID)
	ctx, cancel := context.WithTimeout(ctx, i.cfg.FlushOpTimeout)
	defer cancel()
	err := i.flushChunks(ctx, fp, labels, chunks)
	if err != nil {
		return err
	}

	instance.streamsMtx.Lock()
	for _, chunk := range chunks {
		chunk.flushed = time.Now()
	}
	instance.streamsMtx.Unlock()
	return nil
}

func (i *Ingester) collectChunksToFlush(instance *instance, fp model.Fingerprint, immediate bool) ([]*chunkDesc, []client.LabelPair) {
	instance.streamsMtx.Lock()
	defer instance.streamsMtx.Unlock()

	stream, ok := instance.streams[fp]
	if !ok {
		return nil, nil
	}

	var result []*chunkDesc
	for j := range stream.chunks {
		if immediate || i.shouldFlushChunk(&stream.chunks[j]) {
			result = append(result, &stream.chunks[j])
		}
	}
	return result, stream.labels
}

func (i *Ingester) shouldFlushChunk(chunk *chunkDesc) bool {
	if !chunk.flushed.IsZero() {
		return false
	}

	// Append should close the chunk when the a new one is added.
	if chunk.closed {
		return true
	}

	if time.Since(chunk.lastUpdated) > i.cfg.MaxChunkIdle {
		chunk.closed = true
		return true
	}

	return false
}

func (i *Ingester) removeFlushedChunks(instance *instance, stream *stream) {
	now := time.Now()

	for len(stream.chunks) > 0 {
		if stream.chunks[0].flushed.IsZero() || now.Sub(stream.chunks[0].flushed) < i.cfg.RetainPeriod {
			break
		}

		stream.chunks[0].chunk = nil // erase reference so the chunk can be garbage-collected
		stream.chunks = stream.chunks[1:]
	}

	if len(stream.chunks) == 0 {
		delete(instance.streams, stream.fp)
		instance.index.Delete(stream.labels, stream.fp)
		instance.streamsRemovedTotal.Inc()
	}
}

func (i *Ingester) flushChunks(ctx context.Context, fp model.Fingerprint, labelPairs []client.LabelPair, cs []*chunkDesc) error {
	userID, err := user.ExtractOrgID(ctx)
	if err != nil {
		return err
	}

	metric := fromLabelPairs(labelPairs)
	metric[nameLabel] = logsValue

	wireChunks := make([]chunk.Chunk, 0, len(cs))
	for _, c := range cs {
		firstTime, lastTime := c.chunk.Bounds()
		wireChunks = append(wireChunks, chunk.NewChunk(
			userID, fp, metric,
			chunkenc.NewFacade(c.chunk),
			model.TimeFromUnixNano(firstTime.UnixNano()),
			model.TimeFromUnixNano(lastTime.UnixNano()),
		))
	}

	return i.store.Put(ctx, wireChunks)
}

func fromLabelPairs(ls []client.LabelPair) model.Metric {
	m := make(model.Metric, len(ls))
	for _, l := range ls {
		m[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return m
}
<commit_msg>Explicty encode the chunk.<commit_after>package ingester

import (
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"

	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/common/model"
	"github.com/weaveworks/common/user"

	"github.com/cortexproject/cortex/pkg/chunk"
	"github.com/cortexproject/cortex/pkg/ingester/client"
	"github.com/cortexproject/cortex/pkg/util"
	"github.com/grafana/loki/pkg/chunkenc"
)

const (
	// Backoff for retrying 'immediate' flushes. Only counts for queue
	// position, not wallclock time.
	flushBackoff = 1 * time.Second

	nameLabel = model.LabelName("__name__")
	logsValue = model.LabelValue("logs")
)

// Flush triggers a flush of all the chunks and closes the flush queues.
// Called from the Lifecycler as part of the ingester shutdown.
func (i *Ingester) Flush() {
	i.sweepUsers(true)

	// Close the flush queues, to unblock waiting workers.
	for _, flushQueue := range i.flushQueues {
		flushQueue.Close()
	}

	i.flushQueuesDone.Wait()
}

// FlushHandler triggers a flush of all in memory chunks.  Mainly used for
// local testing.
func (i *Ingester) FlushHandler(w http.ResponseWriter, _ *http.Request) {
	i.sweepUsers(true)
	w.WriteHeader(http.StatusNoContent)
}

type flushOp struct {
	from      model.Time
	userID    string
	fp        model.Fingerprint
	immediate bool
}

func (o *flushOp) Key() string {
	return fmt.Sprintf("%s-%s-%v", o.userID, o.fp, o.immediate)
}

func (o *flushOp) Priority() int64 {
	return -int64(o.from)
}

// sweepUsers periodically schedules series for flushing and garbage collects users with no series
func (i *Ingester) sweepUsers(immediate bool) {
	instances := i.getInstances()

	for _, instance := range instances {
		i.sweepInstance(instance, immediate)
	}
}

func (i *Ingester) sweepInstance(instance *instance, immediate bool) {
	instance.streamsMtx.Lock()
	defer instance.streamsMtx.Unlock()

	for _, stream := range instance.streams {
		i.sweepStream(instance, stream, immediate)
		i.removeFlushedChunks(instance, stream)
	}
}

func (i *Ingester) sweepStream(instance *instance, stream *stream, immediate bool) {
	if len(stream.chunks) == 0 {
		return
	}

	lastChunk := stream.chunks[len(stream.chunks)-1]
	if len(stream.chunks) == 1 && time.Since(lastChunk.lastUpdated) < i.cfg.MaxChunkIdle && !immediate {
		return
	}

	flushQueueIndex := int(uint64(stream.fp) % uint64(i.cfg.ConcurrentFlushes))
	firstTime, _ := stream.chunks[0].chunk.Bounds()
	i.flushQueues[flushQueueIndex].Enqueue(&flushOp{
		model.TimeFromUnixNano(firstTime.UnixNano()), instance.instanceID,
		stream.fp, immediate,
	})
}

func (i *Ingester) flushLoop(j int) {
	defer func() {
		level.Debug(util.Logger).Log("msg", "Ingester.flushLoop() exited")
		i.flushQueuesDone.Done()
	}()

	for {
		o := i.flushQueues[j].Dequeue()
		if o == nil {
			return
		}
		op := o.(*flushOp)

		level.Debug(util.Logger).Log("msg", "flushing stream", "userid", op.userID, "fp", op.fp, "immediate", op.immediate)

		err := i.flushUserSeries(op.userID, op.fp, op.immediate)
		if err != nil {
			level.Error(util.WithUserID(op.userID, util.Logger)).Log("msg", "failed to flush user", "err", err)
		}

		// If we're exiting & we failed to flush, put the failed operation
		// back in the queue at a later point.
		if op.immediate && err != nil {
			op.from = op.from.Add(flushBackoff)
			i.flushQueues[j].Enqueue(op)
		}
	}
}

func (i *Ingester) flushUserSeries(userID string, fp model.Fingerprint, immediate bool) error {
	instance, ok := i.getInstanceByID(userID)
	if !ok {
		return nil
	}

	chunks, labels := i.collectChunksToFlush(instance, fp, immediate)
	if len(chunks) < 1 {
		return nil
	}

	ctx := user.InjectOrgID(context.Background(), userID)
	ctx, cancel := context.WithTimeout(ctx, i.cfg.FlushOpTimeout)
	defer cancel()
	err := i.flushChunks(ctx, fp, labels, chunks)
	if err != nil {
		return err
	}

	instance.streamsMtx.Lock()
	for _, chunk := range chunks {
		chunk.flushed = time.Now()
	}
	instance.streamsMtx.Unlock()
	return nil
}

func (i *Ingester) collectChunksToFlush(instance *instance, fp model.Fingerprint, immediate bool) ([]*chunkDesc, []client.LabelPair) {
	instance.streamsMtx.Lock()
	defer instance.streamsMtx.Unlock()

	stream, ok := instance.streams[fp]
	if !ok {
		return nil, nil
	}

	var result []*chunkDesc
	for j := range stream.chunks {
		if immediate || i.shouldFlushChunk(&stream.chunks[j]) {
			result = append(result, &stream.chunks[j])
		}
	}
	return result, stream.labels
}

func (i *Ingester) shouldFlushChunk(chunk *chunkDesc) bool {
	if !chunk.flushed.IsZero() {
		return false
	}

	// Append should close the chunk when the a new one is added.
	if chunk.closed {
		return true
	}

	if time.Since(chunk.lastUpdated) > i.cfg.MaxChunkIdle {
		chunk.closed = true
		return true
	}

	return false
}

func (i *Ingester) removeFlushedChunks(instance *instance, stream *stream) {
	now := time.Now()

	for len(stream.chunks) > 0 {
		if stream.chunks[0].flushed.IsZero() || now.Sub(stream.chunks[0].flushed) < i.cfg.RetainPeriod {
			break
		}

		stream.chunks[0].chunk = nil // erase reference so the chunk can be garbage-collected
		stream.chunks = stream.chunks[1:]
	}

	if len(stream.chunks) == 0 {
		delete(instance.streams, stream.fp)
		instance.index.Delete(stream.labels, stream.fp)
		instance.streamsRemovedTotal.Inc()
	}
}

func (i *Ingester) flushChunks(ctx context.Context, fp model.Fingerprint, labelPairs []client.LabelPair, cs []*chunkDesc) error {
	userID, err := user.ExtractOrgID(ctx)
	if err != nil {
		return err
	}

	metric := fromLabelPairs(labelPairs)
	metric[nameLabel] = logsValue

	wireChunks := make([]chunk.Chunk, 0, len(cs))
	for _, c := range cs {
		firstTime, lastTime := c.chunk.Bounds()
		c := chunk.NewChunk(
			userID, fp, metric,
			chunkenc.NewFacade(c.chunk),
			model.TimeFromUnixNano(firstTime.UnixNano()),
			model.TimeFromUnixNano(lastTime.UnixNano()),
		)

		if err := c.Encode(); err != nil {
			return err
		}
		wireChunks = append(wireChunks, c)
	}

	return i.store.Put(ctx, wireChunks)
}

func fromLabelPairs(ls []client.LabelPair) model.Metric {
	m := make(model.Metric, len(ls))
	for _, l := range ls {
		m[model.LabelName(l.Name)] = model.LabelValue(l.Value)
	}
	return m
}
<|endoftext|>