<commit_before>package common

import (
	"errors"
	"fmt"
	"log"
	"os"
	"sync"
	"time"
)

const taskIDSize = 20

var (
	Tracker           = &TaskTracker{tasks: map[string]*Task{}}
	TaskStatusUnknown = &TaskStatus{Status: StatusUnknown}
)

func MaintenanceChecker(file string, interval time.Duration) {
	go Tracker.MaintenanceChecker(file, interval)
}

func NewTask(name string, executor TaskExecutor) *Task {
	task := &Task{Tracker: Tracker, Executor: executor}
	task.Status = StatusInit
	task.StatusTime = time.Now()
	task.Name = name
	task.Description = executor.Description()
	task.Request = executor.Request()
	return task
}

type TaskTracker struct {
	sync.RWMutex
	ResultDuration time.Duration
	Maintenance    bool
	tasks          map[string]*Task
}

type Task struct {
	sync.RWMutex
	TaskStatus
	Err      error
	Tracker  *TaskTracker
	ID       string
	Executor TaskExecutor
	Request  interface{}
	Result   interface{}
}

type TaskExecutor interface {
	Request() interface{}
	Result() interface{}
	Description() string
	Execute(t *Task) error
	Authorize() error
}

type TaskMaintenanceExecutor interface {
	AllowDuringMaintenance() bool
}

func createTaskID() string {
	return CreateRandomID(taskIDSize)
}

func (t *TaskTracker) ListIDs(typ string) []string {
	t.Lock()
	ids := make([]string, len(t.tasks))
	i := 0
	for id, task := range t.tasks {
		if task.Name != typ {
			continue
		}
		ids[i] = id
		i++
	}
	t.Unlock()
	return ids
}

func (t *TaskTracker) SetMaintenance(on bool) {
	t.Lock()
	t.Maintenance = on
	t.Unlock()
}

func (t *TaskTracker) UnderMaintenance() bool {
	t.RLock()
	maint := t.Maintenance
	t.RUnlock()
	return maint
}

func (t *TaskTracker) MaintenanceChecker(file string, interval time.Duration) {
	for {
		if _, err := os.Stat(file); err == nil {
			// maintenance file exists
			if !t.UnderMaintenance() {
				log.Println("Begin Maintenance")
				t.SetMaintenance(true)
			}
		} else {
			// maintenance file doesn't exist or there is an error looking for it
			if t.UnderMaintenance() {
				log.Println("End Maintenance")
				t.SetMaintenance(false)
			}
		}
		time.Sleep(interval)
	}
}

func (t *TaskTracker) Idle(checkTask *Task) bool {
	idle := true
	t.RLock()
	for _, task := range t.tasks {
		if task != checkTask && !task.Done {
			idle = false
			break
		}
	}
	t.RUnlock()
	return idle
}

func (t *TaskTracker) ReserveTaskID(task *Task) string {
	t.Lock()
	requestID := createTaskID()
	for _, present := t.tasks[requestID]; present; _, present = t.tasks[requestID] {
		requestID = createTaskID()
	}
	t.tasks[requestID] = task // reserve request id
	t.Unlock()
	task.Lock()
	task.ID = requestID
	task.Unlock()
	return requestID
}

func (t *TaskTracker) ReleaseTaskID(id string) {
	t.Lock()
	delete(t.tasks, id)
	t.Unlock()
}

func (t *TaskTracker) Status(id string) (*TaskStatus, error) {
	t.RLock()
	task := t.tasks[id]
	t.RUnlock()
	if task != nil {
		task.RLock()
		status := task.CopyTaskStatus()
		err := task.Err
		task.RUnlock()
		return status, err
	}
	return TaskStatusUnknown, errors.New("Unknown Task Status")
}

func (t *TaskTracker) Result(id string) interface{} {
	t.RLock()
	task := t.tasks[id]
	t.RUnlock()
	if task != nil {
		task.RLock()
		result := task.Result
		task.RUnlock()
		return result
	}
	return nil
}

func (t *Task) Authorize() error {
	return t.Executor.Authorize()
}

func (t *Task) Run() error {
	if t.Tracker.UnderMaintenance() {
		executor, ok := t.Executor.(TaskMaintenanceExecutor)
		if !ok || !executor.AllowDuringMaintenance() {
			return t.End(errors.New("Under Maintenance"), false)
		}
	}
	t.Tracker.ReserveTaskID(t)
	t.Log("Begin %s", t.Description)
	t.Lock()
	t.StartTime = time.Now()
	t.Unlock()
	err := t.Executor.Authorize()
	if err != nil {
		return t.End(err, false)
	}
	return t.End(t.Executor.Execute(t), false)
}

func (t *Task) RunAsync(r *AsyncReply) error {
	if t.Tracker.UnderMaintenance() {
		executor, ok := t.Executor.(TaskMaintenanceExecutor)
		if !ok || !executor.AllowDuringMaintenance() {
			return t.End(errors.New("Under Maintenance"), false)
		}
	}
	t.Tracker.ReserveTaskID(t)
	t.RLock()
	r.ID = t.ID
	t.RUnlock()
	go func() error {
		t.Log("Begin %s", t.Description)
		t.Lock()
		t.StartTime = time.Now()
		t.Unlock()
		err := t.Executor.Authorize()
		if err != nil {
			return t.End(err, true)
		}
		t.End(t.Executor.Execute(t), true)
		return nil
	}()
	return nil
}

func (t *Task) End(err error, async bool) error {
	logString := fmt.Sprintf("End %s", t.Description)
	t.Lock()
	t.Result = t.Executor.Result()
	t.EndTime = time.Now()
	t.StatusTime = t.EndTime
	if err == nil {
		t.Status = StatusDone
		t.Done = true
	} else {
		t.Status = StatusError
		t.Err = err
		t.Done = true
		logString += fmt.Sprintf(" - Error: %s", err.Error())
	}
	t.Unlock()
	t.Log(logString)
	if async {
		time.AfterFunc(t.Tracker.ResultDuration, func() {
			// keep result around for 30 min in case someone wants to check on it
			t.Tracker.ReleaseTaskID(t.ID)
		})
	} else {
		t.Tracker.ReleaseTaskID(t.ID)
	}
	return err
}

func (t *Task) Log(format string, args ...interface{}) {
	t.RLock()
	log.Printf("[RPC]["+t.Name+"]["+t.ID+"] "+format, args...)
	t.RUnlock()
}

func (t *Task) LogStatus(format string, args ...interface{}) {
	t.Log(format, args...)
	t.Lock()
	t.StatusTime = time.Now()
	t.Status = fmt.Sprintf(format, args...)
	t.Unlock()
}

func (t *Task) AddWarning(warn string) {
	t.Lock()
	if t.Warnings == nil {
		t.Warnings = []string{warn}
	} else {
		t.Warnings = append(t.Warnings, warn)
	}
	t.Unlock()
	t.Log("WARNING: %s", warn)
}

type TaskStatus struct {
	Name        string
	Description string
	Status      string
	Warnings    []string
	Done        bool
	StartTime   time.Time
	StatusTime  time.Time
	EndTime     time.Time
}

func (t *TaskStatus) Map() map[string]interface{} {
	return map[string]interface{}{
		"Name":        t.Name,
		"Description": t.Description,
		"Status":      t.Status,
		"Warnings":    t.Warnings,
		"Done":        t.Done,
		"StartTime":   t.StartTime,
		"StatusTime":  t.StatusTime,
		"EndTime":     t.EndTime,
	}
}

func (t *TaskStatus) String() string {
	return fmt.Sprintf(`%s
Description : %s
Status      : %s
Warnings    : %v
Done        : %t
StartTime   : %s
StatusTime  : %s
EndTime     : %s`, t.Name, t.Description, t.Status, t.Warnings, t.Done, t.StartTime, t.StatusTime,
		t.EndTime)
}

func (t *TaskStatus) CopyTaskStatus() *TaskStatus {
	return &TaskStatus{t.Name, t.Description, t.Status, t.Warnings, t.Done, t.StartTime, t.StatusTime,
		t.EndTime}
}
<commit_msg>allow full whitelist for ListIDs<commit_after>package common

import (
	"errors"
	"fmt"
	"log"
	"os"
	"sync"
	"time"
)

const taskIDSize = 20

var (
	Tracker           = &TaskTracker{tasks: map[string]*Task{}}
	TaskStatusUnknown = &TaskStatus{Status: StatusUnknown}
)

func MaintenanceChecker(file string, interval time.Duration) {
	go Tracker.MaintenanceChecker(file, interval)
}

func NewTask(name string, executor TaskExecutor) *Task {
	task := &Task{Tracker: Tracker, Executor: executor}
	task.Status = StatusInit
	task.StatusTime = time.Now()
	task.Name = name
	task.Description = executor.Description()
	task.Request = executor.Request()
	return task
}

type TaskTracker struct {
	sync.RWMutex
	ResultDuration time.Duration
	Maintenance    bool
	tasks          map[string]*Task
}

type Task struct {
	sync.RWMutex
	TaskStatus
	Err      error
	Tracker  *TaskTracker
	ID       string
	Executor TaskExecutor
	Request  interface{}
	Result   interface{}
}

type TaskExecutor interface {
	Request() interface{}
	Result() interface{}
	Description() string
	Execute(t *Task) error
	Authorize() error
}

type TaskMaintenanceExecutor interface {
	AllowDuringMaintenance() bool
}

func createTaskID() string {
	return CreateRandomID(taskIDSize)
}

func (t *TaskTracker) ListIDs(types []string) []string {
	typesMap = make(map[string]bool, len(types))
	for _, typ := range types {
		typesMap[typ] = true
	}
	ids := []string{}
	t.Lock()
	for id, task := range t.tasks {
		if typesMap[task.Name] {
			ids = append(ids, id)
		}
	}
	t.Unlock()
	return ids
}

func (t *TaskTracker) SetMaintenance(on bool) {
	t.Lock()
	t.Maintenance = on
	t.Unlock()
}

func (t *TaskTracker) UnderMaintenance() bool {
	t.RLock()
	maint := t.Maintenance
	t.RUnlock()
	return maint
}

func (t *TaskTracker) MaintenanceChecker(file string, interval time.Duration) {
	for {
		if _, err := os.Stat(file); err == nil {
			// maintenance file exists
			if !t.UnderMaintenance() {
				log.Println("Begin Maintenance")
				t.SetMaintenance(true)
			}
		} else {
			// maintenance file doesn't exist or there is an error looking for it
			if t.UnderMaintenance() {
				log.Println("End Maintenance")
				t.SetMaintenance(false)
			}
		}
		time.Sleep(interval)
	}
}

func (t *TaskTracker) Idle(checkTask *Task) bool {
	idle := true
	t.RLock()
	for _, task := range t.tasks {
		if task != checkTask && !task.Done {
			idle = false
			break
		}
	}
	t.RUnlock()
	return idle
}

func (t *TaskTracker) ReserveTaskID(task *Task) string {
	t.Lock()
	requestID := createTaskID()
	for _, present := t.tasks[requestID]; present; _, present = t.tasks[requestID] {
		requestID = createTaskID()
	}
	t.tasks[requestID] = task // reserve request id
	t.Unlock()
	task.Lock()
	task.ID = requestID
	task.Unlock()
	return requestID
}

func (t *TaskTracker) ReleaseTaskID(id string) {
	t.Lock()
	delete(t.tasks, id)
	t.Unlock()
}

func (t *TaskTracker) Status(id string) (*TaskStatus, error) {
	t.RLock()
	task := t.tasks[id]
	t.RUnlock()
	if task != nil {
		task.RLock()
		status := task.CopyTaskStatus()
		err := task.Err
		task.RUnlock()
		return status, err
	}
	return TaskStatusUnknown, errors.New("Unknown Task Status")
}

func (t *TaskTracker) Result(id string) interface{} {
	t.RLock()
	task := t.tasks[id]
	t.RUnlock()
	if task != nil {
		task.RLock()
		result := task.Result
		task.RUnlock()
		return result
	}
	return nil
}

func (t *Task) Authorize() error {
	return t.Executor.Authorize()
}

func (t *Task) Run() error {
	if t.Tracker.UnderMaintenance() {
		executor, ok := t.Executor.(TaskMaintenanceExecutor)
		if !ok || !executor.AllowDuringMaintenance() {
			return t.End(errors.New("Under Maintenance"), false)
		}
	}
	t.Tracker.ReserveTaskID(t)
	t.Log("Begin %s", t.Description)
	t.Lock()
	t.StartTime = time.Now()
	t.Unlock()
	err := t.Executor.Authorize()
	if err != nil {
		return t.End(err, false)
	}
	return t.End(t.Executor.Execute(t), false)
}

func (t *Task) RunAsync(r *AsyncReply) error {
	if t.Tracker.UnderMaintenance() {
		executor, ok := t.Executor.(TaskMaintenanceExecutor)
		if !ok || !executor.AllowDuringMaintenance() {
			return t.End(errors.New("Under Maintenance"), false)
		}
	}
	t.Tracker.ReserveTaskID(t)
	t.RLock()
	r.ID = t.ID
	t.RUnlock()
	go func() error {
		t.Log("Begin %s", t.Description)
		t.Lock()
		t.StartTime = time.Now()
		t.Unlock()
		err := t.Executor.Authorize()
		if err != nil {
			return t.End(err, true)
		}
		t.End(t.Executor.Execute(t), true)
		return nil
	}()
	return nil
}

func (t *Task) End(err error, async bool) error {
	logString := fmt.Sprintf("End %s", t.Description)
	t.Lock()
	t.Result = t.Executor.Result()
	t.EndTime = time.Now()
	t.StatusTime = t.EndTime
	if err == nil {
		t.Status = StatusDone
		t.Done = true
	} else {
		t.Status = StatusError
		t.Err = err
		t.Done = true
		logString += fmt.Sprintf(" - Error: %s", err.Error())
	}
	t.Unlock()
	t.Log(logString)
	if async {
		time.AfterFunc(t.Tracker.ResultDuration, func() {
			// keep result around for 30 min in case someone wants to check on it
			t.Tracker.ReleaseTaskID(t.ID)
		})
	} else {
		t.Tracker.ReleaseTaskID(t.ID)
	}
	return err
}

func (t *Task) Log(format string, args ...interface{}) {
	t.RLock()
	log.Printf("[RPC]["+t.Name+"]["+t.ID+"] "+format, args...)
	t.RUnlock()
}

func (t *Task) LogStatus(format string, args ...interface{}) {
	t.Log(format, args...)
	t.Lock()
	t.StatusTime = time.Now()
	t.Status = fmt.Sprintf(format, args...)
	t.Unlock()
}

func (t *Task) AddWarning(warn string) {
	t.Lock()
	if t.Warnings == nil {
		t.Warnings = []string{warn}
	} else {
		t.Warnings = append(t.Warnings, warn)
	}
	t.Unlock()
	t.Log("WARNING: %s", warn)
}

type TaskStatus struct {
	Name        string
	Description string
	Status      string
	Warnings    []string
	Done        bool
	StartTime   time.Time
	StatusTime  time.Time
	EndTime     time.Time
}

func (t *TaskStatus) Map() map[string]interface{} {
	return map[string]interface{}{
		"Name":        t.Name,
		"Description": t.Description,
		"Status":      t.Status,
		"Warnings":    t.Warnings,
		"Done":        t.Done,
		"StartTime":   t.StartTime,
		"StatusTime":  t.StatusTime,
		"EndTime":     t.EndTime,
	}
}

func (t *TaskStatus) String() string {
	return fmt.Sprintf(`%s
Description : %s
Status      : %s
Warnings    : %v
Done        : %t
StartTime   : %s
StatusTime  : %s
EndTime     : %s`, t.Name, t.Description, t.Status, t.Warnings, t.Done, t.StartTime, t.StatusTime,
		t.EndTime)
}

func (t *TaskStatus) CopyTaskStatus() *TaskStatus {
	return &TaskStatus{t.Name, t.Description, t.Status, t.Warnings, t.Done, t.StartTime, t.StatusTime,
		t.EndTime}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

type Generator struct {
	StartUps        []string
	Words           []string
	consumerKey     string
	consumerSecret  string
	accessKey       string
	accessKeySecret string
}

func init() {
	rand.Seed(time.Now().Unix())
}

func New() (*Generator, error) {
	gen := new(Generator)
	doc, err := goquery.NewDocument("https://www.startupranking.com/top/united-states")
	if err != nil {
		return nil, err
	}

	gen.consumerKey = os.Getenv("CONSUMER_KEY")
	gen.consumerSecret = os.Getenv("CONSUMER_SECRET")
	gen.accessKey = os.Getenv("ACCESS_TOKEN")
	gen.accessKeySecret = os.Getenv("ACCESS_TOKEN_SECRET")

	// Find the startup names
	doc.Find(".name").Each(func(i int, s *goquery.Selection) {
		// For each item found, get the name
		startup := s.Text()
		fmt.Printf("Name %d - %s\n", i, startup)
		gen.StartUps = append(gen.StartUps, startup)
	})
	if len(gen.StartUps) == 0 {
		return nil, fmt.Errorf("could not get startup list")
	}

	// Load all uncountabl words to memory
	gen.Words, err = readLines(`./wordlist/unc.txt`)
	if err != nil {
		return nil, err
	}

	return gen, nil
}

func (g *Generator) Tweet(tweet string) {
	fmt.Println(g.accessKey, g.accessKeySecret, g.consumerKey, g.consumerSecret)
	config := oauth1.NewConfig(g.consumerKey, g.consumerSecret)
	token := oauth1.NewToken(g.accessKey, g.accessKeySecret)
	// http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter client
	client := twitter.NewClient(httpClient)

	_, resp, err := client.Statuses.Update(tweet, nil)
	if err != nil {
		fmt.Println("twitter update error", err)
	}
	fmt.Println(resp)

}

// readLines reads a whole file into memory
// and returns a slice of its lines.
func readLines(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	return lines, scanner.Err()
}
<commit_msg>add error print to generator<commit_after>package main

import (
	"bufio"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

type Generator struct {
	StartUps        []string
	Words           []string
	consumerKey     string
	consumerSecret  string
	accessKey       string
	accessKeySecret string
}

func init() {
	rand.Seed(time.Now().Unix())
}

func New() (*Generator, error) {
	gen := new(Generator)
	doc, err := goquery.NewDocument("https://www.startupranking.com/top/united-states")
	if err != nil {
		return nil, err
	}

	gen.consumerKey = os.Getenv("CONSUMER_KEY")
	gen.consumerSecret = os.Getenv("CONSUMER_SECRET")
	gen.accessKey = os.Getenv("ACCESS_TOKEN")
	gen.accessKeySecret = os.Getenv("ACCESS_TOKEN_SECRET")

	// Find the startup names
	doc.Find(".name").Each(func(i int, s *goquery.Selection) {
		// For each item found, get the name
		startup := s.Text()
		fmt.Printf("Name %d - %s\n", i, startup)
		gen.StartUps = append(gen.StartUps, startup)
	})
	if len(gen.StartUps) == 0 {
		return nil, fmt.Errorf("could not get startup list (error=%v)",err)
	}

	// Load all uncountabl words to memory
	gen.Words, err = readLines(`./wordlist/unc.txt`)
	if err != nil {
		return nil, err
	}

	return gen, nil
}

func (g *Generator) Tweet(tweet string) {
	fmt.Println(g.accessKey, g.accessKeySecret, g.consumerKey, g.consumerSecret)
	config := oauth1.NewConfig(g.consumerKey, g.consumerSecret)
	token := oauth1.NewToken(g.accessKey, g.accessKeySecret)
	// http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter client
	client := twitter.NewClient(httpClient)

	_, resp, err := client.Statuses.Update(tweet, nil)
	if err != nil {
		fmt.Println("twitter update error", err)
	}
	fmt.Println(resp)

}

// readLines reads a whole file into memory
// and returns a slice of its lines.
func readLines(path string) ([]string, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}
	return lines, scanner.Err()
}
<|endoftext|>


<commit_before>//  Copyright (c) 2013 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 (
	"flag"
	"fmt"
	"github.com/prataprc/collatejson"
	"strconv"
)

var options struct {
	floatText string
	intText   string
}

func argParse() {
	//flag.BoolVar(&options.ast, "ast", false, "Show the ast of production")
	//flag.IntVar(&options.seed, "s", seed, "Seed value")
	//flag.IntVar(&options.count, "n", 1, "Generate n combinations")
	//flag.StringVar(&options.outfile, "o", "-", "Specify an output file")
	flag.StringVar(&options.floatText, "f", "", "encode floating point number")
	flag.StringVar(&options.intText, "i", "", "encode integer number")
	flag.Parse()
}

func main() {
	argParse()
	if options.floatText != "" {
		encodeFloat(options.floatText)
	} else if options.intText != "" {
		encodeInt(options.intText)
	}
}

func encodeFloat(text string) {
	if f, err := strconv.ParseFloat(text, 64); err != nil {
		panic(err)
	} else {
		ftext := []byte(strconv.FormatFloat(f, 'e', -1, 64))
		fmt.Printf("Encoding %v: %v\n", f, string(collatejson.EncodeFloat(ftext)))
	}
}

func encodeInt(text string) {
	fmt.Println(string(collatejson.EncodeInt([]byte(text))))
}
<commit_msg>Command line tool encode.go accepts and encodes json strings<commit_after>//  Copyright (c) 2013 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 (
	"flag"
	"fmt"
	"github.com/prataprc/collatejson"
	"strconv"
	"unicode/utf8"
)

var options struct {
	floatText  string
	intText    string
	stringText string
}

func argParse() {
	//flag.BoolVar(&options.ast, "ast", false, "Show the ast of production")
	//flag.IntVar(&options.seed, "s", seed, "Seed value")
	//flag.IntVar(&options.count, "n", 1, "Generate n combinations")
	//flag.StringVar(&options.outfile, "o", "-", "Specify an output file")
	flag.StringVar(&options.floatText, "f", "", "encode floating point number")
	flag.StringVar(&options.intText, "i", "", "encode integer number")
	flag.StringVar(&options.stringText, "s", "", "encode string")
	flag.Parse()
}

func main() {
	argParse()
	if options.floatText != "" {
		encodeFloat(options.floatText)
	} else if options.intText != "" {
		encodeInt(options.intText)
	} else {
		s, i := options.stringText, 0
		for {
			r, c := utf8.DecodeRune([]byte(s[i:]))
			i += c
			fmt.Println(r, c)
			if len(s[i:]) == 0 {
				break
			}
		}
	}
}

func encodeFloat(text string) {
	if f, err := strconv.ParseFloat(text, 64); err != nil {
		panic(err)
	} else {
		ftext := []byte(strconv.FormatFloat(f, 'e', -1, 64))
		fmt.Printf("Encoding %v: %v\n", f, string(collatejson.EncodeFloat(ftext)))
	}
}

func encodeInt(text string) {
	fmt.Println(string(collatejson.EncodeInt([]byte(text))))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"time"

	"github.com/atotto/clipboard"
	"github.com/cheggaaa/pb"
	flag "github.com/ogier/pflag"
)

const (
	service = "https://up.depado.eu/"
	method  = "POST"
)

var (
	bar  *pb.ProgressBar
	name string
)

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

func initBar(f *os.File) {
	fi, err := f.Stat()
	if err != nil {
		fmt.Println("Could not stat", f.Name())
		os.Exit(1)
	}
	bar = pb.New64(fi.Size()).SetUnits(pb.U_BYTES).SetRefreshRate(time.Millisecond * 10)
	bar.ShowPercent = true
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.Start()
}

func main() {
	var err error
	var datasource io.Reader

	var tee bool
	var progress bool
	var clip bool
	var argname string

	flag.BoolVarP(&tee, "tee", "t", false, "Displays stdin to stdout")
	flag.BoolVarP(&progress, "progress", "p", false, "Displays a progress bar")
	flag.BoolVarP(&clip, "clipboard", "c", false, "Copy the returned URL directly to the clipboard (needs xclip or xsel)")
	flag.StringVarP(&argname, "name", "n", "", "Specify the filename you want")
	flag.Parse()
	args := flag.Args()

	if len(args) > 0 {
		f, err := os.Open(args[0])
		check(err)
		defer f.Close()
		name = f.Name()
		datasource = f
		if progress {
			initBar(f)
		}
	} else {
		name = "stdin"
		datasource = os.Stdin
	}
	if tee {
		datasource = io.TeeReader(datasource, os.Stdout)
	}
	if argname != "" {
		name = argname
	}

	r, w := io.Pipe()
	multipartWriter := multipart.NewWriter(w)
	go func() {
		var part io.Writer
		defer w.Close()
		if part, err = multipartWriter.CreateFormFile("file", name); err != nil {
			log.Fatal(err)
		}
		if progress {
			part = io.MultiWriter(part, bar)
		}
		if _, err = io.Copy(part, datasource); err != nil {
			log.Fatal(err)
		}
		if err = multipartWriter.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	resp, err := http.Post(service, multipartWriter.FormDataContentType(), r)
	check(err)
	defer resp.Body.Close()
	ret, err := ioutil.ReadAll(resp.Body)
	check(err)
	if clip {
		clipboard.WriteAll(string(ret))
		fmt.Print("Copied URL to clipboard\n")
	} else {
		fmt.Print(string(ret))
	}
}
<commit_msg>Show upload speed and total uploaded size even for stdin<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"time"

	"github.com/atotto/clipboard"
	"github.com/cheggaaa/pb"
	flag "github.com/ogier/pflag"
)

const (
	service = "https://up.depado.eu/"
	method  = "POST"
)

var (
	bar  *pb.ProgressBar
	name string
)

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

func initBar(f *os.File) {
	fi, err := f.Stat()
	if err != nil {
		fmt.Println("Could not stat", f.Name())
		os.Exit(1)
	}
	bar = pb.New64(fi.Size()).SetUnits(pb.U_BYTES).SetRefreshRate(time.Millisecond * 10)
	bar.ShowPercent = true
	bar.ShowSpeed = true
	bar.ShowTimeLeft = true
	bar.Start()
}

func initUnknownBar() {
	bar = pb.New64(0).SetUnits(pb.U_BYTES).SetRefreshRate(time.Millisecond * 10)
	bar.ShowSpeed = true
	bar.ShowCounters = true
	bar.ShowBar = false
	bar.Start()
}

func main() {
	var err error
	var datasource io.Reader

	var tee bool
	var progress bool
	var clip bool
	var argname string

	flag.BoolVarP(&tee, "tee", "t", false, "Displays stdin to stdout")
	flag.BoolVarP(&progress, "progress", "p", false, "Displays a progress bar")
	flag.BoolVarP(&clip, "clipboard", "c", false, "Copy the returned URL directly to the clipboard (needs xclip or xsel)")
	flag.StringVarP(&argname, "name", "n", "", "Specify the filename you want")
	flag.Parse()
	args := flag.Args()

	if len(args) > 0 {
		f, err := os.Open(args[0])
		check(err)
		defer f.Close()
		name = f.Name()
		datasource = f
		if progress {
			initBar(f)
		}
	} else {
		name = "stdin"
		datasource = os.Stdin
		if progress {
			initUnknownBar()
		}
	}
	if tee {
		datasource = io.TeeReader(datasource, os.Stdout)
	}
	if argname != "" {
		name = argname
	}

	r, w := io.Pipe()
	multipartWriter := multipart.NewWriter(w)
	go func() {
		var part io.Writer
		defer w.Close()
		if part, err = multipartWriter.CreateFormFile("file", name); err != nil {
			log.Fatal(err)
		}
		if progress {
			part = io.MultiWriter(part, bar)
		}
		if _, err = io.Copy(part, datasource); err != nil {
			log.Fatal(err)
		}
		if err = multipartWriter.Close(); err != nil {
			log.Fatal(err)
		}
	}()

	resp, err := http.Post(service, multipartWriter.FormDataContentType(), r)
	check(err)
	defer resp.Body.Close()
	ret, err := ioutil.ReadAll(resp.Body)
	check(err)
	if clip {
		clipboard.WriteAll(string(ret))
		fmt.Print("Copied URL to clipboard\n")
	} else {
		fmt.Print(string(ret))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"google.golang.org/protobuf/compiler/protogen"
)

func GenerateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	filename := file.GeneratedFilenamePrefix + ".http.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	g.P("// comment")
	g.P("package main")

	return g
}
<commit_msg>Added head comment and package<commit_after>package main

import (
	"fmt"

	"google.golang.org/protobuf/compiler/protogen"
)

func GenerateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
	filename := file.GeneratedFilenamePrefix + ".http.go"
	g := gen.NewGeneratedFile(filename, file.GoImportPath)

	g.P("// Code generated by protoc-gen-gohttp. DO NOT EDIT.")
	g.P(fmt.Sprintf("// source: %s.proto", file.GeneratedFilenamePrefix))
	g.P()
	g.P(fmt.Sprintf("package %s", file.GoPackageName))

	return g
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/jmhodges/howsmyssl/tls"
	"strings"
)

type Rating string

const (
	okay       Rating = "Probably Okay"
	improvable Rating = "Improvable"
	bad        Rating = "Bad"
)

type clientInfo struct {
	GivenCipherSuites           []string            `json:"given_cipher_suites"`
	EphemeralKeysSupported      bool                `json:"ephemeral_keys_supported"`       // good if true
	SessionTicketsSupported     bool                `json:"session_ticket_supported"`       // good if true
	TLSCompressionSupported     bool                `json:"tls_compression_supported"`      // bad if true
	UnknownCipherSuiteSupported bool                `json:"unknown_cipher_suite_supported"` // bad if true
	BEASTVuln                   bool                `json:"beast_vuln"`                     // bad if true
	AbleToDetectNMinusOneSplitting bool             `json:"able_to_detect_n_minus_one_splitting"` // neutral
	InsecureCipherSuites        map[string][]string `json:"insecure_cipher_suites"`
	TLSVersion                  string              `json:"tls_version"`
	Rating                      Rating              `json:"rating"`
}

func ClientInfo(c *conn) *clientInfo {
	d := &clientInfo{InsecureCipherSuites: make(map[string][]string)}

	c.handshakeMutex.Lock()
	defer c.handshakeMutex.Unlock()

	for _, ci := range c.st.ClientHello.CipherSuites {
		s, found := allCipherSuites[ci]
		if found {
			if strings.Contains(s, "DHE_") {
				d.EphemeralKeysSupported = true
			}
			if c.HasBeastVulnSuites {
				d.BEASTVuln = !c.NMinusOneRecordSplittingDetected
				d.AbleToDetectNMinusOneSplitting = c.AbleToDetectNMinusOneSplitting
			}
			if fewBitCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], fewBitReason)
			}
			if nullCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], nullReason)
			}
			if nullAuthCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], nullAuthReason)
			}
		} else {
			w, found := weirdNSSSuites[ci]
			if !found {
				d.UnknownCipherSuiteSupported = true
				s = fmt.Sprintf("Some unknown cipher suite: %#x", ci)
			} else {
				s = w
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], weirdNSSReason)
			}
		}
		d.GivenCipherSuites = append(d.GivenCipherSuites, s)
	}
	d.SessionTicketsSupported = c.st.ClientHello.TicketSupported

	for _, cm := range c.st.ClientHello.CompressionMethods {
		if cm != 0x0 {
			d.TLSCompressionSupported = true
			break
		}
	}
	vers := c.st.ClientHello.Vers
	switch vers {
	case tls.VersionSSL30:
		d.TLSVersion = "SSL 3.0"
	case tls.VersionTLS10:
		d.TLSVersion = "TLS 1.0"
	case tls.VersionTLS11:
		d.TLSVersion = "TLS 1.1"
	case tls.VersionTLS12:
		d.TLSVersion = "TLS 1.2"
	default:
		d.TLSVersion = "an unknown version of SSL/TLS"
	}
	d.Rating = okay

	if !d.EphemeralKeysSupported || !d.SessionTicketsSupported || vers == tls.VersionTLS11 {
		d.Rating = improvable
	}

	if d.TLSCompressionSupported ||
		d.UnknownCipherSuiteSupported ||
		d.BEASTVuln ||
		len(d.InsecureCipherSuites) != 0 ||
		vers <= tls.VersionTLS10 {
		d.Rating = bad
	}
	return d
}
<commit_msg>gofmt client_info.go<commit_after>package main

import (
	"fmt"
	"github.com/jmhodges/howsmyssl/tls"
	"strings"
)

type Rating string

const (
	okay       Rating = "Probably Okay"
	improvable Rating = "Improvable"
	bad        Rating = "Bad"
)

type clientInfo struct {
	GivenCipherSuites              []string            `json:"given_cipher_suites"`
	EphemeralKeysSupported         bool                `json:"ephemeral_keys_supported"`             // good if true
	SessionTicketsSupported        bool                `json:"session_ticket_supported"`             // good if true
	TLSCompressionSupported        bool                `json:"tls_compression_supported"`            // bad if true
	UnknownCipherSuiteSupported    bool                `json:"unknown_cipher_suite_supported"`       // bad if true
	BEASTVuln                      bool                `json:"beast_vuln"`                           // bad if true
	AbleToDetectNMinusOneSplitting bool                `json:"able_to_detect_n_minus_one_splitting"` // neutral
	InsecureCipherSuites           map[string][]string `json:"insecure_cipher_suites"`
	TLSVersion                     string              `json:"tls_version"`
	Rating                         Rating              `json:"rating"`
}

func ClientInfo(c *conn) *clientInfo {
	d := &clientInfo{InsecureCipherSuites: make(map[string][]string)}

	c.handshakeMutex.Lock()
	defer c.handshakeMutex.Unlock()

	for _, ci := range c.st.ClientHello.CipherSuites {
		s, found := allCipherSuites[ci]
		if found {
			if strings.Contains(s, "DHE_") {
				d.EphemeralKeysSupported = true
			}
			if c.HasBeastVulnSuites {
				d.BEASTVuln = !c.NMinusOneRecordSplittingDetected
				d.AbleToDetectNMinusOneSplitting = c.AbleToDetectNMinusOneSplitting
			}
			if fewBitCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], fewBitReason)
			}
			if nullCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], nullReason)
			}
			if nullAuthCipherSuites[s] {
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], nullAuthReason)
			}
		} else {
			w, found := weirdNSSSuites[ci]
			if !found {
				d.UnknownCipherSuiteSupported = true
				s = fmt.Sprintf("Some unknown cipher suite: %#x", ci)
			} else {
				s = w
				d.InsecureCipherSuites[s] = append(d.InsecureCipherSuites[s], weirdNSSReason)
			}
		}
		d.GivenCipherSuites = append(d.GivenCipherSuites, s)
	}
	d.SessionTicketsSupported = c.st.ClientHello.TicketSupported

	for _, cm := range c.st.ClientHello.CompressionMethods {
		if cm != 0x0 {
			d.TLSCompressionSupported = true
			break
		}
	}
	vers := c.st.ClientHello.Vers
	switch vers {
	case tls.VersionSSL30:
		d.TLSVersion = "SSL 3.0"
	case tls.VersionTLS10:
		d.TLSVersion = "TLS 1.0"
	case tls.VersionTLS11:
		d.TLSVersion = "TLS 1.1"
	case tls.VersionTLS12:
		d.TLSVersion = "TLS 1.2"
	default:
		d.TLSVersion = "an unknown version of SSL/TLS"
	}
	d.Rating = okay

	if !d.EphemeralKeysSupported || !d.SessionTicketsSupported || vers == tls.VersionTLS11 {
		d.Rating = improvable
	}

	if d.TLSCompressionSupported ||
		d.UnknownCipherSuiteSupported ||
		d.BEASTVuln ||
		len(d.InsecureCipherSuites) != 0 ||
		vers <= tls.VersionTLS10 {
		d.Rating = bad
	}
	return d
}
<|endoftext|>


<commit_before>package eneru

import (
	"bytes"
	"encoding/json"
	"github.com/plimble/tsplitter"
	"github.com/stretchr/testify/suite"
	"net/http"
	"net/http/httptest"
	"testing"
)

type ClientSuite struct {
	suite.Suite
	server *httptest.Server
	client *Client
}

func TestClientSuite(t *testing.T) {
	suite.Run(t, &ClientSuite{})
}

func (t *ClientSuite) SetupSuite() {
	t.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	}))

	t.client, _ = NewClient(t.server.URL, 512)
	t.client.dict = tsplitter.NewFileDict("./dictionary.txt")
}

func (t *ClientSuite) TestSplitString() {
	sampleArrayInt := []int{
		45,
		2124,
	}

	sampleJson := map[string]interface{}{
		"Name":   "พลังงานไฟฟ้า",
		"Detail": "พลังงานเกิดจากแสงอาทิตย์",
		"ISBN":   12321342,
		"Tags": []string{
			"ชีวจิตสุขภาพ",
			"การเมืองที่ทำงาน",
			"งานบ้านออฟฟิตคอนโด",
		},
		"Codes": sampleArrayInt,
	}

	bj, err := json.Marshal(sampleJson)
	t.NoError(err)

	b, err := t.client.splitString(bytes.NewBuffer(bj))
	t.NoError(err)

	var data map[string]interface{}
	err = json.Unmarshal(b.Bytes(), &data)
	t.NoError(err)

	t.Equal("พลังงานไฟฟ้า", data["Name"])
	t.Equal("พลังงาน เกิด จาก แสงอาทิตย์", data["Detail"])

	dataInt := int(data["ISBN"].(float64))
	t.Equal(12321342, dataInt)

	setInt, _ := data["Codes"].([]int)
	for i, n := range setInt {
		t.Equal(sampleArrayInt[i], n)
	}

	resultString := []string{
		"ชีวจิต สุขภาพ",
		"การเมือง ที่ทำงาน",
		"งานบ้าน ออฟฟิต คอนโด",
	}
	setString, _ := data["Tags"].([]interface{})
	for i, n := range setString {
		t.Equal(resultString[i], n)
	}
}
<commit_msg>fix test<commit_after>package eneru

import (
	"bytes"
	"encoding/json"
	"github.com/stretchr/testify/suite"
	"net/http"
	"net/http/httptest"
	"testing"
)

type ClientSuite struct {
	suite.Suite
	server *httptest.Server
	client *Client
}

func TestClientSuite(t *testing.T) {
	suite.Run(t, &ClientSuite{})
}

func (t *ClientSuite) SetupSuite() {
	t.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
	}))

	t.client, _ = NewClient(t.server.URL, 512)
	t.client.tsplitterEnable("./dictionary.txt")
}

func (t *ClientSuite) TestSplitString() {
	sampleArrayInt := []int{
		45,
		2124,
	}

	sampleJson := map[string]interface{}{
		"Name":   "พลังงานไฟฟ้า",
		"Detail": "พลังงานเกิดจากแสงอาทิตย์",
		"ISBN":   12321342,
		"Tags": []string{
			"ชีวจิตสุขภาพ",
			"การเมืองที่ทำงาน",
			"งานบ้านออฟฟิตคอนโด",
		},
		"Codes": sampleArrayInt,
	}

	bj, err := json.Marshal(sampleJson)
	t.NoError(err)

	b, err := t.client.splitString(bytes.NewBuffer(bj))
	t.NoError(err)

	var data map[string]interface{}
	err = json.Unmarshal(b.Bytes(), &data)
	t.NoError(err)

	t.Equal("พลังงานไฟฟ้า", data["Name"])
	t.Equal("พลังงาน เกิด จาก แสงอาทิตย์", data["Detail"])

	dataInt := int(data["ISBN"].(float64))
	t.Equal(12321342, dataInt)

	setInt, _ := data["Codes"].([]int)
	for i, n := range setInt {
		t.Equal(sampleArrayInt[i], n)
	}

	resultString := []string{
		"ชีวจิต สุขภาพ",
		"การเมือง ที่ทำงาน",
		"งานบ้าน ออฟฟิต คอนโด",
	}
	setString, _ := data["Tags"].([]interface{})
	for i, n := range setString {
		t.Equal(resultString[i], n)
	}
}
<|endoftext|>


<commit_before>package tvdb_test

import (
	"testing"

	"github.com/danesparza/tvdb"
)

func TestTVDB_Login_ReturnsToken(t *testing.T) {
	//	Arrange
	request := tvdb.AuthRequest{}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.Login(request)

	//	Assert
	if err != nil {
		t.Errorf("Error logging in: %v", err)
	}

	if response.Token == "" {
		t.Errorf("The token is blank, and shouldn't be")
	} else {
		t.Logf("Got a token back: %v", response.Token)
	}
}

func TestTVDB_SeriesSearch_ReturnsInformation(t *testing.T) {
	//	Arrange
	request := tvdb.SearchRequest{
		Name: "Looney Tunes"}

	//	Act
	client := tvdb.TVDBClient{}
	matches, err := client.SeriesSearch(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(matches) == 0 {
		t.Errorf("There are no matches")
	}

	if matches[0].Id != 72514 {
		t.Errorf("Didn't get the series ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_ReturnsInformation(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 72514}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) == 0 {
		t.Errorf("There are no responses")
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	if response[0].Id != 5657563 {
		t.Errorf("Didn't get the episode ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_ReturnsExpectedCount(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 305288}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) != 8 {
		t.Errorf("8 episodes expected, but got %v instead", len(response))
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	if response[0].Id != 5468124 {
		t.Errorf("Didn't get the episode ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_CanMap(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 72514}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) == 0 {
		t.Errorf("Didn't get any episodes")
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	//	Load up the map
	episodes := make(map[string]*tvdb.EpisodeResponse)
	for _, episode := range response {
		episodes[episode.EpisodeName] = &episode
	}

	t.Logf("Created a map with %v items in it", len(episodes))

	//	Check to see if the episode name exists
	//	and then get its season/episode number:
	episodeToFind := "Upswept Hare"
	if episode, ok := episodes[episodeToFind]; ok {
		t.Logf("Found matching episode: s%ve%v", episode.AiredSeason, episode.AiredEpisodeNumber)
	} else {
		t.Errorf("Didn't find the episode '%v'", episodeToFind)
	}
}
<commit_msg>Fixed bug in map test<commit_after>package tvdb_test

import (
	"testing"

	"github.com/danesparza/tvdb"
)

func TestTVDB_Login_ReturnsToken(t *testing.T) {
	//	Arrange
	request := tvdb.AuthRequest{}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.Login(request)

	//	Assert
	if err != nil {
		t.Errorf("Error logging in: %v", err)
	}

	if response.Token == "" {
		t.Errorf("The token is blank, and shouldn't be")
	} else {
		t.Logf("Got a token back: %v", response.Token)
	}
}

func TestTVDB_SeriesSearch_ReturnsInformation(t *testing.T) {
	//	Arrange
	request := tvdb.SearchRequest{
		Name: "Looney Tunes"}

	//	Act
	client := tvdb.TVDBClient{}
	matches, err := client.SeriesSearch(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(matches) == 0 {
		t.Errorf("There are no matches")
	}

	if matches[0].Id != 72514 {
		t.Errorf("Didn't get the series ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_ReturnsInformation(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 72514}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) == 0 {
		t.Errorf("There are no responses")
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	if response[0].Id != 5657563 {
		t.Errorf("Didn't get the episode ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_ReturnsExpectedCount(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 305288}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) != 8 {
		t.Errorf("8 episodes expected, but got %v instead", len(response))
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	if response[0].Id != 5468124 {
		t.Errorf("Didn't get the episode ID back that we expected")
	}
}

func TestTVDB_EpisodesForSeries_CanMap(t *testing.T) {
	//	Arrange
	request := tvdb.EpisodeRequest{
		SeriesId: 72514}

	//	Act
	client := tvdb.TVDBClient{}
	response, err := client.EpisodesForSeries(request)

	//	Assert
	if err != nil {
		t.Errorf("Error getting search results: %v", err)
	}

	if len(response) == 0 {
		t.Errorf("Didn't get any episodes")
	} else {
		t.Logf("Got %v episodes back", len(response))
	}

	//	Load up the map
	episodes := make(map[string]tvdb.EpisodeResponse)
	for _, episode := range response {
		episodes[episode.EpisodeName] = episode
	}

	t.Logf("Created a map with %v items in it", len(episodes))

	//	Check to see if the episode name exists
	//	and then get its season/episode number:
	episodeToFind := "Upswept Hare"
	if episode, ok := episodes[episodeToFind]; ok {
		if episode.AiredSeason != 1953 || episode.AiredEpisodeNumber != 7 {
			t.Errorf("The episode and season don't match what we expect. Expected s1953e7 - Found: s%ve%v", episode.AiredSeason, episode.AiredEpisodeNumber)
		}
	} else {
		t.Errorf("Didn't find the episode '%v'", episodeToFind)
	}
}
<|endoftext|>


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

// Clock specifies a time of day with resolution to the nearest millisecond.
//
package clock

import (
	"math"
	"time"
)

// Clock specifies a time of day. It complements the existing time.Duration, applying
// that to the time since midnight (on some arbitrary day in some arbitrary timezone).
// The resolution is to the nearest millisecond, unlike time.Duration (which has nanosecond
// resolution).
//
// It is not intended that Clock be used to represent periods greater than 24 hours nor
// negative values. However, for such lengths of time, a fixed 24 hours per day
// is assumed and a modulo operation Mod24 is provided to discard whole multiples of 24 hours.
//
// See https://en.wikipedia.org/wiki/ISO_8601#Times
type Clock int32

// Common durations - second, minute, hour and day.
const (
	// Second is one second; it has a similar meaning to time.Second.
	Second Clock = Clock(time.Second / time.Millisecond)
	//	ClockSecond Clock = Clock(time.Second / time.Millisecond)

	// Minute is one minute; it has a similar meaning to time.Minute.
	Minute Clock = Clock(time.Minute / time.Millisecond)
	//	ClockMinute Clock = Clock(time.Minute / time.Millisecond)

	// Hour is one hour; it has a similar meaning to time.Hour.
	Hour Clock = Clock(time.Hour / time.Millisecond)
	//	ClockHour Clock = Clock(time.Hour / time.Millisecond)

	// Day is a fixed period of 24 hours. This does not take account of daylight savings,
	// so is not fully general.
	Day Clock = Clock(time.Hour * 24 / time.Millisecond)

//	ClockDay Clock = Clock(time.Hour * 24 / time.Millisecond)
)

// Midnight is the zero value of a Clock.
const Midnight Clock = 0

// Noon is at 12pm.
const Noon Clock = Hour * 12

// Undefined is provided because the zero value of a Clock *is* defined (i.e. Midnight).
// So a special value is chosen, which is math.MinInt32.
const Undefined Clock = Clock(math.MinInt32)

// New returns a new Clock with specified hour, minute, second and millisecond.
func New(hour, minute, second, millisec int) Clock {
	hx := Clock(hour) * Hour
	mx := Clock(minute) * Minute
	sx := Clock(second) * Second
	return Clock(hx + mx + sx + Clock(millisec))
}

// NewAt returns a new Clock with specified hour, minute, second and millisecond.
func NewAt(t time.Time) Clock {
	hour, minute, second := t.Clock()
	hx := Clock(hour) * Hour
	mx := Clock(minute) * Minute
	sx := Clock(second) * Second
	ms := Clock(t.Nanosecond() / int(time.Millisecond))
	return Clock(hx + mx + sx + ms)
}

// SinceMidnight returns a new Clock based on a duration since some arbitrary midnight.
func SinceMidnight(d time.Duration) Clock {
	return Clock(d / time.Millisecond)
}

// DurationSinceMidnight convert a clock to a time.Duration since some arbitrary midnight.
func (c Clock) DurationSinceMidnight() time.Duration {
	return time.Duration(c) * time.Millisecond
}

// Add returns a new Clock offset from this clock specified hour, minute, second and millisecond.
// The parameters can be negative.
// If required, use Mod24() to correct any overflow or underflow.
func (c Clock) Add(h, m, s, ms int) Clock {
	hx := Clock(h) * Hour
	mx := Clock(m) * Minute
	sx := Clock(s) * Second
	return c + hx + mx + sx + Clock(ms)
}

// IsInOneDay tests whether a clock time is in the range 0 to 24 hours, inclusive. Inside this
// range, a Clock is generally well-behaved. But outside it, there may be errors due to daylight
// savings. Note that 24:00:00 is included as a special case as per ISO-8601 definition of midnight.
func (c Clock) IsInOneDay() bool {
	return Midnight <= c && c <= Day
}

// IsMidnight tests whether a clock time is midnight. This is shorthand for c.Mod24() == 0.
// For large values, this assumes that every day has 24 hours.
func (c Clock) IsMidnight() bool {
	return c.Mod24() == Midnight
}

// Mod24 calculates the remainder vs 24 hours using Euclidean division, in which the result
// will be less than 24 hours and is never negative. Note that this imposes the assumption that
// every day has 24 hours (not correct when daylight saving changes in any timezone).
//
// https://en.wikipedia.org/wiki/Modulo_operation
func (c Clock) Mod24() Clock {
	if Midnight <= c && c < Day {
		return c
	}
	if c < Midnight {
		q := 1 - c/Day
		m := c + (q * Day)
		if m == Day {
			m = Midnight
		}
		return m
	}
	q := c / Day
	return c - (q * Day)
}

// Days gets the number of whole days represented by the Clock, assuming that each day is a fixed
// 24 hour period. Negative values are treated so that the range -23h59m59s to -1s is fully
// enclosed in a day numbered -1, and so on. This means that the result is zero only for the
// clock range 0s to 23h59m59s, for which IsInOneDay() returns true.
func (c Clock) Days() int {
	if c < Midnight {
		return int(c/Day) - 1
	} else {
		return int(c / Day)
	}
}

// Hours gets the clock-face number of hours (calculated from the modulo time, see Mod24).
func (c Clock) Hours() int {
	return int(clockHours(c.Mod24()))
}

// Minutes gets the clock-face number of minutes (calculated from the modulo time, see Mod24).
// For example, for 22:35 this will return 35.
func (c Clock) Minutes() int {
	return int(clockMinutes(c.Mod24()))
}

// Seconds gets the clock-face number of seconds (calculated from the modulo time, see Mod24).
// For example, for 10:20:30 this will return 30.
func (c Clock) Seconds() int {
	return int(clockSeconds(c.Mod24()))
}

// Millisec gets the clock-face number of milliseconds (calculated from the modulo time, see Mod24).
// For example, for 10:20:30.456 this will return 456.
func (c Clock) Millisec() int {
	return int(clockMillisec(c.Mod24()))
}
<commit_msg>Renamed constants<commit_after>// Copyright 2015 Rick Beton. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Clock specifies a time of day with resolution to the nearest millisecond.
//
package clock

import (
	"math"
	"time"
)

// Clock specifies a time of day. It complements the existing time.Duration, applying
// that to the time since midnight (on some arbitrary day in some arbitrary timezone).
// The resolution is to the nearest millisecond, unlike time.Duration (which has nanosecond
// resolution).
//
// It is not intended that Clock be used to represent periods greater than 24 hours nor
// negative values. However, for such lengths of time, a fixed 24 hours per day
// is assumed and a modulo operation Mod24 is provided to discard whole multiples of 24 hours.
//
// See https://en.wikipedia.org/wiki/ISO_8601#Times
type Clock int32

// Common durations - second, minute, hour and day.
const (
	// Second is one second; it has a similar meaning to time.Second.
	Second      Clock = Clock(time.Second / time.Millisecond)
	ClockSecond Clock = Clock(time.Second / time.Millisecond)

	// Minute is one minute; it has a similar meaning to time.Minute.
	Minute      Clock = Clock(time.Minute / time.Millisecond)
	ClockMinute Clock = Clock(time.Minute / time.Millisecond)

	// Hour is one hour; it has a similar meaning to time.Hour.
	Hour      Clock = Clock(time.Hour / time.Millisecond)
	ClockHour Clock = Clock(time.Hour / time.Millisecond)

	// Day is a fixed period of 24 hours. This does not take account of daylight savings,
	// so is not fully general.
	Day      Clock = Clock(time.Hour * 24 / time.Millisecond)
	ClockDay Clock = Clock(time.Hour * 24 / time.Millisecond)
)

// Midnight is the zero value of a Clock.
const Midnight Clock = 0

// Noon is at 12pm.
const Noon Clock = Hour * 12

// Undefined is provided because the zero value of a Clock *is* defined (i.e. Midnight).
// So a special value is chosen, which is math.MinInt32.
const Undefined Clock = Clock(math.MinInt32)

// New returns a new Clock with specified hour, minute, second and millisecond.
func New(hour, minute, second, millisec int) Clock {
	hx := Clock(hour) * Hour
	mx := Clock(minute) * Minute
	sx := Clock(second) * Second
	return Clock(hx + mx + sx + Clock(millisec))
}

// NewAt returns a new Clock with specified hour, minute, second and millisecond.
func NewAt(t time.Time) Clock {
	hour, minute, second := t.Clock()
	hx := Clock(hour) * Hour
	mx := Clock(minute) * Minute
	sx := Clock(second) * Second
	ms := Clock(t.Nanosecond() / int(time.Millisecond))
	return Clock(hx + mx + sx + ms)
}

// SinceMidnight returns a new Clock based on a duration since some arbitrary midnight.
func SinceMidnight(d time.Duration) Clock {
	return Clock(d / time.Millisecond)
}

// DurationSinceMidnight convert a clock to a time.Duration since some arbitrary midnight.
func (c Clock) DurationSinceMidnight() time.Duration {
	return time.Duration(c) * time.Millisecond
}

// Add returns a new Clock offset from this clock specified hour, minute, second and millisecond.
// The parameters can be negative.
// If required, use Mod24() to correct any overflow or underflow.
func (c Clock) Add(h, m, s, ms int) Clock {
	hx := Clock(h) * Hour
	mx := Clock(m) * Minute
	sx := Clock(s) * Second
	return c + hx + mx + sx + Clock(ms)
}

// IsInOneDay tests whether a clock time is in the range 0 to 24 hours, inclusive. Inside this
// range, a Clock is generally well-behaved. But outside it, there may be errors due to daylight
// savings. Note that 24:00:00 is included as a special case as per ISO-8601 definition of midnight.
func (c Clock) IsInOneDay() bool {
	return Midnight <= c && c <= Day
}

// IsMidnight tests whether a clock time is midnight. This is shorthand for c.Mod24() == 0.
// For large values, this assumes that every day has 24 hours.
func (c Clock) IsMidnight() bool {
	return c.Mod24() == Midnight
}

// Mod24 calculates the remainder vs 24 hours using Euclidean division, in which the result
// will be less than 24 hours and is never negative. Note that this imposes the assumption that
// every day has 24 hours (not correct when daylight saving changes in any timezone).
//
// https://en.wikipedia.org/wiki/Modulo_operation
func (c Clock) Mod24() Clock {
	if Midnight <= c && c < Day {
		return c
	}
	if c < Midnight {
		q := 1 - c/Day
		m := c + (q * Day)
		if m == Day {
			m = Midnight
		}
		return m
	}
	q := c / Day
	return c - (q * Day)
}

// Days gets the number of whole days represented by the Clock, assuming that each day is a fixed
// 24 hour period. Negative values are treated so that the range -23h59m59s to -1s is fully
// enclosed in a day numbered -1, and so on. This means that the result is zero only for the
// clock range 0s to 23h59m59s, for which IsInOneDay() returns true.
func (c Clock) Days() int {
	if c < Midnight {
		return int(c/Day) - 1
	} else {
		return int(c / Day)
	}
}

// Hours gets the clock-face number of hours (calculated from the modulo time, see Mod24).
func (c Clock) Hours() int {
	return int(clockHours(c.Mod24()))
}

// Minutes gets the clock-face number of minutes (calculated from the modulo time, see Mod24).
// For example, for 22:35 this will return 35.
func (c Clock) Minutes() int {
	return int(clockMinutes(c.Mod24()))
}

// Seconds gets the clock-face number of seconds (calculated from the modulo time, see Mod24).
// For example, for 10:20:30 this will return 30.
func (c Clock) Seconds() int {
	return int(clockSeconds(c.Mod24()))
}

// Millisec gets the clock-face number of milliseconds (calculated from the modulo time, see Mod24).
// For example, for 10:20:30.456 this will return 456.
func (c Clock) Millisec() int {
	return int(clockMillisec(c.Mod24()))
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Ceph-CSI Authors.

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

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

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

package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"

	"github.com/ceph/ceph-csi/internal/cephfs"
	"github.com/ceph/ceph-csi/internal/controller"
	"github.com/ceph/ceph-csi/internal/controller/persistentvolume"
	"github.com/ceph/ceph-csi/internal/liveness"
	rbddriver "github.com/ceph/ceph-csi/internal/rbd/driver"
	"github.com/ceph/ceph-csi/internal/util"
	"github.com/ceph/ceph-csi/internal/util/log"

	"k8s.io/klog/v2"
)

const (
	rbdType        = "rbd"
	cephFSType     = "cephfs"
	livenessType   = "liveness"
	controllerType = "controller"

	rbdDefaultName      = "rbd.csi.ceph.com"
	cephFSDefaultName   = "cephfs.csi.ceph.com"
	livenessDefaultName = "liveness.csi.ceph.com"

	pollTime     = 60 // seconds
	probeTimeout = 3  // seconds

	// use default namespace if namespace is not set.
	defaultNS = "default"

	defaultPluginPath  = "/var/lib/kubelet/plugins"
	defaultStagingPath = defaultPluginPath + "/kubernetes.io/csi/pv/"
)

var conf util.Config

func init() {
	// common flags
	flag.StringVar(&conf.Vtype, "type", "", "driver type [rbd|cephfs|liveness|controller]")
	flag.StringVar(&conf.Endpoint, "endpoint", "unix://tmp/csi.sock", "CSI endpoint")
	flag.StringVar(&conf.DriverName, "drivername", "", "name of the driver")
	flag.StringVar(&conf.DriverNamespace, "drivernamespace", defaultNS, "namespace in which driver is deployed")
	flag.StringVar(&conf.NodeID, "nodeid", "", "node id")
	flag.StringVar(&conf.PluginPath, "pluginpath", defaultPluginPath, "plugin path")
	flag.StringVar(&conf.StagingPath, "stagingpath", defaultStagingPath, "staging path")
	flag.StringVar(&conf.InstanceID, "instanceid", "", "Unique ID distinguishing this instance of Ceph CSI among other"+
		" instances, when sharing Ceph clusters across CSI instances for provisioning")
	flag.IntVar(&conf.PidLimit, "pidlimit", 0, "the PID limit to configure through cgroups")
	flag.BoolVar(&conf.IsControllerServer, "controllerserver", false, "start cephcsi controller server")
	flag.BoolVar(&conf.IsNodeServer, "nodeserver", false, "start cephcsi node server")
	flag.StringVar(
		&conf.DomainLabels,
		"domainlabels",
		"",
		"list of kubernetes node labels, that determines the topology"+
			" domain the node belongs to, separated by ','")

	// cephfs related flags
	flag.BoolVar(
		&conf.ForceKernelCephFS,
		"forcecephkernelclient",
		false,
		"enable Ceph Kernel clients on kernel < 4.17 which support quotas")

	// liveness/grpc metrics related flags
	flag.IntVar(&conf.MetricsPort, "metricsport", 8080, "TCP port for liveness/grpc metrics requests")
	flag.StringVar(
		&conf.MetricsPath,
		"metricspath",
		"/metrics",
		"path of prometheus endpoint where metrics will be available")
	flag.DurationVar(&conf.PollTime, "polltime", time.Second*pollTime, "time interval in seconds between each poll")
	flag.DurationVar(&conf.PoolTimeout, "timeout", time.Second*probeTimeout, "probe timeout in seconds")

	flag.BoolVar(&conf.EnableGRPCMetrics, "enablegrpcmetrics", false, "[DEPRECATED] enable grpc metrics")
	flag.StringVar(
		&conf.HistogramOption,
		"histogramoption",
		"0.5,2,6",
		"[DEPRECATED] Histogram option for grpc metrics, should be comma separated value, "+
			"ex:= 0.5,2,6 where start=0.5 factor=2, count=6")

	flag.UintVar(
		&conf.RbdHardMaxCloneDepth,
		"rbdhardmaxclonedepth",
		8,
		"Hard limit for maximum number of nested volume clones that are taken before a flatten occurs")
	flag.UintVar(
		&conf.RbdSoftMaxCloneDepth,
		"rbdsoftmaxclonedepth",
		4,
		"Soft limit for maximum number of nested volume clones that are taken before a flatten occurs")
	flag.UintVar(
		&conf.MaxSnapshotsOnImage,
		"maxsnapshotsonimage",
		450,
		"Maximum number of snapshots allowed on rbd image without flattening")
	flag.UintVar(
		&conf.MinSnapshotsOnImage,
		"minsnapshotsonimage",
		250,
		"Minimum number of snapshots required on rbd image to start flattening")
	flag.BoolVar(&conf.SkipForceFlatten, "skipforceflatten", false,
		"skip image flattening if kernel support mapping of rbd images which has the deep-flatten feature")

	flag.BoolVar(&conf.Version, "version", false, "Print cephcsi version information")
	flag.BoolVar(&conf.EnableProfiling, "enableprofiling", false, "enable go profiling")

	// CSI-Addons configuration
	flag.StringVar(&conf.CSIAddonsEndpoint, "csi-addons-endpoint", "unix://tmp/csi-addons.sock", "CSI-Addons endpoint")

	klog.InitFlags(nil)
	if err := flag.Set("logtostderr", "true"); err != nil {
		klog.Exitf("failed to set logtostderr flag: %v", err)
	}
	flag.Parse()
}

func getDriverName() string {
	// was explicitly passed a driver name
	if conf.DriverName != "" {
		return conf.DriverName
	}
	// select driver name based on volume type
	switch conf.Vtype {
	case rbdType:
		return rbdDefaultName
	case cephFSType:
		return cephFSDefaultName
	case livenessType:
		return livenessDefaultName
	default:
		return ""
	}
}

func main() {
	if conf.Version {
		fmt.Println("Cephcsi Version:", util.DriverVersion)
		fmt.Println("Git Commit:", util.GitCommit)
		fmt.Println("Go Version:", runtime.Version())
		fmt.Println("Compiler:", runtime.Compiler)
		fmt.Printf("Platform: %s/%s\n", runtime.GOOS, runtime.GOARCH)
		if kv, err := util.GetKernelVersion(); err == nil {
			fmt.Println("Kernel:", kv)
		}
		os.Exit(0)
	}
	log.DefaultLog("Driver version: %s and Git version: %s", util.DriverVersion, util.GitCommit)

	if conf.Vtype == "" {
		logAndExit("driver type not specified")
	}

	dname := getDriverName()
	err := util.ValidateDriverName(dname)
	if err != nil {
		logAndExit(err.Error())
	}

	// the driver may need a higher PID limit for handling all concurrent requests
	if conf.PidLimit != 0 {
		currentLimit, pidErr := util.GetPIDLimit()
		if pidErr != nil {
			klog.Errorf("Failed to get the PID limit, can not reconfigure: %v", pidErr)
		} else {
			log.DefaultLog("Initial PID limit is set to %d", currentLimit)
			err = util.SetPIDLimit(conf.PidLimit)
			switch {
			case err != nil:
				klog.Errorf("Failed to set new PID limit to %d: %v", conf.PidLimit, err)
			case conf.PidLimit == -1:
				log.DefaultLog("Reconfigured PID limit to %d (max)", conf.PidLimit)
			default:
				log.DefaultLog("Reconfigured PID limit to %d", conf.PidLimit)
			}
		}
	}

	if conf.EnableGRPCMetrics || conf.Vtype == livenessType {
		// validate metrics endpoint
		conf.MetricsIP = os.Getenv("POD_IP")

		if conf.MetricsIP == "" {
			klog.Warning("missing POD_IP env var defaulting to 0.0.0.0")
			conf.MetricsIP = "0.0.0.0"
		}
		err = util.ValidateURL(&conf)
		if err != nil {
			logAndExit(err.Error())
		}
	}

	if err = util.WriteCephConfig(); err != nil {
		log.FatalLogMsg("failed to write ceph configuration file (%v)", err)
	}

	log.DefaultLog("Starting driver type: %v with name: %v", conf.Vtype, dname)
	switch conf.Vtype {
	case rbdType:
		validateCloneDepthFlag(&conf)
		validateMaxSnaphostFlag(&conf)
		driver := rbddriver.NewDriver()
		driver.Run(&conf)

	case cephFSType:
		driver := cephfs.NewDriver()
		driver.Run(&conf)

	case livenessType:
		liveness.Run(&conf)

	case controllerType:
		cfg := controller.Config{
			DriverName: dname,
			Namespace:  conf.DriverNamespace,
		}
		// initialize all controllers before starting.
		initControllers()
		err = controller.Start(cfg)
		if err != nil {
			logAndExit(err.Error())
		}
	}

	os.Exit(0)
}

// initControllers will initialize all the controllers.
func initControllers() {
	// Add list of controller here.
	persistentvolume.Init()
}

func validateCloneDepthFlag(conf *util.Config) {
	// keeping hardlimit to 14 as max to avoid max image depth
	if conf.RbdHardMaxCloneDepth == 0 || conf.RbdHardMaxCloneDepth > 14 {
		logAndExit("rbdhardmaxclonedepth flag value should be between 1 and 14")
	}

	if conf.RbdSoftMaxCloneDepth > conf.RbdHardMaxCloneDepth {
		logAndExit("rbdsoftmaxclonedepth flag value should not be greater than rbdhardmaxclonedepth")
	}
}

func validateMaxSnaphostFlag(conf *util.Config) {
	// maximum number of snapshots on an image are 510 [1] and 16 images in
	// a parent/child chain [2],keeping snapshot limit to 500 to avoid issues.
	// [1] https://github.com/torvalds/linux/blob/master/drivers/block/rbd.c#L98
	// [2] https://github.com/torvalds/linux/blob/master/drivers/block/rbd.c#L92
	if conf.MaxSnapshotsOnImage == 0 || conf.MaxSnapshotsOnImage > 500 {
		logAndExit("maxsnapshotsonimage flag value should be between 1 and 500")
	}

	if conf.MinSnapshotsOnImage > conf.MaxSnapshotsOnImage {
		logAndExit("minsnapshotsonimage flag value should be less than maxsnapshotsonimage")
	}
}

func logAndExit(msg string) {
	klog.Errorln(msg)
	os.Exit(1)
}
<commit_msg>deploy: fix default URL for --csi-addons-endpoint option<commit_after>/*
Copyright 2019 The Ceph-CSI Authors.

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

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

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

package main

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"

	"github.com/ceph/ceph-csi/internal/cephfs"
	"github.com/ceph/ceph-csi/internal/controller"
	"github.com/ceph/ceph-csi/internal/controller/persistentvolume"
	"github.com/ceph/ceph-csi/internal/liveness"
	rbddriver "github.com/ceph/ceph-csi/internal/rbd/driver"
	"github.com/ceph/ceph-csi/internal/util"
	"github.com/ceph/ceph-csi/internal/util/log"

	"k8s.io/klog/v2"
)

const (
	rbdType        = "rbd"
	cephFSType     = "cephfs"
	livenessType   = "liveness"
	controllerType = "controller"

	rbdDefaultName      = "rbd.csi.ceph.com"
	cephFSDefaultName   = "cephfs.csi.ceph.com"
	livenessDefaultName = "liveness.csi.ceph.com"

	pollTime     = 60 // seconds
	probeTimeout = 3  // seconds

	// use default namespace if namespace is not set.
	defaultNS = "default"

	defaultPluginPath  = "/var/lib/kubelet/plugins"
	defaultStagingPath = defaultPluginPath + "/kubernetes.io/csi/pv/"
)

var conf util.Config

func init() {
	// common flags
	flag.StringVar(&conf.Vtype, "type", "", "driver type [rbd|cephfs|liveness|controller]")
	flag.StringVar(&conf.Endpoint, "endpoint", "unix:///tmp/csi.sock", "CSI endpoint")
	flag.StringVar(&conf.DriverName, "drivername", "", "name of the driver")
	flag.StringVar(&conf.DriverNamespace, "drivernamespace", defaultNS, "namespace in which driver is deployed")
	flag.StringVar(&conf.NodeID, "nodeid", "", "node id")
	flag.StringVar(&conf.PluginPath, "pluginpath", defaultPluginPath, "plugin path")
	flag.StringVar(&conf.StagingPath, "stagingpath", defaultStagingPath, "staging path")
	flag.StringVar(&conf.InstanceID, "instanceid", "", "Unique ID distinguishing this instance of Ceph CSI among other"+
		" instances, when sharing Ceph clusters across CSI instances for provisioning")
	flag.IntVar(&conf.PidLimit, "pidlimit", 0, "the PID limit to configure through cgroups")
	flag.BoolVar(&conf.IsControllerServer, "controllerserver", false, "start cephcsi controller server")
	flag.BoolVar(&conf.IsNodeServer, "nodeserver", false, "start cephcsi node server")
	flag.StringVar(
		&conf.DomainLabels,
		"domainlabels",
		"",
		"list of kubernetes node labels, that determines the topology"+
			" domain the node belongs to, separated by ','")

	// cephfs related flags
	flag.BoolVar(
		&conf.ForceKernelCephFS,
		"forcecephkernelclient",
		false,
		"enable Ceph Kernel clients on kernel < 4.17 which support quotas")

	// liveness/grpc metrics related flags
	flag.IntVar(&conf.MetricsPort, "metricsport", 8080, "TCP port for liveness/grpc metrics requests")
	flag.StringVar(
		&conf.MetricsPath,
		"metricspath",
		"/metrics",
		"path of prometheus endpoint where metrics will be available")
	flag.DurationVar(&conf.PollTime, "polltime", time.Second*pollTime, "time interval in seconds between each poll")
	flag.DurationVar(&conf.PoolTimeout, "timeout", time.Second*probeTimeout, "probe timeout in seconds")

	flag.BoolVar(&conf.EnableGRPCMetrics, "enablegrpcmetrics", false, "[DEPRECATED] enable grpc metrics")
	flag.StringVar(
		&conf.HistogramOption,
		"histogramoption",
		"0.5,2,6",
		"[DEPRECATED] Histogram option for grpc metrics, should be comma separated value, "+
			"ex:= 0.5,2,6 where start=0.5 factor=2, count=6")

	flag.UintVar(
		&conf.RbdHardMaxCloneDepth,
		"rbdhardmaxclonedepth",
		8,
		"Hard limit for maximum number of nested volume clones that are taken before a flatten occurs")
	flag.UintVar(
		&conf.RbdSoftMaxCloneDepth,
		"rbdsoftmaxclonedepth",
		4,
		"Soft limit for maximum number of nested volume clones that are taken before a flatten occurs")
	flag.UintVar(
		&conf.MaxSnapshotsOnImage,
		"maxsnapshotsonimage",
		450,
		"Maximum number of snapshots allowed on rbd image without flattening")
	flag.UintVar(
		&conf.MinSnapshotsOnImage,
		"minsnapshotsonimage",
		250,
		"Minimum number of snapshots required on rbd image to start flattening")
	flag.BoolVar(&conf.SkipForceFlatten, "skipforceflatten", false,
		"skip image flattening if kernel support mapping of rbd images which has the deep-flatten feature")

	flag.BoolVar(&conf.Version, "version", false, "Print cephcsi version information")
	flag.BoolVar(&conf.EnableProfiling, "enableprofiling", false, "enable go profiling")

	// CSI-Addons configuration
	flag.StringVar(&conf.CSIAddonsEndpoint, "csi-addons-endpoint", "unix:///tmp/csi-addons.sock", "CSI-Addons endpoint")

	klog.InitFlags(nil)
	if err := flag.Set("logtostderr", "true"); err != nil {
		klog.Exitf("failed to set logtostderr flag: %v", err)
	}
	flag.Parse()
}

func getDriverName() string {
	// was explicitly passed a driver name
	if conf.DriverName != "" {
		return conf.DriverName
	}
	// select driver name based on volume type
	switch conf.Vtype {
	case rbdType:
		return rbdDefaultName
	case cephFSType:
		return cephFSDefaultName
	case livenessType:
		return livenessDefaultName
	default:
		return ""
	}
}

func main() {
	if conf.Version {
		fmt.Println("Cephcsi Version:", util.DriverVersion)
		fmt.Println("Git Commit:", util.GitCommit)
		fmt.Println("Go Version:", runtime.Version())
		fmt.Println("Compiler:", runtime.Compiler)
		fmt.Printf("Platform: %s/%s\n", runtime.GOOS, runtime.GOARCH)
		if kv, err := util.GetKernelVersion(); err == nil {
			fmt.Println("Kernel:", kv)
		}
		os.Exit(0)
	}
	log.DefaultLog("Driver version: %s and Git version: %s", util.DriverVersion, util.GitCommit)

	if conf.Vtype == "" {
		logAndExit("driver type not specified")
	}

	dname := getDriverName()
	err := util.ValidateDriverName(dname)
	if err != nil {
		logAndExit(err.Error())
	}

	// the driver may need a higher PID limit for handling all concurrent requests
	if conf.PidLimit != 0 {
		currentLimit, pidErr := util.GetPIDLimit()
		if pidErr != nil {
			klog.Errorf("Failed to get the PID limit, can not reconfigure: %v", pidErr)
		} else {
			log.DefaultLog("Initial PID limit is set to %d", currentLimit)
			err = util.SetPIDLimit(conf.PidLimit)
			switch {
			case err != nil:
				klog.Errorf("Failed to set new PID limit to %d: %v", conf.PidLimit, err)
			case conf.PidLimit == -1:
				log.DefaultLog("Reconfigured PID limit to %d (max)", conf.PidLimit)
			default:
				log.DefaultLog("Reconfigured PID limit to %d", conf.PidLimit)
			}
		}
	}

	if conf.EnableGRPCMetrics || conf.Vtype == livenessType {
		// validate metrics endpoint
		conf.MetricsIP = os.Getenv("POD_IP")

		if conf.MetricsIP == "" {
			klog.Warning("missing POD_IP env var defaulting to 0.0.0.0")
			conf.MetricsIP = "0.0.0.0"
		}
		err = util.ValidateURL(&conf)
		if err != nil {
			logAndExit(err.Error())
		}
	}

	if err = util.WriteCephConfig(); err != nil {
		log.FatalLogMsg("failed to write ceph configuration file (%v)", err)
	}

	log.DefaultLog("Starting driver type: %v with name: %v", conf.Vtype, dname)
	switch conf.Vtype {
	case rbdType:
		validateCloneDepthFlag(&conf)
		validateMaxSnaphostFlag(&conf)
		driver := rbddriver.NewDriver()
		driver.Run(&conf)

	case cephFSType:
		driver := cephfs.NewDriver()
		driver.Run(&conf)

	case livenessType:
		liveness.Run(&conf)

	case controllerType:
		cfg := controller.Config{
			DriverName: dname,
			Namespace:  conf.DriverNamespace,
		}
		// initialize all controllers before starting.
		initControllers()
		err = controller.Start(cfg)
		if err != nil {
			logAndExit(err.Error())
		}
	}

	os.Exit(0)
}

// initControllers will initialize all the controllers.
func initControllers() {
	// Add list of controller here.
	persistentvolume.Init()
}

func validateCloneDepthFlag(conf *util.Config) {
	// keeping hardlimit to 14 as max to avoid max image depth
	if conf.RbdHardMaxCloneDepth == 0 || conf.RbdHardMaxCloneDepth > 14 {
		logAndExit("rbdhardmaxclonedepth flag value should be between 1 and 14")
	}

	if conf.RbdSoftMaxCloneDepth > conf.RbdHardMaxCloneDepth {
		logAndExit("rbdsoftmaxclonedepth flag value should not be greater than rbdhardmaxclonedepth")
	}
}

func validateMaxSnaphostFlag(conf *util.Config) {
	// maximum number of snapshots on an image are 510 [1] and 16 images in
	// a parent/child chain [2],keeping snapshot limit to 500 to avoid issues.
	// [1] https://github.com/torvalds/linux/blob/master/drivers/block/rbd.c#L98
	// [2] https://github.com/torvalds/linux/blob/master/drivers/block/rbd.c#L92
	if conf.MaxSnapshotsOnImage == 0 || conf.MaxSnapshotsOnImage > 500 {
		logAndExit("maxsnapshotsonimage flag value should be between 1 and 500")
	}

	if conf.MinSnapshotsOnImage > conf.MaxSnapshotsOnImage {
		logAndExit("minsnapshotsonimage flag value should be less than maxsnapshotsonimage")
	}
}

func logAndExit(msg string) {
	klog.Errorln(msg)
	os.Exit(1)
}
<|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 (
	"fmt"
	"sync"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/blance"
	"github.com/couchbaselabs/cbgt"
)

type rebalancer struct {
	version       string
	cfg           cbgt.Cfg
	server        string
	nodesAll      []string          // Array of node UUID's.
	nodesToAdd    []string          // Array of node UUID's.
	nodesToRemove []string          // Array of node UUID's.
	nodeWeights   map[string]int    // Keyed by node UUID.
	nodeHierarchy map[string]string // Keyed by node UUID.

	begIndexDefs    *cbgt.IndexDefs
	begNodeDefs     *cbgt.NodeDefs
	begPlanPIndexes *cbgt.PlanPIndexes

	m sync.Mutex // Protects the mutatable fields that follow.

	cas uint64

	endPlanPIndexes *cbgt.PlanPIndexes

	o *blance.Orchestrator

	// Map of index -> partition -> node -> stateOp.
	currStates map[string]map[string]map[string]stateOp
}

type stateOp struct {
	state string
	op    string // May be "" for unknown or no in-flight op.
}

// runRebalancer implements the "master, central planner (MCP)"
// rebalance workflow.
func runRebalancer(version string, cfg cbgt.Cfg, server string) (
	// TODO: Need to ensure that all nodes are up, especially those
	// that haven't been removed yet.
	//
	// TODO: Need timeouts on moves.
	changed bool, err error) {
	if cfg == nil { // Can occur during testing.
		return false, nil
	}

	uuid := "" // We don't have a uuid, as we're not a node.

	begIndexDefs, begNodeDefs, begPlanPIndexes, cas, err :=
		cbgt.PlannerGetPlan(cfg, version, uuid)
	if err != nil {
		return false, err
	}

	nodesAll, nodesToAdd, nodesToRemove,
		nodeWeights, nodeHierarchy :=
		cbgt.CalcNodesLayout(begIndexDefs, begNodeDefs, begPlanPIndexes)

	log.Printf("runRebalancer: nodesAll: %#v", nodesAll)
	log.Printf("runRebalancer: nodesToAdd: %#v", nodesToAdd)
	log.Printf("runRebalancer: nodesToRemove: %#v", nodesToRemove)
	log.Printf("runRebalancer: nodeWeights: %#v", nodeWeights)
	log.Printf("runRebalancer: nodeHierarchy: %#v", nodeHierarchy)
	log.Printf("runRebalancer: begIndexDefs: %#v", begIndexDefs)
	log.Printf("runRebalancer: begNodeDefs: %#v", begNodeDefs)
	log.Printf("runRebalancer: begPlanPIndexes: %#v, cas: %v",
		begPlanPIndexes, cas)

	r := &rebalancer{
		version:         version,
		cfg:             cfg,
		server:          server,
		nodesAll:        nodesAll,
		nodesToAdd:      nodesToAdd,
		nodesToRemove:   nodesToRemove,
		nodeWeights:     nodeWeights,
		nodeHierarchy:   nodeHierarchy,
		begIndexDefs:    begIndexDefs,
		begNodeDefs:     begNodeDefs,
		begPlanPIndexes: begPlanPIndexes,
		endPlanPIndexes: cbgt.NewPlanPIndexes(version),
		cas:             cas,
		currStates:      map[string]map[string]map[string]stateOp{},
	}

	// TODO: Prepopulate currStates so that we can double-check that
	// our state transitions(assignPartition) are valid.

	return r.run()
}

// The run method rebalances each index, one at a time.
func (r *rebalancer) run() (bool, error) {
	changedAny := false

	for _, indexDef := range r.begIndexDefs.IndexDefs {
		changed, err := r.runIndex(indexDef)
		if err != nil {
			log.Printf("run: indexDef.Name: %s, err: %#v",
				indexDef.Name, err)
		}

		changedAny = changedAny || changed
	}

	return changedAny, nil
}

// The runIndex method rebalances a single index.
func (r *rebalancer) runIndex(indexDef *cbgt.IndexDef) (
	changed bool, err error) {
	log.Printf(" runIndex: indexDef.Name: %s", indexDef.Name)

	r.m.Lock()
	if cbgt.CasePlanFrozen(indexDef, r.begPlanPIndexes, r.endPlanPIndexes) {
		r.m.Unlock()

		log.Printf("  plan frozen: indexDef.Name: %s,"+
			" cloned previous plan", indexDef.Name)

		return false, nil
	}
	r.m.Unlock()

	partitionModel, begMap, endMap, err := r.calcBegEndMaps(indexDef)
	if err != nil {
		return false, err
	}

	assignPartitionFunc := func(stopCh chan struct{},
		partition, node, state, op string) error {
		return r.assignPartition(stopCh,
			indexDef.Name, partition, node, state, op)
	}

	partitionStateFunc := func(stopCh chan struct{},
		partition string, node string) (
		state string, pct float32, err error) {
		return r.partitionState(stopCh,
			indexDef.Name, partition, node)
	}

	o, err := blance.OrchestrateMoves(
		partitionModel,
		blance.OrchestratorOptions{}, // TODO.
		r.nodesAll,
		begMap,
		endMap,
		assignPartitionFunc,
		partitionStateFunc,
		blance.LowestWeightPartitionMoveForNode) // TODO: concurrency.
	if err != nil {
		return false, err
	}

	r.m.Lock()
	r.o = o
	r.m.Unlock()

	numProgress := 0
	errs := []error(nil)

	for progress := range o.ProgressCh() {
		numProgress++

		log.Printf("   numProgress: %d,"+
			" indexDef.Name: %s, progress: %#v",
			numProgress, indexDef.Name, progress)

		errs = append(errs, progress.Errors...)
	}

	o.Stop()

	// TDOO: Check that the plan in the cfg should match our endMap...
	//
	// _, err = cbgt.CfgSetPlanPIndexes(cfg, planPIndexesFFwd, cas)
	// if err != nil {
	//     return false, fmt.Errorf("mcp: could not save new plan,"+
	//     " perhaps a concurrent planner won, cas: %d, err: %v",
	//     cas, err)
	// }

	if len(errs) > 0 {
		return true, errs[0] // TODO: Propogate errs better.
	}

	return true, err // TODO: compute proper change response.
}

func (r *rebalancer) calcBegEndMaps(indexDef *cbgt.IndexDef) (
	partitionModel blance.PartitionModel,
	begMap blance.PartitionMap,
	endMap blance.PartitionMap,
	err error) {
	r.m.Lock()
	defer r.m.Unlock()

	// The endPlanPIndexesForIndex is a working data structure that's
	// mutated as calcBegEndMaps progresses.
	endPlanPIndexesForIndex, err := cbgt.SplitIndexDefIntoPlanPIndexes(
		indexDef, r.server, r.endPlanPIndexes)
	if err != nil {
		log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
			" could not SplitIndexDefIntoPlanPIndexes,"+
			" indexDef: %#v, server: %s, err: %v",
			indexDef.Name, indexDef, r.server, err)

		return partitionModel, begMap, endMap, err
	}

	// Invoke blance to assign the endPlanPIndexesForIndex to nodes.
	warnings := cbgt.BlancePlanPIndexes(indexDef,
		endPlanPIndexesForIndex, r.begPlanPIndexes,
		r.nodesAll, r.nodesToAdd, r.nodesToRemove,
		r.nodeWeights, r.nodeHierarchy)

	r.endPlanPIndexes.Warnings[indexDef.Name] = warnings

	// TODO: handle blance ffwd plan warnings.

	for _, warning := range warnings {
		log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
			" BlancePlanPIndexes warning: %q, indexDef: %#v",
			indexDef.Name, warning, indexDef)
	}

	log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
		" endPlanPIndexes: %#v", indexDef.Name, r.endPlanPIndexes)

	partitionModel, _ = cbgt.BlancePartitionModel(indexDef)

	begMap = cbgt.BlanceMap(endPlanPIndexesForIndex, r.begPlanPIndexes)
	endMap = cbgt.BlanceMap(endPlanPIndexesForIndex, r.endPlanPIndexes)

	return partitionModel, begMap, endMap, nil
}

// --------------------------------------------------------

func (r *rebalancer) assignPartition(stopCh chan struct{},
	index, partition, node, state, op string) error {
	log.Printf("  assignPartitionFunc: index: %s,"+
		" partition: %s, node: %s, state: %s, op: %s",
		index, partition, node, state, op)

	var err error

	r.m.Lock()

	// Update currStates to the assigned index/partition/node/state.
	partitions, exists := r.currStates[index]
	if !exists || partitions == nil {
		partitions = map[string]map[string]stateOp{}
		r.currStates[index] = partitions
	}

	nodes, exists := partitions[partition]
	if !exists || nodes == nil {
		nodes = map[string]stateOp{}
		partitions[partition] = nodes
	}

	if op == "add" {
		if stateOp, exists := nodes[node]; exists && stateOp.state != "" {
			r.m.Unlock()

			return fmt.Errorf("assignPartition:"+
				" op was add when exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				" stateOp: %#v,"+
				index, partition, node, state, op, stateOp)
		}
	} else {
		if stateOp, exists := nodes[node]; !exists || stateOp.state == "" {
			r.m.Unlock()

			return fmt.Errorf("assignPartition:"+
				" op was non-add when not exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				" stateOp: %#v,"+
				index, partition, node, state, op, stateOp)
		}
	}

	nodes[node] = stateOp{state, op}

	r.m.Unlock()

	indexDefs, err := cbgt.PlannerGetIndexDefs(r.cfg, r.version)
	if err != nil {
		return err
	}

	indexDef := indexDefs.IndexDefs[index]
	if indexDef == nil {
		return fmt.Errorf("assignPartition: no indexDef,"+
			" index: %s, partition: %s, node: %s, state: %s, op: %s",
			index, partition, node, state, op)
	}

	planPIndexes, cas, err :=
		cbgt.PlannerGetPlanPIndexes(r.cfg, r.version)
	if err != nil {
		return err
	}

	planPIndex := planPIndexes.PlanPIndexes[partition]
	if planPIndex == nil {
		r.m.Lock()
		endPlanPIndex := r.endPlanPIndexes.PlanPIndexes[partition]
		if endPlanPIndex != nil {
			p := *endPlanPIndex // Copy.
			planPIndex = &p
			planPIndex.Nodes = nil
			planPIndexes.PlanPIndexes[partition] = planPIndex
		}
		r.m.Unlock()
	}

	if planPIndex == nil {
		return fmt.Errorf("assignPartition: no planPIndex,"+
			" index: %s, partition: %s, node: %s, state: %s, op: %s",
			index, partition, node, state, op)
	}

	if planPIndex.Nodes == nil {
		planPIndex.Nodes = make(map[string]*cbgt.PlanPIndexNode)
	}

	planPIndex.UUID = cbgt.NewUUID()

	canRead := true
	canWrite := true
	nodePlanParam := cbgt.GetNodePlanParam(
		indexDef.PlanParams.NodePlanParams, node,
		indexDef.Name, partition)
	if nodePlanParam != nil {
		canRead = nodePlanParam.CanRead
		canWrite = nodePlanParam.CanWrite
	}

	priority := 0
	if state == "replica" {
		priority = len(planPIndex.Nodes)
	}

	if op == "add" {
		if planPIndex.Nodes[node] != nil {
			return fmt.Errorf("assignPartition: planPIndex already exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s,"+
				" planPIndex: %#v",
				index, partition, node, state, op, planPIndex)
		}

		// TODO: Need to shift the other node priorities around?
		planPIndex.Nodes[node] = &cbgt.PlanPIndexNode{
			CanRead:  canRead,
			CanWrite: canWrite,
			Priority: priority,
		}
	} else {
		if planPIndex.Nodes[node] == nil {
			return fmt.Errorf("assignPartition: planPIndex missing,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				index, partition, node, state, op)
		}

		if op == "del" {
			// TODO: Need to shift the other node priorities around?
			delete(planPIndex.Nodes, node)
		} else {
			// TODO: Need to shift the other node priorities around?
			planPIndex.Nodes[node] = &cbgt.PlanPIndexNode{
				CanRead:  canRead,
				CanWrite: canWrite,
				Priority: priority,
			}
		}
	}

	// TODO: stopCh handling.

	planPIndexes.UUID = cbgt.NewUUID()

	_, err = cbgt.CfgSetPlanPIndexes(r.cfg, planPIndexes, cas)
	if err != nil {
		return fmt.Errorf("assignPartition: update plan,"+
			" perhaps a concurrent planner won, cas: %d, err: %v",
			cas, err)
	}

	return nil
}

func (r *rebalancer) partitionState(stopCh chan struct{},
	index, partition, node string) (
	state string, pct float32, err error) {
	log.Printf("  partitionStateFunc: index: %s,"+
		" partition: %s, node: %s", index, partition, node)

	var stateOp stateOp

	r.m.Lock()
	if r.currStates[index] != nil &&
		r.currStates[index][partition] != nil {
		stateOp = r.currStates[index][partition][node]
	}
	r.m.Unlock()

	// TODO: real state & pct, with stopCh handling.

	return stateOp.state, 1.0, nil
}
<commit_msg>only send lastProgress errors<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 (
	"fmt"
	"sync"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/blance"
	"github.com/couchbaselabs/cbgt"
)

type rebalancer struct {
	version       string
	cfg           cbgt.Cfg
	server        string
	nodesAll      []string          // Array of node UUID's.
	nodesToAdd    []string          // Array of node UUID's.
	nodesToRemove []string          // Array of node UUID's.
	nodeWeights   map[string]int    // Keyed by node UUID.
	nodeHierarchy map[string]string // Keyed by node UUID.

	begIndexDefs    *cbgt.IndexDefs
	begNodeDefs     *cbgt.NodeDefs
	begPlanPIndexes *cbgt.PlanPIndexes

	m sync.Mutex // Protects the mutatable fields that follow.

	cas uint64

	endPlanPIndexes *cbgt.PlanPIndexes

	o *blance.Orchestrator

	// Map of index -> partition -> node -> stateOp.
	currStates map[string]map[string]map[string]stateOp
}

type stateOp struct {
	state string
	op    string // May be "" for unknown or no in-flight op.
}

// runRebalancer implements the "master, central planner (MCP)"
// rebalance workflow.
func runRebalancer(version string, cfg cbgt.Cfg, server string) (
	// TODO: Need to ensure that all nodes are up, especially those
	// that haven't been removed yet.
	//
	// TODO: Need timeouts on moves.
	changed bool, err error) {
	if cfg == nil { // Can occur during testing.
		return false, nil
	}

	uuid := "" // We don't have a uuid, as we're not a node.

	begIndexDefs, begNodeDefs, begPlanPIndexes, cas, err :=
		cbgt.PlannerGetPlan(cfg, version, uuid)
	if err != nil {
		return false, err
	}

	nodesAll, nodesToAdd, nodesToRemove,
		nodeWeights, nodeHierarchy :=
		cbgt.CalcNodesLayout(begIndexDefs, begNodeDefs, begPlanPIndexes)

	log.Printf("runRebalancer: nodesAll: %#v", nodesAll)
	log.Printf("runRebalancer: nodesToAdd: %#v", nodesToAdd)
	log.Printf("runRebalancer: nodesToRemove: %#v", nodesToRemove)
	log.Printf("runRebalancer: nodeWeights: %#v", nodeWeights)
	log.Printf("runRebalancer: nodeHierarchy: %#v", nodeHierarchy)
	log.Printf("runRebalancer: begIndexDefs: %#v", begIndexDefs)
	log.Printf("runRebalancer: begNodeDefs: %#v", begNodeDefs)
	log.Printf("runRebalancer: begPlanPIndexes: %#v, cas: %v",
		begPlanPIndexes, cas)

	r := &rebalancer{
		version:         version,
		cfg:             cfg,
		server:          server,
		nodesAll:        nodesAll,
		nodesToAdd:      nodesToAdd,
		nodesToRemove:   nodesToRemove,
		nodeWeights:     nodeWeights,
		nodeHierarchy:   nodeHierarchy,
		begIndexDefs:    begIndexDefs,
		begNodeDefs:     begNodeDefs,
		begPlanPIndexes: begPlanPIndexes,
		endPlanPIndexes: cbgt.NewPlanPIndexes(version),
		cas:             cas,
		currStates:      map[string]map[string]map[string]stateOp{},
	}

	// TODO: Prepopulate currStates so that we can double-check that
	// our state transitions(assignPartition) are valid.

	return r.run()
}

// The run method rebalances each index, one at a time.
func (r *rebalancer) run() (bool, error) {
	changedAny := false

	for _, indexDef := range r.begIndexDefs.IndexDefs {
		changed, err := r.runIndex(indexDef)
		if err != nil {
			log.Printf("run: indexDef.Name: %s, err: %#v",
				indexDef.Name, err)
		}

		changedAny = changedAny || changed
	}

	return changedAny, nil
}

// The runIndex method rebalances a single index.
func (r *rebalancer) runIndex(indexDef *cbgt.IndexDef) (
	changed bool, err error) {
	log.Printf(" runIndex: indexDef.Name: %s", indexDef.Name)

	r.m.Lock()
	if cbgt.CasePlanFrozen(indexDef, r.begPlanPIndexes, r.endPlanPIndexes) {
		r.m.Unlock()

		log.Printf("  plan frozen: indexDef.Name: %s,"+
			" cloned previous plan", indexDef.Name)

		return false, nil
	}
	r.m.Unlock()

	partitionModel, begMap, endMap, err := r.calcBegEndMaps(indexDef)
	if err != nil {
		return false, err
	}

	assignPartitionFunc := func(stopCh chan struct{},
		partition, node, state, op string) error {
		return r.assignPartition(stopCh,
			indexDef.Name, partition, node, state, op)
	}

	partitionStateFunc := func(stopCh chan struct{},
		partition string, node string) (
		state string, pct float32, err error) {
		return r.partitionState(stopCh,
			indexDef.Name, partition, node)
	}

	o, err := blance.OrchestrateMoves(
		partitionModel,
		blance.OrchestratorOptions{}, // TODO.
		r.nodesAll,
		begMap,
		endMap,
		assignPartitionFunc,
		partitionStateFunc,
		blance.LowestWeightPartitionMoveForNode) // TODO: concurrency.
	if err != nil {
		return false, err
	}

	r.m.Lock()
	r.o = o
	r.m.Unlock()

	numProgress := 0
	var lastProgress blance.OrchestratorProgress

	for progress := range o.ProgressCh() {
		numProgress++
		lastProgress = progress

		log.Printf("   numProgress: %d,"+
			" indexDef.Name: %s, progress: %#v",
			numProgress, indexDef.Name, progress)
	}

	o.Stop()

	// TDOO: Check that the plan in the cfg should match our endMap...
	//
	// _, err = cbgt.CfgSetPlanPIndexes(cfg, planPIndexesFFwd, cas)
	// if err != nil {
	//     return false, fmt.Errorf("mcp: could not save new plan,"+
	//     " perhaps a concurrent planner won, cas: %d, err: %v",
	//     cas, err)
	// }

	if len(lastProgress.Errors) > 0 {
		// TODO: Propogate errors better.
		return true, lastProgress.Errors[0]
	}

	return true, nil // TODO: compute proper change response.
}

func (r *rebalancer) calcBegEndMaps(indexDef *cbgt.IndexDef) (
	partitionModel blance.PartitionModel,
	begMap blance.PartitionMap,
	endMap blance.PartitionMap,
	err error) {
	r.m.Lock()
	defer r.m.Unlock()

	// The endPlanPIndexesForIndex is a working data structure that's
	// mutated as calcBegEndMaps progresses.
	endPlanPIndexesForIndex, err := cbgt.SplitIndexDefIntoPlanPIndexes(
		indexDef, r.server, r.endPlanPIndexes)
	if err != nil {
		log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
			" could not SplitIndexDefIntoPlanPIndexes,"+
			" indexDef: %#v, server: %s, err: %v",
			indexDef.Name, indexDef, r.server, err)

		return partitionModel, begMap, endMap, err
	}

	// Invoke blance to assign the endPlanPIndexesForIndex to nodes.
	warnings := cbgt.BlancePlanPIndexes(indexDef,
		endPlanPIndexesForIndex, r.begPlanPIndexes,
		r.nodesAll, r.nodesToAdd, r.nodesToRemove,
		r.nodeWeights, r.nodeHierarchy)

	r.endPlanPIndexes.Warnings[indexDef.Name] = warnings

	// TODO: handle blance ffwd plan warnings.

	for _, warning := range warnings {
		log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
			" BlancePlanPIndexes warning: %q, indexDef: %#v",
			indexDef.Name, warning, indexDef)
	}

	log.Printf("  calcBegEndMaps: indexDef.Name: %s,"+
		" endPlanPIndexes: %#v", indexDef.Name, r.endPlanPIndexes)

	partitionModel, _ = cbgt.BlancePartitionModel(indexDef)

	begMap = cbgt.BlanceMap(endPlanPIndexesForIndex, r.begPlanPIndexes)
	endMap = cbgt.BlanceMap(endPlanPIndexesForIndex, r.endPlanPIndexes)

	return partitionModel, begMap, endMap, nil
}

// --------------------------------------------------------

func (r *rebalancer) assignPartition(stopCh chan struct{},
	index, partition, node, state, op string) error {
	log.Printf("  assignPartitionFunc: index: %s,"+
		" partition: %s, node: %s, state: %s, op: %s",
		index, partition, node, state, op)

	var err error

	r.m.Lock()

	// Update currStates to the assigned index/partition/node/state.
	partitions, exists := r.currStates[index]
	if !exists || partitions == nil {
		partitions = map[string]map[string]stateOp{}
		r.currStates[index] = partitions
	}

	nodes, exists := partitions[partition]
	if !exists || nodes == nil {
		nodes = map[string]stateOp{}
		partitions[partition] = nodes
	}

	if op == "add" {
		if stateOp, exists := nodes[node]; exists && stateOp.state != "" {
			r.m.Unlock()

			return fmt.Errorf("assignPartition:"+
				" op was add when exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				" stateOp: %#v,"+
				index, partition, node, state, op, stateOp)
		}
	} else {
		if stateOp, exists := nodes[node]; !exists || stateOp.state == "" {
			r.m.Unlock()

			return fmt.Errorf("assignPartition:"+
				" op was non-add when not exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				" stateOp: %#v,"+
				index, partition, node, state, op, stateOp)
		}
	}

	nodes[node] = stateOp{state, op}

	r.m.Unlock()

	indexDefs, err := cbgt.PlannerGetIndexDefs(r.cfg, r.version)
	if err != nil {
		return err
	}

	indexDef := indexDefs.IndexDefs[index]
	if indexDef == nil {
		return fmt.Errorf("assignPartition: no indexDef,"+
			" index: %s, partition: %s, node: %s, state: %s, op: %s",
			index, partition, node, state, op)
	}

	planPIndexes, cas, err :=
		cbgt.PlannerGetPlanPIndexes(r.cfg, r.version)
	if err != nil {
		return err
	}

	planPIndex := planPIndexes.PlanPIndexes[partition]
	if planPIndex == nil {
		r.m.Lock()
		endPlanPIndex := r.endPlanPIndexes.PlanPIndexes[partition]
		if endPlanPIndex != nil {
			p := *endPlanPIndex // Copy.
			planPIndex = &p
			planPIndex.Nodes = nil
			planPIndexes.PlanPIndexes[partition] = planPIndex
		}
		r.m.Unlock()
	}

	if planPIndex == nil {
		return fmt.Errorf("assignPartition: no planPIndex,"+
			" index: %s, partition: %s, node: %s, state: %s, op: %s",
			index, partition, node, state, op)
	}

	if planPIndex.Nodes == nil {
		planPIndex.Nodes = make(map[string]*cbgt.PlanPIndexNode)
	}

	planPIndex.UUID = cbgt.NewUUID()

	canRead := true
	canWrite := true
	nodePlanParam := cbgt.GetNodePlanParam(
		indexDef.PlanParams.NodePlanParams, node,
		indexDef.Name, partition)
	if nodePlanParam != nil {
		canRead = nodePlanParam.CanRead
		canWrite = nodePlanParam.CanWrite
	}

	priority := 0
	if state == "replica" {
		priority = len(planPIndex.Nodes)
	}

	if op == "add" {
		if planPIndex.Nodes[node] != nil {
			return fmt.Errorf("assignPartition: planPIndex already exists,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s,"+
				" planPIndex: %#v",
				index, partition, node, state, op, planPIndex)
		}

		// TODO: Need to shift the other node priorities around?
		planPIndex.Nodes[node] = &cbgt.PlanPIndexNode{
			CanRead:  canRead,
			CanWrite: canWrite,
			Priority: priority,
		}
	} else {
		if planPIndex.Nodes[node] == nil {
			return fmt.Errorf("assignPartition: planPIndex missing,"+
				" index: %s, partition: %s, node: %s, state: %s, op: %s"+
				index, partition, node, state, op)
		}

		if op == "del" {
			// TODO: Need to shift the other node priorities around?
			delete(planPIndex.Nodes, node)
		} else {
			// TODO: Need to shift the other node priorities around?
			planPIndex.Nodes[node] = &cbgt.PlanPIndexNode{
				CanRead:  canRead,
				CanWrite: canWrite,
				Priority: priority,
			}
		}
	}

	// TODO: stopCh handling.

	planPIndexes.UUID = cbgt.NewUUID()

	_, err = cbgt.CfgSetPlanPIndexes(r.cfg, planPIndexes, cas)
	if err != nil {
		return fmt.Errorf("assignPartition: update plan,"+
			" perhaps a concurrent planner won, cas: %d, err: %v",
			cas, err)
	}

	return nil
}

func (r *rebalancer) partitionState(stopCh chan struct{},
	index, partition, node string) (
	state string, pct float32, err error) {
	log.Printf("  partitionStateFunc: index: %s,"+
		" partition: %s, node: %s", index, partition, node)

	var stateOp stateOp

	r.m.Lock()
	if r.currStates[index] != nil &&
		r.currStates[index][partition] != nil {
		stateOp = r.currStates[index][partition][node]
	}
	r.m.Unlock()

	// TODO: real state & pct, with stopCh handling.

	return stateOp.state, 1.0, nil
}
<|endoftext|>