<commit_before>//
// Treasure Data API client for Go
//
// Copyright (C) 2014 Treasure Data, Inc.
//
// 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 td_client

import (
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"strconv"
	"time"
)

// ListTablesResultElement represents an item of the result of ListTables API
// call
type ListTablesResultElement struct {
	Id                   int
	Name                 string
	Type                 string
	Count                int
	CreatedAt            time.Time
	UpdatedAt            time.Time
	LastImport           time.Time
	LastLogTimestamp     time.Time
	EstimatedStorageSize int
	Schema               []interface{}
	ExpireDays           int
	PrimaryKey           string
	PrimaryKeyType       string
}

var showTableSchema = map[string]interface{}{
	"id":                     0,
	"name":                   "",
	"type":                   Optional{"", "?"},
	"count":                  Optional{0, 0},
	"created_at":             time.Time{},
	"updated_at":             time.Time{},
	"counter_updated_at":     Optional{time.Time{}, time.Time{}},
	"last_log_timestamp":     Optional{time.Time{}, time.Time{}},
	"delete_protected":       false,
	"estimated_storage_size": 0,
	"schema":                 Optional{EmbeddedJSON([]interface{}{}), nil},
	"expire_days":            Optional{0, 0},
	"primary_key":            Optional{"", ""},
	"primary_key_type":       Optional{"", ""},
	"include_v":              false,
}

// ListTablesResult is a collection of ListTablesResultElement
type ListTablesResult []ListTablesResultElement

var listTablesSchema = map[string]interface{}{
	"database": "",
	"tables": []map[string]interface{}{
		map[string]interface{}{
			"id":                     0,
			"name":                   "",
			"type":                   Optional{"", "?"},
			"count":                  Optional{0, 0},
			"created_at":             time.Time{},
			"updated_at":             time.Time{},
			"counter_updated_at":     Optional{time.Time{}, time.Time{}},
			"last_log_timestamp":     Optional{time.Time{}, time.Time{}},
			"delete_protected":       false,
			"estimated_storage_size": 0,
			"schema":                 Optional{EmbeddedJSON([]interface{}{}), nil},
			"expire_days":            Optional{0, 0},
			"primary_key":            Optional{"", ""},
			"primary_key_type":       Optional{"", ""},
			"include_v":              false,
		},
	},
}

var deleteTableSchema = map[string]interface{}{
	"table":    "",
	"database": "",
	"type":     Optional{"", "?"},
}

func (client *TDClient) ShowTable(db, table string) (*ListTablesResultElement, error) {
	resp, err := client.get(fmt.Sprintf("/v3/table/show/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, client.buildError(resp, -1, "Show table failed", nil)
	}
	js, err := client.checkedJson(resp, showTableSchema)
	if err != nil {
		return nil, err
	}
	return &ListTablesResultElement{
		Id:                   js["id"].(int),
		Name:                 js["name"].(string),
		Type:                 js["type"].(string),
		Count:                js["count"].(int),
		CreatedAt:            js["created_at"].(time.Time),
		UpdatedAt:            js["updated_at"].(time.Time),
		LastImport:           js["counter_updated_at"].(time.Time),
		LastLogTimestamp:     js["last_log_timestamp"].(time.Time),
		EstimatedStorageSize: js["estimated_storage_size"].(int),
		Schema:               js["schema"].([]interface{}),
		ExpireDays:           js["expire_days"].(int),
		PrimaryKey:           js["primary_key"].(string),
		PrimaryKeyType:       js["primary_key_type"].(string),
	}, nil
}

func (client *TDClient) ListTables(db string) (*ListTablesResult, error) {
	resp, err := client.get(fmt.Sprintf("/v3/table/list/%s", url.QueryEscape(db)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, client.buildError(resp, -1, "List tables failed", nil)
	}
	js, err := client.checkedJson(resp, listTablesSchema)
	if err != nil {
		return nil, err
	}
	tables := js["tables"].([]map[string]interface{})
	retval := make(ListTablesResult, len(tables))
	for i, v := range tables {
		retval[i] = ListTablesResultElement{
			Id:                   v["id"].(int),
			Name:                 v["name"].(string),
			Type:                 v["type"].(string),
			Count:                v["count"].(int),
			CreatedAt:            v["created_at"].(time.Time),
			UpdatedAt:            v["updated_at"].(time.Time),
			LastImport:           v["counter_updated_at"].(time.Time),
			LastLogTimestamp:     v["last_log_timestamp"].(time.Time),
			EstimatedStorageSize: v["estimated_storage_size"].(int),
			Schema:               v["schema"].([]interface{}),
			ExpireDays:           v["expire_days"].(int),
			PrimaryKey:           v["primary_key"].(string),
			PrimaryKeyType:       v["primary_key_type"].(string),
		}
	}
	return &retval, nil
}

func (client *TDClient) createTable(db string, table string, type_ string, params map[string]string) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/create/%s/%s/%s", url.QueryEscape(db), url.QueryEscape(table), url.QueryEscape(type_)), dictToValues(params))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, fmt.Sprintf("Create %s table failed", type_), nil)
	}
	return nil
}

func (client *TDClient) CreateItemTable(db string, table string, primaryKey string, primaryKeyType string) error {
	return client.createTable(
		db, table, "item",
		map[string]string{
			"primary_key":      primaryKey,
			"primary_key_type": primaryKeyType,
		},
	)
}

func (client *TDClient) CreateLogTable(db string, table string) error {
	return client.createTable(db, table, "log", nil)
}

func (client *TDClient) SwapTable(db string, table1 string, table2 string) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/swap/%s/%s/%s", url.QueryEscape(db), url.QueryEscape(table1), url.QueryEscape(table2)), nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Swap tables failed", nil)
	}
	return nil
}

func (client *TDClient) UpdateSchema(db string, table string, schema []interface{}) error {
	jsStr, err := json.Marshal(schema)
	if err != nil {
		return err
	}
	resp, err := client.post(fmt.Sprintf("/v3/table/update-schema/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), url.Values{"schema": {string(jsStr)}})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Update schema failed", nil)
	}
	return nil
}

func (client *TDClient) UpdateExpire(db string, table string, expireDays int) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/update/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), url.Values{"expire_days": {strconv.Itoa(expireDays)}})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Update expire failed", nil)
	}
	return nil
}

func (client *TDClient) DeleteTable(db string, table string) (string, error) {
	resp, err := client.post(fmt.Sprintf("/v3/table/delete/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", client.buildError(resp, -1, "Delete table failed", nil)
	}
	js, err := client.checkedJson(resp, deleteTableSchema)
	if err != nil {
		return "", err
	}
	return js["type"].(string), err
}

func (client *TDClient) Tail(db string, table string, count int, to time.Time, from time.Time, reader func(interface{}) error) error {
	params := url.Values{}
	if count > 0 {
		params.Set("count", strconv.Itoa(count))
	}
	if !to.IsZero() {
		params.Set("to", to.UTC().Format(TDAPIDateTime))
	}
	if !from.IsZero() {
		params.Set("from", from.UTC().Format(TDAPIDateTime))
	}
	resp, err := client.post(fmt.Sprintf("/v3/table/tail/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), params)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Tail failed", nil)
	}
	dec := client.getMessagePackDecoder(resp.Body)
	for {
		v := (interface{})(nil)
		err := dec.Decode(&v)
		if err != nil {
			if err == io.EOF {
				break
			}
			return client.buildError(resp, -1, "Invalid MessagePack stream", nil)
		}
		err = reader(v)
		if err != nil {
			return client.buildError(resp, -1, "Reader returned error status", err)
		}
	}
	return nil
}
<commit_msg>Add `IncludeV` to `ListTablesResultElement`<commit_after>//
// Treasure Data API client for Go
//
// Copyright (C) 2014 Treasure Data, Inc.
//
// 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 td_client

import (
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"strconv"
	"time"
)

// ListTablesResultElement represents an item of the result of ListTables API
// call
type ListTablesResultElement struct {
	Id                   int
	Name                 string
	Type                 string
	Count                int
	CreatedAt            time.Time
	UpdatedAt            time.Time
	LastImport           time.Time
	LastLogTimestamp     time.Time
	EstimatedStorageSize int
	Schema               []interface{}
	ExpireDays           int
	PrimaryKey           string
	PrimaryKeyType       string
	IncludeV             bool
}

var showTableSchema = map[string]interface{}{
	"id":                     0,
	"name":                   "",
	"type":                   Optional{"", "?"},
	"count":                  Optional{0, 0},
	"created_at":             time.Time{},
	"updated_at":             time.Time{},
	"counter_updated_at":     Optional{time.Time{}, time.Time{}},
	"last_log_timestamp":     Optional{time.Time{}, time.Time{}},
	"delete_protected":       false,
	"estimated_storage_size": 0,
	"schema":                 Optional{EmbeddedJSON([]interface{}{}), nil},
	"expire_days":            Optional{0, 0},
	"primary_key":            Optional{"", ""},
	"primary_key_type":       Optional{"", ""},
	"include_v":              false,
}

// ListTablesResult is a collection of ListTablesResultElement
type ListTablesResult []ListTablesResultElement

var listTablesSchema = map[string]interface{}{
	"database": "",
	"tables": []map[string]interface{}{
		map[string]interface{}{
			"id":                     0,
			"name":                   "",
			"type":                   Optional{"", "?"},
			"count":                  Optional{0, 0},
			"created_at":             time.Time{},
			"updated_at":             time.Time{},
			"counter_updated_at":     Optional{time.Time{}, time.Time{}},
			"last_log_timestamp":     Optional{time.Time{}, time.Time{}},
			"delete_protected":       false,
			"estimated_storage_size": 0,
			"schema":                 Optional{EmbeddedJSON([]interface{}{}), nil},
			"expire_days":            Optional{0, 0},
			"primary_key":            Optional{"", ""},
			"primary_key_type":       Optional{"", ""},
			"include_v":              false,
		},
	},
}

var deleteTableSchema = map[string]interface{}{
	"table":    "",
	"database": "",
	"type":     Optional{"", "?"},
}

func (client *TDClient) ShowTable(db, table string) (*ListTablesResultElement, error) {
	resp, err := client.get(fmt.Sprintf("/v3/table/show/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, client.buildError(resp, -1, "Show table failed", nil)
	}
	js, err := client.checkedJson(resp, showTableSchema)
	if err != nil {
		return nil, err
	}
	return &ListTablesResultElement{
		Id:                   js["id"].(int),
		Name:                 js["name"].(string),
		Type:                 js["type"].(string),
		Count:                js["count"].(int),
		CreatedAt:            js["created_at"].(time.Time),
		UpdatedAt:            js["updated_at"].(time.Time),
		LastImport:           js["counter_updated_at"].(time.Time),
		LastLogTimestamp:     js["last_log_timestamp"].(time.Time),
		EstimatedStorageSize: js["estimated_storage_size"].(int),
		Schema:               js["schema"].([]interface{}),
		ExpireDays:           js["expire_days"].(int),
		PrimaryKey:           js["primary_key"].(string),
		PrimaryKeyType:       js["primary_key_type"].(string),
		IncludeV:             js["include_v"].(bool),
	}, nil
}

func (client *TDClient) ListTables(db string) (*ListTablesResult, error) {
	resp, err := client.get(fmt.Sprintf("/v3/table/list/%s", url.QueryEscape(db)), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, client.buildError(resp, -1, "List tables failed", nil)
	}
	js, err := client.checkedJson(resp, listTablesSchema)
	if err != nil {
		return nil, err
	}
	tables := js["tables"].([]map[string]interface{})
	retval := make(ListTablesResult, len(tables))
	for i, v := range tables {
		retval[i] = ListTablesResultElement{
			Id:                   v["id"].(int),
			Name:                 v["name"].(string),
			Type:                 v["type"].(string),
			Count:                v["count"].(int),
			CreatedAt:            v["created_at"].(time.Time),
			UpdatedAt:            v["updated_at"].(time.Time),
			LastImport:           v["counter_updated_at"].(time.Time),
			LastLogTimestamp:     v["last_log_timestamp"].(time.Time),
			EstimatedStorageSize: v["estimated_storage_size"].(int),
			Schema:               v["schema"].([]interface{}),
			ExpireDays:           v["expire_days"].(int),
			PrimaryKey:           v["primary_key"].(string),
			PrimaryKeyType:       v["primary_key_type"].(string),
			IncludeV:             v["include_v"].(bool),
		}
	}
	return &retval, nil
}

func (client *TDClient) createTable(db string, table string, type_ string, params map[string]string) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/create/%s/%s/%s", url.QueryEscape(db), url.QueryEscape(table), url.QueryEscape(type_)), dictToValues(params))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, fmt.Sprintf("Create %s table failed", type_), nil)
	}
	return nil
}

func (client *TDClient) CreateItemTable(db string, table string, primaryKey string, primaryKeyType string) error {
	return client.createTable(
		db, table, "item",
		map[string]string{
			"primary_key":      primaryKey,
			"primary_key_type": primaryKeyType,
		},
	)
}

func (client *TDClient) CreateLogTable(db string, table string) error {
	return client.createTable(db, table, "log", nil)
}

func (client *TDClient) SwapTable(db string, table1 string, table2 string) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/swap/%s/%s/%s", url.QueryEscape(db), url.QueryEscape(table1), url.QueryEscape(table2)), nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Swap tables failed", nil)
	}
	return nil
}

func (client *TDClient) UpdateSchema(db string, table string, schema []interface{}) error {
	jsStr, err := json.Marshal(schema)
	if err != nil {
		return err
	}
	resp, err := client.post(fmt.Sprintf("/v3/table/update-schema/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), url.Values{"schema": {string(jsStr)}})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Update schema failed", nil)
	}
	return nil
}

func (client *TDClient) UpdateExpire(db string, table string, expireDays int) error {
	resp, err := client.post(fmt.Sprintf("/v3/table/update/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), url.Values{"expire_days": {strconv.Itoa(expireDays)}})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Update expire failed", nil)
	}
	return nil
}

func (client *TDClient) DeleteTable(db string, table string) (string, error) {
	resp, err := client.post(fmt.Sprintf("/v3/table/delete/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), nil)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", client.buildError(resp, -1, "Delete table failed", nil)
	}
	js, err := client.checkedJson(resp, deleteTableSchema)
	if err != nil {
		return "", err
	}
	return js["type"].(string), err
}

func (client *TDClient) Tail(db string, table string, count int, to time.Time, from time.Time, reader func(interface{}) error) error {
	params := url.Values{}
	if count > 0 {
		params.Set("count", strconv.Itoa(count))
	}
	if !to.IsZero() {
		params.Set("to", to.UTC().Format(TDAPIDateTime))
	}
	if !from.IsZero() {
		params.Set("from", from.UTC().Format(TDAPIDateTime))
	}
	resp, err := client.post(fmt.Sprintf("/v3/table/tail/%s/%s", url.QueryEscape(db), url.QueryEscape(table)), params)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return client.buildError(resp, -1, "Tail failed", nil)
	}
	dec := client.getMessagePackDecoder(resp.Body)
	for {
		v := (interface{})(nil)
		err := dec.Decode(&v)
		if err != nil {
			if err == io.EOF {
				break
			}
			return client.buildError(resp, -1, "Invalid MessagePack stream", nil)
		}
		err = reader(v)
		if err != nil {
			return client.buildError(resp, -1, "Reader returned error status", err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Command apidiff determines whether two versions of a package are compatible
package main

import (
	"flag"
	"fmt"
	"go/token"
	"go/types"
	"os"

	"golang.org/x/exp/apidiff"
	"golang.org/x/tools/go/gcexportdata"
	"golang.org/x/tools/go/packages"
)

var (
	exportDataOutfile = flag.String("w", "", "file for export data")
	incompatibleOnly  = flag.Bool("incompatible", false, "display only incompatible changes")
)

func main() {
	flag.Usage = func() {
		w := flag.CommandLine.Output()
		fmt.Fprintf(w, "usage:\n")
		fmt.Fprintf(w, "apidiff OLD NEW\n")
		fmt.Fprintf(w, "   compares OLD and NEW package APIs\n")
		fmt.Fprintf(w, "   where OLD and NEW are either import paths or files of export data\n")
		fmt.Fprintf(w, "apidiff -w FILE IMPORT_PATH\n")
		fmt.Fprintf(w, "   writes export data of the package at IMPORT_PATH to FILE\n")
		fmt.Fprintf(w, "   NOTE: In a GOPATH-less environment, this option consults the\n")
		fmt.Fprintf(w, "   module cache by default, unless used in the directory that\n")
		fmt.Fprintf(w, "   contains the go.mod module definition that IMPORT_PATH belongs\n")
		fmt.Fprintf(w, "   to. In most cases users want the latter behavior, so be sure\n")
		fmt.Fprintf(w, "   to cd to the exact directory which contains the module\n")
		fmt.Fprintf(w, "   definition of IMPORT_PATH.\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if *exportDataOutfile != "" {
		if len(flag.Args()) != 1 {
			flag.Usage()
		}
		pkg := mustLoadPackage(flag.Arg(0))
		if err := writeExportData(pkg, *exportDataOutfile); err != nil {
			die("writing export data: %v", err)
		}
	} else {
		if len(flag.Args()) != 2 {
			flag.Usage()
		}
		oldpkg := mustLoadOrRead(flag.Arg(0))
		newpkg := mustLoadOrRead(flag.Arg(1))

		report := apidiff.Changes(oldpkg, newpkg)
		var err error
		if *incompatibleOnly {
			err = report.TextIncompatible(os.Stdout, false)
		} else {
			err = report.Text(os.Stdout)
		}
		if err != nil {
			die("writing report: %v", err)
		}
	}
}

func mustLoadOrRead(importPathOrFile string) *types.Package {
	fileInfo, err := os.Stat(importPathOrFile)
	if err == nil && fileInfo.Mode().IsRegular() {
		pkg, err := readExportData(importPathOrFile)
		if err != nil {
			die("reading export data from %s: %v", importPathOrFile, err)
		}
		return pkg
	} else {
		return mustLoadPackage(importPathOrFile).Types
	}
}

func mustLoadPackage(importPath string) *packages.Package {
	pkg, err := loadPackage(importPath)
	if err != nil {
		die("loading %s: %v", importPath, err)
	}
	return pkg
}

func loadPackage(importPath string) (*packages.Package, error) {
	cfg := &packages.Config{Mode: packages.LoadTypes}
	pkgs, err := packages.Load(cfg, importPath)
	if err != nil {
		return nil, err
	}
	if len(pkgs) == 0 {
		return nil, fmt.Errorf("found no packages for import %s", importPath)
	}
	if len(pkgs[0].Errors) > 0 {
		return nil, pkgs[0].Errors[0]
	}
	return pkgs[0], nil
}

func readExportData(filename string) (*types.Package, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return gcexportdata.Read(f, token.NewFileSet(), map[string]*types.Package{}, filename)
}

func writeExportData(pkg *packages.Package, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	err1 := gcexportdata.Write(f, pkg.Fset, pkg.Types)
	err2 := f.Close()
	if err1 != nil {
		return err1
	}
	return err2
}

func die(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	os.Exit(1)
}
<commit_msg>cmd/apidiff: exit on bad input<commit_after>// Command apidiff determines whether two versions of a package are compatible
package main

import (
	"flag"
	"fmt"
	"go/token"
	"go/types"
	"os"

	"golang.org/x/exp/apidiff"
	"golang.org/x/tools/go/gcexportdata"
	"golang.org/x/tools/go/packages"
)

var (
	exportDataOutfile = flag.String("w", "", "file for export data")
	incompatibleOnly  = flag.Bool("incompatible", false, "display only incompatible changes")
)

func main() {
	flag.Usage = func() {
		w := flag.CommandLine.Output()
		fmt.Fprintf(w, "usage:\n")
		fmt.Fprintf(w, "apidiff OLD NEW\n")
		fmt.Fprintf(w, "   compares OLD and NEW package APIs\n")
		fmt.Fprintf(w, "   where OLD and NEW are either import paths or files of export data\n")
		fmt.Fprintf(w, "apidiff -w FILE IMPORT_PATH\n")
		fmt.Fprintf(w, "   writes export data of the package at IMPORT_PATH to FILE\n")
		fmt.Fprintf(w, "   NOTE: In a GOPATH-less environment, this option consults the\n")
		fmt.Fprintf(w, "   module cache by default, unless used in the directory that\n")
		fmt.Fprintf(w, "   contains the go.mod module definition that IMPORT_PATH belongs\n")
		fmt.Fprintf(w, "   to. In most cases users want the latter behavior, so be sure\n")
		fmt.Fprintf(w, "   to cd to the exact directory which contains the module\n")
		fmt.Fprintf(w, "   definition of IMPORT_PATH.\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if *exportDataOutfile != "" {
		if len(flag.Args()) != 1 {
			flag.Usage()
			os.Exit(2)
		}
		pkg := mustLoadPackage(flag.Arg(0))
		if err := writeExportData(pkg, *exportDataOutfile); err != nil {
			die("writing export data: %v", err)
		}
	} else {
		if len(flag.Args()) != 2 {
			flag.Usage()
			os.Exit(2)
		}
		oldpkg := mustLoadOrRead(flag.Arg(0))
		newpkg := mustLoadOrRead(flag.Arg(1))

		report := apidiff.Changes(oldpkg, newpkg)
		var err error
		if *incompatibleOnly {
			err = report.TextIncompatible(os.Stdout, false)
		} else {
			err = report.Text(os.Stdout)
		}
		if err != nil {
			die("writing report: %v", err)
		}
	}
}

func mustLoadOrRead(importPathOrFile string) *types.Package {
	fileInfo, err := os.Stat(importPathOrFile)
	if err == nil && fileInfo.Mode().IsRegular() {
		pkg, err := readExportData(importPathOrFile)
		if err != nil {
			die("reading export data from %s: %v", importPathOrFile, err)
		}
		return pkg
	} else {
		return mustLoadPackage(importPathOrFile).Types
	}
}

func mustLoadPackage(importPath string) *packages.Package {
	pkg, err := loadPackage(importPath)
	if err != nil {
		die("loading %s: %v", importPath, err)
	}
	return pkg
}

func loadPackage(importPath string) (*packages.Package, error) {
	cfg := &packages.Config{Mode: packages.LoadTypes}
	pkgs, err := packages.Load(cfg, importPath)
	if err != nil {
		return nil, err
	}
	if len(pkgs) == 0 {
		return nil, fmt.Errorf("found no packages for import %s", importPath)
	}
	if len(pkgs[0].Errors) > 0 {
		return nil, pkgs[0].Errors[0]
	}
	return pkgs[0], nil
}

func readExportData(filename string) (*types.Package, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return gcexportdata.Read(f, token.NewFileSet(), map[string]*types.Package{}, filename)
}

func writeExportData(pkg *packages.Package, filename string) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	err1 := gcexportdata.Write(f, pkg.Fset, pkg.Types)
	err2 := f.Close()
	if err1 != nil {
		return err1
	}
	return err2
}

func die(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/google/subcommands"
	"github.com/jhillyerd/inbucket/rest/client"
)

type matchCmd struct {
	mailbox string
	output  string
	outFunc func(headers []*client.MessageHeader) error
	delete  bool
	// match criteria
	from    regexFlag
	subject regexFlag
	to      regexFlag
	maxAge  time.Duration
}

func (*matchCmd) Name() string {
	return "match"
}

func (*matchCmd) Synopsis() string {
	return "output messages matching criteria"
}

func (*matchCmd) Usage() string {
	return `match [flags] <mailbox>:
	output messages matching all specified criteria
	exit status will be 1 if no matches were found, otherwise 0
`
}

func (m *matchCmd) SetFlags(f *flag.FlagSet) {
	f.StringVar(&m.output, "output", "id", "output format: id, json, or mbox")
	f.BoolVar(&m.delete, "delete", false, "delete matched messages after output")
	f.Var(&m.from, "from", "From header matching regexp")
	f.Var(&m.subject, "subject", "Subject header matching regexp")
	f.Var(&m.to, "to", "To header matching regexp (must match one)")
	f.DurationVar(
		&m.maxAge, "maxage", 0,
		"Matches must have been received in this time frame (ex: \"10s\", \"5m\")")
}

func (m *matchCmd) Execute(
	_ context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	mailbox := f.Arg(0)
	if mailbox == "" {
		return usage("mailbox required")
	}
	// Select output function
	switch m.output {
	case "id":
		m.outFunc = outputID
	case "json":
		m.outFunc = outputJSON
	case "mbox":
		m.outFunc = outputMbox
	default:
		return usage("unknown output type: " + m.output)
	}
	// Setup REST client
	c, err := client.New(baseURL())
	if err != nil {
		return fatal("Couldn't build client", err)
	}
	// Get list
	headers, err := c.ListMailbox(mailbox)
	if err != nil {
		return fatal("List REST call failed", err)
	}
	// Find matches
	matches := make([]*client.MessageHeader, 0, len(headers))
	for _, h := range headers {
		if m.match(h) {
			matches = append(matches, h)
		}
	}
	// Return error status if no matches
	if len(matches) == 0 {
		return subcommands.ExitFailure
	}
	// Output matches
	err = m.outFunc(matches)
	if err != nil {
		return fatal("Error", err)
	}
	if m.delete {
		// Delete matches
		for _, h := range matches {
			err = h.Delete()
			if err != nil {
				return fatal("Delete REST call failed", err)
			}
		}
	}
	return subcommands.ExitSuccess
}

// match returns true if header matches all defined criteria
func (m *matchCmd) match(header *client.MessageHeader) bool {
	if m.maxAge > 0 {
		if time.Since(header.Date) > m.maxAge {
			return false
		}
	}
	if m.subject.Defined() {
		if !m.subject.MatchString(header.Subject) {
			return false
		}
	}
	if m.from.Defined() {
		if !m.from.MatchString(header.From) {
			return false
		}
	}
	if m.to.Defined() {
		match := false
		for _, to := range header.To {
			if m.to.MatchString(to) {
				match = true
				break
			}
		}
		if !match {
			return false
		}
	}
	return true
}

func outputID(headers []*client.MessageHeader) error {
	for _, h := range headers {
		fmt.Println(h.ID)
	}
	return nil
}

func outputJSON(headers []*client.MessageHeader) error {
	jsonEncoder := json.NewEncoder(os.Stdout)
	jsonEncoder.SetEscapeHTML(false)
	jsonEncoder.SetIndent("", "  ")
	return jsonEncoder.Encode(headers)
}
<commit_msg>Address matching should only apply to address, not name<commit_after>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/mail"
	"os"
	"time"

	"github.com/google/subcommands"
	"github.com/jhillyerd/inbucket/rest/client"
)

type matchCmd struct {
	mailbox string
	output  string
	outFunc func(headers []*client.MessageHeader) error
	delete  bool
	// match criteria
	from    regexFlag
	subject regexFlag
	to      regexFlag
	maxAge  time.Duration
}

func (*matchCmd) Name() string {
	return "match"
}

func (*matchCmd) Synopsis() string {
	return "output messages matching criteria"
}

func (*matchCmd) Usage() string {
	return `match [flags] <mailbox>:
	output messages matching all specified criteria
	exit status will be 1 if no matches were found, otherwise 0
`
}

func (m *matchCmd) SetFlags(f *flag.FlagSet) {
	f.StringVar(&m.output, "output", "id", "output format: id, json, or mbox")
	f.BoolVar(&m.delete, "delete", false, "delete matched messages after output")
	f.Var(&m.from, "from", "From header matching regexp (address, not name)")
	f.Var(&m.subject, "subject", "Subject header matching regexp")
	f.Var(&m.to, "to", "To header matching regexp (must match 1+ to address)")
	f.DurationVar(
		&m.maxAge, "maxage", 0,
		"Matches must have been received in this time frame (ex: \"10s\", \"5m\")")
}

func (m *matchCmd) Execute(
	_ context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
	mailbox := f.Arg(0)
	if mailbox == "" {
		return usage("mailbox required")
	}
	// Select output function
	switch m.output {
	case "id":
		m.outFunc = outputID
	case "json":
		m.outFunc = outputJSON
	case "mbox":
		m.outFunc = outputMbox
	default:
		return usage("unknown output type: " + m.output)
	}
	// Setup REST client
	c, err := client.New(baseURL())
	if err != nil {
		return fatal("Couldn't build client", err)
	}
	// Get list
	headers, err := c.ListMailbox(mailbox)
	if err != nil {
		return fatal("List REST call failed", err)
	}
	// Find matches
	matches := make([]*client.MessageHeader, 0, len(headers))
	for _, h := range headers {
		if m.match(h) {
			matches = append(matches, h)
		}
	}
	// Return error status if no matches
	if len(matches) == 0 {
		return subcommands.ExitFailure
	}
	// Output matches
	err = m.outFunc(matches)
	if err != nil {
		return fatal("Error", err)
	}
	if m.delete {
		// Delete matches
		for _, h := range matches {
			err = h.Delete()
			if err != nil {
				return fatal("Delete REST call failed", err)
			}
		}
	}
	return subcommands.ExitSuccess
}

// match returns true if header matches all defined criteria
func (m *matchCmd) match(header *client.MessageHeader) bool {
	if m.maxAge > 0 {
		if time.Since(header.Date) > m.maxAge {
			return false
		}
	}
	if m.subject.Defined() {
		if !m.subject.MatchString(header.Subject) {
			return false
		}
	}
	if m.from.Defined() {
		from := header.From
		addr, err := mail.ParseAddress(from)
		if err == nil {
			// Parsed successfully
			from = addr.Address
		}
		if !m.from.MatchString(from) {
			return false
		}
	}
	if m.to.Defined() {
		match := false
		for _, to := range header.To {
			addr, err := mail.ParseAddress(to)
			if err == nil {
				// Parsed successfully
				to = addr.Address
			}
			if m.to.MatchString(to) {
				match = true
				break
			}
		}
		if !match {
			return false
		}
	}
	return true
}

func outputID(headers []*client.MessageHeader) error {
	for _, h := range headers {
		fmt.Println(h.ID)
	}
	return nil
}

func outputJSON(headers []*client.MessageHeader) error {
	jsonEncoder := json.NewEncoder(os.Stdout)
	jsonEncoder.SetEscapeHTML(false)
	jsonEncoder.SetIndent("", "  ")
	return jsonEncoder.Encode(headers)
}
<|endoftext|>


<commit_before>/*
dokugen is a simple command line utility that exposes many of the basic functions of the
sudoku package. It's able to generate puzzles (with difficutly) and solve provided puzzles.
Run with -h to see help on how to use it.
*/
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"github.com/gosuri/uiprogress"
	"github.com/jkomoros/sudoku"
	"github.com/jkomoros/sudoku/sdkconverter"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
)

//TODO: let people pass in a filename to export to.

const STORED_PUZZLES_DIRECTORY = ".puzzles"

type appOptions struct {
	GENERATE            bool
	HELP                bool
	PUZZLE_TO_SOLVE     string
	NUM                 int
	PRINT_STATS         bool
	WALKTHROUGH         bool
	RAW_SYMMETRY        string
	RAW_DIFFICULTY      string
	SYMMETRY            sudoku.SymmetryType
	SYMMETRY_PROPORTION float64
	MIN_FILLED_CELLS    int
	MIN_DIFFICULTY      float64
	MAX_DIFFICULTY      float64
	NO_CACHE            bool
	PUZZLE_FORMAT       string
	OUTPUT_CSV          bool
	CONVERTER           sdkconverter.SudokuPuzzleConverter
}

var difficultyRanges map[string]struct {
	low, high float64
}

func init() {
	//grid.Difficulty can make use of a number of processes simultaneously.
	runtime.GOMAXPROCS(6)

	difficultyRanges = map[string]struct{ low, high float64 }{
		"gentle": {0.0, 0.3},
		"easy":   {0.3, 0.6},
		"medium": {0.6, 0.7},
		"tough":  {0.7, 1.0},
	}
}

func getOptions() *appOptions {

	var options appOptions

	flag.BoolVar(&options.GENERATE, "g", false, "if true, will generate a puzzle.")
	flag.BoolVar(&options.HELP, "h", false, "If provided, will print help and exit.")
	flag.IntVar(&options.NUM, "n", 1, "Number of things to generate")
	flag.BoolVar(&options.PRINT_STATS, "p", false, "If provided, will print stats.")
	flag.StringVar(&options.PUZZLE_TO_SOLVE, "s", "", "If provided, will solve the puzzle at the given filename and print solution.")
	flag.BoolVar(&options.WALKTHROUGH, "w", false, "If provided, will print out a walkthrough to solve the provided puzzle.")
	flag.StringVar(&options.RAW_SYMMETRY, "y", "vertical", "Valid values: 'none', 'both', 'horizontal', 'vertical")
	flag.Float64Var(&options.SYMMETRY_PROPORTION, "r", 0.7, "What proportion of cells should be filled according to symmetry")
	flag.IntVar(&options.MIN_FILLED_CELLS, "min-filled-cells", 0, "The minimum number of cells that should be filled in the generated puzzles.")
	flag.Float64Var(&options.MIN_DIFFICULTY, "min", 0.0, "Minimum difficulty for generated puzzle")
	flag.Float64Var(&options.MAX_DIFFICULTY, "max", 1.0, "Maximum difficulty for generated puzzle")
	flag.BoolVar(&options.NO_CACHE, "no-cache", false, "If provided, will not vend generated puzzles from the cache of previously generated puzzles.")
	//TODO: the format should also be how we interpret loads, too.
	flag.StringVar(&options.PUZZLE_FORMAT, "format", "sdk", "Which format to export puzzles from. Defaults to 'sdk'")
	flag.BoolVar(&options.OUTPUT_CSV, "csv", false, "Output the results in CSV.")
	flag.StringVar(&options.RAW_DIFFICULTY, "d", "", "difficulty, one of {gentle, easy, medium, tough}")
	flag.Parse()

	options.RAW_SYMMETRY = strings.ToLower(options.RAW_SYMMETRY)
	switch options.RAW_SYMMETRY {
	case "none":
		options.SYMMETRY = sudoku.SYMMETRY_NONE
	case "both":
		options.SYMMETRY = sudoku.SYMMETRY_BOTH
	case "horizontal":
		options.SYMMETRY = sudoku.SYMMETRY_HORIZONTAL
	case "vertical":
		options.SYMMETRY = sudoku.SYMMETRY_VERTICAL
	default:
		log.Fatal("Unknown symmetry flag: ", options.RAW_SYMMETRY)
	}

	options.RAW_DIFFICULTY = strings.ToLower(options.RAW_DIFFICULTY)
	if options.RAW_DIFFICULTY != "" {
		vals, ok := difficultyRanges[options.RAW_DIFFICULTY]
		if !ok {
			log.Fatal("Invalid difficulty option:", options.RAW_DIFFICULTY)
		}
		options.MIN_DIFFICULTY = vals.low
		options.MAX_DIFFICULTY = vals.high
		log.Println("Using difficulty max:", strconv.FormatFloat(vals.high, 'f', -1, 64), "min:", strconv.FormatFloat(vals.low, 'f', -1, 64))
	}

	options.CONVERTER = sdkconverter.Converters[options.PUZZLE_FORMAT]

	if options.CONVERTER == nil {
		log.Fatal("Invalid format option:", options.PUZZLE_FORMAT)
	}

	return &options
}

func main() {

	//TODO: figure out how to test this.

	options := getOptions()

	output := os.Stdout

	if options.HELP {
		flag.PrintDefaults()
		return
	}

	var grid *sudoku.Grid

	var csvWriter *csv.Writer
	var csvRec []string

	if options.OUTPUT_CSV {
		csvWriter = csv.NewWriter(output)
	}

	var bar *uiprogress.Bar

	//TODO: do more useful / explanatory printing here.
	if options.NUM > 1 {
		uiprogress.DefaultProgress.Out = os.Stderr
		uiprogress.Start()
		bar = uiprogress.AddBar(options.NUM).PrependElapsed().AppendCompleted()
	}

	for i := 0; i < options.NUM; i++ {

		if options.OUTPUT_CSV {
			csvRec = nil
		}

		//TODO: allow the type of symmetry to be configured.
		if options.GENERATE {
			grid = generatePuzzle(options.MIN_DIFFICULTY, options.MAX_DIFFICULTY, options.SYMMETRY, options.SYMMETRY_PROPORTION, options.MIN_FILLED_CELLS, options.NO_CACHE)
			//TODO: factor out all of this double-printing.
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))
			}
		} else if options.PUZZLE_TO_SOLVE != "" {
			//TODO: detect if the load failed.
			grid = sudoku.NewGrid()

			data, err := ioutil.ReadFile(options.PUZZLE_TO_SOLVE)

			if err != nil {
				log.Fatalln("Read error for specified file:", err)
			}

			//TODO: shouldn't a load method have a way to say the string provided is invalid?
			options.CONVERTER.Load(grid, string(data))
		}

		if grid == nil {
			//No grid to do anything with.
			log.Fatalln("No grid loaded.")
		}

		//TODO: use of this option leads to a busy loop somewhere... Is it related to the generate-multiple-and-difficulty hang?

		var directions *sudoku.SolveDirections

		if options.WALKTHROUGH || options.PRINT_STATS {
			directions = grid.HumanSolution(nil)
			if len(directions.Steps) == 0 {
				//We couldn't solve it. Let's check and see if the puzzle is well formed.
				if grid.HasMultipleSolutions() {
					//TODO: figure out why guesses wouldn't be used here effectively.
					log.Println("The puzzle had multiple solutions; that means it's not well-formed")
				}
			}
		}

		if options.WALKTHROUGH {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, directions.Walkthrough())
			} else {
				fmt.Fprintln(output, directions.Walkthrough())
			}
		}
		if options.PRINT_STATS {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, strconv.FormatFloat(grid.Difficulty(), 'f', -1, 64))
				//We won't print out the directions.Stats() like we do for just printing to stdout,
				//because that's mostly noise in this format.
			} else {
				fmt.Fprintln(output, grid.Difficulty())
				//TODO: consider actually printing out the Signals stats (with a Stats method on signals)
				fmt.Fprintln(output, strings.Join(directions.Stats(), "\n"))
			}
		}
		if options.PUZZLE_TO_SOLVE != "" {
			grid.Solve()
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))

			}
		}

		if options.OUTPUT_CSV {
			csvWriter.Write(csvRec)
		}

		if options.PUZZLE_TO_SOLVE != "" {
			//If we're asked to solve, n could only be 1 anyway.
			return
		}
		grid.Done()
		if options.NUM > 1 {
			bar.Incr()
		}
	}
	if options.OUTPUT_CSV {
		csvWriter.Flush()
	}

}

func puzzleDirectoryParts(symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) []string {
	return []string{
		STORED_PUZZLES_DIRECTORY,
		"SYM_TYPE_" + strconv.Itoa(int(symmetryType)),
		"SYM_PERCENTAGE_" + strconv.FormatFloat(symmetryPercentage, 'f', -1, 64),
		"MIN_FILED_CELLS_" + strconv.Itoa(minFilledCells),
	}
}

func storePuzzle(grid *sudoku.Grid, difficulty float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) bool {
	//TODO: we should include a hashed version of our difficulty weights file so we don't cache ones with old weights.
	directoryParts := puzzleDirectoryParts(symmetryType, symmetryPercentage, minFilledCells)

	fileNamePart := strconv.FormatFloat(difficulty, 'f', -1, 64) + ".sdk"

	pathSoFar := ""

	for i, part := range directoryParts {
		if i == 0 {
			pathSoFar = part
		} else {
			pathSoFar = filepath.Join(pathSoFar, part)
		}
		if _, err := os.Stat(pathSoFar); os.IsNotExist(err) {
			//need to create it.
			os.Mkdir(pathSoFar, 0700)
		}
	}

	fileName := filepath.Join(pathSoFar, fileNamePart)

	file, err := os.Create(fileName)

	if err != nil {
		log.Println(err)
		return false
	}

	defer file.Close()

	puzzleText := grid.DataString()

	n, err := io.WriteString(file, puzzleText)

	if err != nil {
		log.Println(err)
		return false
	} else {
		if n < len(puzzleText) {
			log.Println("Didn't write full file, only wrote", n, "bytes of", len(puzzleText))
			return false
		}
	}
	return true
}

func vendPuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) *sudoku.Grid {

	directory := filepath.Join(puzzleDirectoryParts(symmetryType, symmetryPercentage, minFilledCells)...)

	if files, err := ioutil.ReadDir(directory); os.IsNotExist(err) {
		//The directory doesn't exist.
		return nil
	} else {
		//OK, the directory exists, now see which puzzles are there and if any fit. If one does, vend it and delete the file.
		for _, file := range files {
			//See what this actually returns.
			filenameParts := strings.Split(file.Name(), ".")

			//Remember: there's a dot in the filename due to the float seperator.
			//TODO: shouldn't "sdk" be in a constant somewhere?
			if len(filenameParts) != 3 || filenameParts[2] != "sdk" {
				continue
			}

			difficulty, err := strconv.ParseFloat(strings.Join(filenameParts[0:2], "."), 64)
			if err != nil {
				continue
			}

			if min <= difficulty && difficulty <= max {
				//Found a puzzle!
				grid := sudoku.NewGrid()
				fullFileName := filepath.Join(directory, file.Name())
				grid.LoadFromFile(fullFileName)
				os.Remove(fullFileName)
				return grid
			}
		}
	}
	return nil
}

func generatePuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int, skipCache bool) *sudoku.Grid {
	var result *sudoku.Grid

	if !skipCache {
		result = vendPuzzle(min, max, symmetryType, symmetryPercentage, minFilledCells)

		if result != nil {
			log.Println("Vending a puzzle from the cache.")
			return result
		}
	}

	options := sudoku.GenerationOptions{
		Symmetry:           symmetryType,
		SymmetryPercentage: symmetryPercentage,
		MinFilledCells:     minFilledCells,
	}

	//We'll have to generate one ourselves.
	count := 0
	for {
		//The first time we don't bother saying what number attemp it is, because if the first run is likely to generate a useable puzzle it's just noise.
		if count != 0 {
			log.Println("Attempt", count, "at generating puzzle.")
		}

		result = sudoku.GenerateGrid(&options)

		difficulty := result.Difficulty()

		if difficulty >= min && difficulty <= max {
			return result
		}

		log.Println("Rejecting grid of difficulty", difficulty)
		if storePuzzle(result, difficulty, symmetryType, symmetryPercentage, minFilledCells) {
			log.Println("Stored the puzzle for future use.")
		}

		count++
	}
	return nil
}
<commit_msg>Factored out almost everything in main into process method, which does the real meat of the work<commit_after>/*
dokugen is a simple command line utility that exposes many of the basic functions of the
sudoku package. It's able to generate puzzles (with difficutly) and solve provided puzzles.
Run with -h to see help on how to use it.
*/
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"github.com/gosuri/uiprogress"
	"github.com/jkomoros/sudoku"
	"github.com/jkomoros/sudoku/sdkconverter"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
)

//TODO: let people pass in a filename to export to.

const STORED_PUZZLES_DIRECTORY = ".puzzles"

type appOptions struct {
	GENERATE            bool
	HELP                bool
	PUZZLE_TO_SOLVE     string
	NUM                 int
	PRINT_STATS         bool
	WALKTHROUGH         bool
	RAW_SYMMETRY        string
	RAW_DIFFICULTY      string
	SYMMETRY            sudoku.SymmetryType
	SYMMETRY_PROPORTION float64
	MIN_FILLED_CELLS    int
	MIN_DIFFICULTY      float64
	MAX_DIFFICULTY      float64
	NO_CACHE            bool
	PUZZLE_FORMAT       string
	OUTPUT_CSV          bool
	CONVERTER           sdkconverter.SudokuPuzzleConverter
}

var difficultyRanges map[string]struct {
	low, high float64
}

func init() {
	//grid.Difficulty can make use of a number of processes simultaneously.
	runtime.GOMAXPROCS(6)

	difficultyRanges = map[string]struct{ low, high float64 }{
		"gentle": {0.0, 0.3},
		"easy":   {0.3, 0.6},
		"medium": {0.6, 0.7},
		"tough":  {0.7, 1.0},
	}
}

func getOptions() *appOptions {

	var options appOptions

	flag.BoolVar(&options.GENERATE, "g", false, "if true, will generate a puzzle.")
	flag.BoolVar(&options.HELP, "h", false, "If provided, will print help and exit.")
	flag.IntVar(&options.NUM, "n", 1, "Number of things to generate")
	flag.BoolVar(&options.PRINT_STATS, "p", false, "If provided, will print stats.")
	flag.StringVar(&options.PUZZLE_TO_SOLVE, "s", "", "If provided, will solve the puzzle at the given filename and print solution.")
	flag.BoolVar(&options.WALKTHROUGH, "w", false, "If provided, will print out a walkthrough to solve the provided puzzle.")
	flag.StringVar(&options.RAW_SYMMETRY, "y", "vertical", "Valid values: 'none', 'both', 'horizontal', 'vertical")
	flag.Float64Var(&options.SYMMETRY_PROPORTION, "r", 0.7, "What proportion of cells should be filled according to symmetry")
	flag.IntVar(&options.MIN_FILLED_CELLS, "min-filled-cells", 0, "The minimum number of cells that should be filled in the generated puzzles.")
	flag.Float64Var(&options.MIN_DIFFICULTY, "min", 0.0, "Minimum difficulty for generated puzzle")
	flag.Float64Var(&options.MAX_DIFFICULTY, "max", 1.0, "Maximum difficulty for generated puzzle")
	flag.BoolVar(&options.NO_CACHE, "no-cache", false, "If provided, will not vend generated puzzles from the cache of previously generated puzzles.")
	//TODO: the format should also be how we interpret loads, too.
	flag.StringVar(&options.PUZZLE_FORMAT, "format", "sdk", "Which format to export puzzles from. Defaults to 'sdk'")
	flag.BoolVar(&options.OUTPUT_CSV, "csv", false, "Output the results in CSV.")
	flag.StringVar(&options.RAW_DIFFICULTY, "d", "", "difficulty, one of {gentle, easy, medium, tough}")
	flag.Parse()

	options.RAW_SYMMETRY = strings.ToLower(options.RAW_SYMMETRY)
	switch options.RAW_SYMMETRY {
	case "none":
		options.SYMMETRY = sudoku.SYMMETRY_NONE
	case "both":
		options.SYMMETRY = sudoku.SYMMETRY_BOTH
	case "horizontal":
		options.SYMMETRY = sudoku.SYMMETRY_HORIZONTAL
	case "vertical":
		options.SYMMETRY = sudoku.SYMMETRY_VERTICAL
	default:
		log.Fatal("Unknown symmetry flag: ", options.RAW_SYMMETRY)
	}

	options.RAW_DIFFICULTY = strings.ToLower(options.RAW_DIFFICULTY)
	if options.RAW_DIFFICULTY != "" {
		vals, ok := difficultyRanges[options.RAW_DIFFICULTY]
		if !ok {
			log.Fatal("Invalid difficulty option:", options.RAW_DIFFICULTY)
		}
		options.MIN_DIFFICULTY = vals.low
		options.MAX_DIFFICULTY = vals.high
		log.Println("Using difficulty max:", strconv.FormatFloat(vals.high, 'f', -1, 64), "min:", strconv.FormatFloat(vals.low, 'f', -1, 64))
	}

	options.CONVERTER = sdkconverter.Converters[options.PUZZLE_FORMAT]

	if options.CONVERTER == nil {
		log.Fatal("Invalid format option:", options.PUZZLE_FORMAT)
	}

	return &options
}

func main() {

	//TODO: figure out how to test this.
	process(getOptions())
}

func process(options *appOptions) {
	output := os.Stdout

	if options.HELP {
		flag.PrintDefaults()
		return
	}

	var grid *sudoku.Grid

	var csvWriter *csv.Writer
	var csvRec []string

	if options.OUTPUT_CSV {
		csvWriter = csv.NewWriter(output)
	}

	var bar *uiprogress.Bar

	//TODO: do more useful / explanatory printing here.
	if options.NUM > 1 {
		uiprogress.DefaultProgress.Out = os.Stderr
		uiprogress.Start()
		bar = uiprogress.AddBar(options.NUM).PrependElapsed().AppendCompleted()
	}

	for i := 0; i < options.NUM; i++ {

		if options.OUTPUT_CSV {
			csvRec = nil
		}

		//TODO: allow the type of symmetry to be configured.
		if options.GENERATE {
			grid = generatePuzzle(options.MIN_DIFFICULTY, options.MAX_DIFFICULTY, options.SYMMETRY, options.SYMMETRY_PROPORTION, options.MIN_FILLED_CELLS, options.NO_CACHE)
			//TODO: factor out all of this double-printing.
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))
			}
		} else if options.PUZZLE_TO_SOLVE != "" {
			//TODO: detect if the load failed.
			grid = sudoku.NewGrid()

			data, err := ioutil.ReadFile(options.PUZZLE_TO_SOLVE)

			if err != nil {
				log.Fatalln("Read error for specified file:", err)
			}

			//TODO: shouldn't a load method have a way to say the string provided is invalid?
			options.CONVERTER.Load(grid, string(data))
		}

		if grid == nil {
			//No grid to do anything with.
			log.Fatalln("No grid loaded.")
		}

		//TODO: use of this option leads to a busy loop somewhere... Is it related to the generate-multiple-and-difficulty hang?

		var directions *sudoku.SolveDirections

		if options.WALKTHROUGH || options.PRINT_STATS {
			directions = grid.HumanSolution(nil)
			if len(directions.Steps) == 0 {
				//We couldn't solve it. Let's check and see if the puzzle is well formed.
				if grid.HasMultipleSolutions() {
					//TODO: figure out why guesses wouldn't be used here effectively.
					log.Println("The puzzle had multiple solutions; that means it's not well-formed")
				}
			}
		}

		if options.WALKTHROUGH {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, directions.Walkthrough())
			} else {
				fmt.Fprintln(output, directions.Walkthrough())
			}
		}
		if options.PRINT_STATS {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, strconv.FormatFloat(grid.Difficulty(), 'f', -1, 64))
				//We won't print out the directions.Stats() like we do for just printing to stdout,
				//because that's mostly noise in this format.
			} else {
				fmt.Fprintln(output, grid.Difficulty())
				//TODO: consider actually printing out the Signals stats (with a Stats method on signals)
				fmt.Fprintln(output, strings.Join(directions.Stats(), "\n"))
			}
		}
		if options.PUZZLE_TO_SOLVE != "" {
			grid.Solve()
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))

			}
		}

		if options.OUTPUT_CSV {
			csvWriter.Write(csvRec)
		}

		if options.PUZZLE_TO_SOLVE != "" {
			//If we're asked to solve, n could only be 1 anyway.
			return
		}
		grid.Done()
		if options.NUM > 1 {
			bar.Incr()
		}
	}
	if options.OUTPUT_CSV {
		csvWriter.Flush()
	}
}

func puzzleDirectoryParts(symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) []string {
	return []string{
		STORED_PUZZLES_DIRECTORY,
		"SYM_TYPE_" + strconv.Itoa(int(symmetryType)),
		"SYM_PERCENTAGE_" + strconv.FormatFloat(symmetryPercentage, 'f', -1, 64),
		"MIN_FILED_CELLS_" + strconv.Itoa(minFilledCells),
	}
}

func storePuzzle(grid *sudoku.Grid, difficulty float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) bool {
	//TODO: we should include a hashed version of our difficulty weights file so we don't cache ones with old weights.
	directoryParts := puzzleDirectoryParts(symmetryType, symmetryPercentage, minFilledCells)

	fileNamePart := strconv.FormatFloat(difficulty, 'f', -1, 64) + ".sdk"

	pathSoFar := ""

	for i, part := range directoryParts {
		if i == 0 {
			pathSoFar = part
		} else {
			pathSoFar = filepath.Join(pathSoFar, part)
		}
		if _, err := os.Stat(pathSoFar); os.IsNotExist(err) {
			//need to create it.
			os.Mkdir(pathSoFar, 0700)
		}
	}

	fileName := filepath.Join(pathSoFar, fileNamePart)

	file, err := os.Create(fileName)

	if err != nil {
		log.Println(err)
		return false
	}

	defer file.Close()

	puzzleText := grid.DataString()

	n, err := io.WriteString(file, puzzleText)

	if err != nil {
		log.Println(err)
		return false
	} else {
		if n < len(puzzleText) {
			log.Println("Didn't write full file, only wrote", n, "bytes of", len(puzzleText))
			return false
		}
	}
	return true
}

func vendPuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int) *sudoku.Grid {

	directory := filepath.Join(puzzleDirectoryParts(symmetryType, symmetryPercentage, minFilledCells)...)

	if files, err := ioutil.ReadDir(directory); os.IsNotExist(err) {
		//The directory doesn't exist.
		return nil
	} else {
		//OK, the directory exists, now see which puzzles are there and if any fit. If one does, vend it and delete the file.
		for _, file := range files {
			//See what this actually returns.
			filenameParts := strings.Split(file.Name(), ".")

			//Remember: there's a dot in the filename due to the float seperator.
			//TODO: shouldn't "sdk" be in a constant somewhere?
			if len(filenameParts) != 3 || filenameParts[2] != "sdk" {
				continue
			}

			difficulty, err := strconv.ParseFloat(strings.Join(filenameParts[0:2], "."), 64)
			if err != nil {
				continue
			}

			if min <= difficulty && difficulty <= max {
				//Found a puzzle!
				grid := sudoku.NewGrid()
				fullFileName := filepath.Join(directory, file.Name())
				grid.LoadFromFile(fullFileName)
				os.Remove(fullFileName)
				return grid
			}
		}
	}
	return nil
}

func generatePuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, minFilledCells int, skipCache bool) *sudoku.Grid {
	var result *sudoku.Grid

	if !skipCache {
		result = vendPuzzle(min, max, symmetryType, symmetryPercentage, minFilledCells)

		if result != nil {
			log.Println("Vending a puzzle from the cache.")
			return result
		}
	}

	options := sudoku.GenerationOptions{
		Symmetry:           symmetryType,
		SymmetryPercentage: symmetryPercentage,
		MinFilledCells:     minFilledCells,
	}

	//We'll have to generate one ourselves.
	count := 0
	for {
		//The first time we don't bother saying what number attemp it is, because if the first run is likely to generate a useable puzzle it's just noise.
		if count != 0 {
			log.Println("Attempt", count, "at generating puzzle.")
		}

		result = sudoku.GenerateGrid(&options)

		difficulty := result.Difficulty()

		if difficulty >= min && difficulty <= max {
			return result
		}

		log.Println("Rejecting grid of difficulty", difficulty)
		if storePuzzle(result, difficulty, symmetryType, symmetryPercentage, minFilledCells) {
			log.Println("Stored the puzzle for future use.")
		}

		count++
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Google Inc.
//
// 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 main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"path/filepath"

	"github.com/golang/protobuf/proto"
	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/data/pack"
	"github.com/google/gapid/core/log"
)

type unpackVerb struct{ UnpackFlags }

func init() {
	verb := &unpackVerb{}
	app.AddVerb(&app.Verb{
		Name:      "unpack",
		ShortHelp: "Displays the raw protos in a protopack file",
		Action:    verb,
	})
}

func (verb *unpackVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one protopack file expected, got %d", flags.NArg())
		return nil
	}

	filepath, err := filepath.Abs(flags.Arg(0))
	ctx = log.V{"filepath": filepath}.Bind(ctx)
	if err != nil {
		return log.Err(ctx, err, "Could not find capture file")
	}

	r, err := os.Open(filepath)
	if err != nil {
		return err
	}
	defer r.Close()

	return pack.Read(ctx, r, unpacker{})
}

type unpacker struct{}

func (unpacker) BeginGroup(ctx context.Context, msg proto.Message, id uint64) error {
	log.I(ctx, "BeginGroup(msg: %v, id: %v)", msgString(msg), id)
	return nil
}
func (unpacker) BeginChildGroup(ctx context.Context, msg proto.Message, id, parentID uint64) error {
	log.I(ctx, "BeginChildGroup(msg: %v, id: %v, parentID: %v)", msgString(msg), id, parentID)
	return nil
}
func (unpacker) EndGroup(ctx context.Context, id uint64) error {
	log.I(ctx, "EndGroup(id: %v)", id)
	return nil
}
func (unpacker) Object(ctx context.Context, msg proto.Message) error {
	log.I(ctx, "Object(msg: %v)", msgString(msg))
	return nil
}
func (unpacker) ChildObject(ctx context.Context, msg proto.Message, parentID uint64) error {
	log.I(ctx, "ChildObject(msg: %v, parentID: %v)", msgString(msg), parentID)
	return nil
}

func msgString(msg proto.Message) string {
	if msg, ok := msg.(*pack.Dynamic); ok {
		return fmt.Sprintf("%+v", msg)
	}
	return fmt.Sprintf("%T{%+v}", msg, msg)
}
<commit_msg>cmd/gapit/unpack: Truncate stupid long messages.<commit_after>// Copyright (C) 2017 Google Inc.
//
// 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 main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"path/filepath"

	"github.com/golang/protobuf/proto"
	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/data/pack"
	"github.com/google/gapid/core/log"
)

type unpackVerb struct{ UnpackFlags }

func init() {
	verb := &unpackVerb{}
	app.AddVerb(&app.Verb{
		Name:      "unpack",
		ShortHelp: "Displays the raw protos in a protopack file",
		Action:    verb,
	})
}

func (verb *unpackVerb) Run(ctx context.Context, flags flag.FlagSet) error {
	if flags.NArg() != 1 {
		app.Usage(ctx, "Exactly one protopack file expected, got %d", flags.NArg())
		return nil
	}

	filepath, err := filepath.Abs(flags.Arg(0))
	ctx = log.V{"filepath": filepath}.Bind(ctx)
	if err != nil {
		return log.Err(ctx, err, "Could not find capture file")
	}

	r, err := os.Open(filepath)
	if err != nil {
		return err
	}
	defer r.Close()

	return pack.Read(ctx, r, unpacker{})
}

type unpacker struct{}

func (unpacker) BeginGroup(ctx context.Context, msg proto.Message, id uint64) error {
	log.I(ctx, "BeginGroup(msg: %v, id: %v)", msgString(msg), id)
	return nil
}
func (unpacker) BeginChildGroup(ctx context.Context, msg proto.Message, id, parentID uint64) error {
	log.I(ctx, "BeginChildGroup(msg: %v, id: %v, parentID: %v)", msgString(msg), id, parentID)
	return nil
}
func (unpacker) EndGroup(ctx context.Context, id uint64) error {
	log.I(ctx, "EndGroup(id: %v)", id)
	return nil
}
func (unpacker) Object(ctx context.Context, msg proto.Message) error {
	log.I(ctx, "Object(msg: %v)", msgString(msg))
	return nil
}
func (unpacker) ChildObject(ctx context.Context, msg proto.Message, parentID uint64) error {
	log.I(ctx, "ChildObject(msg: %v, parentID: %v)", msgString(msg), parentID)
	return nil
}

func msgString(msg proto.Message) string {
	var str string
	switch msg := msg.(type) {
	case *pack.Dynamic:
		str = fmt.Sprintf("%+v", msg)
	default:
		str = fmt.Sprintf("%T{%+v}", msg, msg)
	}
	if len(str) > 100 {
		str = str[:97] + "..." // TODO: Consider unicode.
	}
	return str
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 Couchbase, Inc.
//  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 main

import (
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"testing"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
)

func TestRunRebalancer(t *testing.T) {
	testDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(testDir)

	log.Printf("testDir: %s", testDir)

	tests := []struct {
		label      string
		ops        string // Space separated "+a", "-x".
		params     map[string]string
		expNodes   string // Space separated list of nodes ("a"..."v").
		expIndexes string // Space separated list of indxes ("x"..."z").
		expChanged bool
		expErr     bool
	}{
		{"1st node",
			"+a", nil,
			"a",
			"",
			false, true,
		},
		{"add 1st index x",
			"+x", nil,
			"a",
			"x",
			true, false,
		},
		{"add 2nd node b",
			"+b", nil,
			"a b",
			"x",
			true, false,
		},
	}

	cfg := cbgt.NewCfgMem()

	mgrs := map[string]*cbgt.Manager{}

	var mgr0 *cbgt.Manager

	server := "."

	waitUntilEmptyCfgEvents := func(ch chan cbgt.CfgEvent) {
		for {
			select {
			case <-ch:
			default:
				return
			}
		}
	}

	cfgEventsNodeDefsWanted := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.NODE_DEFS_WANTED, cfgEventsNodeDefsWanted)

	waitUntilEmptyCfgEventsNodeDefsWanted := func() {
		waitUntilEmptyCfgEvents(cfgEventsNodeDefsWanted)
	}

	cfgEventsIndexDefs := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.INDEX_DEFS_KEY, cfgEventsIndexDefs)

	waitUntilEmptyCfgEventsIndexDefs := func() {
		waitUntilEmptyCfgEvents(cfgEventsIndexDefs)
	}

	for testi, test := range tests {
		log.Printf("testi: %d, label: %q", testi, test.label)

		for opi, op := range strings.Split(test.ops, " ") {
			log.Printf(" opi: %d, op: %s", opi, op)

			name := op[1:2]

			isIndexOp := name >= "x"
			if isIndexOp {
				indexName := name
				log.Printf(" indexOp: %s, indexName: %s", op[0:1], indexName)

				sourceType := "primary"
				if test.params["sourceType"] != "" {
					sourceType = test.params["sourceType"]
				}
				if test.params[indexName+".sourceType"] != "" {
					sourceType = test.params[indexName+".sourceType"]
				}

				sourceName := "default"
				if test.params["sourceName"] != "" {
					sourceName = test.params["sourceName"]
				}
				if test.params[indexName+".sourceName"] != "" {
					sourceName = test.params[indexName+".sourceName"]
				}

				sourceUUID := ""
				if test.params["sourceUUID"] != "" {
					sourceUUID = test.params["sourceUUID"]
				}
				if test.params[indexName+".sourceUUID"] != "" {
					sourceUUID = test.params[indexName+".sourceUUID"]
				}

				sourceParams := ""
				if test.params["sourceParams"] != "" {
					sourceParams = test.params["sourceParams"]
				}
				if test.params[indexName+".sourceParams"] != "" {
					sourceParams = test.params[indexName+".sourceParams"]
				}

				indexType := "blackhole"
				if test.params["indexType"] != "" {
					indexType = test.params["indexType"]
				}
				if test.params[indexName+".indexType"] != "" {
					indexType = test.params[indexName+".indexType"]
				}

				indexParams := ""
				if test.params["indexParams"] != "" {
					indexParams = test.params["indexParams"]
				}
				if test.params[indexName+".indexParams"] != "" {
					indexParams = test.params[indexName+".indexParams"]
				}

				prevIndexUUID := ""
				if test.params["prevIndexUUID"] != "" {
					prevIndexUUID = test.params["prevIndexUUID"]
				}
				if test.params[indexName+".prevIndexUUID"] != "" {
					prevIndexUUID = test.params[indexName+".prevIndexUUID"]
				}

				planParams := cbgt.PlanParams{}

				waitUntilEmptyCfgEventsIndexDefs()

				err := mgr0.CreateIndex(
					sourceType, sourceName, sourceUUID, sourceParams,
					indexType, indexName, indexParams,
					planParams,
					prevIndexUUID)
				if err != nil {
					t.Errorf("expected no err, got: %#v", err)
				}

				waitUntilEmptyCfgEventsIndexDefs()
			} else { // It's a node op.
				nodeName := name
				log.Printf(" nodeOp: %s, nodeName: %s", op[0:1], nodeName)

				register := "wanted"
				if op[0:1] == "-" {
					register = "unknown"
				}
				if test.params["register"] != "" {
					register = test.params["register"]
				}
				if test.params[nodeName+".register"] != "" {
					register = test.params[nodeName+".register"]
				}

				if mgrs[nodeName] != nil {
					mgrs[nodeName].Stop()
					delete(mgrs, nodeName)
				}

				waitUntilEmptyCfgEventsNodeDefsWanted()

				mgr, err := startNodeManager(testDir, cfg,
					name, register, test.params, server)
				if err != nil || mgr == nil {
					t.Errorf("expected no err, got: %#v", err)
				}
				if mgr0 == nil {
					mgr0 = mgr
				}

				if register != "unknown" {
					mgrs[nodeName] = mgr
				}

				mgr.Kick("kick")

				waitUntilEmptyCfgEventsNodeDefsWanted()
			}
		}

		changed, err := runRebalancer(cbgt.VERSION, cfg, ".")
		if changed != test.expChanged {
			t.Errorf("testi: %d, label: %q,"+
				" expChanged: %v, but got: %v",
				testi, test.label,
				test.expChanged, changed)
		}
		if (test.expErr && err == nil) ||
			(!test.expErr && err != nil) {
			t.Errorf("testi: %d, label: %q,"+
				" expErr: %v, but got: %v",
				testi, test.label,
				test.expErr, err)
		}
	}
}

func startNodeManager(testDir string, cfg cbgt.Cfg, node, register string,
	params map[string]string, server string) (
	mgr *cbgt.Manager, err error) {
	uuid := node
	if params["uuid"] != "" {
		uuid = params["uuid"]
	}
	if params[node+".uuid"] != "" {
		uuid = params[node+".uuid"]
	}

	// No planner in tags because mcp provides the planner.
	tags := []string{"feed", "pindex", "janitor", "queryer"}
	if params["tags"] != "" {
		tags = strings.Split(params["tags"], ",")
	}
	if params[node+".tags"] != "" {
		tags = strings.Split(params[node+".tags"], ",")
	}

	container := ""
	if params["container"] != "" {
		container = params["container"]
	}
	if params[node+".container"] != "" {
		container = params[node+".container"]
	}

	weight := 1
	if params["weight"] != "" {
		weight, err = strconv.Atoi(params["weight"])
	}
	if params[node+".weight"] != "" {
		weight, err = strconv.Atoi(params[node+".weight"])
	}
	if weight < 1 {
		weight = 1
	}

	extras := ""

	bindHttp := node

	dataDir := testDir + string(os.PathSeparator) + node

	os.MkdirAll(dataDir, 0700)

	meh := cbgt.ManagerEventHandlers(nil)

	mgr = cbgt.NewManager(cbgt.VERSION, cfg, uuid,
		tags, container, weight, extras,
		bindHttp, dataDir, server, meh)

	err = mgr.Start(register)
	if err != nil {
		mgr.Stop()

		return nil, err
	}

	return mgr, nil
}
<commit_msg>testing with 4 partitions in mcp_test<commit_after>//  Copyright (c) 2015 Couchbase, Inc.
//  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 main

import (
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"testing"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
)

func TestRunRebalancer(t *testing.T) {
	testDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(testDir)

	log.Printf("testDir: %s", testDir)

	tests := []struct {
		label      string
		ops        string // Space separated "+a", "-x".
		params     map[string]string
		expNodes   string // Space separated list of nodes ("a"..."v").
		expIndexes string // Space separated list of indxes ("x"..."z").
		expChanged bool
		expErr     bool
	}{
		{"1st node",
			"+a", nil,
			"a",
			"",
			false, true,
		},
		{"add 1st index x",
			"+x", nil,
			"a",
			"x",
			true, false,
		},
		{"add 2nd node b",
			"+b", nil,
			"a b",
			"x",
			true, false,
		},
	}

	cfg := cbgt.NewCfgMem()

	mgrs := map[string]*cbgt.Manager{}

	var mgr0 *cbgt.Manager

	server := "."

	waitUntilEmptyCfgEvents := func(ch chan cbgt.CfgEvent) {
		for {
			select {
			case <-ch:
			default:
				return
			}
		}
	}

	cfgEventsNodeDefsWanted := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.NODE_DEFS_WANTED, cfgEventsNodeDefsWanted)

	waitUntilEmptyCfgEventsNodeDefsWanted := func() {
		waitUntilEmptyCfgEvents(cfgEventsNodeDefsWanted)
	}

	cfgEventsIndexDefs := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.INDEX_DEFS_KEY, cfgEventsIndexDefs)

	waitUntilEmptyCfgEventsIndexDefs := func() {
		waitUntilEmptyCfgEvents(cfgEventsIndexDefs)
	}

	for testi, test := range tests {
		log.Printf("testi: %d, label: %q", testi, test.label)

		for opi, op := range strings.Split(test.ops, " ") {
			log.Printf(" opi: %d, op: %s", opi, op)

			name := op[1:2]

			isIndexOp := name >= "x"
			if isIndexOp {
				indexName := name
				log.Printf(" indexOp: %s, indexName: %s", op[0:1], indexName)

				sourceType := "primary"
				if test.params["sourceType"] != "" {
					sourceType = test.params["sourceType"]
				}
				if test.params[indexName+".sourceType"] != "" {
					sourceType = test.params[indexName+".sourceType"]
				}

				sourceName := "default"
				if test.params["sourceName"] != "" {
					sourceName = test.params["sourceName"]
				}
				if test.params[indexName+".sourceName"] != "" {
					sourceName = test.params[indexName+".sourceName"]
				}

				sourceUUID := ""
				if test.params["sourceUUID"] != "" {
					sourceUUID = test.params["sourceUUID"]
				}
				if test.params[indexName+".sourceUUID"] != "" {
					sourceUUID = test.params[indexName+".sourceUUID"]
				}

				sourceParams := `{"numPartitions":4}`
				if test.params["sourceParams"] != "" {
					sourceParams = test.params["sourceParams"]
				}
				if test.params[indexName+".sourceParams"] != "" {
					sourceParams = test.params[indexName+".sourceParams"]
				}

				indexType := "blackhole"
				if test.params["indexType"] != "" {
					indexType = test.params["indexType"]
				}
				if test.params[indexName+".indexType"] != "" {
					indexType = test.params[indexName+".indexType"]
				}

				indexParams := ""
				if test.params["indexParams"] != "" {
					indexParams = test.params["indexParams"]
				}
				if test.params[indexName+".indexParams"] != "" {
					indexParams = test.params[indexName+".indexParams"]
				}

				prevIndexUUID := ""
				if test.params["prevIndexUUID"] != "" {
					prevIndexUUID = test.params["prevIndexUUID"]
				}
				if test.params[indexName+".prevIndexUUID"] != "" {
					prevIndexUUID = test.params[indexName+".prevIndexUUID"]
				}

				planParams := cbgt.PlanParams{
					MaxPartitionsPerPIndex: 1,
				}

				waitUntilEmptyCfgEventsIndexDefs()

				err := mgr0.CreateIndex(
					sourceType, sourceName, sourceUUID, sourceParams,
					indexType, indexName, indexParams,
					planParams,
					prevIndexUUID)
				if err != nil {
					t.Errorf("expected no err, got: %#v", err)
				}

				waitUntilEmptyCfgEventsIndexDefs()
			} else { // It's a node op.
				nodeName := name
				log.Printf(" nodeOp: %s, nodeName: %s", op[0:1], nodeName)

				register := "wanted"
				if op[0:1] == "-" {
					register = "unknown"
				}
				if test.params["register"] != "" {
					register = test.params["register"]
				}
				if test.params[nodeName+".register"] != "" {
					register = test.params[nodeName+".register"]
				}

				if mgrs[nodeName] != nil {
					mgrs[nodeName].Stop()
					delete(mgrs, nodeName)
				}

				waitUntilEmptyCfgEventsNodeDefsWanted()

				mgr, err := startNodeManager(testDir, cfg,
					name, register, test.params, server)
				if err != nil || mgr == nil {
					t.Errorf("expected no err, got: %#v", err)
				}
				if mgr0 == nil {
					mgr0 = mgr
				}

				if register != "unknown" {
					mgrs[nodeName] = mgr
				}

				mgr.Kick("kick")

				waitUntilEmptyCfgEventsNodeDefsWanted()
			}
		}

		changed, err := runRebalancer(cbgt.VERSION, cfg, ".")
		if changed != test.expChanged {
			t.Errorf("testi: %d, label: %q,"+
				" expChanged: %v, but got: %v",
				testi, test.label,
				test.expChanged, changed)
		}
		if (test.expErr && err == nil) ||
			(!test.expErr && err != nil) {
			t.Errorf("testi: %d, label: %q,"+
				" expErr: %v, but got: %v",
				testi, test.label,
				test.expErr, err)
		}
	}
}

func startNodeManager(testDir string, cfg cbgt.Cfg, node, register string,
	params map[string]string, server string) (
	mgr *cbgt.Manager, err error) {
	uuid := node
	if params["uuid"] != "" {
		uuid = params["uuid"]
	}
	if params[node+".uuid"] != "" {
		uuid = params[node+".uuid"]
	}

	// No planner in tags because mcp provides the planner.
	tags := []string{"feed", "pindex", "janitor", "queryer"}
	if params["tags"] != "" {
		tags = strings.Split(params["tags"], ",")
	}
	if params[node+".tags"] != "" {
		tags = strings.Split(params[node+".tags"], ",")
	}

	container := ""
	if params["container"] != "" {
		container = params["container"]
	}
	if params[node+".container"] != "" {
		container = params[node+".container"]
	}

	weight := 1
	if params["weight"] != "" {
		weight, err = strconv.Atoi(params["weight"])
	}
	if params[node+".weight"] != "" {
		weight, err = strconv.Atoi(params[node+".weight"])
	}
	if weight < 1 {
		weight = 1
	}

	extras := ""

	bindHttp := node

	dataDir := testDir + string(os.PathSeparator) + node

	os.MkdirAll(dataDir, 0700)

	meh := cbgt.ManagerEventHandlers(nil)

	mgr = cbgt.NewManager(cbgt.VERSION, cfg, uuid,
		tags, container, weight, extras,
		bindHttp, dataDir, server, meh)

	err = mgr.Start(register)
	if err != nil {
		mgr.Stop()

		return nil, err
	}

	return mgr, nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestNewRootCmd(t *testing.T) {
	cmd := NewRootCmd()
	assert.Equal(t, "mockery", cmd.Name())
}
<commit_msg>test: add test for env var configurations<commit_after>package cmd

import (
	"fmt"
	"os"
	"testing"

	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/vektra/mockery/v2/pkg/config"
)

func TestNewRootCmd(t *testing.T) {
	cmd := NewRootCmd()
	assert.Equal(t, "mockery", cmd.Name())
}

func TestConfigEnvFlags(t *testing.T) {
	expected := config.Config{
		Config:               "my_file.yaml",
		Name:                 "SomeInterface",
		Print:                true,
		Output:               "/some/dir",
		Outpkg:               "some/package",
		Packageprefix:        "prefix_",
		Dir:                  "dir/to/search",
		Recursive:            true,
		All:                  true,
		InPackage:            true,
		TestOnly:             true,
		Case:                 "underscore",
		Note:                 "// this is a test",
		Cpuprofile:           "test.pprof",
		Version:              true,
		KeepTree:             true,
		BuildTags:            "test mock",
		FileName:             "my-file.go",
		StructName:           "Interface1",
		LogLevel:             "warn",
		SrcPkg:               "some/other/package",
		DryRun:               true,
		DisableVersionString: true,
		BoilerplateFile:      "some/file",
		UnrollVariadic:       false,
		Exported:             true,
		WithExpecter:         true,
	}

	env(t, "CONFIG", expected.Config)
	env(t, "NAME", expected.Name)
	env(t, "PRINT", fmt.Sprint(expected.Print))
	env(t, "OUTPUT", expected.Output)
	env(t, "OUTPKG", expected.Outpkg)
	env(t, "PACKAGEPREFIX", expected.Packageprefix)
	env(t, "DIR", expected.Dir)
	env(t, "RECURSIVE", fmt.Sprint(expected.Recursive))
	env(t, "ALL", fmt.Sprint(expected.All))
	env(t, "INPACKAGE", fmt.Sprint(expected.InPackage))
	env(t, "TESTONLY", fmt.Sprint(expected.TestOnly))
	env(t, "CASE", expected.Case)
	env(t, "NOTE", expected.Note)
	env(t, "CPUPROFILE", expected.Cpuprofile)
	env(t, "VERSION", fmt.Sprint(expected.Version))
	env(t, "QUIET", fmt.Sprint(expected.Quiet))
	env(t, "KEEPTREE", fmt.Sprint(expected.KeepTree))
	env(t, "TAGS", expected.BuildTags)
	env(t, "FILENAME", expected.FileName)
	env(t, "STRUCTNAME", expected.StructName)
	env(t, "LOG_LEVEL", expected.LogLevel)
	env(t, "SRCPKG", expected.SrcPkg)
	env(t, "DRY_RUN", fmt.Sprint(expected.DryRun))
	env(t, "DISABLE_VERSION_STRING", fmt.Sprint(expected.DisableVersionString))
	env(t, "BOILERPLATE_FILE", expected.BoilerplateFile)
	env(t, "UNROLL_VARIADIC", fmt.Sprint(expected.UnrollVariadic))
	env(t, "EXPORTED", fmt.Sprint(expected.Exported))
	env(t, "WITH_EXPECTER", fmt.Sprint(expected.WithExpecter))

	initConfig()

	app, err := GetRootAppFromViper(viper.GetViper())
	require.NoError(t, err)

	assert.Equal(t, expected, app.Config)
}

func env(t *testing.T, key, value string) {
	key = "MOCKERY_" + key
	t.Cleanup(func() { os.Unsetenv(key) })
	os.Setenv(key, value)
}
<|endoftext|>


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

// Monitor is a tool designed to occasionally poll an active cluster and save
// the status to disk. The monitor program will exit if the status of the
// cluster can not be determined.

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/retry"
	"github.com/cockroachdb/cockroach/util/stop"
)

const (
	// urlPath is the http path of the status endpoint at each monitored address.
	urlPath = "health"
)

var interval = flag.Duration("interval", 10*time.Second, "Interval in which to poll the status of each monitored address.")
var addrs = flag.String("addrs", ":26257", "Comma-separated list of host:port addresses to monitor.")
var insecure = flag.Bool("insecure", false, "True if using an insecure connection.")
var user = flag.String("user", security.RootUser, "User used to connect to the cluster.")
var certs = flag.String("certs", "certs", "Directory containing RSA key and x509 certs. This flag is required if --insecure=false.")

var retryOptions = retry.Options{
	InitialBackoff: 100 * time.Millisecond,
	MaxRetries:     10,
	Multiplier:     2,
}

type statusMonitor struct {
	addr       string
	url        string
	httpClient *http.Client
	file       *io.Writer
}

func newStatusMonitor(context *base.Context, addr string) (*statusMonitor, error) {
	monitor := &statusMonitor{
		addr: addr,
	}
	var err error
	monitor.httpClient, err = context.GetHTTPClient()
	if err != nil {
		return nil, err
	}
	monitor.url = fmt.Sprintf("%s://%s/%s", context.HTTPRequestScheme(), monitor.addr, urlPath)
	return monitor, nil
}

func (m *statusMonitor) queryStatus() error {
	var queryErr error
	for r := retry.Start(retryOptions); r.Next(); {
		if log.V(1) && queryErr != nil {
			log.Infof("retrying after error: %s", queryErr)
		}

		// Construct a new HTTP GET Request.
		req, err := http.NewRequest("GET", m.url, nil)
		if err != nil {
			queryErr = fmt.Errorf("could not create http request for %s: %s", m.url, err)
			// Break immediately, this is not recoverable.
			break
		}
		req.Header.Set(util.AcceptHeader, util.JSONContentType)

		// Execute request.
		resp, err := m.httpClient.Do(req)
		if err != nil {
			queryErr = fmt.Errorf("could not GET %s - %s", m.url, err)
			continue
		}
		defer resp.Body.Close()

		// Read and verify body of response.
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			queryErr = fmt.Errorf("could not ready body for %s - %s", m.url, err)
			continue
		}
		if resp.StatusCode != http.StatusOK {
			queryErr = fmt.Errorf("could not GET %s - statuscode: %d - body: %s", m.url, resp.StatusCode, body)
			continue
		}
		returnedContentType := resp.Header.Get(util.ContentTypeHeader)
		if returnedContentType != util.JSONContentType {
			queryErr = fmt.Errorf("unexpected content type: %v", returnedContentType)
			continue
		}
		return nil
	}
	return queryErr
}

func main() {
	flag.Parse()
	parsedAddrs := strings.Split(*addrs, ",")

	ctx := base.Context{Insecure: *insecure, Certs: *certs, User: *user}

	startTime := time.Now()
	stopper := stop.NewStopper()
	for _, addr := range parsedAddrs {
		client, err := newStatusMonitor(&ctx, addr)
		if err != nil {
			log.Errorf("error creating client: %s", err)
			return
		}
		log.Infof("Monitoring Status URL: %s", client.url)
		stopper.RunWorker(func() {
			timer := time.Tick(*interval)
			for {
				select {
				case <-stopper.ShouldStop():
					return
				case <-timer:
					elapsed := time.Since(startTime)
					if err := client.queryStatus(); err != nil {
						log.Warningf("Could not get status from url %s. Time since monitor started %s.", client.url, elapsed)
						stopper.Stop()
						continue
					}
					log.Infof("Got status from url %s. Time since start: %s", client.url, elapsed)
				}
			}
		})
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, os.Kill, syscall.SIGTERM)
	// Block until a termination signal is received, or the stopper is closed by
	// an error in one of the client routines.
	select {
	case <-stopper.ShouldStop():
		log.Infof("Monitor stopped by error...")
	case <-signalCh:
		log.Infof("Stopping status monitor...")
		stopper.Stop()
	}
}
<commit_msg>Exit with non-zero code<commit_after>// Copyright 2015 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Bram Gruneir

// Monitor is a tool designed to occasionally poll an active cluster and save
// the status to disk. The monitor program will exit if the status of the
// cluster can not be determined.

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/cockroachdb/cockroach/base"
	"github.com/cockroachdb/cockroach/security"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/retry"
	"github.com/cockroachdb/cockroach/util/stop"
)

const (
	// urlPath is the http path of the status endpoint at each monitored address.
	urlPath = "health"
)

var interval = flag.Duration("interval", 10*time.Second, "Interval in which to poll the status of each monitored address.")
var addrs = flag.String("addrs", ":26257", "Comma-separated list of host:port addresses to monitor.")
var insecure = flag.Bool("insecure", false, "True if using an insecure connection.")
var user = flag.String("user", security.RootUser, "User used to connect to the cluster.")
var certs = flag.String("certs", "certs", "Directory containing RSA key and x509 certs. This flag is required if --insecure=false.")

var retryOptions = retry.Options{
	InitialBackoff: 100 * time.Millisecond,
	MaxRetries:     10,
	Multiplier:     2,
}

type statusMonitor struct {
	addr       string
	url        string
	httpClient *http.Client
	file       *io.Writer
}

func newStatusMonitor(context *base.Context, addr string) (*statusMonitor, error) {
	monitor := &statusMonitor{
		addr: addr,
	}
	var err error
	monitor.httpClient, err = context.GetHTTPClient()
	if err != nil {
		return nil, err
	}
	monitor.url = fmt.Sprintf("%s://%s/%s", context.HTTPRequestScheme(), monitor.addr, urlPath)
	return monitor, nil
}

func (m *statusMonitor) queryStatus() error {
	var queryErr error
	for r := retry.Start(retryOptions); r.Next(); {
		if log.V(1) && queryErr != nil {
			log.Infof("retrying after error: %s", queryErr)
		}

		// Construct a new HTTP GET Request.
		req, err := http.NewRequest("GET", m.url, nil)
		if err != nil {
			queryErr = fmt.Errorf("could not create http request for %s: %s", m.url, err)
			// Break immediately, this is not recoverable.
			break
		}
		req.Header.Set(util.AcceptHeader, util.JSONContentType)

		// Execute request.
		resp, err := m.httpClient.Do(req)
		if err != nil {
			queryErr = fmt.Errorf("could not GET %s - %s", m.url, err)
			continue
		}
		defer resp.Body.Close()

		// Read and verify body of response.
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			queryErr = fmt.Errorf("could not ready body for %s - %s", m.url, err)
			continue
		}
		if resp.StatusCode != http.StatusOK {
			queryErr = fmt.Errorf("could not GET %s - statuscode: %d - body: %s", m.url, resp.StatusCode, body)
			continue
		}
		returnedContentType := resp.Header.Get(util.ContentTypeHeader)
		if returnedContentType != util.JSONContentType {
			queryErr = fmt.Errorf("unexpected content type: %v", returnedContentType)
			continue
		}
		return nil
	}
	return queryErr
}

func main() {
	flag.Parse()
	parsedAddrs := strings.Split(*addrs, ",")

	ctx := base.Context{Insecure: *insecure, Certs: *certs, User: *user}

	startTime := time.Now()
	stopper := stop.NewStopper()
	for _, addr := range parsedAddrs {
		client, err := newStatusMonitor(&ctx, addr)
		if err != nil {
			log.Errorf("error creating client: %s", err)
			return
		}
		log.Infof("Monitoring Status URL: %s", client.url)
		stopper.RunWorker(func() {
			timer := time.Tick(*interval)
			for {
				select {
				case <-stopper.ShouldStop():
					return
				case <-timer:
					elapsed := time.Since(startTime)
					if err := client.queryStatus(); err != nil {
						log.Warningf("Could not get status from url %s. Time since monitor started %s.", client.url, elapsed)
						stopper.Stop()
						continue
					}
					log.Infof("Got status from url %s. Time since start: %s", client.url, elapsed)
				}
			}
		})
	}

	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, os.Kill, syscall.SIGTERM)
	// Block until a termination signal is received, or the stopper is closed by
	// an error in one of the client routines.
	select {
	case <-stopper.ShouldStop():
		log.Infof("Monitor stopped by error...")
		os.Exit(1)
	case <-signalCh:
		log.Infof("Stopping status monitor...")
		stopper.Stop()
	}
}
<|endoftext|>


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

package main

import (
	"context"
	"flag"
	logpkg "log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/exporter/stackdriver"
	gops "github.com/google/gops/agent"
	"github.com/neovim/go-client/nvim/plugin"
	"github.com/pkg/errors"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
	"go.uber.org/multierr"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"golang.org/x/exp/errors/fmt"
	"golang.org/x/sync/errgroup"

	"github.com/zchee/nvim-go/pkg/autocmd"
	"github.com/zchee/nvim-go/pkg/buildctxt"
	"github.com/zchee/nvim-go/pkg/command"
	"github.com/zchee/nvim-go/pkg/config"
	"github.com/zchee/nvim-go/pkg/logger"
	"github.com/zchee/nvim-go/pkg/nctx"
	"github.com/zchee/nvim-go/pkg/server"
	"github.com/zchee/nvim-go/pkg/version"
)

// flags
var (
	fVersion    = flag.Bool("version", false, "Show the version information.")
	pluginHost  = flag.String("manifest", "", "Write plugin manifest for `host` to stdout")
	vimFilePath = flag.String("location", "", "Manifest is automatically written to `.vim file`")
)

func init() {
	flag.Parse()
	logpkg.SetPrefix("nvim-go: ")
}

func main() {
	if *fVersion {
		fmt.Printf("%s:\n  version: %s\n", nctx.AppName, version.Version)
		return
	}

	ctx, cancel := context.WithCancel(Context())
	defer cancel()

	if *pluginHost != "" {
		os.Unsetenv("NVIM_GO_DEBUG")               // disable zap output
		ctx = logger.NewContext(ctx, zap.NewNop()) // avoid nil panic on logger.FromContext

		fn := func(p *plugin.Plugin) error {
			return func(ctx context.Context, p *plugin.Plugin) error {
				bctxt := buildctxt.NewContext()
				c := command.Register(ctx, p, bctxt)
				autocmd.Register(ctx, p, bctxt, c)
				return nil
			}(ctx, p)
		}
		if err := Plugin(fn); err != nil {
			logpkg.Fatal(err)
		}
		return
	}

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	sighupFn := func() {}
	sigintFn := func() {
		logpkg.Println("Start shutdown gracefully")
		cancel()
	}
	go signalHandler(sigc, sighupFn, sigintFn)

	errc := make(chan error, 1)
	go func() {
		defer close(errc)
		errc <- startServer(ctx)
	}()

	select {
	case <-ctx.Done():
	case err := <-errc:
		if err != nil {
			logpkg.Fatal(err)
		}
	}
	logpkg.Println("shutdown nvim-go server")
}

func signalHandler(ch <-chan os.Signal, sighupFn, sigintFn func()) {
	for {
		select {
		case sig := <-ch:
			switch sig {
			case syscall.SIGHUP:
				logpkg.Printf("catch signal %s", sig)
				sighupFn()
			case syscall.SIGINT, syscall.SIGTERM:
				logpkg.Printf("catch signal %s", sig)
				sigintFn()
			}
		}
	}
}

func startServer(ctx context.Context) (errs error) {
	env := config.Process()

	var lv zapcore.Level
	if err := lv.UnmarshalText([]byte(env.LogLevel)); err != nil {
		return fmt.Errorf("failed to parse log level: %s, err: %v", env.LogLevel, err)
	}
	log, undo := logger.NewRedirectZapLogger(lv)
	defer undo()
	ctx = logger.NewContext(ctx, log)

	// Open socket for using gops to get stacktraces of the daemon.
	if err := gops.Listen(gops.Options{ConfigDir: "/tmp/gops", ShutdownCleanup: true}); err != nil {
		return fmt.Errorf("unable to start gops: %s", err)
	}
	log.Info("starting gops agent")

	if gcpProjectID := env.GCPProjectID; gcpProjectID != "" {
		// OpenCensus tracing with Stackdriver exporter
		sdOpts := stackdriver.Options{
			ProjectID: gcpProjectID,
			OnError: func(err error) {
				errs = multierr.Append(errs, fmt.Errorf("stackdriver.Exporter: %v", err))
			},
			MetricPrefix: nctx.AppName,
			Context:      ctx,
		}
		sd, err := stackdriver.NewExporter(sdOpts)
		if err != nil {
			logpkg.Fatalf("failed to create stackdriver exporter: %v", err)
		}
		defer sd.Flush()
		trace.RegisterExporter(sd)
		view.RegisterExporter(sd)
		log.Info("opencensus", zap.String("trace", "enabled Stackdriver exporter"))

		// Stackdriver Profiler
		profConf := profiler.Config{
			Service:        nctx.AppName,
			ServiceVersion: version.Tag,
			MutexProfiling: true,
			ProjectID:      gcpProjectID,
		}
		if err := profiler.Start(profConf); err != nil {
			logpkg.Fatalf("failed to start stackdriver profiler: %v", err)
		}
		log.Info("stackdriver", zap.String("profiler", "enabled Stackdriver profiler"))

		trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
		var span *trace.Span
		ctx, span = trace.StartSpan(ctx, "main") // start root span
		defer span.End()
	}

	fn := func(p *plugin.Plugin) error {
		return func(ctx context.Context, p *plugin.Plugin) error {
			log := logger.FromContext(ctx).Named("main")
			ctx = logger.NewContext(ctx, log)

			bctxt := buildctxt.NewContext()
			cmd := command.Register(ctx, p, bctxt)
			autocmd.Register(ctx, p, bctxt, cmd)

			// switch to unix socket rpc-connection
			if n, err := server.Dial(ctx); err == nil {
				p.Nvim = n
			}

			return nil
		}(ctx, p)
	}

	eg := new(errgroup.Group)
	eg.Go(func() error {
		return Plugin(fn)
	})
	eg.Go(func() error {
		return subscribeServer(ctx)
	})

	log.Info(fmt.Sprintf("starting %s server", nctx.AppName), zap.Object("env", env))
	if err := eg.Wait(); err != nil {
		log.Fatal("occurred error", zap.Error(err))
	}

	return errs
}

func subscribeServer(ctx context.Context) error {
	log := logger.FromContext(ctx).Named("child")
	ctx = logger.NewContext(ctx, log)

	s, err := server.NewServer(ctx)
	if err != nil {
		return errors.Wrap(err, "failed to create NewServer")
	}
	go s.Serve()

	s.Nvim.Subscribe(nctx.Method)

	select {
	case <-ctx.Done():
		if err := s.Close(); err != nil {
			log.Fatal("Close", zap.Error(err))
		}
		return nil
	}
}
<commit_msg>cmd/nvim-go: remove gops/agent<commit_after>// Copyright 2016 The nvim-go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"flag"
	logpkg "log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/exporter/stackdriver"
	"github.com/neovim/go-client/nvim/plugin"
	"github.com/pkg/errors"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/trace"
	"go.uber.org/multierr"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"golang.org/x/exp/errors/fmt"
	"golang.org/x/sync/errgroup"

	"github.com/zchee/nvim-go/pkg/autocmd"
	"github.com/zchee/nvim-go/pkg/buildctxt"
	"github.com/zchee/nvim-go/pkg/command"
	"github.com/zchee/nvim-go/pkg/config"
	"github.com/zchee/nvim-go/pkg/logger"
	"github.com/zchee/nvim-go/pkg/nctx"
	"github.com/zchee/nvim-go/pkg/server"
	"github.com/zchee/nvim-go/pkg/version"
)

// flags
var (
	fVersion    = flag.Bool("version", false, "Show the version information.")
	pluginHost  = flag.String("manifest", "", "Write plugin manifest for `host` to stdout")
	vimFilePath = flag.String("location", "", "Manifest is automatically written to `.vim file`")
)

func init() {
	flag.Parse()
	logpkg.SetPrefix("nvim-go: ")
}

func main() {
	if *fVersion {
		fmt.Printf("%s:\n  version: %s\n", nctx.AppName, version.Version)
		return
	}

	ctx, cancel := context.WithCancel(Context())
	defer cancel()

	if *pluginHost != "" {
		os.Unsetenv("NVIM_GO_DEBUG")               // disable zap output
		ctx = logger.NewContext(ctx, zap.NewNop()) // avoid nil panic on logger.FromContext

		fn := func(p *plugin.Plugin) error {
			return func(ctx context.Context, p *plugin.Plugin) error {
				bctxt := buildctxt.NewContext()
				c := command.Register(ctx, p, bctxt)
				autocmd.Register(ctx, p, bctxt, c)
				return nil
			}(ctx, p)
		}
		if err := Plugin(fn); err != nil {
			logpkg.Fatal(err)
		}
		return
	}

	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL)
	sighupFn := func() {}
	sigintFn := func() {
		logpkg.Println("Start shutdown gracefully")
		cancel()
	}
	go signalHandler(sigc, sighupFn, sigintFn)

	errc := make(chan error, 1)
	go func() {
		defer close(errc)
		errc <- startServer(ctx)
	}()

	select {
	case <-ctx.Done():
	case err := <-errc:
		if err != nil {
			logpkg.Fatal(err)
		}
	}
	logpkg.Println("shutdown nvim-go server")
}

func signalHandler(ch <-chan os.Signal, sighupFn, sigintFn func()) {
	for {
		select {
		case sig := <-ch:
			switch sig {
			case syscall.SIGHUP:
				logpkg.Printf("catch signal %s", sig)
				sighupFn()
			case syscall.SIGINT, syscall.SIGTERM:
				logpkg.Printf("catch signal %s", sig)
				sigintFn()
			}
		}
	}
}

func startServer(ctx context.Context) (errs error) {
	env := config.Process()

	var lv zapcore.Level
	if err := lv.UnmarshalText([]byte(env.LogLevel)); err != nil {
		return fmt.Errorf("failed to parse log level: %s, err: %v", env.LogLevel, err)
	}
	log, undo := logger.NewRedirectZapLogger(lv)
	defer undo()
	ctx = logger.NewContext(ctx, log)

	if gcpProjectID, ok := config.HasGCPProjectID(); ok {
		// OpenCensus tracing with Stackdriver exporter
		sdOpts := stackdriver.Options{
			ProjectID: gcpProjectID,
			OnError: func(err error) {
				errs = multierr.Append(errs, fmt.Errorf("stackdriver.Exporter: %v", err))
			},
			MetricPrefix: nctx.AppName,
			Context:      ctx,
		}
		sd, err := stackdriver.NewExporter(sdOpts)
		if err != nil {
			logpkg.Fatalf("failed to create stackdriver exporter: %v", err)
		}
		defer sd.Flush()
		trace.RegisterExporter(sd)
		view.RegisterExporter(sd)
		log.Info("opencensus", zap.String("trace", "enabled Stackdriver exporter"))

		// Stackdriver Profiler
		profConf := profiler.Config{
			Service:        nctx.AppName,
			ServiceVersion: version.Tag,
			MutexProfiling: true,
			ProjectID:      gcpProjectID,
		}
		if err := profiler.Start(profConf); err != nil {
			logpkg.Fatalf("failed to start stackdriver profiler: %v", err)
		}
		log.Info("stackdriver", zap.String("profiler", "enabled Stackdriver profiler"))

		trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
		var span *trace.Span
		ctx, span = trace.StartSpan(ctx, "main") // start root span
		defer span.End()
	}

	fn := func(p *plugin.Plugin) error {
		return func(ctx context.Context, p *plugin.Plugin) error {
			log := logger.FromContext(ctx).Named("main")
			ctx = logger.NewContext(ctx, log)

			bctxt := buildctxt.NewContext()
			cmd := command.Register(ctx, p, bctxt)
			autocmd.Register(ctx, p, bctxt, cmd)

			// switch to unix socket rpc-connection
			if n, err := server.Dial(ctx); err == nil {
				p.Nvim = n
			}

			return nil
		}(ctx, p)
	}

	eg := new(errgroup.Group)
	eg.Go(func() error {
		return Plugin(fn)
	})
	eg.Go(func() error {
		return subscribeServer(ctx)
	})

	log.Info(fmt.Sprintf("starting %s server", nctx.AppName), zap.Object("env", env))
	if err := eg.Wait(); err != nil {
		log.Fatal("occurred error", zap.Error(err))
	}

	return errs
}

func subscribeServer(ctx context.Context) error {
	log := logger.FromContext(ctx).Named("child")
	ctx = logger.NewContext(ctx, log)

	s, err := server.NewServer(ctx)
	if err != nil {
		return errors.Wrap(err, "failed to create NewServer")
	}
	go s.Serve()

	s.Nvim.Subscribe(nctx.Method)

	select {
	case <-ctx.Done():
		if err := s.Close(); err != nil {
			log.Fatal("Close", zap.Error(err))
		}
		return nil
	}
}
<|endoftext|>