<commit_before>package process

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"strconv"

	"kego.io/kerr"
)

type fileType string

const (
	F_MAIN         fileType = "main"
	F_TYPES                 = "types"
	F_GLOBALS               = "globals"
	F_CMD_MAIN              = "cmd_main"
	F_CMD_TYPES             = "cmd_types"
	F_CMD_VALIDATE          = "cmd_validate"
)

var generatorUpdateFlag = flag.Bool("u", false, "Update: update all import packages e.g. go get -u")
var generatorPathFlag = flag.String("p", "", "Package: full package path e.g. github.com/foo/bar")
var generatorRecursiveFlag = flag.Bool("r", false, "Recursive: scan subdirectories for objects")
var generatorVerboseFlag = flag.Bool("v", false, "Verbose")

func Initialise() (dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string, err error) {

	if !flag.Parsed() {
		flag.Parse()
	}

	update = *generatorUpdateFlag
	path = *generatorPathFlag
	recursive = *generatorRecursiveFlag
	verbose = *generatorVerboseFlag

	dir, err = os.Getwd()
	if err != nil {
		err = kerr.New("OKOLXAMBSJ", err, "process.Initialise", "os.Getwd")
		return
	}

	if path == "" {
		path, err = getPackage(dir, os.Getenv("GOPATH"))
		if err != nil {
			err = kerr.New("PSRAWHQCPV", err, "process.Initialise", "getPackage")
			return
		}
	}

	aliases, err = ScanForAliases(dir, recursive, path)
	if err != nil {
		err = kerr.New("IAAETYCHSW", err, "process.Initialise", "ScanForImports")
		return
	}

	return

}

func KegoCmd(dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {
	if err := GenerateAndRunCmd(F_CMD_MAIN, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	if err := GenerateAndRunCmd(F_CMD_TYPES, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	if err := GenerateAndRunCmd(F_CMD_VALIDATE, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	return nil
}

// GenerateFiles generates the source code from templates and writes
// the files to the correct folders.
//
// file == F_MAIN: the generated.go in the root of the package.
//
// file == F_TYPES: the generated.go containing advanced type information
// in the "types" sub package. Note that to generate this file, we need to
// have the main generated.go compiled in, so we generate a temporary
// command and run it with "go run".
//
// file == F_CMD_TYPES: this is the temporary command that we create in order to
// generate the types source file
//
// file == F_CMD_VALIDATE: this is the temporary command that we create in order
// to run the validation
//
func GenerateFiles(file fileType, dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {

	if verbose {
		fmt.Println("Generating", file)
	}

	outputDir := dir
	if file == F_TYPES {
		outputDir = filepath.Join(dir, "types")
	}

	ignoreUnknownTypes := true
	if file == F_TYPES {
		ignoreUnknownTypes = false
	}

	if file == F_MAIN || file == F_TYPES {
		// If type == F_GLOBALS, we have already generated and imported the types, so
		// there is no need to scan.
		if err := ScanForTypes(dir, ignoreUnknownTypes, recursive, path, aliases); err != nil {
			return kerr.New("XYIUHERDHE", err, "process.GenerateFiles", "ScanForTypes")
		}
	} else {
		// However, we need to scan for the globals.
		if err := ScanForGlobals(dir, recursive, path, aliases); err != nil {
			return kerr.New("JQLAQVKLAN", err, "process.GenerateFiles", "ScanForGlobals")
		}
	}

	source, err := Generate(file, path, aliases)
	if err != nil {
		return kerr.New("XFNESBLBTQ", err, "process.GenerateFiles", "Generate")
	}

	// We only backup in the system types because they are the only
	// generated files you'll ever need to roll back
	backup := path == "kego.io/system" || path == "kego.io/system/types"

	filename := "generated.go"
	if file == F_GLOBALS {
		filename = "globals.go"
	}

	if err = save(outputDir, source, filename, backup); err != nil {
		return kerr.New("UONJTTSTWW", err, "process.GenerateFiles", "save")
	}

	return nil
}

// This creates a temporary folder in the package, in which the go source
// for a command is generated. This command is then compiled and run with
// "go run". When run, this command generates the extra types data in
// the "types" subpackage.
func GenerateAndRunCmd(file fileType, dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {

	if verbose {
		fmt.Println("Generating", file)
	}

	source, err := Generate(file, path, aliases)
	if err != nil {
		return kerr.New("SPRFABSRWK", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "Generate")
	}

	outputDir, err := ioutil.TempDir(dir, "temporary")
	if err != nil {
		return kerr.New("HWOPVXYMCT", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "ioutil.TempDir")
	}
	defer os.RemoveAll(outputDir)
	outputName := "generated_cmd.go"
	outputPath := filepath.Join(outputDir, outputName)

	if err = save(outputDir, source, outputName, false); err != nil {
		return kerr.New("FRLCYFOWCJ", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "save")
	}

	if file == F_CMD_VALIDATE {
		cmd := exec.Command("go", "build", "-o", "validate", outputPath)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return kerr.New("OEPAEEYKIS", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "cmd.Run (go build):\n%s", out)
		}
		if verbose {
			fmt.Print(string(out))
		}
	}

	command := ""
	params := []string{}

	if file == F_CMD_VALIDATE {
		command = "./validate"
	} else {
		command = "go"
		params = []string{"run", outputPath}
	}

	if *generatorPathFlag != "" {
		params = append(params, fmt.Sprintf("-p=%s", strconv.Quote(*generatorPathFlag)))
	}
	if update {
		params = append(params, "-u")
	}
	if recursive {
		params = append(params, "-r")
	}
	if verbose {
		params = append(params, "-v")
	}
	cmd := exec.Command(command, params...)

	if verbose {
		fmt.Println("Running", file)
	}

	out, err := cmd.CombinedOutput()
	if err != nil {
		return kerr.New("UDDSSMQRHA", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "cmd.Run: \n%s", out)
	}

	if verbose {
		fmt.Print(string(out))
	}

	return nil
}

func save(dir string, contents []byte, name string, backup bool) error {

	if _, err := os.Stat(dir); os.IsNotExist(err) {
		if err = os.MkdirAll(dir, 0777); err != nil {
			return kerr.New("BPGOUIYPXO", err, "process.save", "os.MkdirAll")
		}
	}

	file := filepath.Join(dir, name)

	if backup {
		backupPath := filepath.Join(dir, fmt.Sprintf("%s.backup", name))
		if _, err := os.Stat(backupPath); err == nil {
			os.Remove(backupPath)
		}
		if _, err := os.Stat(file); err == nil {
			os.Rename(file, backupPath)
		}
	} else {
		os.Remove(file)
	}

	if len(contents) == 0 {
		return nil
	}

	output, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return kerr.New("NWLWHSGJWP", err, "process.save", "os.OpenFile (could not open output file)")
	}
	defer output.Close()

	if _, err := output.Write(contents); err != nil {
		return kerr.New("FBMGPRWQBL", err, "process.save", "output.Write")
	}

	if err := output.Sync(); err != nil {
		return kerr.New("EGFNTMNKFX", err, "process.save", "output.Sync")
	}

	return nil
}

func getPackage(dir string, gopathEnv string) (string, error) {
	gopaths := filepath.SplitList(gopathEnv)
	var savedError error
	for _, gopath := range gopaths {
		if strings.HasPrefix(dir, gopath) {
			gosrc := fmt.Sprintf("%s/src", gopath)
			relpath, err := filepath.Rel(gosrc, dir)
			if err != nil {
				savedError = err
				continue
			}
			if relpath == "" {
				continue
			}
			return relpath, nil
		}
	}
	if savedError != nil {
		return "", savedError
	}
	return "", kerr.New("CXOETFPTGM", nil, "process.getPackage", "Package not found for %s", dir)
}
<commit_msg>Added go get to aliases<commit_after>package process

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"strconv"

	"kego.io/kerr"
)

type fileType string

const (
	F_MAIN         fileType = "main"
	F_TYPES                 = "types"
	F_GLOBALS               = "globals"
	F_CMD_MAIN              = "cmd_main"
	F_CMD_TYPES             = "cmd_types"
	F_CMD_VALIDATE          = "cmd_validate"
)

var generatorUpdateFlag = flag.Bool("u", false, "Update: update all import packages e.g. go get -u")
var generatorPathFlag = flag.String("p", "", "Package: full package path e.g. github.com/foo/bar")
var generatorRecursiveFlag = flag.Bool("r", false, "Recursive: scan subdirectories for objects")
var generatorVerboseFlag = flag.Bool("v", false, "Verbose")

func Initialise() (dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string, err error) {

	if !flag.Parsed() {
		flag.Parse()
	}

	update = *generatorUpdateFlag
	path = *generatorPathFlag
	recursive = *generatorRecursiveFlag
	verbose = *generatorVerboseFlag

	dir, err = os.Getwd()
	if err != nil {
		err = kerr.New("OKOLXAMBSJ", err, "process.Initialise", "os.Getwd")
		return
	}

	if path == "" {
		path, err = getPackage(dir, os.Getenv("GOPATH"))
		if err != nil {
			err = kerr.New("PSRAWHQCPV", err, "process.Initialise", "getPackage")
			return
		}
	}

	aliases, err = ScanForAliases(dir, recursive, path)
	if err != nil {
		err = kerr.New("IAAETYCHSW", err, "process.Initialise", "ScanForImports")
		return
	}

	for p, _ := range aliases {
		params := []string{}
		if update {
			params = []string{"get", "-u", p}
		} else {
			params = []string{"get", p}
		}
		var out []byte
		if out, err = exec.Command("go", params...).CombinedOutput(); err != nil {
			err = kerr.New("HHKSTQMAKG", err, "process.Initialize", "go get command: %s", out)
			return
		}
	}

	return

}

func KegoCmd(dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {
	if err := GenerateAndRunCmd(F_CMD_MAIN, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	if err := GenerateAndRunCmd(F_CMD_TYPES, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	if err := GenerateAndRunCmd(F_CMD_VALIDATE, dir, update, recursive, verbose, path, aliases); err != nil {
		return err
	}
	return nil
}

// GenerateFiles generates the source code from templates and writes
// the files to the correct folders.
//
// file == F_MAIN: the generated.go in the root of the package.
//
// file == F_TYPES: the generated.go containing advanced type information
// in the "types" sub package. Note that to generate this file, we need to
// have the main generated.go compiled in, so we generate a temporary
// command and run it with "go run".
//
// file == F_CMD_TYPES: this is the temporary command that we create in order to
// generate the types source file
//
// file == F_CMD_VALIDATE: this is the temporary command that we create in order
// to run the validation
//
func GenerateFiles(file fileType, dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {

	if verbose {
		fmt.Println("Generating", file)
	}

	outputDir := dir
	if file == F_TYPES {
		outputDir = filepath.Join(dir, "types")
	}

	ignoreUnknownTypes := true
	if file == F_TYPES {
		ignoreUnknownTypes = false
	}

	if file == F_MAIN || file == F_TYPES {
		// If type == F_GLOBALS, we have already generated and imported the types, so
		// there is no need to scan.
		if err := ScanForTypes(dir, ignoreUnknownTypes, recursive, path, aliases); err != nil {
			return kerr.New("XYIUHERDHE", err, "process.GenerateFiles", "ScanForTypes")
		}
	} else {
		// However, we need to scan for the globals.
		if err := ScanForGlobals(dir, recursive, path, aliases); err != nil {
			return kerr.New("JQLAQVKLAN", err, "process.GenerateFiles", "ScanForGlobals")
		}
	}

	source, err := Generate(file, path, aliases)
	if err != nil {
		return kerr.New("XFNESBLBTQ", err, "process.GenerateFiles", "Generate")
	}

	// We only backup in the system types because they are the only
	// generated files you'll ever need to roll back
	backup := path == "kego.io/system" || path == "kego.io/system/types"

	filename := "generated.go"
	if file == F_GLOBALS {
		filename = "globals.go"
	}

	if err = save(outputDir, source, filename, backup); err != nil {
		return kerr.New("UONJTTSTWW", err, "process.GenerateFiles", "save")
	}

	return nil
}

// This creates a temporary folder in the package, in which the go source
// for a command is generated. This command is then compiled and run with
// "go run". When run, this command generates the extra types data in
// the "types" subpackage.
func GenerateAndRunCmd(file fileType, dir string, update bool, recursive bool, verbose bool, path string, aliases map[string]string) error {

	if verbose {
		fmt.Println("Generating", file)
	}

	source, err := Generate(file, path, aliases)
	if err != nil {
		return kerr.New("SPRFABSRWK", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "Generate")
	}

	outputDir, err := ioutil.TempDir(dir, "temporary")
	if err != nil {
		return kerr.New("HWOPVXYMCT", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "ioutil.TempDir")
	}
	defer os.RemoveAll(outputDir)
	outputName := "generated_cmd.go"
	outputPath := filepath.Join(outputDir, outputName)

	if err = save(outputDir, source, outputName, false); err != nil {
		return kerr.New("FRLCYFOWCJ", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "save")
	}

	if file == F_CMD_VALIDATE {
		cmd := exec.Command("go", "build", "-o", "validate", outputPath)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return kerr.New("OEPAEEYKIS", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "cmd.Run (go build):\n%s", out)
		}
		if verbose {
			fmt.Print(string(out))
		}
	}

	command := ""
	params := []string{}

	if file == F_CMD_VALIDATE {
		command = "./validate"
	} else {
		command = "go"
		params = []string{"run", outputPath}
	}

	if *generatorPathFlag != "" {
		params = append(params, fmt.Sprintf("-p=%s", strconv.Quote(*generatorPathFlag)))
	}
	if update {
		params = append(params, "-u")
	}
	if recursive {
		params = append(params, "-r")
	}
	if verbose {
		params = append(params, "-v")
	}
	cmd := exec.Command(command, params...)

	if verbose {
		fmt.Println("Running", file)
	}

	out, err := cmd.CombinedOutput()
	if err != nil {
		return kerr.New("UDDSSMQRHA", err, fmt.Sprintf("process.GenerateAndRunCmd %s", file), "cmd.Run: \n%s", out)
	}

	if verbose {
		fmt.Print(string(out))
	}

	return nil
}

func save(dir string, contents []byte, name string, backup bool) error {

	if _, err := os.Stat(dir); os.IsNotExist(err) {
		if err = os.MkdirAll(dir, 0777); err != nil {
			return kerr.New("BPGOUIYPXO", err, "process.save", "os.MkdirAll")
		}
	}

	file := filepath.Join(dir, name)

	if backup {
		backupPath := filepath.Join(dir, fmt.Sprintf("%s.backup", name))
		if _, err := os.Stat(backupPath); err == nil {
			os.Remove(backupPath)
		}
		if _, err := os.Stat(file); err == nil {
			os.Rename(file, backupPath)
		}
	} else {
		os.Remove(file)
	}

	if len(contents) == 0 {
		return nil
	}

	output, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return kerr.New("NWLWHSGJWP", err, "process.save", "os.OpenFile (could not open output file)")
	}
	defer output.Close()

	if _, err := output.Write(contents); err != nil {
		return kerr.New("FBMGPRWQBL", err, "process.save", "output.Write")
	}

	if err := output.Sync(); err != nil {
		return kerr.New("EGFNTMNKFX", err, "process.save", "output.Sync")
	}

	return nil
}

func getPackage(dir string, gopathEnv string) (string, error) {
	gopaths := filepath.SplitList(gopathEnv)
	var savedError error
	for _, gopath := range gopaths {
		if strings.HasPrefix(dir, gopath) {
			gosrc := fmt.Sprintf("%s/src", gopath)
			relpath, err := filepath.Rel(gosrc, dir)
			if err != nil {
				savedError = err
				continue
			}
			if relpath == "" {
				continue
			}
			return relpath, nil
		}
	}
	if savedError != nil {
		return "", savedError
	}
	return "", kerr.New("CXOETFPTGM", nil, "process.getPackage", "Package not found for %s", dir)
}
<|endoftext|>


<commit_before>// Package lint provides abstractions on top of go/analysis.
package lint

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"strconv"
	"strings"

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

type Analyzer struct {
	// The analyzer's documentation. Unlike go/analysis.Analyzer.Doc,
	// this field is structured, providing access to severity, options
	// etc.
	Doc      *Documentation
	Analyzer *analysis.Analyzer
}

func (a *Analyzer) initialize() {
	a.Analyzer.Doc = a.Doc.String()
	if a.Analyzer.Flags.Usage == nil {
		fs := flag.NewFlagSet("", flag.PanicOnError)
		fs.Var(newVersionFlag(), "go", "Target Go version")
		a.Analyzer.Flags = *fs
	}
}

func InitializeAnalyzers(docs map[string]*Documentation, analyzers map[string]*analysis.Analyzer) []*Analyzer {
	out := make([]*Analyzer, 0, len(analyzers))
	for k, v := range analyzers {
		v.Name = k
		a := &Analyzer{
			Doc:      docs[k],
			Analyzer: v,
		}
		a.initialize()
		out = append(out, a)
	}
	return out
}

type Severity int

const (
	SeverityNone Severity = iota
	SeverityError
	SeverityDeprecated
	SeverityWarning
	SeverityInfo
	SeverityHint
)

type MergeStrategy int

const (
	MergeIfAny MergeStrategy = iota
	MergeIfAll
)

type RawDocumentation struct {
	Title      string
	Text       string
	Before     string
	After      string
	Since      string
	NonDefault bool
	Options    []string
	Severity   Severity
	MergeIf    MergeStrategy
}

type Documentation struct {
	Title string
	Text  string

	TitleMarkdown string
	TextMarkdown  string

	Before     string
	After      string
	Since      string
	NonDefault bool
	Options    []string
	Severity   Severity
	MergeIf    MergeStrategy
}

func Markdownify(m map[string]*RawDocumentation) map[string]*Documentation {
	out := make(map[string]*Documentation, len(m))
	for k, v := range m {
		out[k] = &Documentation{
			Title: strings.TrimSpace(stripMarkdown(v.Title)),
			Text:  strings.TrimSpace(stripMarkdown(v.Text)),

			TitleMarkdown: strings.TrimSpace(toMarkdown(v.Title)),
			TextMarkdown:  strings.TrimSpace(toMarkdown(v.Text)),

			Before:     strings.TrimSpace(v.Before),
			After:      strings.TrimSpace(v.After),
			Since:      v.Since,
			NonDefault: v.NonDefault,
			Options:    v.Options,
			Severity:   v.Severity,
			MergeIf:    v.MergeIf,
		}
	}
	return out
}

func toMarkdown(s string) string {
	return strings.NewReplacer(`\'`, "`", `\"`, "`").Replace(s)
}

func stripMarkdown(s string) string {
	return strings.NewReplacer(`\'`, "", `\"`, "'").Replace(s)
}

func (doc *Documentation) Format(metadata bool) string {
	return doc.format(false, metadata)
}

func (doc *Documentation) FormatMarkdown(metadata bool) string {
	return doc.format(true, metadata)
}

func (doc *Documentation) format(markdown bool, metadata bool) string {
	b := &strings.Builder{}
	if markdown {
		fmt.Fprintf(b, "%s\n\n", doc.TitleMarkdown)
		if doc.Text != "" {
			fmt.Fprintf(b, "%s\n\n", doc.TextMarkdown)
		}
	} else {
		fmt.Fprintf(b, "%s\n\n", doc.Title)
		if doc.Text != "" {
			fmt.Fprintf(b, "%s\n\n", doc.Text)
		}
	}

	if doc.Before != "" {
		fmt.Fprintln(b, "Before:")
		fmt.Fprintln(b, "")
		for _, line := range strings.Split(doc.Before, "\n") {
			fmt.Fprint(b, "    ", line, "\n")
		}
		fmt.Fprintln(b, "")
		fmt.Fprintln(b, "After:")
		fmt.Fprintln(b, "")
		for _, line := range strings.Split(doc.After, "\n") {
			fmt.Fprint(b, "    ", line, "\n")
		}
		fmt.Fprintln(b, "")
	}

	if metadata {
		fmt.Fprint(b, "Available since\n    ")
		if doc.Since == "" {
			fmt.Fprint(b, "unreleased")
		} else {
			fmt.Fprintf(b, "%s", doc.Since)
		}
		if doc.NonDefault {
			fmt.Fprint(b, ", non-default")
		}
		fmt.Fprint(b, "\n")
		if len(doc.Options) > 0 {
			fmt.Fprintf(b, "\nOptions\n")
			for _, opt := range doc.Options {
				fmt.Fprintf(b, "    %s", opt)
			}
			fmt.Fprint(b, "\n")
		}
	}

	return b.String()
}

func (doc *Documentation) String() string {
	return doc.Format(true)
}

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 fmt.Errorf("invalid Go version: %q", s)
	}
	if s[0] != '1' {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	if s[1] != '.' {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	i, err := strconv.Atoi(s[2:])
	if err != nil {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	*v = VersionFlag(i)
	return nil
}

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

// ExhaustiveTypeSwitch panics when called. It can be used to ensure
// that type switches are exhaustive.
func ExhaustiveTypeSwitch(v interface{}) {
	panic(fmt.Sprintf("internal error: unhandled case %T", v))
}

// A directive is a comment of the form '//lint:<command>
// [arguments...]'. It represents instructions to the static analysis
// tool.
type Directive struct {
	Command   string
	Arguments []string
	Directive *ast.Comment
	Node      ast.Node
}

func parseDirective(s string) (cmd string, args []string) {
	if !strings.HasPrefix(s, "//lint:") {
		return "", nil
	}
	s = strings.TrimPrefix(s, "//lint:")
	fields := strings.Split(s, " ")
	return fields[0], fields[1:]
}

func ParseDirectives(files []*ast.File, fset *token.FileSet) []Directive {
	var dirs []Directive
	for _, f := range files {
		// OPT(dh): in our old code, we skip all the comment map work if we
		// couldn't find any directives, benchmark if that's actually
		// worth doing
		cm := ast.NewCommentMap(fset, f, f.Comments)
		for node, cgs := range cm {
			for _, cg := range cgs {
				for _, c := range cg.List {
					if !strings.HasPrefix(c.Text, "//lint:") {
						continue
					}
					cmd, args := parseDirective(c.Text)
					d := Directive{
						Command:   cmd,
						Arguments: args,
						Directive: c,
						Node:      node,
					}
					dirs = append(dirs, d)
				}
			}
		}
	}
	return dirs
}
<commit_msg>analysis/lint: write some documentation<commit_after>// Package lint provides abstractions on top of go/analysis.
// These abstractions add extra information to analyzes, such as structured documentation and severities.
package lint

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/token"
	"strconv"
	"strings"

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

// Analyzer wraps a go/analysis.Analyzer and provides structured documentation.
type Analyzer struct {
	// The analyzer's documentation. Unlike go/analysis.Analyzer.Doc,
	// this field is structured, providing access to severity, options
	// etc.
	Doc      *Documentation
	Analyzer *analysis.Analyzer
}

func (a *Analyzer) initialize() {
	a.Analyzer.Doc = a.Doc.String()
	if a.Analyzer.Flags.Usage == nil {
		fs := flag.NewFlagSet("", flag.PanicOnError)
		fs.Var(newVersionFlag(), "go", "Target Go version")
		a.Analyzer.Flags = *fs
	}
}

// InitializeAnalyzers takes a map of documentation and a map of go/analysis.Analyzers and returns a slice of Analyzers.
// The map keys are the analyzer names.
func InitializeAnalyzers(docs map[string]*Documentation, analyzers map[string]*analysis.Analyzer) []*Analyzer {
	out := make([]*Analyzer, 0, len(analyzers))
	for k, v := range analyzers {
		v.Name = k
		a := &Analyzer{
			Doc:      docs[k],
			Analyzer: v,
		}
		a.initialize()
		out = append(out, a)
	}
	return out
}

// Severity describes the severity of diagnostics reported by an analyzer.
type Severity int

const (
	SeverityNone Severity = iota
	SeverityError
	SeverityDeprecated
	SeverityWarning
	SeverityInfo
	SeverityHint
)

// MergeStrategy sets how merge mode should behave for diagnostics of an analyzer.
type MergeStrategy int

const (
	MergeIfAny MergeStrategy = iota
	MergeIfAll
)

type RawDocumentation struct {
	Title      string
	Text       string
	Before     string
	After      string
	Since      string
	NonDefault bool
	Options    []string
	Severity   Severity
	MergeIf    MergeStrategy
}

type Documentation struct {
	Title string
	Text  string

	TitleMarkdown string
	TextMarkdown  string

	Before     string
	After      string
	Since      string
	NonDefault bool
	Options    []string
	Severity   Severity
	MergeIf    MergeStrategy
}

func Markdownify(m map[string]*RawDocumentation) map[string]*Documentation {
	out := make(map[string]*Documentation, len(m))
	for k, v := range m {
		out[k] = &Documentation{
			Title: strings.TrimSpace(stripMarkdown(v.Title)),
			Text:  strings.TrimSpace(stripMarkdown(v.Text)),

			TitleMarkdown: strings.TrimSpace(toMarkdown(v.Title)),
			TextMarkdown:  strings.TrimSpace(toMarkdown(v.Text)),

			Before:     strings.TrimSpace(v.Before),
			After:      strings.TrimSpace(v.After),
			Since:      v.Since,
			NonDefault: v.NonDefault,
			Options:    v.Options,
			Severity:   v.Severity,
			MergeIf:    v.MergeIf,
		}
	}
	return out
}

func toMarkdown(s string) string {
	return strings.NewReplacer(`\'`, "`", `\"`, "`").Replace(s)
}

func stripMarkdown(s string) string {
	return strings.NewReplacer(`\'`, "", `\"`, "'").Replace(s)
}

func (doc *Documentation) Format(metadata bool) string {
	return doc.format(false, metadata)
}

func (doc *Documentation) FormatMarkdown(metadata bool) string {
	return doc.format(true, metadata)
}

func (doc *Documentation) format(markdown bool, metadata bool) string {
	b := &strings.Builder{}
	if markdown {
		fmt.Fprintf(b, "%s\n\n", doc.TitleMarkdown)
		if doc.Text != "" {
			fmt.Fprintf(b, "%s\n\n", doc.TextMarkdown)
		}
	} else {
		fmt.Fprintf(b, "%s\n\n", doc.Title)
		if doc.Text != "" {
			fmt.Fprintf(b, "%s\n\n", doc.Text)
		}
	}

	if doc.Before != "" {
		fmt.Fprintln(b, "Before:")
		fmt.Fprintln(b, "")
		for _, line := range strings.Split(doc.Before, "\n") {
			fmt.Fprint(b, "    ", line, "\n")
		}
		fmt.Fprintln(b, "")
		fmt.Fprintln(b, "After:")
		fmt.Fprintln(b, "")
		for _, line := range strings.Split(doc.After, "\n") {
			fmt.Fprint(b, "    ", line, "\n")
		}
		fmt.Fprintln(b, "")
	}

	if metadata {
		fmt.Fprint(b, "Available since\n    ")
		if doc.Since == "" {
			fmt.Fprint(b, "unreleased")
		} else {
			fmt.Fprintf(b, "%s", doc.Since)
		}
		if doc.NonDefault {
			fmt.Fprint(b, ", non-default")
		}
		fmt.Fprint(b, "\n")
		if len(doc.Options) > 0 {
			fmt.Fprintf(b, "\nOptions\n")
			for _, opt := range doc.Options {
				fmt.Fprintf(b, "    %s", opt)
			}
			fmt.Fprint(b, "\n")
		}
	}

	return b.String()
}

func (doc *Documentation) String() string {
	return doc.Format(true)
}

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 fmt.Errorf("invalid Go version: %q", s)
	}
	if s[0] != '1' {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	if s[1] != '.' {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	i, err := strconv.Atoi(s[2:])
	if err != nil {
		return fmt.Errorf("invalid Go version: %q", s)
	}
	*v = VersionFlag(i)
	return nil
}

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

// ExhaustiveTypeSwitch panics when called. It can be used to ensure
// that type switches are exhaustive.
func ExhaustiveTypeSwitch(v interface{}) {
	panic(fmt.Sprintf("internal error: unhandled case %T", v))
}

// A directive is a comment of the form '//lint:<command>
// [arguments...]'. It represents instructions to the static analysis
// tool.
type Directive struct {
	Command   string
	Arguments []string
	Directive *ast.Comment
	Node      ast.Node
}

func parseDirective(s string) (cmd string, args []string) {
	if !strings.HasPrefix(s, "//lint:") {
		return "", nil
	}
	s = strings.TrimPrefix(s, "//lint:")
	fields := strings.Split(s, " ")
	return fields[0], fields[1:]
}

// ParseDirectives extracts all directives from a list of Go files.
func ParseDirectives(files []*ast.File, fset *token.FileSet) []Directive {
	var dirs []Directive
	for _, f := range files {
		// OPT(dh): in our old code, we skip all the comment map work if we
		// couldn't find any directives, benchmark if that's actually
		// worth doing
		cm := ast.NewCommentMap(fset, f, f.Comments)
		for node, cgs := range cm {
			for _, cg := range cgs {
				for _, c := range cg.List {
					if !strings.HasPrefix(c.Text, "//lint:") {
						continue
					}
					cmd, args := parseDirective(c.Text)
					d := Directive{
						Command:   cmd,
						Arguments: args,
						Directive: c,
						Node:      node,
					}
					dirs = append(dirs, d)
				}
			}
		}
	}
	return dirs
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package cpumanager

import (
	"fmt"
	"sort"

	"k8s.io/klog/v2"

	"k8s.io/kubernetes/pkg/kubelet/cm/cpumanager/topology"
	"k8s.io/kubernetes/pkg/kubelet/cm/cpuset"
)

type cpuAccumulator struct {
	topo          *topology.CPUTopology
	details       topology.CPUDetails
	numCPUsNeeded int
	result        cpuset.CPUSet
}

func newCPUAccumulator(topo *topology.CPUTopology, availableCPUs cpuset.CPUSet, numCPUs int) *cpuAccumulator {
	return &cpuAccumulator{
		topo:          topo,
		details:       topo.CPUDetails.KeepOnly(availableCPUs),
		numCPUsNeeded: numCPUs,
		result:        cpuset.NewCPUSet(),
	}
}

// Returns true if the supplied NUMANode is fully available in `topoDetails`.
func (a *cpuAccumulator) isNUMANodeFree(numaID int) bool {
	return a.details.CPUsInNUMANodes(numaID).Size() == a.topo.CPUDetails.CPUsInNUMANodes(numaID).Size()
}

// Returns true if the supplied socket is fully available in `topoDetails`.
func (a *cpuAccumulator) isSocketFree(socketID int) bool {
	return a.details.CPUsInSockets(socketID).Size() == a.topo.CPUsPerSocket()
}

// Returns true if the supplied core is fully available in `topoDetails`.
func (a *cpuAccumulator) isCoreFree(coreID int) bool {
	return a.details.CPUsInCores(coreID).Size() == a.topo.CPUsPerCore()
}

// Returns free NUMA Node IDs as a slice sorted by sortAvailableNUMANodes().
func (a *cpuAccumulator) freeNUMANodes() []int {
	free := []int{}
	for _, numa := range a.sortAvailableNUMANodes() {
		if a.isNUMANodeFree(numa) {
			free = append(free, numa)
		}
	}
	return free
}

// Returns free socket IDs as a slice sorted by sortAvailableSockets().
func (a *cpuAccumulator) freeSockets() []int {
	free := []int{}
	for _, socket := range a.sortAvailableSockets() {
		if a.isSocketFree(socket) {
			free = append(free, socket)
		}
	}
	return free
}

// Returns free core IDs as a slice sorted by sortAvailableCores().
func (a *cpuAccumulator) freeCores() []int {
	free := []int{}
	for _, core := range a.sortAvailableCores() {
		if a.isCoreFree(core) {
			free = append(free, core)
		}
	}
	return free
}

// Returns free CPU IDs as a slice sorted by sortAvailableCPUs().
func (a *cpuAccumulator) freeCPUs() []int {
	return a.sortAvailableCPUs()
}

// Sorts the provided list of NUMA nodes/sockets/cores/cpus referenced in 'ids'
// by the number of available CPUs contained within them (smallest to largest).
// The 'getCPU()' paramater defines the function that should be called to
// retrieve the list of available CPUs for the type being referenced. If two
// NUMA nodes/sockets/cores/cpus have the same number of available CPUs, they
// are sorted in ascending order by their id.
func (a *cpuAccumulator) sort(ids []int, getCPUs func(ids ...int) cpuset.CPUSet) {
	sort.Slice(ids,
		func(i, j int) bool {
			iCPUs := getCPUs(ids[i])
			jCPUs := getCPUs(ids[j])
			if iCPUs.Size() < jCPUs.Size() {
				return true
			}
			if iCPUs.Size() > jCPUs.Size() {
				return false
			}
			return ids[i] < ids[j]
		})
}

// Sort all NUMA nodes with free CPUs:
// - If NUMA nodes are higher than sockets in the memory hierarchy then sort
//   them directly using the sort() algorithm defined above.
// - Otherwise sort them:
//   - First by socket using sortAvailableSockets().
//   - Then within each socket, using the sort() algorithm defined above.
func (a *cpuAccumulator) sortAvailableNUMANodes() []int {
	// If NUMA nodes are equal or higher in the memory hierarchy than sockets
	if a.topo.NumSockets >= a.topo.NumNUMANodes {
		numas := a.details.NUMANodes().ToSliceNoSort()
		a.sort(numas, a.details.CPUsInNUMANodes)
		return numas
	}

	// Otherwise each socket has multiple NUMA nodes
	var result []int
	for _, socket := range a.sortAvailableSockets() {
		numas := a.details.NUMANodesInSockets(socket).ToSliceNoSort()
		a.sort(numas, a.details.CPUsInNUMANodes)
		result = append(result, numas...)
	}
	return result
}

// Sort all sockets with free CPUs:
// - If sockets are higher than NUMA nodes in the memory hierarchy then sort
//   them directly using the sort() algorithm defined above.
// - Otherwise sort them:
//   - First by NUMA node using sortAvailableNUMANodes().
//   - Then within each NUMA node, using the sort() algorithm defined above.
func (a *cpuAccumulator) sortAvailableSockets() []int {
	// If sockets are higher than NUMA nodes in the memory hierarchy
	if a.topo.NumNUMANodes >= a.topo.NumSockets {
		sockets := a.details.Sockets().ToSliceNoSort()
		a.sort(sockets, a.details.CPUsInSockets)
		return sockets
	}

	// Otherwise each NUMA Node has multiple sockets
	var result []int
	for _, numa := range a.sortAvailableNUMANodes() {
		sockets := a.details.SocketsInNUMANodes(numa).ToSliceNoSort()
		a.sort(sockets, a.details.CPUsInSockets)
		result = append(result, sockets...)
	}
	return result
}

// Sort all cores with free CPUs:
// - First by socket (or NUMA node) using sortAvailableSockets() (or sortAvailableNUMANodes()).
// - Then within each socket or NUMA node, using the sort() algorithm defined above.
func (a *cpuAccumulator) sortAvailableCores() []int {
	// If NUMA nodes are higher in the memory hierarchy than sockets, then
	// cores sit directly below sockets in the memory hierarchy.
	if a.topo.NumSockets >= a.topo.NumNUMANodes {
		var result []int
		for _, socket := range a.sortAvailableSockets() {
			cores := a.details.CoresInSockets(socket).ToSliceNoSort()
			a.sort(cores, a.details.CPUsInCores)
			result = append(result, cores...)
		}
		return result
	}

	// Otherwise they sit directly below NUMA nodes.
	var result []int
	for _, numa := range a.sortAvailableNUMANodes() {
		cores := a.details.CoresInNUMANodes(numa).ToSliceNoSort()
		a.sort(cores, a.details.CPUsInCores)
		result = append(result, cores...)
	}
	return result
}

// Sort all available CPUs:
// - First by core using sortAvailableCores().
// - Then within each core, using the sort() algorithm defined above.
func (a *cpuAccumulator) sortAvailableCPUs() []int {
	var result []int
	for _, core := range a.sortAvailableCores() {
		cpus := a.details.CPUsInCores(core).ToSliceNoSort()
		sort.Ints(cpus)
		result = append(result, cpus...)
	}
	return result
}

func (a *cpuAccumulator) take(cpus cpuset.CPUSet) {
	a.result = a.result.Union(cpus)
	a.details = a.details.KeepOnly(a.details.CPUs().Difference(a.result))
	a.numCPUsNeeded -= cpus.Size()
}

func (a *cpuAccumulator) takeFullNUMANodes() {
	for _, numa := range a.freeNUMANodes() {
		cpusInNUMANode := a.topo.CPUDetails.CPUsInNUMANodes(numa)
		if !a.needs(cpusInNUMANode.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullNUMANodes: claiming NUMA node", "numa", numa)
		a.take(cpusInNUMANode)
	}
}

func (a *cpuAccumulator) takeFullSockets() {
	for _, socket := range a.freeSockets() {
		cpusInSocket := a.topo.CPUDetails.CPUsInSockets(socket)
		if !a.needs(cpusInSocket.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullSockets: claiming socket", "socket", socket)
		a.take(cpusInSocket)
	}
}

func (a *cpuAccumulator) takeFullCores() {
	for _, core := range a.freeCores() {
		cpusInCore := a.topo.CPUDetails.CPUsInCores(core)
		if !a.needs(cpusInCore.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullCores: claiming core", "core", core)
		a.take(cpusInCore)
	}
}

func (a *cpuAccumulator) takeRemainingCPUs() {
	for _, cpu := range a.sortAvailableCPUs() {
		klog.V(4).InfoS("takeRemainingCPUs: claiming CPU", "cpu", cpu)
		a.take(cpuset.NewCPUSet(cpu))
		if a.isSatisfied() {
			return
		}
	}
}

func (a *cpuAccumulator) needs(n int) bool {
	return a.numCPUsNeeded >= n
}

func (a *cpuAccumulator) isSatisfied() bool {
	return a.numCPUsNeeded < 1
}

func (a *cpuAccumulator) isFailed() bool {
	return a.numCPUsNeeded > a.details.CPUs().Size()
}

func takeByTopology(topo *topology.CPUTopology, availableCPUs cpuset.CPUSet, numCPUs int) (cpuset.CPUSet, error) {
	acc := newCPUAccumulator(topo, availableCPUs, numCPUs)
	if acc.isSatisfied() {
		return acc.result, nil
	}
	if acc.isFailed() {
		return cpuset.NewCPUSet(), fmt.Errorf("not enough cpus available to satisfy request")
	}

	// Algorithm: topology-aware best-fit
	// 1. Acquire whole NUMA nodes and sockets, if available and the container
	//    requires at least a NUMA node or socket's-worth of CPUs. If NUMA
	//    Nodes map to 1 or more sockets, pull from NUMA nodes first.
	//    Otherwise pull from sockets first.
	if acc.topo.NumSockets >= acc.topo.NumNUMANodes {
		acc.takeFullNUMANodes()
		if acc.isSatisfied() {
			return acc.result, nil
		}

		acc.takeFullSockets()
		if acc.isSatisfied() {
			return acc.result, nil
		}
	} else {
		acc.takeFullSockets()
		if acc.isSatisfied() {
			return acc.result, nil
		}

		acc.takeFullNUMANodes()
		if acc.isSatisfied() {
			return acc.result, nil
		}
	}

	// 2. Acquire whole cores, if available and the container requires at least
	//    a core's-worth of CPUs.
	acc.takeFullCores()
	if acc.isSatisfied() {
		return acc.result, nil
	}

	// 3. Acquire single threads, preferring to fill partially-allocated cores
	//    on the same sockets as the whole cores we have already taken in this
	//    allocation.
	acc.takeRemainingCPUs()
	if acc.isSatisfied() {
		return acc.result, nil
	}

	return cpuset.NewCPUSet(), fmt.Errorf("failed to allocate cpus")
}
<commit_msg>Abstract out whether NUMA or Sockets come first in the memory hierarchy<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package cpumanager

import (
	"fmt"
	"sort"

	"k8s.io/klog/v2"

	"k8s.io/kubernetes/pkg/kubelet/cm/cpumanager/topology"
	"k8s.io/kubernetes/pkg/kubelet/cm/cpuset"
)

type numaOrSocketsFirstFuncs interface {
	takeFullFirstLevel()
	takeFullSecondLevel()
	sortAvailableNUMANodes() []int
	sortAvailableSockets() []int
	sortAvailableCores() []int
}

type numaFirst struct{ acc *cpuAccumulator }
type socketsFirst struct{ acc *cpuAccumulator }

var _ numaOrSocketsFirstFuncs = (*numaFirst)(nil)
var _ numaOrSocketsFirstFuncs = (*socketsFirst)(nil)

// If NUMA nodes are higher in the memory hierarchy than sockets, then we take
// from the set of NUMA Nodes as the first level.
func (n *numaFirst) takeFullFirstLevel() {
	n.acc.takeFullNUMANodes()
}

// If NUMA nodes are higher in the memory hierarchy than sockets, then we take
// from the set of sockets as the second level.
func (n *numaFirst) takeFullSecondLevel() {
	n.acc.takeFullSockets()
}

// If NUMA nodes are higher in the memory hierarchy than sockets, then just
// sort the NUMA nodes directly, and return them.
func (n *numaFirst) sortAvailableNUMANodes() []int {
	numas := n.acc.details.NUMANodes().ToSliceNoSort()
	n.acc.sort(numas, n.acc.details.CPUsInNUMANodes)
	return numas
}

// If NUMA nodes are higher in the memory hierarchy than sockets, then we need
// to pull the set of sockets out of each sorted NUMA node, and accumulate the
// partial order across them.
func (n *numaFirst) sortAvailableSockets() []int {
	var result []int
	for _, numa := range n.sortAvailableNUMANodes() {
		sockets := n.acc.details.SocketsInNUMANodes(numa).ToSliceNoSort()
		n.acc.sort(sockets, n.acc.details.CPUsInSockets)
		result = append(result, sockets...)
	}
	return result
}

// If NUMA nodes are higher in the memory hierarchy than sockets, then
// cores sit directly below sockets in the memory hierarchy.
func (n *numaFirst) sortAvailableCores() []int {
	var result []int
	for _, socket := range n.acc.sortAvailableSockets() {
		cores := n.acc.details.CoresInSockets(socket).ToSliceNoSort()
		n.acc.sort(cores, n.acc.details.CPUsInCores)
		result = append(result, cores...)
	}
	return result
}

// If sockets are higher in the memory hierarchy than NUMA nodes, then we take
// from the set of sockets as the first level.
func (s *socketsFirst) takeFullFirstLevel() {
	s.acc.takeFullSockets()
}

// If sockets are higher in the memory hierarchy than NUMA nodes, then we take
// from the set of NUMA Nodes as the second level.
func (s *socketsFirst) takeFullSecondLevel() {
	s.acc.takeFullNUMANodes()
}

// If sockets are higher in the memory hierarchy than NUMA nodes, then we need
// to pull the set of NUMA nodes out of each sorted Socket, and accumulate the
// partial order across them.
func (s *socketsFirst) sortAvailableNUMANodes() []int {
	var result []int
	for _, socket := range s.sortAvailableSockets() {
		numas := s.acc.details.NUMANodesInSockets(socket).ToSliceNoSort()
		s.acc.sort(numas, s.acc.details.CPUsInNUMANodes)
		result = append(result, numas...)
	}
	return result
}

// If sockets are higher in the memory hierarchy than NUMA nodes, then just
// sort the sockets directly, and return them.
func (s *socketsFirst) sortAvailableSockets() []int {
	sockets := s.acc.details.Sockets().ToSliceNoSort()
	s.acc.sort(sockets, s.acc.details.CPUsInSockets)
	return sockets
}

// If sockets are higher in the memory hierarchy than NUMA nodes, then cores
// sit directly below NUMA Nodes in the memory hierarchy.
func (s *socketsFirst) sortAvailableCores() []int {
	var result []int
	for _, numa := range s.acc.sortAvailableNUMANodes() {
		cores := s.acc.details.CoresInNUMANodes(numa).ToSliceNoSort()
		s.acc.sort(cores, s.acc.details.CPUsInCores)
		result = append(result, cores...)
	}
	return result
}

type cpuAccumulator struct {
	topo               *topology.CPUTopology
	details            topology.CPUDetails
	numCPUsNeeded      int
	result             cpuset.CPUSet
	numaOrSocketsFirst numaOrSocketsFirstFuncs
}

func newCPUAccumulator(topo *topology.CPUTopology, availableCPUs cpuset.CPUSet, numCPUs int) *cpuAccumulator {
	acc := &cpuAccumulator{
		topo:          topo,
		details:       topo.CPUDetails.KeepOnly(availableCPUs),
		numCPUsNeeded: numCPUs,
		result:        cpuset.NewCPUSet(),
	}

	if topo.NumSockets >= topo.NumNUMANodes {
		acc.numaOrSocketsFirst = &numaFirst{acc}
	} else {
		acc.numaOrSocketsFirst = &socketsFirst{acc}
	}

	return acc
}

// Returns true if the supplied NUMANode is fully available in `topoDetails`.
func (a *cpuAccumulator) isNUMANodeFree(numaID int) bool {
	return a.details.CPUsInNUMANodes(numaID).Size() == a.topo.CPUDetails.CPUsInNUMANodes(numaID).Size()
}

// Returns true if the supplied socket is fully available in `topoDetails`.
func (a *cpuAccumulator) isSocketFree(socketID int) bool {
	return a.details.CPUsInSockets(socketID).Size() == a.topo.CPUsPerSocket()
}

// Returns true if the supplied core is fully available in `topoDetails`.
func (a *cpuAccumulator) isCoreFree(coreID int) bool {
	return a.details.CPUsInCores(coreID).Size() == a.topo.CPUsPerCore()
}

// Returns free NUMA Node IDs as a slice sorted by sortAvailableNUMANodes().
func (a *cpuAccumulator) freeNUMANodes() []int {
	free := []int{}
	for _, numa := range a.sortAvailableNUMANodes() {
		if a.isNUMANodeFree(numa) {
			free = append(free, numa)
		}
	}
	return free
}

// Returns free socket IDs as a slice sorted by sortAvailableSockets().
func (a *cpuAccumulator) freeSockets() []int {
	free := []int{}
	for _, socket := range a.sortAvailableSockets() {
		if a.isSocketFree(socket) {
			free = append(free, socket)
		}
	}
	return free
}

// Returns free core IDs as a slice sorted by sortAvailableCores().
func (a *cpuAccumulator) freeCores() []int {
	free := []int{}
	for _, core := range a.sortAvailableCores() {
		if a.isCoreFree(core) {
			free = append(free, core)
		}
	}
	return free
}

// Returns free CPU IDs as a slice sorted by sortAvailableCPUs().
func (a *cpuAccumulator) freeCPUs() []int {
	return a.sortAvailableCPUs()
}

// Sorts the provided list of NUMA nodes/sockets/cores/cpus referenced in 'ids'
// by the number of available CPUs contained within them (smallest to largest).
// The 'getCPU()' paramater defines the function that should be called to
// retrieve the list of available CPUs for the type being referenced. If two
// NUMA nodes/sockets/cores/cpus have the same number of available CPUs, they
// are sorted in ascending order by their id.
func (a *cpuAccumulator) sort(ids []int, getCPUs func(ids ...int) cpuset.CPUSet) {
	sort.Slice(ids,
		func(i, j int) bool {
			iCPUs := getCPUs(ids[i])
			jCPUs := getCPUs(ids[j])
			if iCPUs.Size() < jCPUs.Size() {
				return true
			}
			if iCPUs.Size() > jCPUs.Size() {
				return false
			}
			return ids[i] < ids[j]
		})
}

// Sort all NUMA nodes with free CPUs.
func (a *cpuAccumulator) sortAvailableNUMANodes() []int {
	return a.numaOrSocketsFirst.sortAvailableNUMANodes()
}

// Sort all sockets with free CPUs.
func (a *cpuAccumulator) sortAvailableSockets() []int {
	return a.numaOrSocketsFirst.sortAvailableSockets()
}

// Sort all cores with free CPUs:
func (a *cpuAccumulator) sortAvailableCores() []int {
	return a.numaOrSocketsFirst.sortAvailableCores()
}

// Sort all available CPUs:
// - First by core using sortAvailableCores().
// - Then within each core, using the sort() algorithm defined above.
func (a *cpuAccumulator) sortAvailableCPUs() []int {
	var result []int
	for _, core := range a.sortAvailableCores() {
		cpus := a.details.CPUsInCores(core).ToSliceNoSort()
		sort.Ints(cpus)
		result = append(result, cpus...)
	}
	return result
}

func (a *cpuAccumulator) take(cpus cpuset.CPUSet) {
	a.result = a.result.Union(cpus)
	a.details = a.details.KeepOnly(a.details.CPUs().Difference(a.result))
	a.numCPUsNeeded -= cpus.Size()
}

func (a *cpuAccumulator) takeFullNUMANodes() {
	for _, numa := range a.freeNUMANodes() {
		cpusInNUMANode := a.topo.CPUDetails.CPUsInNUMANodes(numa)
		if !a.needs(cpusInNUMANode.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullNUMANodes: claiming NUMA node", "numa", numa)
		a.take(cpusInNUMANode)
	}
}

func (a *cpuAccumulator) takeFullSockets() {
	for _, socket := range a.freeSockets() {
		cpusInSocket := a.topo.CPUDetails.CPUsInSockets(socket)
		if !a.needs(cpusInSocket.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullSockets: claiming socket", "socket", socket)
		a.take(cpusInSocket)
	}
}

func (a *cpuAccumulator) takeFullCores() {
	for _, core := range a.freeCores() {
		cpusInCore := a.topo.CPUDetails.CPUsInCores(core)
		if !a.needs(cpusInCore.Size()) {
			continue
		}
		klog.V(4).InfoS("takeFullCores: claiming core", "core", core)
		a.take(cpusInCore)
	}
}

func (a *cpuAccumulator) takeRemainingCPUs() {
	for _, cpu := range a.sortAvailableCPUs() {
		klog.V(4).InfoS("takeRemainingCPUs: claiming CPU", "cpu", cpu)
		a.take(cpuset.NewCPUSet(cpu))
		if a.isSatisfied() {
			return
		}
	}
}

func (a *cpuAccumulator) needs(n int) bool {
	return a.numCPUsNeeded >= n
}

func (a *cpuAccumulator) isSatisfied() bool {
	return a.numCPUsNeeded < 1
}

func (a *cpuAccumulator) isFailed() bool {
	return a.numCPUsNeeded > a.details.CPUs().Size()
}

func takeByTopology(topo *topology.CPUTopology, availableCPUs cpuset.CPUSet, numCPUs int) (cpuset.CPUSet, error) {
	acc := newCPUAccumulator(topo, availableCPUs, numCPUs)
	if acc.isSatisfied() {
		return acc.result, nil
	}
	if acc.isFailed() {
		return cpuset.NewCPUSet(), fmt.Errorf("not enough cpus available to satisfy request")
	}

	// Algorithm: topology-aware best-fit
	// 1. Acquire whole NUMA nodes and sockets, if available and the container
	//    requires at least a NUMA node or socket's-worth of CPUs. If NUMA
	//    Nodes map to 1 or more sockets, pull from NUMA nodes first.
	//    Otherwise pull from sockets first.
	acc.numaOrSocketsFirst.takeFullFirstLevel()
	if acc.isSatisfied() {
		return acc.result, nil
	}
	acc.numaOrSocketsFirst.takeFullSecondLevel()
	if acc.isSatisfied() {
		return acc.result, nil
	}

	// 2. Acquire whole cores, if available and the container requires at least
	//    a core's-worth of CPUs.
	acc.takeFullCores()
	if acc.isSatisfied() {
		return acc.result, nil
	}

	// 3. Acquire single threads, preferring to fill partially-allocated cores
	//    on the same sockets as the whole cores we have already taken in this
	//    allocation.
	acc.takeRemainingCPUs()
	if acc.isSatisfied() {
		return acc.result, nil
	}

	return cpuset.NewCPUSet(), fmt.Errorf("failed to allocate cpus")
}
<|endoftext|>


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

package memory

var (
	memset func(b []byte, c byte)
)

// Set assigns the value c to every element of the slice buf.
func Set(buf []byte, c byte) {
	memset(buf, c)
}

// memory_memset_go reference implementation
func memory_memset_go(buf []byte, c byte) {
	for i := 0; i < len(buf); i++ {
		buf[i] = c
	}
}
<commit_msg>ARROW-6966: [Go] Set a default memset for when the platform doesn't set one<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package memory

var (
	memset func(b []byte, c byte) = memory_memset_go
)

// Set assigns the value c to every element of the slice buf.
func Set(buf []byte, c byte) {
	memset(buf, c)
}

// memory_memset_go reference implementation
func memory_memset_go(buf []byte, c byte) {
	for i := 0; i < len(buf); i++ {
		buf[i] = c
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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 bitlocker provides functionality for managing Bitlocker.
package bitlocker

import (
	"fmt"

	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/iamacarpet/go-win64api"
)

var (
	// Test Helpers
	funcBackup       = winapi.BackupBitLockerRecoveryKeys
	funcRecoveryInfo = winapi.GetBitLockerRecoveryInfo
)

// BackupToAD backs up Bitlocker recovery keys to Active Directory.
func BackupToAD() error {
	infos, err := funcRecoveryInfo()
	if err != nil {
		return err
	}
	volIDs := []string{}
	for _, i := range infos {
		if i.ConversionStatus != 1 {
			logger.Warningf("Skipping volume %s due to conversion status (%d).", i.DriveLetter, i.ConversionStatus)
			continue
		}
		logger.Infof("Backing up Bitlocker recovery password for drive %q.", i.DriveLetter)
		volIDs = append(volIDs, i.PersistentVolumeID)
	}
	return funcBackup(volIDs)
}

type wmi struct {
	intf *ole.IDispatch
	svc  *ole.IDispatch
}

func (w *wmi) connect() error {
	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		return fmt.Errorf("unable to create initial object, %w", err)
	}
	defer unknown.Release()
	w.intf, err = unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return fmt.Errorf("unable to create initial object, %w", err)
	}
	serviceRaw, err := oleutil.CallMethod(w.intf, "ConnectServer", nil, `\\.\ROOT\CIMV2\Security\MicrosoftVolumeEncryption`)
	if err != nil {
		return fmt.Errorf("permission denied: %w", err)
	}
	w.svc = serviceRaw.ToIDispatch()
	return nil
}

func (w *wmi) close() {
	w.svc.Release()
	w.intf.Release()
}

const (
	// Encryption Methods
	// https://docs.microsoft.com/en-us/windows/win32/secprov/getencryptionmethod-win32-encryptablevolume
	None int32 = iota
	AES128WithDiffuser
	AES256WithDiffuser
	AES128
	AES256
	HardwareEncryption
	XtsAES128
	XtsAES256

	// Encryption Flags
	// https://docs.microsoft.com/en-us/windows/win32/secprov/encrypt-win32-encryptablevolume
	EncryptDataOnly    int32 = 0x00000001
	EncryptDemandWipe  int32 = 0x00000002
	EncryptSynchronous int32 = 0x00010000

	// Error Codes
	ERROR_IO_DEVICE        int32 = -2147023779
	FVE_E_BOOTABLE_CDDVD   int32 = -2144272336
	FVE_E_PROTECTOR_EXISTS int32 = -2144272335
)

func encryptErrHandler(val int32) error {
	switch val {
	case ERROR_IO_DEVICE:
		return fmt.Errorf("an I/O error has occurred during encryption; the device may need to be reset")
	case FVE_E_BOOTABLE_CDDVD:
		return fmt.Errorf("BitLocker Drive Encryption detected bootable media (CD or DVD) in the computer. " +
			"Remove the media and restart the computer before configuring BitLocker.")
	case FVE_E_PROTECTOR_EXISTS:
		return fmt.Errorf("key protector cannot be added; only one key protector of this type is allowed for this drive")
	default:
		return fmt.Errorf("error code returned during encryption: %d", val)
	}
}

// EncryptWithTPM encrypts the drive with Bitlocker using TPM key protection.
//
// Example: bitlocker.EncryptWithTPM("c:", bitlocker.XtsAES256, bitlocker.EncryptDataOnly)
func EncryptWithTPM(driveLetter string, method int32, flags int32) error {
	ole.CoInitialize(0)
	defer ole.CoUninitialize()
	w := &wmi{}
	if err := w.connect(); err != nil {
		return fmt.Errorf("wmi.Connect: %w", err)
	}
	defer w.close()
	raw, err := oleutil.CallMethod(w.svc, "ExecQuery",
		"SELECT * FROM Win32_EncryptableVolume WHERE DriveLetter = '"+driveLetter+"'")
	if err != nil {
		return fmt.Errorf("ExecQuery: %w", err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	itemRaw, err := oleutil.CallMethod(result, "ItemIndex", 0)
	if err != nil {
		return fmt.Errorf("failed to fetch result row while processing BitLocker info: %w", err)
	}
	item := itemRaw.ToIDispatch()
	defer item.Release()

	// https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(item, "ProtectKeyWithTPM", nil, nil, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("error calling ProtectKeyWithTPM(%s): %w", driveLetter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return encryptErrHandler(val)
	}

	resultRaw, err = oleutil.CallMethod(item, "Encrypt", method, flags)
	if err != nil {
		return fmt.Errorf("error calling Encrypt(%s): %w", driveLetter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return encryptErrHandler(val)
	}
	return nil
}
<commit_msg>Add a bit more decoration to return code error handling.<commit_after>// Copyright 2021 Google LLC
//
// 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 bitlocker provides functionality for managing Bitlocker.
package bitlocker

import (
	"fmt"

	"github.com/google/logger"
	"github.com/go-ole/go-ole"
	"github.com/go-ole/go-ole/oleutil"
	"github.com/iamacarpet/go-win64api"
)

var (
	// Test Helpers
	funcBackup       = winapi.BackupBitLockerRecoveryKeys
	funcRecoveryInfo = winapi.GetBitLockerRecoveryInfo
)

// BackupToAD backs up Bitlocker recovery keys to Active Directory.
func BackupToAD() error {
	infos, err := funcRecoveryInfo()
	if err != nil {
		return err
	}
	volIDs := []string{}
	for _, i := range infos {
		if i.ConversionStatus != 1 {
			logger.Warningf("Skipping volume %s due to conversion status (%d).", i.DriveLetter, i.ConversionStatus)
			continue
		}
		logger.Infof("Backing up Bitlocker recovery password for drive %q.", i.DriveLetter)
		volIDs = append(volIDs, i.PersistentVolumeID)
	}
	return funcBackup(volIDs)
}

type wmi struct {
	intf *ole.IDispatch
	svc  *ole.IDispatch
}

func (w *wmi) connect() error {
	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
	if err != nil {
		return fmt.Errorf("unable to create initial object, %w", err)
	}
	defer unknown.Release()
	w.intf, err = unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return fmt.Errorf("unable to create initial object, %w", err)
	}
	serviceRaw, err := oleutil.CallMethod(w.intf, "ConnectServer", nil, `\\.\ROOT\CIMV2\Security\MicrosoftVolumeEncryption`)
	if err != nil {
		return fmt.Errorf("permission denied: %w", err)
	}
	w.svc = serviceRaw.ToIDispatch()
	return nil
}

func (w *wmi) close() {
	w.svc.Release()
	w.intf.Release()
}

const (
	// Encryption Methods
	// https://docs.microsoft.com/en-us/windows/win32/secprov/getencryptionmethod-win32-encryptablevolume
	None int32 = iota
	AES128WithDiffuser
	AES256WithDiffuser
	AES128
	AES256
	HardwareEncryption
	XtsAES128
	XtsAES256

	// Encryption Flags
	// https://docs.microsoft.com/en-us/windows/win32/secprov/encrypt-win32-encryptablevolume
	EncryptDataOnly    int32 = 0x00000001
	EncryptDemandWipe  int32 = 0x00000002
	EncryptSynchronous int32 = 0x00010000

	// Error Codes
	ERROR_IO_DEVICE        int32 = -2147023779
	FVE_E_BOOTABLE_CDDVD   int32 = -2144272336
	FVE_E_PROTECTOR_EXISTS int32 = -2144272335
)

func encryptErrHandler(val int32) error {
	switch val {
	case ERROR_IO_DEVICE:
		return fmt.Errorf("an I/O error has occurred during encryption; the device may need to be reset")
	case FVE_E_BOOTABLE_CDDVD:
		return fmt.Errorf("BitLocker Drive Encryption detected bootable media (CD or DVD) in the computer. " +
			"Remove the media and restart the computer before configuring BitLocker.")
	case FVE_E_PROTECTOR_EXISTS:
		return fmt.Errorf("key protector cannot be added; only one key protector of this type is allowed for this drive")
	default:
		return fmt.Errorf("error code returned during encryption: %d", val)
	}
}

// EncryptWithTPM encrypts the drive with Bitlocker using TPM key protection.
//
// Example: bitlocker.EncryptWithTPM("c:", bitlocker.XtsAES256, bitlocker.EncryptDataOnly)
func EncryptWithTPM(driveLetter string, method int32, flags int32) error {
	ole.CoInitialize(0)
	defer ole.CoUninitialize()
	w := &wmi{}
	if err := w.connect(); err != nil {
		return fmt.Errorf("wmi.Connect: %w", err)
	}
	defer w.close()
	raw, err := oleutil.CallMethod(w.svc, "ExecQuery",
		"SELECT * FROM Win32_EncryptableVolume WHERE DriveLetter = '"+driveLetter+"'")
	if err != nil {
		return fmt.Errorf("ExecQuery: %w", err)
	}
	result := raw.ToIDispatch()
	defer result.Release()

	itemRaw, err := oleutil.CallMethod(result, "ItemIndex", 0)
	if err != nil {
		return fmt.Errorf("failed to fetch result row while processing BitLocker info: %w", err)
	}
	item := itemRaw.ToIDispatch()
	defer item.Release()

	// https://docs.microsoft.com/en-us/windows/win32/secprov/protectkeywithtpm-win32-encryptablevolume
	var volumeKeyProtectorID ole.VARIANT
	ole.VariantInit(&volumeKeyProtectorID)
	resultRaw, err := oleutil.CallMethod(item, "ProtectKeyWithTPM", nil, nil, &volumeKeyProtectorID)
	if err != nil {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", driveLetter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("ProtectKeyWithTPM(%s): %w", driveLetter, encryptErrHandler(val))
	}

	resultRaw, err = oleutil.CallMethod(item, "Encrypt", method, flags)
	if err != nil {
		return fmt.Errorf("Encrypt(%s): %w", driveLetter, err)
	} else if val, ok := resultRaw.Value().(int32); val != 0 || !ok {
		return fmt.Errorf("Encrypt(%s): %w", driveLetter, encryptErrHandler(val))
	}
	return nil
}
<|endoftext|>


<commit_before>package search

import (
	"context"
	"fmt"
	"sync"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type store struct {
	sync.Mutex
	globals.Contextified
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// store keeps an encrypted index of chat messages for all conversations
// to enable full inbox search locally.  Index data for each conversation
// is stored in an encrypted leveldb in the form:
// (uid,convID) -> {
//  Index: {
//     token: { msgID,... },
//     ...
//  },
//  Alias: {
//    alias: { token,... },
//    ...
//  },
//  Metadata: chat1.ConversationIndexMetadata
//}
// NOTE: as a performance optimization we may want to split the metadata
// from the index itself so we can quickly operate on the metadata
// separately from the index and have less bytes to encrypt/decrypt on
// reads (metadata only contains only msg ids and no user content).
func newStore(g *globals.Context) *store {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return storage.GetSecretBoxKey(ctx, g.ExternalG(), storage.DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	return &store{
		Contextified: globals.NewContextified(g),
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "Search.store", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
	}
}

func (s *store) dbKey(convID chat1.ConversationID, uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatIndex,
		Key: fmt.Sprintf("idx:%s:%s", uid, convID),
	}
}

func (s *store) getLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (ret *chat1.ConversationIndex, err error) {
	defer func() {
		// return a blank index
		if err == nil && ret == nil {
			ret = &chat1.ConversationIndex{}
			ret.Index = map[string]map[chat1.MessageID]chat1.EmptyStruct{}
			ret.Alias = map[string]map[string]chat1.EmptyStruct{}
			ret.Metadata.SeenIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		if err != nil {
			if derr := s.deleteLocked(ctx, convID, uid); derr != nil {
				s.Debug(ctx, "unable to delete: %v", derr)
			}
		}
	}()
	dbKey := s.dbKey(convID, uid)
	var entry chat1.ConversationIndex
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		return nil, err
	}
	if entry.Metadata.Version != IndexVersion {
		// drop the whole index for this conv
		err = s.deleteLocked(ctx, convID, uid)
		return nil, err
	}
	return &entry, nil
}

func (s *store) putLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID, entry *chat1.ConversationIndex) error {
	if entry == nil {
		return nil
	}
	dbKey := s.dbKey(convID, uid)
	entry.Metadata.Version = IndexVersion
	return s.encryptedDB.Put(ctx, dbKey, *entry)
}

func (s *store) deleteLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) error {
	dbKey := s.dbKey(convID, uid)
	return s.encryptedDB.Delete(ctx, dbKey)
}

func (s *store) getConvIndex(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (entry *chat1.ConversationIndex, err error) {
	s.Lock()
	defer s.Unlock()
	return s.getLocked(ctx, convID, uid)
}

// addTokensLocked add the given tokens to the index under the given message
// id, when ingesting EDIT messages the msgID is of the superseded msg but the
// tokens are from the EDIT itself.
func (s *store) addTokensLocked(entry *chat1.ConversationIndex, tokens tokenMap, msgID chat1.MessageID) {
	for token, aliases := range tokens {
		msgIDs, ok := entry.Index[token]
		if !ok {
			msgIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		msgIDs[msgID] = chat1.EmptyStruct{}
		entry.Index[token] = msgIDs
		for alias := range aliases {
			atoken, ok := entry.Alias[alias]
			if !ok {
				atoken = map[string]chat1.EmptyStruct{}
			}
			atoken[token] = chat1.EmptyStruct{}
			entry.Alias[alias] = atoken
		}
	}
}

func (s *store) addMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	tokens := tokensFromMsg(msg)
	s.addTokensLocked(entry, tokens, msg.GetMessageID())
}

func (s *store) removeMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	// find the msgID that the index stores
	var msgID chat1.MessageID
	switch msg.GetMessageType() {
	case chat1.MessageType_EDIT, chat1.MessageType_ATTACHMENTUPLOADED:
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil || len(superIDs) != 1 {
			return
		}
		msgID = superIDs[0]
	default:
		msgID = msg.GetMessageID()
	}

	for token, aliases := range tokensFromMsg(msg) {
		msgIDs := entry.Index[token]
		delete(msgIDs, msgID)
		if len(msgIDs) == 0 {
			delete(entry.Index, token)
		}
		for alias := range aliases {
			for atoken := range entry.Alias[alias] {
				_, ok := entry.Index[atoken]
				if !ok {
					delete(entry.Alias[alias], atoken)
				}
			}
			if len(entry.Alias[alias]) == 0 {
				delete(entry.Alias, alias)
			}
		}
	}
}

func (s *store) add(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	fetchSupersededMsgs := func(msg chat1.MessageUnboxed) []chat1.MessageUnboxed {
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil {
			s.Debug(ctx, "unable to get supersedes: %v", err)
			return nil
		}
		reason := chat1.GetThreadReason_INDEXED_SEARCH
		supersededMsgs, err := s.G().ChatHelper.GetMessages(ctx, uid, convID, superIDs,
			false /* resolveSupersedes*/, &reason)
		if err != nil {
			// Log but ignore error
			s.Debug(ctx, "unable to get fetch messages: %v", err)
			return nil
		}
		return supersededMsgs
	}

	for _, msg := range msgs {
		seenIDs := entry.Metadata.SeenIDs
		// Don't add if we've seen
		if _, ok := seenIDs[msg.GetMessageID()]; ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		// NOTE DELETE and DELETEHISTORY are handled through calls to `remove`,
		// other messages will be added if there is any content that can be
		// indexed.
		switch msg.GetMessageType() {
		case chat1.MessageType_ATTACHMENTUPLOADED:
			supersededMsgs := fetchSupersededMsgs(msg)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.addMsgLocked(entry, sm)
			}
		case chat1.MessageType_EDIT:
			tokens := tokensFromMsg(msg)
			supersededMsgs := fetchSupersededMsgs(msg)
			// remove the original message text and replace it with the edited
			// contents (using the original id in the index)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.removeMsgLocked(entry, sm)
				s.addTokensLocked(entry, tokens, sm.GetMessageID())
			}
		default:
			s.addMsgLocked(entry, msg)
		}
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}

// Remove tokenizes the message content and updates/removes index keys for each token.
func (s *store) remove(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	s.Lock()
	defer s.Unlock()

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	seenIDs := entry.Metadata.SeenIDs
	for _, msg := range msgs {
		// Don't remove if we haven't seen
		if _, ok := entry.Metadata.SeenIDs[msg.GetMessageID()]; !ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		s.removeMsgLocked(entry, msg)
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}
<commit_msg>use lock tab (#17059)<commit_after>package search

import (
	"context"
	"fmt"

	"github.com/keybase/client/go/chat/globals"
	"github.com/keybase/client/go/chat/storage"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/encrypteddb"
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type store struct {
	lockTab *libkb.LockTable
	globals.Contextified
	utils.DebugLabeler
	encryptedDB *encrypteddb.EncryptedDB
}

// store keeps an encrypted index of chat messages for all conversations
// to enable full inbox search locally.  Index data for each conversation
// is stored in an encrypted leveldb in the form:
// (uid,convID) -> {
//  Index: {
//     token: { msgID,... },
//     ...
//  },
//  Alias: {
//    alias: { token,... },
//    ...
//  },
//  Metadata: chat1.ConversationIndexMetadata
//}
// NOTE: as a performance optimization we may want to split the metadata
// from the index itself so we can quickly operate on the metadata
// separately from the index and have less bytes to encrypt/decrypt on
// reads (metadata only contains only msg ids and no user content).
func newStore(g *globals.Context) *store {
	keyFn := func(ctx context.Context) ([32]byte, error) {
		return storage.GetSecretBoxKey(ctx, g.ExternalG(), storage.DefaultSecretUI)
	}
	dbFn := func(g *libkb.GlobalContext) *libkb.JSONLocalDb {
		return g.LocalChatDb
	}
	return &store{
		Contextified: globals.NewContextified(g),
		DebugLabeler: utils.NewDebugLabeler(g.GetLog(), "Search.store", false),
		encryptedDB:  encrypteddb.New(g.ExternalG(), dbFn, keyFn),
		lockTab:      &libkb.LockTable{},
	}
}

func (s *store) dbKey(convID chat1.ConversationID, uid gregor1.UID) libkb.DbKey {
	return libkb.DbKey{
		Typ: libkb.DBChatIndex,
		Key: fmt.Sprintf("idx:%s:%s", uid, convID),
	}
}

func (s *store) getLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (ret *chat1.ConversationIndex, err error) {
	defer func() {
		// return a blank index
		if err == nil && ret == nil {
			ret = &chat1.ConversationIndex{}
			ret.Index = map[string]map[chat1.MessageID]chat1.EmptyStruct{}
			ret.Alias = map[string]map[string]chat1.EmptyStruct{}
			ret.Metadata.SeenIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		if err != nil {
			if derr := s.deleteLocked(ctx, convID, uid); derr != nil {
				s.Debug(ctx, "unable to delete: %v", derr)
			}
		}
	}()
	dbKey := s.dbKey(convID, uid)
	var entry chat1.ConversationIndex
	found, err := s.encryptedDB.Get(ctx, dbKey, &entry)
	if err != nil || !found {
		return nil, err
	}
	if entry.Metadata.Version != IndexVersion {
		// drop the whole index for this conv
		err = s.deleteLocked(ctx, convID, uid)
		return nil, err
	}
	return &entry, nil
}

func (s *store) putLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID, entry *chat1.ConversationIndex) error {
	if entry == nil {
		return nil
	}
	dbKey := s.dbKey(convID, uid)
	entry.Metadata.Version = IndexVersion
	return s.encryptedDB.Put(ctx, dbKey, *entry)
}

func (s *store) deleteLocked(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) error {
	dbKey := s.dbKey(convID, uid)
	return s.encryptedDB.Delete(ctx, dbKey)
}

func (s *store) getConvIndex(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID) (entry *chat1.ConversationIndex, err error) {
	lock := s.lockTab.AcquireOnName(ctx, s.G(), convID.String())
	defer lock.Release(ctx)
	return s.getLocked(ctx, convID, uid)
}

// addTokensLocked add the given tokens to the index under the given message
// id, when ingesting EDIT messages the msgID is of the superseded msg but the
// tokens are from the EDIT itself.
func (s *store) addTokensLocked(entry *chat1.ConversationIndex, tokens tokenMap, msgID chat1.MessageID) {
	for token, aliases := range tokens {
		msgIDs, ok := entry.Index[token]
		if !ok {
			msgIDs = map[chat1.MessageID]chat1.EmptyStruct{}
		}
		msgIDs[msgID] = chat1.EmptyStruct{}
		entry.Index[token] = msgIDs
		for alias := range aliases {
			atoken, ok := entry.Alias[alias]
			if !ok {
				atoken = map[string]chat1.EmptyStruct{}
			}
			atoken[token] = chat1.EmptyStruct{}
			entry.Alias[alias] = atoken
		}
	}
}

func (s *store) addMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	tokens := tokensFromMsg(msg)
	s.addTokensLocked(entry, tokens, msg.GetMessageID())
}

func (s *store) removeMsgLocked(entry *chat1.ConversationIndex, msg chat1.MessageUnboxed) {
	// find the msgID that the index stores
	var msgID chat1.MessageID
	switch msg.GetMessageType() {
	case chat1.MessageType_EDIT, chat1.MessageType_ATTACHMENTUPLOADED:
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil || len(superIDs) != 1 {
			return
		}
		msgID = superIDs[0]
	default:
		msgID = msg.GetMessageID()
	}

	for token, aliases := range tokensFromMsg(msg) {
		msgIDs := entry.Index[token]
		delete(msgIDs, msgID)
		if len(msgIDs) == 0 {
			delete(entry.Index, token)
		}
		for alias := range aliases {
			for atoken := range entry.Alias[alias] {
				_, ok := entry.Index[atoken]
				if !ok {
					delete(entry.Alias[alias], atoken)
				}
			}
			if len(entry.Alias[alias]) == 0 {
				delete(entry.Alias, alias)
			}
		}
	}
}

func (s *store) add(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	lock := s.lockTab.AcquireOnName(ctx, s.G(), convID.String())
	defer lock.Release(ctx)

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	fetchSupersededMsgs := func(msg chat1.MessageUnboxed) []chat1.MessageUnboxed {
		superIDs, err := utils.GetSupersedes(msg)
		if err != nil {
			s.Debug(ctx, "unable to get supersedes: %v", err)
			return nil
		}
		reason := chat1.GetThreadReason_INDEXED_SEARCH
		supersededMsgs, err := s.G().ChatHelper.GetMessages(ctx, uid, convID, superIDs,
			false /* resolveSupersedes*/, &reason)
		if err != nil {
			// Log but ignore error
			s.Debug(ctx, "unable to get fetch messages: %v", err)
			return nil
		}
		return supersededMsgs
	}

	for _, msg := range msgs {
		seenIDs := entry.Metadata.SeenIDs
		// Don't add if we've seen
		if _, ok := seenIDs[msg.GetMessageID()]; ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		// NOTE DELETE and DELETEHISTORY are handled through calls to `remove`,
		// other messages will be added if there is any content that can be
		// indexed.
		switch msg.GetMessageType() {
		case chat1.MessageType_ATTACHMENTUPLOADED:
			supersededMsgs := fetchSupersededMsgs(msg)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.addMsgLocked(entry, sm)
			}
		case chat1.MessageType_EDIT:
			tokens := tokensFromMsg(msg)
			supersededMsgs := fetchSupersededMsgs(msg)
			// remove the original message text and replace it with the edited
			// contents (using the original id in the index)
			for _, sm := range supersededMsgs {
				seenIDs[sm.GetMessageID()] = chat1.EmptyStruct{}
				s.removeMsgLocked(entry, sm)
				s.addTokensLocked(entry, tokens, sm.GetMessageID())
			}
		default:
			s.addMsgLocked(entry, msg)
		}
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}

// Remove tokenizes the message content and updates/removes index keys for each token.
func (s *store) remove(ctx context.Context, convID chat1.ConversationID, uid gregor1.UID,
	msgs []chat1.MessageUnboxed) (err error) {
	lock := s.lockTab.AcquireOnName(ctx, s.G(), convID.String())
	defer lock.Release(ctx)

	entry, err := s.getLocked(ctx, convID, uid)
	if err != nil {
		return err
	}

	seenIDs := entry.Metadata.SeenIDs
	for _, msg := range msgs {
		// Don't remove if we haven't seen
		if _, ok := entry.Metadata.SeenIDs[msg.GetMessageID()]; !ok {
			continue
		}
		seenIDs[msg.GetMessageID()] = chat1.EmptyStruct{}
		s.removeMsgLocked(entry, msg)
	}
	err = s.putLocked(ctx, convID, uid, entry)
	return err
}
<|endoftext|>


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

// Package oppobloom implements a filter data structure that may report false
// negatives but no false positives.
package oppobloom

import (
	"bytes"
	"crypto/md5"
	"errors"
	"hash"
	"math"
	"sync/atomic"
	"unsafe"
)

type Filter struct {
	array []*[]byte
	sizeMask uint32
}

var ErrSizeTooLarge = errors.New("oppobloom: size given to large to round to a power of 2")
var ErrSizeTooSmall = errors.New("oppobloom: filter cannot have a zero or negative size")
var MaxFilterSize = 1 << 30

func NewFilter(size int) (*Filter, error) {
	if size > MaxFilterSize {
		return nil, ErrSizeTooLarge
	}
	if size <= 0 {
		return nil, ErrSizeTooSmall
	}
	// round to the next largest power of two
	size = int(math.Pow(2, math.Ceil(math.Log2(float64(size)))))
	slice := make([]*[]byte, size)
	sizeMask := uint32(size - 1)
	return &Filter{slice, sizeMask}, nil
}

func (f *Filter) Contains(id []byte) bool {
	h := md5UintHash{md5.New()}
	h.Write(id)
	uindex := h.Sum32() & f.sizeMask
	index := int32(uindex)
	oldId := getAndSet(f.array, index, id)
	return bytes.Equal(oldId, id)
}

func (f *Filter) Size() int {
	return len(f.array)
}

type md5UintHash struct {
	hash.Hash // a hack with knowledge of how md5 works
}

func (m md5UintHash) Sum32() uint32 {
	sum := m.Sum(nil)
	x := uint32(sum[0])
	for _, val := range sum[1:3] {
		x = x << 3
		x += uint32(val)
	}
	return x
}

// Returns the id that was in the slice at the given index after putting the
// new id in the slice at that index, atomically.
func getAndSet(arr []*[]byte, index int32, id []byte) []byte {
	indexPtr := (*unsafe.Pointer)(unsafe.Pointer(&arr[index]))
	idUnsafe := unsafe.Pointer(&id)
	var oldId []byte
	for {
		oldIdUnsafe := atomic.LoadPointer(indexPtr)
		if atomic.CompareAndSwapPointer(indexPtr, oldIdUnsafe, idUnsafe) {
			oldIdPtr := (*[]byte)(oldIdUnsafe)
			if oldIdPtr != nil {
				oldId = *oldIdPtr
			}
			break
		}
	}
	return oldId
}
<commit_msg>Fix typo.<commit_after>// Copyright 2012 Jeff Hodges. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package oppobloom implements a filter data structure that may report false
// negatives but no false positives.
package oppobloom

import (
	"bytes"
	"crypto/md5"
	"errors"
	"hash"
	"math"
	"sync/atomic"
	"unsafe"
)

type Filter struct {
	array []*[]byte
	sizeMask uint32
}

var ErrSizeTooLarge = errors.New("oppobloom: size given too large to round to a power of 2")
var ErrSizeTooSmall = errors.New("oppobloom: filter cannot have a zero or negative size")
var MaxFilterSize = 1 << 30

func NewFilter(size int) (*Filter, error) {
	if size > MaxFilterSize {
		return nil, ErrSizeTooLarge
	}
	if size <= 0 {
		return nil, ErrSizeTooSmall
	}
	// round to the next largest power of two
	size = int(math.Pow(2, math.Ceil(math.Log2(float64(size)))))
	slice := make([]*[]byte, size)
	sizeMask := uint32(size - 1)
	return &Filter{slice, sizeMask}, nil
}

func (f *Filter) Contains(id []byte) bool {
	h := md5UintHash{md5.New()}
	h.Write(id)
	uindex := h.Sum32() & f.sizeMask
	index := int32(uindex)
	oldId := getAndSet(f.array, index, id)
	return bytes.Equal(oldId, id)
}

func (f *Filter) Size() int {
	return len(f.array)
}

type md5UintHash struct {
	hash.Hash // a hack with knowledge of how md5 works
}

func (m md5UintHash) Sum32() uint32 {
	sum := m.Sum(nil)
	x := uint32(sum[0])
	for _, val := range sum[1:3] {
		x = x << 3
		x += uint32(val)
	}
	return x
}

// Returns the id that was in the slice at the given index after putting the
// new id in the slice at that index, atomically.
func getAndSet(arr []*[]byte, index int32, id []byte) []byte {
	indexPtr := (*unsafe.Pointer)(unsafe.Pointer(&arr[index]))
	idUnsafe := unsafe.Pointer(&id)
	var oldId []byte
	for {
		oldIdUnsafe := atomic.LoadPointer(indexPtr)
		if atomic.CompareAndSwapPointer(indexPtr, oldIdUnsafe, idUnsafe) {
			oldIdPtr := (*[]byte)(oldIdUnsafe)
			if oldIdPtr != nil {
				oldId = *oldIdPtr
			}
			break
		}
	}
	return oldId
}
<|endoftext|>


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

// RuleSetGroup API support - Fetch, Create, Update, Delete, and Search
// See: https://login.circonus.com/resources/api/calls/rule_set_group

package api

import (
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"

	"github.com/circonus-labs/circonus-gometrics/api/config"
)

// RuleSetGroupRule defines a rulesetGroup rule
type RuleSetGroupRule struct {
	Criteria          string `json:"criteria"`
	Severity          uint   `json:"severity"`
	Value             string `json:"value"`
	WindowingDuration uint   `json:"windowing_duration,omitempty"`
	WindowingFunction string `json:"windowing_function,omitempty"`
	Wait              uint   `json:"wait,omitempty"`
}

// RuleSetGroupFormula defines a formula for raising alerts
type RuleSetGroupFormula struct {
	Expression    string `json:"expression"`
	RaiseSeverity uint   `json:"raise_severity"`
	Wait          uint   `json:"wait"`
}

// RuleSetGroupCondition defines conditions for raising alerts
type RuleSetGroupCondition struct {
	MatchingSeverities []string `json:"matching_serverities"`
	RuleSetCID         string   `json:"rule_set"`
}

// RuleSetGroup defines a ruleset group. See https://login.circonus.com/resources/api/calls/rule_set_group for more information.
type RuleSetGroup struct {
	CID               string                  `json:"_cid,omitempty"`
	ContactGroups     map[uint8][]string      `json:"contact_groups"`
	Formulas          []RuleSetGroupFormula   `json:"formulas"`
	Name              string                  `json:"name"`
	RuleSetConditions []RuleSetGroupCondition `json:"rule_set_conditions"`
	Tags              []string                `json:"tags"`
}

// NewRuleSetGroup returns a new RuleSetGroup (with defaults, if applicable)
func NewRuleSetGroup() *RuleSetGroup {
	return &RuleSetGroup{}
}

// FetchRuleSetGroup retrieves rule set group with passed cid.
func (a *API) FetchRuleSetGroup(cid CIDType) (*RuleSetGroup, error) {
	if cid == nil || *cid == "" {
		return nil, fmt.Errorf("Invalid rule set group CID [none]")
	}

	groupCID := string(*cid)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	result, err := a.Get(groupCID)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] fetch rule set group, received JSON: %s", string(result))
	}

	rulesetGroup := &RuleSetGroup{}
	if err := json.Unmarshal(result, rulesetGroup); err != nil {
		return nil, err
	}

	return rulesetGroup, nil
}

// FetchRuleSetGroups retrieves all rule set groups available to API Token.
func (a *API) FetchRuleSetGroups() (*[]RuleSetGroup, error) {
	result, err := a.Get(config.RuleSetGroupPrefix)
	if err != nil {
		return nil, err
	}

	var rulesetGroups []RuleSetGroup
	if err := json.Unmarshal(result, &rulesetGroups); err != nil {
		return nil, err
	}

	return &rulesetGroups, nil
}

// UpdateRuleSetGroup updates passed rule set group.
func (a *API) UpdateRuleSetGroup(cfg *RuleSetGroup) (*RuleSetGroup, error) {
	if cfg == nil {
		return nil, fmt.Errorf("Invalid rule set group config [nil]")
	}

	groupCID := string(cfg.CID)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	jsonCfg, err := json.Marshal(cfg)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] update rule set group, sending JSON: %s", string(jsonCfg))
	}

	result, err := a.Put(groupCID, jsonCfg)
	if err != nil {
		return nil, err
	}

	groups := &RuleSetGroup{}
	if err := json.Unmarshal(result, groups); err != nil {
		return nil, err
	}

	return groups, nil
}

// CreateRuleSetGroup creates a new rule set group.
func (a *API) CreateRuleSetGroup(cfg *RuleSetGroup) (*RuleSetGroup, error) {
	if cfg == nil {
		return nil, fmt.Errorf("Invalid rule set group config [nil]")
	}

	jsonCfg, err := json.Marshal(cfg)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] create rule set group, sending JSON: %s", string(jsonCfg))
	}

	result, err := a.Post(config.RuleSetGroupPrefix, jsonCfg)
	if err != nil {
		return nil, err
	}

	group := &RuleSetGroup{}
	if err := json.Unmarshal(result, group); err != nil {
		return nil, err
	}

	return group, nil
}

// DeleteRuleSetGroup deletes passed rule set group.
func (a *API) DeleteRuleSetGroup(cfg *RuleSetGroup) (bool, error) {
	if cfg == nil {
		return false, fmt.Errorf("Invalid rule set group config [nil]")
	}
	return a.DeleteRuleSetGroupByCID(CIDType(&cfg.CID))
}

// DeleteRuleSetGroupByCID deletes rule set group wiht passed cid.
func (a *API) DeleteRuleSetGroupByCID(cid CIDType) (bool, error) {
	if cid == nil || *cid == "" {
		return false, fmt.Errorf("Invalid rule set group CID [none]")
	}

	groupCID := string(*cid)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return false, err
	}
	if !matched {
		return false, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	_, err = a.Delete(groupCID)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SearchRuleSetGroups returns rule set groups matching the
// specified search query and/or filter. If nil is passed for
// both parameters all rule set groups will be returned.
func (a *API) SearchRuleSetGroups(searchCriteria *SearchQueryType, filterCriteria *SearchFilterType) (*[]RuleSetGroup, error) {
	q := url.Values{}

	if searchCriteria != nil && *searchCriteria != "" {
		q.Set("search", string(*searchCriteria))
	}

	if filterCriteria != nil && len(*filterCriteria) > 0 {
		for filter, criteria := range *filterCriteria {
			for _, val := range criteria {
				q.Add(filter, val)
			}
		}
	}

	if q.Encode() == "" {
		return a.FetchRuleSetGroups()
	}

	reqURL := url.URL{
		Path:     config.RuleSetGroupPrefix,
		RawQuery: q.Encode(),
	}

	result, err := a.Get(reqURL.String())
	if err != nil {
		return nil, fmt.Errorf("[ERROR] API call error %+v", err)
	}

	var groups []RuleSetGroup
	if err := json.Unmarshal(result, &groups); err != nil {
		return nil, err
	}

	return &groups, nil
}
<commit_msg>upd: document/update struct member types to reflect what is received from api<commit_after>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// RuleSetGroup API support - Fetch, Create, Update, Delete, and Search
// See: https://login.circonus.com/resources/api/calls/rule_set_group

package api

import (
	"encoding/json"
	"fmt"
	"net/url"
	"regexp"

	"github.com/circonus-labs/circonus-gometrics/api/config"
)

// RuleSetGroupFormula defines a formula for raising alerts
type RuleSetGroupFormula struct {
	Expression    interface{} `json:"expression"`     // string or uint BUG doc: string, api: string or numeric
	RaiseSeverity uint        `json:"raise_severity"` // uint
	Wait          uint        `json:"wait"`           // uint
}

// RuleSetGroupCondition defines conditions for raising alerts
type RuleSetGroupCondition struct {
	MatchingSeverities []string `json:"matching_serverities"` // [] len >= 1
	RuleSetCID         string   `json:"rule_set"`             // string
}

// RuleSetGroup defines a ruleset group. See https://login.circonus.com/resources/api/calls/rule_set_group for more information.
type RuleSetGroup struct {
	CID               string                  `json:"_cid,omitempty"`      // string
	ContactGroups     map[uint8][]string      `json:"contact_groups"`      // [] len == 5
	Formulas          []RuleSetGroupFormula   `json:"formulas"`            // [] len >= 0
	Name              string                  `json:"name"`                // string
	RuleSetConditions []RuleSetGroupCondition `json:"rule_set_conditions"` // [] len >= 1
	Tags              []string                `json:"tags"`                // [] len >= 0
}

// NewRuleSetGroup returns a new RuleSetGroup (with defaults, if applicable)
func NewRuleSetGroup() *RuleSetGroup {
	return &RuleSetGroup{}
}

// FetchRuleSetGroup retrieves rule set group with passed cid.
func (a *API) FetchRuleSetGroup(cid CIDType) (*RuleSetGroup, error) {
	if cid == nil || *cid == "" {
		return nil, fmt.Errorf("Invalid rule set group CID [none]")
	}

	groupCID := string(*cid)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	result, err := a.Get(groupCID)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] fetch rule set group, received JSON: %s", string(result))
	}

	rulesetGroup := &RuleSetGroup{}
	if err := json.Unmarshal(result, rulesetGroup); err != nil {
		return nil, err
	}

	return rulesetGroup, nil
}

// FetchRuleSetGroups retrieves all rule set groups available to API Token.
func (a *API) FetchRuleSetGroups() (*[]RuleSetGroup, error) {
	result, err := a.Get(config.RuleSetGroupPrefix)
	if err != nil {
		return nil, err
	}

	var rulesetGroups []RuleSetGroup
	if err := json.Unmarshal(result, &rulesetGroups); err != nil {
		return nil, err
	}

	return &rulesetGroups, nil
}

// UpdateRuleSetGroup updates passed rule set group.
func (a *API) UpdateRuleSetGroup(cfg *RuleSetGroup) (*RuleSetGroup, error) {
	if cfg == nil {
		return nil, fmt.Errorf("Invalid rule set group config [nil]")
	}

	groupCID := string(cfg.CID)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return nil, err
	}
	if !matched {
		return nil, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	jsonCfg, err := json.Marshal(cfg)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] update rule set group, sending JSON: %s", string(jsonCfg))
	}

	result, err := a.Put(groupCID, jsonCfg)
	if err != nil {
		return nil, err
	}

	groups := &RuleSetGroup{}
	if err := json.Unmarshal(result, groups); err != nil {
		return nil, err
	}

	return groups, nil
}

// CreateRuleSetGroup creates a new rule set group.
func (a *API) CreateRuleSetGroup(cfg *RuleSetGroup) (*RuleSetGroup, error) {
	if cfg == nil {
		return nil, fmt.Errorf("Invalid rule set group config [nil]")
	}

	jsonCfg, err := json.Marshal(cfg)
	if err != nil {
		return nil, err
	}

	if a.Debug {
		a.Log.Printf("[DEBUG] create rule set group, sending JSON: %s", string(jsonCfg))
	}

	result, err := a.Post(config.RuleSetGroupPrefix, jsonCfg)
	if err != nil {
		return nil, err
	}

	group := &RuleSetGroup{}
	if err := json.Unmarshal(result, group); err != nil {
		return nil, err
	}

	return group, nil
}

// DeleteRuleSetGroup deletes passed rule set group.
func (a *API) DeleteRuleSetGroup(cfg *RuleSetGroup) (bool, error) {
	if cfg == nil {
		return false, fmt.Errorf("Invalid rule set group config [nil]")
	}
	return a.DeleteRuleSetGroupByCID(CIDType(&cfg.CID))
}

// DeleteRuleSetGroupByCID deletes rule set group wiht passed cid.
func (a *API) DeleteRuleSetGroupByCID(cid CIDType) (bool, error) {
	if cid == nil || *cid == "" {
		return false, fmt.Errorf("Invalid rule set group CID [none]")
	}

	groupCID := string(*cid)

	matched, err := regexp.MatchString(config.RuleSetGroupCIDRegex, groupCID)
	if err != nil {
		return false, err
	}
	if !matched {
		return false, fmt.Errorf("Invalid rule set group CID [%s]", groupCID)
	}

	_, err = a.Delete(groupCID)
	if err != nil {
		return false, err
	}

	return true, nil
}

// SearchRuleSetGroups returns rule set groups matching the
// specified search query and/or filter. If nil is passed for
// both parameters all rule set groups will be returned.
func (a *API) SearchRuleSetGroups(searchCriteria *SearchQueryType, filterCriteria *SearchFilterType) (*[]RuleSetGroup, error) {
	q := url.Values{}

	if searchCriteria != nil && *searchCriteria != "" {
		q.Set("search", string(*searchCriteria))
	}

	if filterCriteria != nil && len(*filterCriteria) > 0 {
		for filter, criteria := range *filterCriteria {
			for _, val := range criteria {
				q.Add(filter, val)
			}
		}
	}

	if q.Encode() == "" {
		return a.FetchRuleSetGroups()
	}

	reqURL := url.URL{
		Path:     config.RuleSetGroupPrefix,
		RawQuery: q.Encode(),
	}

	result, err := a.Get(reqURL.String())
	if err != nil {
		return nil, fmt.Errorf("[ERROR] API call error %+v", err)
	}

	var groups []RuleSetGroup
	if err := json.Unmarshal(result, &groups); err != nil {
		return nil, err
	}

	return &groups, nil
}
<|endoftext|>


<commit_before>package server

import (
	"net/http"

	"github.com/gin-gonic/gin"
)

// Version of IronFunctions
var Version = "0.0.50"

func handleVersion(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"version": Version})
}
<commit_msg>functions: 0.0.51 release [skip ci]<commit_after>package server

import (
	"net/http"

	"github.com/gin-gonic/gin"
)

// Version of IronFunctions
var Version = "0.0.51"

func handleVersion(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"version": Version})
}
<|endoftext|>