<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"
)

type logger struct {
	ticket   int
	hostname string
	buf      *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, l.hostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+time.Since(T_START).String()+"\n",
		)
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		fmt.Fprintf(loggerOut,
			"execute: remotely: host: %s stdout: [%s]\n",
			hostname,
			output.Stdout)

		fmt.Fprintf(loggerErr,
			"execute: remotely: host: %s stderr: [%s]\n",
			hostname,
			output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	slaveAddress    string
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(
			p.logger,
			"adding slave: %s ok: %s",
			slave.Address,
			ok)
	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.slaveAddress)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	flag_slave_address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, "concurrency limit: %d", *flag_jobs)
	fmt.Fprintf(logger, "slaves: %s", *flag_slaves)

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.slaveAddress = *flag_slave_address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	if *flag_slave == false {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as master\n")
		mainMaster(&p)
	} else {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as slave on: %s\n", p.slaveAddress)
		mainSlave(&p)
	}
}
<commit_msg>adding prometheus metrics<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime/debug"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type logger struct {
	ticket   int
	hostname string
	buf      *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, l.hostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func CheckFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e)
		panic(e)
	}
	return e
}

func CheckNotFatal(e error) error {
	if e != nil {
		debug.PrintStack()
		log.Println("CHECK: ERROR:", e, e.Error())
	}
	return e
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	p.StatNumCommandsStart.Inc()
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		dt := time.Since(T_START)
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+dt.String()+"\n",
		)
		if err == nil {
			p.StatNumCommandsDone.Inc()
			p.StatCommandLatency.Observe(dt.Seconds())
		}
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		fmt.Fprintf(loggerOut,
			"execute: remotely: host: %s stdout: [%s]\n",
			hostname,
			output.Stdout)

		fmt.Fprintf(loggerErr,
			"execute: remotely: host: %s stderr: [%s]\n",
			hostname,
			output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	slaveAddress    string

	StatNumCommandsStart prometheus.Counter
	StatNumCommandsDone  prometheus.Counter
	StatCommandLatency   prometheus.Summary
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(
			p.logger,
			"adding slave: %s ok: %s",
			slave.Address,
			ok)
	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.slaveAddress)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	flag_slave_address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	flag_slave_metrics_address := flag.String(
		"metrics_address",
		"localhost:9011",
		"slave metric address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, "concurrency limit: %d", *flag_jobs)
	fmt.Fprintf(logger, "slaves: %s", *flag_slaves)

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.slaveAddress = *flag_slave_address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	// stats:
	p.StatNumCommandsStart = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "num_commands_start",
			Help: "num received"})
	err := prometheus.Register(p.StatNumCommandsStart)
	CheckFatal(err)

	p.StatNumCommandsDone = prometheus.NewCounter(
		prometheus.CounterOpts{
			Name: "num_commands_done",
			Help: "num received"})
	err = prometheus.Register(p.StatNumCommandsDone)
	CheckFatal(err)

	p.StatCommandLatency = prometheus.NewSummary(prometheus.SummaryOpts{
		Name: "commands_latency",
		Help: "commands latency stat",
	})
	err = prometheus.Register(p.StatCommandLatency)
	CheckFatal(err)

	if *flag_slave == false {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as master\n")
		mainMaster(&p)
	} else {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		// metrics
		metricsHandler := promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{})
		http.HandleFunc("/metrics", func(c http.ResponseWriter, req *http.Request) {
			metricsHandler.ServeHTTP(c, req)
		})

		http.HandleFunc("/",
			func(c http.ResponseWriter, req *http.Request) {
				io.WriteString(c, fmt.Sprintf("parallel %s", time.Now().String()))
			})

		go http.ListenAndServe(*flag_slave_metrics_address, nil)

		fmt.Fprintf(logger, "running as slave on: %s\n", p.slaveAddress)
		mainSlave(&p)
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"log"
	"io/ioutil"
	"os"
	"fmt"
	"archive/zip"
	"github.com/gosuri/uilive"
	"time"
	"strings"
	"github.com/fatih/color"
	"path/filepath"
	"gopkg.in/yaml.v2"
)

var (
	allResFiles  map[string]bool
	updatedFilesMap map[string]bool
	distFileMap map[string]bool
	addedFilesMap map[string]bool
)

//Entry point of the validate command
func Validate(updateLocation, distributionLocation string, logsEnabled bool) {

	if (!logsEnabled) {
		log.SetOutput(ioutil.Discard)
	} else {
		log.Println("Logs enabled")
	}
	log.Println("validate command called")

	//Initialize variables
	initialize()

	//Update location should be a zip file
	log.Println("Update Loc: " + updateLocation)
	if !isAZipFile(updateLocation) {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE]: Update file should be a zip file.")
		color.Unset()
		os.Exit(1)
	}
	//Check whether the update location exists
	updateLocationExists := checkFile(updateLocation)
	if updateLocationExists {
		log.Println("Update location exists.")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE]: Update location does not exist. Enter a valid file location.")
		color.Unset()
		os.Exit(1)
	}

	log.Println("Reading update zip...")
	readUpdateZip(updateLocation, logsEnabled)
	log.Println("Update zip successfully read.")
	log.Println("Entries in update zip: ", updatedFilesMap)

	log.Println("Distribution Loc: " + distributionLocation)
	//Check whether the distribution is a zip or a directory
	if isAZipFile(distributionLocation) {
		//Check whether the distribution zip exists
		zipFileExists := checkFile(distributionLocation)
		if zipFileExists {
			log.Println("Distribution location exists.")
			readDistZip(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE]: Distribution zip does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	} else {
		//Check whether the distribution location exists
		distributionLocationExists := checkDir(distributionLocation)
		if distributionLocationExists {
			log.Println("Distribution location exists.")
			readDistDir(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE]: Distribution location does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	}
	//Validate files
	validate()
}

//This initializes the variables
func initialize() {
	allResFiles = make(map[string]bool)
	allResFiles[_LICENSE_FILE_NAME] = true
	allResFiles[_NOT_A_CONTRIBUTION_FILE_NAME] = true
	allResFiles[_README_FILE_NAME] = true
	allResFiles[_UPDATE_DESCRIPTOR_FILE_NAME] = true
	allResFiles[_INSTRUCTIONS_FILE_NAME] = true

	updatedFilesMap = make(map[string]bool)
	distFileMap = make(map[string]bool)
	addedFilesMap = make(map[string]bool)
}

//This method validates the files
func validate() {
	//Iterate through all the files in the update. All files should be in the distribution unless they are newly
	// added files
	for updateLoc, _ := range updatedFilesMap {
		log.Println("Checking location: ", updateLoc)
		//Check whether the distribution has a file with the same name
		_, found := distFileMap[updateLoc]
		//If there is a file
		if found {
			log.Println(updateLoc, "found in distFileMap")
		} else {
			//If there is no file
			log.Println(updateLoc, "not found in distFileMap")
			//Check whether it is a newly added file
			_, found := addedFilesMap[updateLoc]
			//if it is a newly added file
			if found {
				log.Println(updateLoc, "found in addedFilesMap")
			} else {
				//If it is not a newly added file, print an error
				log.Println(updateLoc, "not found in addedFilesMap")
				log.Println("addedFilesMap: ", addedFilesMap)
				color.Set(color.FgRed)
				fmt.Println("[FAILURE]:", updateLoc, "not found in distribution and it is not a " +
				"newly added file.")
				fmt.Println("If it is a new file, please add an entry in", _UPDATE_DESCRIPTOR_FILE_NAME,
					"file.")
				fmt.Println("\nValidation FAILED\n")
				color.Unset()
				os.Exit(1)
			}
		}
	}
	color.Set(color.FgGreen)
	fmt.Println("\n[INFO] Validation SUCCESSFUL\n")
	color.Unset()
}

//This function reads the files of the given update zip
func readUpdateZip(zipLocation string, logsEnabled bool) {
	log.Println("Zip file reading started: ", zipLocation)

	updateName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(updateName, string(os.PathSeparator)); lastIndex > -1 {
		updateName = updateName[lastIndex + 1:]
	}
	log.Println("Update name: ", updateName)

	//Check whether the update name has the required prefix
	if !strings.HasPrefix(updateName, _UPDATE_NAME_PREFIX) {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Update file does not have", _UPDATE_NAME_PREFIX, "prefix")
		color.Unset()
		os.Exit(1)
	} else {
		log.Println("Update file does have", _UPDATE_NAME_PREFIX, "prefix")
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a new writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in the zip
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from update zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		//Every file should be in a root folder. Check for the os.PathSeparator character to identify this
		index := strings.Index(file.Name, string(os.PathSeparator))
		if index == -1 {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE] Update zip file should have a root folder called", updateName)
			color.Unset()
			os.Exit(1)
		} else {
			rootFolder := file.Name[:index]
			log.Println("RootFolder:", rootFolder)
			if rootFolder != updateName {
				color.Set(color.FgRed)
				fmt.Println("[FAILURE]", file.Name, "should be in", updateName, "root folder. But it " +
				"is in ", rootFolder, "folder")
				color.Unset()
				os.Exit(1)
			}
		}
		//Check whether the file is a resource file
		_, found := allResFiles[file.FileInfo().Name()]
		//If it is not a resource file
		if !found {
			//It should be in a carbon.home folder
			containsCarbonHome := strings.Contains(file.Name, _CARBON_HOME)
			if (!containsCarbonHome) {
				color.Set(color.FgRed)
				fmt.Println("[FAILURE] '" + file.Name + "' is not a known resource file. " +
				"It should be in '" + updateName + string(os.PathSeparator) + _CARBON_HOME +
				string(os.PathSeparator) + "' folder")
				color.Unset()
				os.Exit(1)
			}
			log.Println("Have a", _CARBON_HOME, "folder")

			temp := strings.TrimPrefix(file.Name, updateName + string(os.PathSeparator) + _CARBON_HOME)
			log.Println("Entry: ", temp)
			updatedFilesMap[temp] = true
		} else {
			//If the file is a resource file, delete the entry from allResFiles. This map is later used
			// to track missing resource files
			log.Println(file.FileInfo().Name(), "was found in resource map")
			delete(allResFiles, file.FileInfo().Name())
			log.Println(file.FileInfo().Name(), "was removed from the map")
			//If the file is update-descriptor.yaml file, we need to read the newly added files.
			// Otherwise there will be no match for these files and validation will be failed
			if file.FileInfo().Name() == _UPDATE_DESCRIPTOR_FILE_NAME {
				//Open the file
				yamlFile, err := file.Open()
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("[FAILURE] Error occurred while reading the", _UPDATE_DESCRIPTOR_FILE_NAME, "file:",
						err)
					color.Unset()
				}
				//Get the byte array
				data, err := ioutil.ReadAll(yamlFile)
				if err != nil {
					log.Fatal(err)
				}
				descriptor := update_descriptor{}
				//Get the values
				err = yaml.Unmarshal(data, &descriptor)
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("[FAILURE] Error occurred while unmarshalling the yaml:", err)
					color.Unset()
				}
				log.Println("descriptor:", descriptor)
				//Add all files to addedFilesMap
				for _, addedFile := range descriptor.File_changes.Added_files {
					addedFilesMap[addedFile] = true
				}
			}
		}
	}
	//Stop the writer
	writer.Stop()

	//Delete instructions.txt file if it is left in the map because it is optional
	delete(allResFiles, _INSTRUCTIONS_FILE_NAME)
	log.Println("Resource map:", allResFiles)
	log.Println(updatedFilesMap)

	//At this point, the size of the allResFiles should be zero. If one or more files are not found, that means
	// that some required files are missing
	if (len(allResFiles) != 0) {
		//Print the missing files
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Following resource file(s) were not found in the update zip: ")
		for key, _ := range allResFiles {
			fmt.Println("\t", "-", key)
		}
		color.Unset()
		os.Exit(1)
	}

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] All files not read from zip file")
		color.Unset()
		os.Exit(1)
	}
}

//This function reads the files of the given distribution zip
func readDistZip(zipLocation string, logsEnabled bool) {
	log.Println("Zip file reading started: ", zipLocation)

	//Get the distribution name
	distName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(distName, string(os.PathSeparator)); lastIndex > -1 {
		distName = distName[lastIndex + 1:]
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		//Get the relative path in the zip
		temp := strings.TrimPrefix(file.Name, distName)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true

	}
	//Stop the writer
	writer.Stop()

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] All files not read from zip file")
		color.Unset()
		os.Exit(1)
	}
}

func readDistDir(distributionLocation string, logsEnabled bool) {
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()
	fileCount := 0
	//Start the walk
	err := filepath.Walk(distributionLocation, func(path string, fileInfo os.FileInfo, err error) error {
		fileCount++;
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution directory: %d files read\n", fileCount)
			time.Sleep(time.Millisecond * 2)
		}
		log.Println("Walking: ", path)
		//Get the relative path
		temp := strings.TrimPrefix(path, distributionLocation)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true
		return nil
	})
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading the zip file: ", err)
		os.Exit(1)
		color.Unset()
	}
	log.Println("Total files read:", fileCount)
	//stop the writer
	writer.Stop()
}
<commit_msg>Validation not working properly in windows issue fixed<commit_after>package cmd

import (
	"log"
	"io/ioutil"
	"os"
	"fmt"
	"archive/zip"
	"github.com/gosuri/uilive"
	"time"
	"strings"
	"github.com/fatih/color"
	"path/filepath"
	"gopkg.in/yaml.v2"
)

var (
	allResFiles  map[string]bool
	updatedFilesMap map[string]bool
	distFileMap map[string]bool
	addedFilesMap map[string]bool
)

//Entry point of the validate command
func Validate(updateLocation, distributionLocation string, logsEnabled bool) {

	if (!logsEnabled) {
		log.SetOutput(ioutil.Discard)
	} else {
		log.Println("Logs enabled")
	}
	log.Println("validate command called")

	//Initialize variables
	initialize()

	//Update location should be a zip file
	log.Println("Update Loc: " + updateLocation)
	if !isAZipFile(updateLocation) {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE]: Update file should be a zip file.")
		color.Unset()
		os.Exit(1)
	}
	//Check whether the update location exists
	updateLocationExists := checkFile(updateLocation)
	if updateLocationExists {
		log.Println("Update location exists.")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE]: Update location does not exist. Enter a valid file location.")
		color.Unset()
		os.Exit(1)
	}

	log.Println("Reading update zip...")
	readUpdateZip(updateLocation, logsEnabled)
	log.Println("Update zip successfully read.")
	log.Println("Entries in update zip: ", updatedFilesMap)

	log.Println("Distribution Loc: " + distributionLocation)
	//Check whether the distribution is a zip or a directory
	if isAZipFile(distributionLocation) {
		//Check whether the distribution zip exists
		zipFileExists := checkFile(distributionLocation)
		if zipFileExists {
			log.Println("Distribution location exists.")
			readDistZip(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE]: Distribution zip does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	} else {
		//Check whether the distribution location exists
		distributionLocationExists := checkDir(distributionLocation)
		if distributionLocationExists {
			log.Println("Distribution location exists.")
			readDistDir(distributionLocation, logsEnabled)
		} else {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE]: Distribution location does not exist. Enter a valid location.")
			color.Unset()
			os.Exit(1)
		}
	}
	//Validate files
	validate()
}

//This initializes the variables
func initialize() {
	allResFiles = make(map[string]bool)
	allResFiles[_LICENSE_FILE_NAME] = true
	allResFiles[_NOT_A_CONTRIBUTION_FILE_NAME] = true
	allResFiles[_README_FILE_NAME] = true
	allResFiles[_UPDATE_DESCRIPTOR_FILE_NAME] = true
	allResFiles[_INSTRUCTIONS_FILE_NAME] = true

	updatedFilesMap = make(map[string]bool)
	distFileMap = make(map[string]bool)
	addedFilesMap = make(map[string]bool)
}

//This method validates the files
func validate() {
	//Iterate through all the files in the update. All files should be in the distribution unless they are newly
	// added files
	for updateLoc, _ := range updatedFilesMap {
		log.Println("Checking location: ", updateLoc)
		//Check whether the distribution has a file with the same name
		_, found := distFileMap[updateLoc]
		//If there is a file
		if found {
			log.Println(updateLoc, "found in distFileMap")
		} else {
			//If there is no file
			log.Println(updateLoc, "not found in distFileMap")
			//Check whether it is a newly added file
			_, found := addedFilesMap[updateLoc]
			//if it is a newly added file
			if found {
				log.Println(updateLoc, "found in addedFilesMap")
			} else {
				//If it is not a newly added file, print an error
				log.Println(updateLoc, "not found in addedFilesMap")
				log.Println("addedFilesMap: ", addedFilesMap)
				color.Set(color.FgRed)
				fmt.Println("[FAILURE]:", updateLoc, "not found in distribution and it is not a " +
				"newly added file.")
				fmt.Println("If it is a new file, please add an entry in", _UPDATE_DESCRIPTOR_FILE_NAME,
					"file.")
				fmt.Println("\nValidation FAILED\n")
				color.Unset()
				os.Exit(1)
			}
		}
	}
	color.Set(color.FgGreen)
	fmt.Println("\n[INFO] Validation SUCCESSFUL\n")
	color.Unset()
}

//This function reads the files of the given update zip
func readUpdateZip(zipLocation string, logsEnabled bool) {
	log.Println("Zip file reading started: ", zipLocation)

	updateName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(updateName, string(os.PathSeparator)); lastIndex > -1 {
		updateName = updateName[lastIndex + 1:]
	}
	log.Println("Update name: ", updateName)

	//Check whether the update name has the required prefix
	if !strings.HasPrefix(updateName, _UPDATE_NAME_PREFIX) {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Update file does not have", _UPDATE_NAME_PREFIX, "prefix")
		color.Unset()
		os.Exit(1)
	} else {
		log.Println("Update file does have", _UPDATE_NAME_PREFIX, "prefix")
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a new writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in the zip
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from update zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		//Every file should be in a root folder. Check for the os.PathSeparator character to identify this
		index := strings.Index(file.Name, "/")//string(os.PathSeparator) removed because it does not work
		// properly in windows
		if index == -1 {
			color.Set(color.FgRed)
			fmt.Println("[FAILURE] Update zip file should have a root folder called", updateName)
			color.Unset()
			os.Exit(1)
		} else {
			rootFolder := file.Name[:index]
			log.Println("RootFolder:", rootFolder)
			if rootFolder != updateName {
				color.Set(color.FgRed)
				fmt.Println("[FAILURE]", file.Name, "should be in", updateName, "root folder. But it " +
				"is in ", rootFolder, "folder")
				color.Unset()
				os.Exit(1)
			}
		}
		//Check whether the file is a resource file
		_, found := allResFiles[file.FileInfo().Name()]
		//If it is not a resource file
		if !found {
			//It should be in a carbon.home folder
			containsCarbonHome := strings.Contains(file.Name, _CARBON_HOME)
			if (!containsCarbonHome) {
				color.Set(color.FgRed)
				//string(os.PathSeparator) removed because it does not work properly in windows
				fmt.Println("[FAILURE] '" + file.Name + "' is not a known resource file. " +
				"It should be in '" + updateName + "/" + _CARBON_HOME + "/" + "' folder")
				color.Unset()
				os.Exit(1)
			}
			log.Println("Have a", _CARBON_HOME, "folder")
			//string(os.PathSeparator) removed because it does not work properly in windows
			temp := strings.TrimPrefix(file.Name, updateName + "/" + _CARBON_HOME)
			log.Println("Entry: ", temp)
			updatedFilesMap[temp] = true
		} else {
			//If the file is a resource file, delete the entry from allResFiles. This map is later used
			// to track missing resource files
			log.Println(file.FileInfo().Name(), "was found in resource map")
			delete(allResFiles, file.FileInfo().Name())
			log.Println(file.FileInfo().Name(), "was removed from the map")
			//If the file is update-descriptor.yaml file, we need to read the newly added files.
			// Otherwise there will be no match for these files and validation will be failed
			if file.FileInfo().Name() == _UPDATE_DESCRIPTOR_FILE_NAME {
				//Open the file
				yamlFile, err := file.Open()
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("[FAILURE] Error occurred while reading the", _UPDATE_DESCRIPTOR_FILE_NAME, "file:",
						err)
					color.Unset()
				}
				//Get the byte array
				data, err := ioutil.ReadAll(yamlFile)
				if err != nil {
					log.Fatal(err)
				}
				descriptor := update_descriptor{}
				//Get the values
				err = yaml.Unmarshal(data, &descriptor)
				if err != nil {
					color.Set(color.FgRed)
					fmt.Println("[FAILURE] Error occurred while unmarshalling the yaml:", err)
					color.Unset()
				}
				log.Println("descriptor:", descriptor)
				//Add all files to addedFilesMap
				for _, addedFile := range descriptor.File_changes.Added_files {
					addedFilesMap[addedFile] = true
				}
			}
		}
	}
	//Stop the writer
	writer.Stop()

	//Delete instructions.txt file if it is left in the map because it is optional
	delete(allResFiles, _INSTRUCTIONS_FILE_NAME)
	log.Println("Resource map:", allResFiles)
	log.Println(updatedFilesMap)

	//At this point, the size of the allResFiles should be zero. If one or more files are not found, that means
	// that some required files are missing
	if (len(allResFiles) != 0) {
		//Print the missing files
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Following resource file(s) were not found in the update zip: ")
		for key, _ := range allResFiles {
			fmt.Println("\t", "-", key)
		}
		color.Unset()
		os.Exit(1)
	}

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] All files not read from zip file")
		color.Unset()
		os.Exit(1)
	}
}

//This function reads the files of the given distribution zip
func readDistZip(zipLocation string, logsEnabled bool) {
	log.Println("Zip file reading started: ", zipLocation)

	//Get the distribution name
	distName := strings.TrimSuffix(zipLocation, ".zip")
	if lastIndex := strings.LastIndex(distName, string(os.PathSeparator)); lastIndex > -1 {
		distName = distName[lastIndex + 1:]
	}

	// Create a reader out of the zip archive
	zipReader, err := zip.OpenReader(zipLocation)
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading zip:", err)
		color.Unset()
		log.Fatal(err)
	}
	defer zipReader.Close()

	totalFiles := len(zipReader.Reader.File)
	log.Println("File count in zip: ", totalFiles)

	fileCount := 0
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()

	// Iterate through each file/dir found in
	for _, file := range zipReader.Reader.File {
		fileCount++
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution zip: (%d/%d)\n", fileCount, totalFiles)
			time.Sleep(time.Millisecond * 2)
		}

		log.Println("Checking file: ", file.Name)

		//Get the relative path in the zip
		temp := strings.TrimPrefix(file.Name, distName)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true

	}
	//Stop the writer
	writer.Stop()

	//Check whether all files are read
	log.Println("Zip file reading finished")
	log.Println("Total files read: ", fileCount)
	if totalFiles == fileCount {
		log.Println("All files read")
	} else {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] All files not read from zip file")
		color.Unset()
		os.Exit(1)
	}
}

func readDistDir(distributionLocation string, logsEnabled bool) {
	//Create a writer to show the progress
	writer := uilive.New()
	//start listening for updates and render
	writer.Start()
	fileCount := 0
	//Start the walk
	err := filepath.Walk(distributionLocation, func(path string, fileInfo os.FileInfo, err error) error {
		fileCount++;
		if (!logsEnabled) {
			fmt.Fprintf(writer, "Reading files from distribution directory: %d files read\n", fileCount)
			time.Sleep(time.Millisecond * 2)
		}
		log.Println("Walking: ", path)
		//Get the relative path
		temp := strings.TrimPrefix(path, distributionLocation)
		log.Println("Entry: ", temp)
		//Add to the map
		distFileMap[temp] = true
		return nil
	})
	if err != nil {
		color.Set(color.FgRed)
		fmt.Println("[FAILURE] Error occurred while reading the zip file: ", err)
		os.Exit(1)
		color.Unset()
	}
	log.Println("Total files read:", fileCount)
	//stop the writer
	writer.Stop()
}
<|endoftext|>


<commit_before>/**
 * Copyright (C) 2015 Red Hat, 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 cmds

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/fabric8io/gofabric8/client"
	"github.com/fabric8io/gofabric8/util"
	"github.com/spf13/cobra"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/resource"
	k8sclient "k8s.io/kubernetes/pkg/client/unversioned"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

const (
	sshCommandFlag = "ssh-command"
)

func NewCmdVolumes(f *cmdutil.Factory) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "volumes",
		Short: "Creates a persisent volume for any pending persistance volume claims",
		Long:  `Creates a persisent volume for any pending persistance volume claims`,
		PreRun: func(cmd *cobra.Command, args []string) {
			showBanner()
		},
		Run: func(cmd *cobra.Command, args []string) {
			c, _ := client.NewClient(f)
			ns, _, err := f.DefaultNamespace()
			if err != nil {
				util.Fatal("No default namespace")
			} else {

				found, pendingClaimNames := findPendingPVS(c, ns)
				if found {
					createPV(c, ns, pendingClaimNames, cmd)
				}
			}
		},
	}
	cmd.PersistentFlags().String(sshCommandFlag, "", "the ssh command to run commands inside the VM of the single node cluster")
	return cmd
}

func findPendingPVS(c *k8sclient.Client, ns string) (bool, []string) {

	pvcs, err := c.PersistentVolumeClaims(ns).List(api.ListOptions{})
	if err != nil {
		util.Infof("Failed to find any PersistentVolumeClaims, %s in namespace %s\n", err, ns)
	}

	if pvcs != nil {
		items := pvcs.Items
		pendingClaimNames := make([]string, 0, len(items))
		for _, item := range items {
			status := item.Status.Phase
			if status == "Pending" || status == "Lost" {
				pendingClaimNames = append(pendingClaimNames, item.ObjectMeta.Name)
			}
		}
		if len(pendingClaimNames) > 0 {
			return true, pendingClaimNames
		}
	}
	return false, nil
}

func createPV(c *k8sclient.Client, ns string, pvcNames []string, cmd *cobra.Command) (Result, error) {

	for _, pvcName := range pvcNames {
		hostPath := "/data/" + pvcName
		pvs := c.PersistentVolumes()
		rc, err := pvs.List(api.ListOptions{})
		if err != nil {
			util.Errorf("Failed to load PersistentVolumes with error %v\n", err)
		}
		items := rc.Items
		for _, volume := range items {
			vname := volume.ObjectMeta.Name
			if vname == pvcName {
				util.Infof("Already created PersistentVolumes for %s\n", pvcName)
			}
		}

		err = configureHostPathVolume(c, ns, hostPath, cmd)
		if err != nil {
			util.Errorf("Failed to configure the host path %s with error %v\n", hostPath, err)
		}

		// lets create a new PV
		util.Infof("PersistentVolume name %s will be created on host path %s\n", pvcName, hostPath)
		pv := api.PersistentVolume{
			ObjectMeta: api.ObjectMeta{
				Name: pvcName,
			},
			Spec: api.PersistentVolumeSpec{
				Capacity: api.ResourceList{
					api.ResourceName(api.ResourceStorage): resource.MustParse("5Gi"),
				},
				AccessModes: []api.PersistentVolumeAccessMode{api.ReadWriteOnce},
				PersistentVolumeSource: api.PersistentVolumeSource{
					HostPath: &api.HostPathVolumeSource{Path: hostPath},
				},
			},
		}

		_, err = pvs.Create(&pv)
		if err != nil {
			util.Errorf("Failed to create PersistentVolume %s at %s with error %v\n", pvcName, hostPath, err)
		}
	}

	return Success, nil
}

// if we are on minikube or minishift lets try to create the
// hostPath folders with relaxed persmissions
func configureHostPathVolume(c *k8sclient.Client, ns string, hostPath string, corbaCmd *cobra.Command) error {
	cli := ""
	flag := corbaCmd.Flags().Lookup(sshCommandFlag)
	if flag != nil {
		cli = flag.Value.String()
	}

	if len(cli) == 0 {
		nodes, err := c.Nodes().List(api.ListOptions{})
		if err != nil {
			return err
		}
		if len(nodes.Items) == 1 {
			node := nodes.Items[0]
			if node.Name == minikubeNodeName || node.Name == minishiftNodeName {
				// lets figure out which one we are
				// TODO there's no obvious annotation yet to know
				// if it was created via minikube or minishift
				// so lets look at the images
				cli = "minikube"

				for _, image := range node.Status.Images {
					for _, imageName := range image.Names {
						if strings.HasPrefix(imageName, "openshift/origin-pod:") {
							cli = "minishift"
							break
						}
					}

				}
			}
		}
	}
	if len(cli) == 0 {
		// lets default to using vagrant if we have a Vagrantfile
		if _, err := os.Stat("Vagrantfile"); err == nil {
			cli = "vagrant"
		}
	}
	if len(cli) == 0 {
		return nil
	}

	shellCommands := []string{
		fmt.Sprintf("sudo mkdir -p %s", hostPath),
		fmt.Sprintf("sudo chmod 777 %s", hostPath),
		fmt.Sprintf("echo hostPath is setup correctly at: %s", hostPath),
	}
	util.Infof("About to modify host paths on the VM via the command: %s\n", cli)

	for _, shellCmd := range shellCommands {
		args := []string{"ssh", fmt.Sprintf("/bin/sh -c '%s'", shellCmd)}
		cmd := exec.Command(cli, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		var waitStatus syscall.WaitStatus
		if err := cmd.Run(); err != nil {
			printErr(err)
			if exitError, ok := err.(*exec.ExitError); ok {
				waitStatus = exitError.Sys().(syscall.WaitStatus)
				printStatus(waitStatus.ExitStatus())
			}
			return err
		} else {
			waitStatus = cmd.ProcessState.Sys().(syscall.WaitStatus)
			printStatus(waitStatus.ExitStatus())
		}
	}
	return nil
}
<commit_msg>minikube vm name can be boot2docker see fabric8io/gofabric8#152<commit_after>/**
 * Copyright (C) 2015 Red Hat, 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 cmds

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/fabric8io/gofabric8/client"
	"github.com/fabric8io/gofabric8/util"
	"github.com/spf13/cobra"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/resource"
	k8sclient "k8s.io/kubernetes/pkg/client/unversioned"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

const (
	sshCommandFlag = "ssh-command"
)

func NewCmdVolumes(f *cmdutil.Factory) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "volumes",
		Short: "Creates a persisent volume for any pending persistance volume claims",
		Long:  `Creates a persisent volume for any pending persistance volume claims`,
		PreRun: func(cmd *cobra.Command, args []string) {
			showBanner()
		},
		Run: func(cmd *cobra.Command, args []string) {
			c, _ := client.NewClient(f)
			ns, _, err := f.DefaultNamespace()
			if err != nil {
				util.Fatal("No default namespace")
			} else {

				found, pendingClaimNames := findPendingPVS(c, ns)
				if found {
					createPV(c, ns, pendingClaimNames, cmd)
				}
			}
		},
	}
	cmd.PersistentFlags().String(sshCommandFlag, "", "the ssh command to run commands inside the VM of the single node cluster")
	return cmd
}

func findPendingPVS(c *k8sclient.Client, ns string) (bool, []string) {

	pvcs, err := c.PersistentVolumeClaims(ns).List(api.ListOptions{})
	if err != nil {
		util.Infof("Failed to find any PersistentVolumeClaims, %s in namespace %s\n", err, ns)
	}

	if pvcs != nil {
		items := pvcs.Items
		pendingClaimNames := make([]string, 0, len(items))
		for _, item := range items {
			status := item.Status.Phase
			if status == "Pending" || status == "Lost" {
				pendingClaimNames = append(pendingClaimNames, item.ObjectMeta.Name)
			}
		}
		if len(pendingClaimNames) > 0 {
			return true, pendingClaimNames
		}
	}
	return false, nil
}

func createPV(c *k8sclient.Client, ns string, pvcNames []string, cmd *cobra.Command) (Result, error) {

	for _, pvcName := range pvcNames {
		hostPath := "/data/" + pvcName
		pvs := c.PersistentVolumes()
		rc, err := pvs.List(api.ListOptions{})
		if err != nil {
			util.Errorf("Failed to load PersistentVolumes with error %v\n", err)
		}
		items := rc.Items
		for _, volume := range items {
			vname := volume.ObjectMeta.Name
			if vname == pvcName {
				util.Infof("Already created PersistentVolumes for %s\n", pvcName)
			}
		}

		err = configureHostPathVolume(c, ns, hostPath, cmd)
		if err != nil {
			util.Errorf("Failed to configure the host path %s with error %v\n", hostPath, err)
		}

		// lets create a new PV
		util.Infof("PersistentVolume name %s will be created on host path %s\n", pvcName, hostPath)
		pv := api.PersistentVolume{
			ObjectMeta: api.ObjectMeta{
				Name: pvcName,
			},
			Spec: api.PersistentVolumeSpec{
				Capacity: api.ResourceList{
					api.ResourceName(api.ResourceStorage): resource.MustParse("5Gi"),
				},
				AccessModes: []api.PersistentVolumeAccessMode{api.ReadWriteOnce},
				PersistentVolumeSource: api.PersistentVolumeSource{
					HostPath: &api.HostPathVolumeSource{Path: hostPath},
				},
			},
		}

		_, err = pvs.Create(&pv)
		if err != nil {
			util.Errorf("Failed to create PersistentVolume %s at %s with error %v\n", pvcName, hostPath, err)
		}
	}

	return Success, nil
}

// if we are on minikube or minishift lets try to create the
// hostPath folders with relaxed persmissions
func configureHostPathVolume(c *k8sclient.Client, ns string, hostPath string, corbaCmd *cobra.Command) error {
	cli := ""
	flag := corbaCmd.Flags().Lookup(sshCommandFlag)
	if flag != nil {
		cli = flag.Value.String()
	}

	if len(cli) == 0 {
		nodes, err := c.Nodes().List(api.ListOptions{})
		if err != nil {
			return err
		}
		if len(nodes.Items) == 1 {
			node := nodes.Items[0]
			if node.Name == minikubeNodeName || node.Name == minishiftNodeName || node.Name == boot2docker {
				// lets figure out which one we are
				// TODO there's no obvious annotation yet to know
				// if it was created via minikube or minishift
				// so lets look at the images
				cli = "minikube"

				for _, image := range node.Status.Images {
					for _, imageName := range image.Names {
						if strings.HasPrefix(imageName, "openshift/origin-pod:") {
							cli = "minishift"
							break
						}
					}

				}
			}
		}
	}
	if len(cli) == 0 {
		// lets default to using vagrant if we have a Vagrantfile
		if _, err := os.Stat("Vagrantfile"); err == nil {
			cli = "vagrant"
		}
	}
	if len(cli) == 0 {
		return nil
	}

	shellCommands := []string{
		fmt.Sprintf("sudo mkdir -p %s", hostPath),
		fmt.Sprintf("sudo chmod 777 %s", hostPath),
		fmt.Sprintf("echo hostPath is setup correctly at: %s", hostPath),
	}
	util.Infof("About to modify host paths on the VM via the command: %s\n", cli)

	for _, shellCmd := range shellCommands {
		args := []string{"ssh", fmt.Sprintf("/bin/sh -c '%s'", shellCmd)}
		cmd := exec.Command(cli, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		var waitStatus syscall.WaitStatus
		if err := cmd.Run(); err != nil {
			printErr(err)
			if exitError, ok := err.(*exec.ExitError); ok {
				waitStatus = exitError.Sys().(syscall.WaitStatus)
				printStatus(waitStatus.ExitStatus())
			}
			return err
		} else {
			waitStatus = cmd.ProcessState.Sys().(syscall.WaitStatus)
			printStatus(waitStatus.ExitStatus())
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012-2020 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

import "io"

// encWriter abstracts writing to a byte array or to an io.Writer.
type encWriter interface {
	writeb([]byte)
	writestr(string)
	writeqstr(string) // write string wrapped in quotes ie "..."
	writen1(byte)

	// add convenience functions for writing 2,4
	writen2(byte, byte)
	writen4(byte, byte, byte, byte)

	end()
}

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

// bufioEncWriter
type bufioEncWriter struct {
	w io.Writer

	buf []byte

	n int

	b [16]byte // scratch buffer and padding (cache-aligned)
}

func (z *bufioEncWriter) reset(w io.Writer, bufsize int, blist *bytesFreelist) {
	z.w = w
	z.n = 0
	if bufsize <= 0 {
		bufsize = defEncByteBufSize
	}
	// bufsize must be >= 8, to accomodate writen methods (where n <= 8)
	if bufsize <= 8 {
		bufsize = 8
	}
	if cap(z.buf) < bufsize {
		if len(z.buf) > 0 && &z.buf[0] != &z.b[0] {
			blist.put(z.buf)
		}
		if len(z.b) > bufsize {
			z.buf = z.b[:]
		} else {
			z.buf = blist.get(bufsize)
		}
	}
	z.buf = z.buf[:cap(z.buf)]
}

func (z *bufioEncWriter) flushErr() (err error) {
	n, err := z.w.Write(z.buf[:z.n])
	z.n -= n
	if z.n > 0 && err == nil {
		err = io.ErrShortWrite
	}
	if n > 0 && z.n > 0 {
		copy(z.buf, z.buf[n:z.n+n])
	}
	return err
}

func (z *bufioEncWriter) flush() {
	halt.onerror(z.flushErr())
}

func (z *bufioEncWriter) writeb(s []byte) {
LOOP:
	a := len(z.buf) - z.n
	if len(s) > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
}

func (z *bufioEncWriter) writestr(s string) {
	// z.writeb(bytesView(s)) // inlined below
LOOP:
	a := len(z.buf) - z.n
	if len(s) > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
}

func (z *bufioEncWriter) writeqstr(s string) {
	// z.writen1('"')
	// z.writestr(s)
	// z.writen1('"')

	if z.n+len(s)+2 > len(z.buf) {
		z.flush()
	}
	z.buf[z.n] = '"'
	z.n++
LOOP:
	a := len(z.buf) - z.n
	if len(s)+1 > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
	z.buf[z.n] = '"'
	z.n++
}

func (z *bufioEncWriter) writen1(b1 byte) {
	if 1 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n] = b1
	z.n++
}

func (z *bufioEncWriter) writen2(b1, b2 byte) {
	if 2 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n+1] = b2
	z.buf[z.n] = b1
	z.n += 2
}

func (z *bufioEncWriter) writen4(b1, b2, b3, b4 byte) {
	if 4 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n+3] = b4
	z.buf[z.n+2] = b3
	z.buf[z.n+1] = b2
	z.buf[z.n] = b1
	z.n += 4
}

func (z *bufioEncWriter) endErr() (err error) {
	if z.n > 0 {
		err = z.flushErr()
	}
	return
}

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

// bytesEncAppender implements encWriter and can write to an byte slice.
type bytesEncAppender struct {
	b   []byte
	out *[]byte
}

func (z *bytesEncAppender) writeb(s []byte) {
	z.b = append(z.b, s...)
}
func (z *bytesEncAppender) writestr(s string) {
	z.b = append(z.b, s...)
}
func (z *bytesEncAppender) writeqstr(s string) {
	z.b = append(append(append(z.b, '"'), s...), '"')

	// z.b = append(z.b, '"')
	// z.b = append(z.b, s...)
	// z.b = append(z.b, '"')
}
func (z *bytesEncAppender) writen1(b1 byte) {
	z.b = append(z.b, b1)
}
func (z *bytesEncAppender) writen2(b1, b2 byte) {
	z.b = append(z.b, b1, b2)
}
func (z *bytesEncAppender) writen4(b1, b2, b3, b4 byte) {
	z.b = append(z.b, b1, b2, b3, b4)
}
func (z *bytesEncAppender) endErr() error {
	*(z.out) = z.b
	return nil
}
func (z *bytesEncAppender) reset(in []byte, out *[]byte) {
	z.b = in[:0]
	z.out = out
}

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

type encWr struct {
	bytes bool // encoding to []byte
	js    bool // is json encoder?
	be    bool // is binary encoder?

	c containerState

	calls uint16

	wb bytesEncAppender
	wf *bufioEncWriter
}

func (z *encWr) writeb(s []byte) {
	if z.bytes {
		z.wb.writeb(s)
	} else {
		z.wf.writeb(s)
	}
}
func (z *encWr) writeqstr(s string) {
	if z.bytes {
		// MARKER: manually inline, else this function is not inlined.
		// Keep in sync with bytesEncWriter.writeqstr
		// z.wb.writeqstr(s)
		z.wb.b = append(append(append(z.wb.b, '"'), s...), '"')
	} else {
		z.wf.writeqstr(s)
	}
}
func (z *encWr) writestr(s string) {
	if z.bytes {
		z.wb.writestr(s)
	} else {
		z.wf.writestr(s)
	}
}
func (z *encWr) writen1(b1 byte) {
	if z.bytes {
		z.wb.writen1(b1)
	} else {
		z.wf.writen1(b1)
	}
}

// MARKER: manually inline bytesEncAppender.writenx methods,
// as calling them causes encWr.writenx methods to not be inlined.
//
// i.e. instead of writing z.wb.writen2(b1, b2), use z.wb.b = append(z.wb.b, b1, b2)

func (z *encWr) writen2(b1, b2 byte) {
	if z.bytes {
		z.wb.b = append(z.wb.b, b1, b2)
	} else {
		z.wf.writen2(b1, b2)
	}
}
func (z *encWr) writen4(b1, b2, b3, b4 byte) {
	if z.bytes {
		z.wb.b = append(z.wb.b, b1, b2, b3, b4)
	} else {
		z.wf.writen4(b1, b2, b3, b4)
	}
}

func (z *encWr) endErr() error {
	if z.bytes {
		return z.wb.endErr()
	}
	return z.wf.endErr()
}

func (z *encWr) end() {
	halt.onerror(z.endErr())
}

var _ encWriter = (*encWr)(nil)
<commit_msg>codec: clean up writer.go<commit_after>// Copyright (c) 2012-2020 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

import "io"

// encWriter abstracts writing to a byte array or to an io.Writer.
type encWriter interface {
	writeb([]byte)
	writestr(string)
	writeqstr(string) // write string wrapped in quotes ie "..."
	writen1(byte)

	// add convenience functions for writing 2,4
	writen2(byte, byte)
	writen4(byte, byte, byte, byte)

	end()
}

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

// bufioEncWriter
type bufioEncWriter struct {
	w io.Writer

	buf []byte

	n int

	b [16]byte // scratch buffer and padding (cache-aligned)
}

func (z *bufioEncWriter) reset(w io.Writer, bufsize int, blist *bytesFreelist) {
	z.w = w
	z.n = 0
	if bufsize <= 0 {
		bufsize = defEncByteBufSize
	}
	// bufsize must be >= 8, to accomodate writen methods (where n <= 8)
	if bufsize <= 8 {
		bufsize = 8
	}
	if cap(z.buf) < bufsize {
		if len(z.buf) > 0 && &z.buf[0] != &z.b[0] {
			blist.put(z.buf)
		}
		if len(z.b) > bufsize {
			z.buf = z.b[:]
		} else {
			z.buf = blist.get(bufsize)
		}
	}
	z.buf = z.buf[:cap(z.buf)]
}

func (z *bufioEncWriter) flushErr() (err error) {
	n, err := z.w.Write(z.buf[:z.n])
	z.n -= n
	if z.n > 0 {
		if err == nil {
			err = io.ErrShortWrite
		}
		if n > 0 {
			copy(z.buf, z.buf[n:z.n+n])
		}
	}
	return err
}

func (z *bufioEncWriter) flush() {
	halt.onerror(z.flushErr())
}

func (z *bufioEncWriter) writeb(s []byte) {
LOOP:
	a := len(z.buf) - z.n
	if len(s) > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
}

func (z *bufioEncWriter) writestr(s string) {
	// z.writeb(bytesView(s)) // inlined below
LOOP:
	a := len(z.buf) - z.n
	if len(s) > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
}

func (z *bufioEncWriter) writeqstr(s string) {
	// z.writen1('"')
	// z.writestr(s)
	// z.writen1('"')

	if z.n+len(s)+2 > len(z.buf) {
		z.flush()
	}
	z.buf[z.n] = '"'
	z.n++
LOOP:
	a := len(z.buf) - z.n
	if len(s)+1 > a {
		z.n += copy(z.buf[z.n:], s[:a])
		s = s[a:]
		z.flush()
		goto LOOP
	}
	z.n += copy(z.buf[z.n:], s)
	z.buf[z.n] = '"'
	z.n++
}

func (z *bufioEncWriter) writen1(b1 byte) {
	if 1 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n] = b1
	z.n++
}
func (z *bufioEncWriter) writen2(b1, b2 byte) {
	if 2 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n+1] = b2
	z.buf[z.n] = b1
	z.n += 2
}
func (z *bufioEncWriter) writen4(b1, b2, b3, b4 byte) {
	if 4 > len(z.buf)-z.n {
		z.flush()
	}
	z.buf[z.n+3] = b4
	z.buf[z.n+2] = b3
	z.buf[z.n+1] = b2
	z.buf[z.n] = b1
	z.n += 4
}

func (z *bufioEncWriter) endErr() (err error) {
	if z.n > 0 {
		err = z.flushErr()
	}
	return
}

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

// bytesEncAppender implements encWriter and can write to an byte slice.
type bytesEncAppender struct {
	b   []byte
	out *[]byte
}

func (z *bytesEncAppender) writeb(s []byte) {
	z.b = append(z.b, s...)
}
func (z *bytesEncAppender) writestr(s string) {
	z.b = append(z.b, s...)
}
func (z *bytesEncAppender) writeqstr(s string) {
	z.b = append(append(append(z.b, '"'), s...), '"')
	// z.b = append(z.b, '"')
	// z.b = append(z.b, s...)
	// z.b = append(z.b, '"')
}
func (z *bytesEncAppender) writen1(b1 byte) {
	z.b = append(z.b, b1)
}
func (z *bytesEncAppender) writen2(b1, b2 byte) {
	z.b = append(z.b, b1, b2)
}
func (z *bytesEncAppender) writen4(b1, b2, b3, b4 byte) {
	z.b = append(z.b, b1, b2, b3, b4)
}
func (z *bytesEncAppender) endErr() error {
	*(z.out) = z.b
	return nil
}
func (z *bytesEncAppender) reset(in []byte, out *[]byte) {
	z.b = in[:0]
	z.out = out
}

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

type encWr struct {
	bytes bool // encoding to []byte
	js    bool // is json encoder?
	be    bool // is binary encoder?

	c containerState

	calls uint16

	wb bytesEncAppender
	wf *bufioEncWriter
}

// MARKER: manually inline bytesEncAppender.writenx/writeqstr methods,
// as calling them causes encWr.writenx/writeqstr methods to not be inlined (cost > 80).
//
// i.e. e.g. instead of writing z.wb.writen2(b1, b2), use z.wb.b = append(z.wb.b, b1, b2)

func (z *encWr) writeb(s []byte) {
	if z.bytes {
		z.wb.writeb(s)
	} else {
		z.wf.writeb(s)
	}
}
func (z *encWr) writeqstr(s string) {
	if z.bytes {
		// MARKER: z.wb.writeqstr(s)
		z.wb.b = append(append(append(z.wb.b, '"'), s...), '"')
	} else {
		z.wf.writeqstr(s)
	}
}
func (z *encWr) writestr(s string) {
	if z.bytes {
		z.wb.writestr(s)
	} else {
		z.wf.writestr(s)
	}
}
func (z *encWr) writen1(b1 byte) {
	if z.bytes {
		z.wb.writen1(b1)
	} else {
		z.wf.writen1(b1)
	}
}

func (z *encWr) writen2(b1, b2 byte) {
	if z.bytes {
		// MARKER: z.wb.writen2(b1, b2)
		z.wb.b = append(z.wb.b, b1, b2)
	} else {
		z.wf.writen2(b1, b2)
	}
}
func (z *encWr) writen4(b1, b2, b3, b4 byte) {
	if z.bytes {
		// MARKER: z.wb.writen4(b1, b2, b3, b4)
		z.wb.b = append(z.wb.b, b1, b2, b3, b4)
	} else {
		z.wf.writen4(b1, b2, b3, b4)
	}
}

func (z *encWr) endErr() error {
	if z.bytes {
		return z.wb.endErr()
	}
	return z.wf.endErr()
}

func (z *encWr) end() {
	halt.onerror(z.endErr())
}

var _ encWriter = (*encWr)(nil)
<|endoftext|>


<commit_before>package codegen

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ssa"

	"github.com/bjwbell/gensimd/simd"
)

var sse2intructions = []Intrinsic{
	{
		Name:            "AddEpi64",
		InstructionName: "PADDQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SubEpi64",
		InstructionName: "PSUBQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "MulEpu32",
		InstructionName: "PMULULQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShufflehiEpi16",
		InstructionName: "PSHUFHW",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShuffleloEpi16",
		InstructionName: "PSHUFLW",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShuffleEpi32",
		InstructionName: "PSHUFD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SlliSi128",
		InstructionName: "PSLLO",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SrliSi128",
		InstructionName: "PSRLO",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "UnpackhiEpi64",
		InstructionName: "PUNPCKHLQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "UnpackloEpi64",
		InstructionName: "PUNPCKLLQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AddPd",
		InstructionName: "ADDPD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AddSd",
		InstructionName: "ADDSD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AndnotPd",
		InstructionName: "ANDNPD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "CmpeqPd",
		InstructionName: "CMPPD",
		VarOps: []VarOp{
			{
				Var:   Variable{},
				Const: 0,
				Op:    Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "CmpeqSd",
		InstructionName: "CMPPD",
		VarOps: []VarOp{
			{
				Var:   Variable{},
				Const: 0,
				Op:    Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 2,
	},
}

func getSSE2(name string) (Intrinsic, bool) {
	for _, sse2instr := range sse2intructions {
		if sse2instr.Name == name {
			return sse2instr, true
		}
	}
	return Intrinsic{}, false
}

func sse2Intrinsic(f *Function, loc ssa.Instruction, call *ssa.Call, intrinsic Intrinsic, args []ssa.Value) string {
	var asm string
	idents := []*identifier{}
	argRegs := []*register{}
	result := f.Ident(call)
	asm += fmt.Sprintf("// BEGIN INTRINSIC %v\n", intrinsic.Name)
	asm += fmt.Sprintf("// BEGIN LOAD ARGS %v\n", intrinsic.Name)
	// load the arguments into registers
	for _, arg := range args {
		ident := f.Ident(arg)
		a, reg, err := f.LoadIdent(loc, ident, 0, ident.size())
		if err != nil {
			ice(fmt.Sprintf("%v", err.Err))
		}
		reg.inUse = true
		asm += a
		idents = append(idents, ident)
		argRegs = append(argRegs, reg)
	}
	asm += fmt.Sprintf("// END LOAD ARGS %v\n", intrinsic.Name)

	// construct the assembly instruction
	var regResult *register
	argIdx := 0
	assembly := fmt.Sprintf("%-9v    ", intrinsic.InstructionName)
	for i, varop := range intrinsic.VarOps {
		flags := varop.Op.Flags
		if flags&Implicit == 1 {
			continue
		}

		if flags&Const != 0 {
			if flags&Imm8 != 0 || flags&Imm32 != 0 {
				asm += "$" + strconv.Itoa(varop.Const)
			} else if flags&ImmF32 != 0 || flags&ImmF64 != 0 {
				asm += ice("unimplemented case")
			} else {
				ice("unexpected case")
			}
			continue
		}
		if flags&In == 0 && flags&Out == 0 {
			continue
		}
		reg := argRegs[argIdx]
		if reflectType(args[argIdx].Type()).String() != varop.Var.VarType.String() {
			got := reflectType(args[argIdx].Type()).String()
			expected := varop.Var.VarType.String()
			msg := fmt.Sprintf("wrong type for argument, Got: %v, Expected: %v", got, expected)
			ice(msg)
		}
		if flags&In != 0 {
			assembly += reg.name + ", "

		}
		if flags&Out != 0 {
			if i == intrinsic.ResultIdxOp {
				regResult = reg
			} else if flags&Out != 0 {
				reg.dirty = true
			}
		}
		argIdx++

	}
	assembly = strings.TrimSuffix(assembly, ", ") + "\n"
	asm += assembly
	if regResult != nil {
		c, err := f.StoreSSE2(loc, regResult, result)
		if err != nil {
			ice(fmt.Sprint(err.Err))
		}
		asm += c
	} else {
		ice("no result from intrinsic")
	}
	for _, reg := range argRegs {
		f.freeReg(reg)
	}
	asm += fmt.Sprintf("// END INTRINSIC %v\n", intrinsic.Name)
	return asm
}
<commit_msg>Intrinsics add immediate support<commit_after>package codegen

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"golang.org/x/tools/go/ssa"

	"github.com/bjwbell/gensimd/simd"
)

var sse2intructions = []Intrinsic{
	{
		Name:            "AddEpi64",
		InstructionName: "PADDQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SubEpi64",
		InstructionName: "PSUBQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "MulEpu32",
		InstructionName: "PMULULQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShufflehiEpi16",
		InstructionName: "PSHUFHW",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShuffleloEpi16",
		InstructionName: "PSHUFLW",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "ShuffleEpi32",
		InstructionName: "PSHUFD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SlliSi128",
		InstructionName: "PSLLO",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "SrliSi128",
		InstructionName: "PSRLO",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "UnpackhiEpi64",
		InstructionName: "PUNPCKHLQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "UnpackloEpi64",
		InstructionName: "PUNPCKLLQ",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128i{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AddPd",
		InstructionName: "ADDPD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AddSd",
		InstructionName: "ADDSD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "AndnotPd",
		InstructionName: "ANDNPD",
		VarOps: []VarOp{
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128d{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "CmpeqPd",
		InstructionName: "CMPPD",
		VarOps: []VarOp{
			{
				Var:   Variable{},
				Const: 0,
				Op:    Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 1,
	},
	{
		Name:            "CmpeqSd",
		InstructionName: "CMPPD",
		VarOps: []VarOp{
			{
				Var:   Variable{},
				Const: 0,
				Op:    Operand{Flags: Imm8, Reg: INVALID_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In, Reg: XMM_REG, NamedReg: REG_INVALID}},
			{
				Var: Variable{reflect.TypeOf(simd.M128{})},
				Op:  Operand{Flags: In | Out, Reg: XMM_REG, NamedReg: REG_INVALID},
			},
		},
		ResultIdxOp: 2,
	},
}

func getSSE2(name string) (Intrinsic, bool) {
	for _, sse2instr := range sse2intructions {
		if sse2instr.Name == name {
			return sse2instr, true
		}
	}
	return Intrinsic{}, false
}

func sse2Intrinsic(f *Function, loc ssa.Instruction, call *ssa.Call, intrinsic Intrinsic, args []ssa.Value) string {
	var asm string
	idents := []*identifier{}
	argRegs := []*register{}
	result := f.Ident(call)
	asm += fmt.Sprintf("// BEGIN INTRINSIC %v\n", intrinsic.Name)
	asm += fmt.Sprintf("// BEGIN LOAD ARGS %v\n", intrinsic.Name)
	// load the arguments into registers
	for _, arg := range args {
		ident := f.Ident(arg)
		a, reg, err := f.LoadIdent(loc, ident, 0, ident.size())
		if err != nil {
			ice(fmt.Sprintf("%v", err.Err))
		}
		reg.inUse = true
		asm += a
		idents = append(idents, ident)
		argRegs = append(argRegs, reg)
	}
	asm += fmt.Sprintf("// END LOAD ARGS %v\n", intrinsic.Name)

	// construct the assembly instruction
	var regResult *register
	argIdx := 0
	assembly := fmt.Sprintf("%-9v    ", intrinsic.InstructionName)
	for i, varop := range intrinsic.VarOps {
		flags := varop.Op.Flags
		if flags&Implicit == 1 {
			continue
		}

		if flags&Const != 0 {
			if flags&Imm8 != 0 || flags&Imm32 != 0 {
				asm += "$" + strconv.Itoa(varop.Const)
			} else if flags&ImmF32 != 0 || flags&ImmF64 != 0 {
				asm += ice("unimplemented case")
			} else {
				ice("unexpected case")
			}
			continue
		}
		if flags&In == 0 && flags&Out == 0 {
			continue
		}
		ident := idents[argIdx]
		reg := argRegs[argIdx]
		if reflectType(args[argIdx].Type()).String() != varop.Var.VarType.String() {
			got := reflectType(args[argIdx].Type()).String()
			expected := varop.Var.VarType.String()
			msg := fmt.Sprintf("wrong type for argument, Got: %v, Expected: %v", got, expected)
			ice(msg)
		}
		if flags&In != 0 {
			if flags&Imm8 != 0 || flags&Imm32 != 0 {
				if ident.cnst == nil {
					panic("Provided argument must be a constant")
				} else {
					assembly += "$" + strconv.Itoa(int(ident.cnst.Uint64()))
				}
			} else if flags&ImmF32 != 0 || flags&ImmF64 != 0 {
				ice("unimplemented case")
			}
			assembly += reg.name + ", "

		}
		if flags&Out != 0 {
			if i == intrinsic.ResultIdxOp {
				regResult = reg
			} else if flags&Out != 0 {
				reg.dirty = true
			}
		}
		argIdx++

	}
	assembly = strings.TrimSuffix(assembly, ", ") + "\n"
	asm += assembly
	if regResult != nil {
		c, err := f.StoreSSE2(loc, regResult, result)
		if err != nil {
			ice(fmt.Sprint(err.Err))
		}
		asm += c
	} else {
		ice("no result from intrinsic")
	}
	for _, reg := range argRegs {
		f.freeReg(reg)
	}
	asm += fmt.Sprintf("// END INTRINSIC %v\n", intrinsic.Name)
	return asm
}
<|endoftext|>


<commit_before>package redis

import (
	"context"
	"fmt"
	"time"

	libredis "github.com/go-redis/redis/v7"
	"github.com/pkg/errors"

	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/common"
)

// Client is an interface thats allows to use a redis cluster or a redis single client seamlessly.
type Client interface {
	Ping() *libredis.StatusCmd
	Get(key string) *libredis.StringCmd
	Set(key string, value interface{}, expiration time.Duration) *libredis.StatusCmd
	Watch(handler func(*libredis.Tx) error, keys ...string) error
	Del(keys ...string) *libredis.IntCmd
	SetNX(key string, value interface{}, expiration time.Duration) *libredis.BoolCmd
	Eval(script string, keys []string, args ...interface{}) *libredis.Cmd
}

// Store is the redis store.
type Store struct {
	// Prefix used for the key.
	Prefix string
	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int
	// client used to communicate with redis server.
	client Client
}

// NewStore returns an instance of redis store with defaults.
func NewStore(client Client) (limiter.Store, error) {
	return NewStoreWithOptions(client, limiter.StoreOptions{
		Prefix:          limiter.DefaultPrefix,
		CleanUpInterval: limiter.DefaultCleanUpInterval,
		MaxRetry:        limiter.DefaultMaxRetry,
	})
}

// NewStoreWithOptions returns an instance of redis store with options.
func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.Store, error) {
	store := &Store{
		client:   client,
		Prefix:   options.Prefix,
		MaxRetry: options.MaxRetry,
	}

	if store.MaxRetry <= 0 {
		store.MaxRetry = 1
	}

	_, err := store.ping()
	if err != nil {
		return nil, err
	}

	return store, nil
}

// Get returns the limit for given identifier.
func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		created, err := store.doSetValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		if created {
			expiration := now.Add(rate.Period)
			lctx = common.GetContextFromState(now, rate, expiration, 1)
			return nil
		}

		count, ttl, err := store.doUpdateValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot get value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Peek returns the limit for given identifier, without modification on current values.
func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {
		count, ttl, err := store.doPeekValue(rtx, key)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot peek value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Reset returns the limit for given identifier which is set to zero.
func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		err := store.doResetValue(rtx, key)
		if err != nil {
			return err
		}

		count := int64(0)
		expiration := now.Add(rate.Period)

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot reset value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// doPeekValue will execute peekValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doPeekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := peekValue(rtx, key)
		if err == nil {
			return count, ttl, nil
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// peekValue will retrieve the counter and its expiration for given key.
func peekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	pipe := rtx.TxPipeline()
	value := pipe.Get(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	count, err := value.Int64()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	return count, ttl, nil
}

// doSetValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doSetValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	for i := 0; i < store.MaxRetry; i++ {
		created, err := setValue(rtx, key, expiration)
		if err == nil {
			return created, nil
		}
	}
	return false, errors.New("retry limit exceeded")
}

// setValue will try to initialize a new counter if given key doesn't exists.
func setValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	value := rtx.SetNX(key, 1, expiration)

	created, err := value.Result()
	if err != nil {
		return false, err
	}

	return created, nil
}

// doUpdateValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doUpdateValue(rtx *libredis.Tx, key string,
	expiration time.Duration) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := updateValue(rtx, key, expiration)
		if err == nil {
			return count, ttl, nil
		}

		// If ttl is negative and there is an error, do not retry an update.
		if ttl < 0 {
			return 0, 0, err
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// updateValue will try to increment the counter identified by given key.
func updateValue(rtx *libredis.Tx, key string, expiration time.Duration) (int64, time.Duration, error) {
	pipe := rtx.TxPipeline()
	value := pipe.Incr(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil {
		return 0, 0, err
	}

	count, err := value.Result()
	if err != nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	// If ttl is -1ms, we have to define key expiration.
	// The PTTL command returns -2 if the key does not exist, and -1 if the key exists, but there is no expiry set.
	// We shouldn't try to set an expiry on a key that doesn't exist.
	if isKeyExpired(ttl) {
		expire := rtx.Expire(key, expiration)

		ok, err := expire.Result()
		if err != nil {
			return count, ttl, err
		}

		if !ok {
			return count, ttl, errors.New("cannot configure timeout on key")
		}
	}

	return count, ttl, nil
}

// doResetValue will execute resetValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doResetValue(rtx *libredis.Tx, key string) error {
	for i := 0; i < store.MaxRetry; i++ {
		err := resetValue(rtx, key)
		if err == nil {
			return nil
		}
	}
	return errors.New("retry limit exceeded")
}

// resetValue will try to reset the counter identified by given key.
func resetValue(rtx *libredis.Tx, key string) error {
	deletion := rtx.Del(key)

	_, err := deletion.Result()
	if err != nil {
		return err
	}

	return nil
}

// ping checks if redis is alive.
func (store *Store) ping() (bool, error) {
	cmd := store.client.Ping()

	pong, err := cmd.Result()
	if err != nil {
		return false, errors.Wrap(err, "limiter: cannot ping redis server")
	}

	return (pong == "PONG"), nil
}

// isKeyExpired returns from PTTL command -2 if the key does not exist, and -1 if the key exists.
// Usually, it should be returned in nanosecond, but some users have reported that it could be in millisecond as well.
// Better safe than sorry: we handle both.
func isKeyExpired(ttl time.Duration) bool {
	switch ttl {
	case -1 * time.Nanosecond, -1 * time.Millisecond:
		return true
	default:
		return false
	}
}
<commit_msg>refactor(redis): Better naming<commit_after>package redis

import (
	"context"
	"fmt"
	"time"

	libredis "github.com/go-redis/redis/v7"
	"github.com/pkg/errors"

	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/common"
)

// Client is an interface thats allows to use a redis cluster or a redis single client seamlessly.
type Client interface {
	Ping() *libredis.StatusCmd
	Get(key string) *libredis.StringCmd
	Set(key string, value interface{}, expiration time.Duration) *libredis.StatusCmd
	Watch(handler func(*libredis.Tx) error, keys ...string) error
	Del(keys ...string) *libredis.IntCmd
	SetNX(key string, value interface{}, expiration time.Duration) *libredis.BoolCmd
	Eval(script string, keys []string, args ...interface{}) *libredis.Cmd
}

// Store is the redis store.
type Store struct {
	// Prefix used for the key.
	Prefix string
	// MaxRetry is the maximum number of retry under race conditions.
	MaxRetry int
	// client used to communicate with redis server.
	client Client
}

// NewStore returns an instance of redis store with defaults.
func NewStore(client Client) (limiter.Store, error) {
	return NewStoreWithOptions(client, limiter.StoreOptions{
		Prefix:          limiter.DefaultPrefix,
		CleanUpInterval: limiter.DefaultCleanUpInterval,
		MaxRetry:        limiter.DefaultMaxRetry,
	})
}

// NewStoreWithOptions returns an instance of redis store with options.
func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.Store, error) {
	store := &Store{
		client:   client,
		Prefix:   options.Prefix,
		MaxRetry: options.MaxRetry,
	}

	if store.MaxRetry <= 0 {
		store.MaxRetry = 1
	}

	_, err := store.ping()
	if err != nil {
		return nil, err
	}

	return store, nil
}

// Get returns the limit for given identifier.
func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		created, err := store.doSetValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		if created {
			expiration := now.Add(rate.Period)
			lctx = common.GetContextFromState(now, rate, expiration, 1)
			return nil
		}

		count, ttl, err := store.doUpdateValue(rtx, key, rate.Period)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot get value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Peek returns the limit for given identifier, without modification on current values.
func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {
		count, ttl, err := store.doPeekValue(rtx, key)
		if err != nil {
			return err
		}

		expiration := now.Add(rate.Period)
		if ttl > 0 {
			expiration = now.Add(ttl)
		}

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot peek value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// Reset returns the limit for given identifier which is set to zero.
func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
	key = fmt.Sprintf("%s:%s", store.Prefix, key)
	now := time.Now()

	lctx := limiter.Context{}
	onWatch := func(rtx *libredis.Tx) error {

		err := store.doResetValue(rtx, key)
		if err != nil {
			return err
		}

		count := int64(0)
		expiration := now.Add(rate.Period)

		lctx = common.GetContextFromState(now, rate, expiration, count)
		return nil
	}

	err := store.client.Watch(onWatch, key)
	if err != nil {
		err = errors.Wrapf(err, "limiter: cannot reset value for %s", key)
		return limiter.Context{}, err
	}

	return lctx, nil
}

// doPeekValue will execute peekValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doPeekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := peekValue(rtx, key)
		if err == nil {
			return count, ttl, nil
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// peekValue will retrieve the counter and its expiration for given key.
func peekValue(rtx *libredis.Tx, key string) (int64, time.Duration, error) {
	pipe := rtx.TxPipeline()
	value := pipe.Get(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	count, err := value.Int64()
	if err != nil && err != libredis.Nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	return count, ttl, nil
}

// doSetValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doSetValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	for i := 0; i < store.MaxRetry; i++ {
		created, err := setValue(rtx, key, expiration)
		if err == nil {
			return created, nil
		}
	}
	return false, errors.New("retry limit exceeded")
}

// setValue will try to initialize a new counter if given key doesn't exists.
func setValue(rtx *libredis.Tx, key string, expiration time.Duration) (bool, error) {
	value := rtx.SetNX(key, 1, expiration)

	created, err := value.Result()
	if err != nil {
		return false, err
	}

	return created, nil
}

// doUpdateValue will execute setValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doUpdateValue(rtx *libredis.Tx, key string,
	expiration time.Duration) (int64, time.Duration, error) {
	for i := 0; i < store.MaxRetry; i++ {
		count, ttl, err := updateValue(rtx, key, expiration)
		if err == nil {
			return count, ttl, nil
		}

		// If ttl is negative and there is an error, do not retry an update.
		if ttl < 0 {
			return 0, 0, err
		}
	}
	return 0, 0, errors.New("retry limit exceeded")
}

// updateValue will try to increment the counter identified by given key.
func updateValue(rtx *libredis.Tx, key string, expiration time.Duration) (int64, time.Duration, error) {
	pipe := rtx.TxPipeline()
	value := pipe.Incr(key)
	expire := pipe.PTTL(key)

	_, err := pipe.Exec()
	if err != nil {
		return 0, 0, err
	}

	count, err := value.Result()
	if err != nil {
		return 0, 0, err
	}

	ttl, err := expire.Result()
	if err != nil {
		return 0, 0, err
	}

	// If ttl is less than zero, we have to define key expiration.
	// The PTTL command returns -2 if the key does not exist, and -1 if the key exists, but there is no expiry set.
	// We shouldn't try to set an expiry on a key that doesn't exist.
	if isExpirationRequired(ttl) {
		expire := rtx.Expire(key, expiration)

		ok, err := expire.Result()
		if err != nil {
			return count, ttl, err
		}

		if !ok {
			return count, ttl, errors.New("cannot configure timeout on key")
		}
	}

	return count, ttl, nil
}

// doResetValue will execute resetValue with a retry mecanism (optimistic locking) until store.MaxRetry is reached.
func (store *Store) doResetValue(rtx *libredis.Tx, key string) error {
	for i := 0; i < store.MaxRetry; i++ {
		err := resetValue(rtx, key)
		if err == nil {
			return nil
		}
	}
	return errors.New("retry limit exceeded")
}

// resetValue will try to reset the counter identified by given key.
func resetValue(rtx *libredis.Tx, key string) error {
	deletion := rtx.Del(key)

	_, err := deletion.Result()
	if err != nil {
		return err
	}

	return nil
}

// ping checks if redis is alive.
func (store *Store) ping() (bool, error) {
	cmd := store.client.Ping()

	pong, err := cmd.Result()
	if err != nil {
		return false, errors.Wrap(err, "limiter: cannot ping redis server")
	}

	return (pong == "PONG"), nil
}

// isExpirationRequired returns if we should set an expiration on a key, using (error) result from PTTL command.
// The error code is -2 if the key does not exist, and -1 if the key exists.
// Usually, it should be returned in nanosecond, but some users have reported that it could be in millisecond as well.
// Better safe than sorry: we handle both.
func isExpirationRequired(ttl time.Duration) bool {
	switch ttl {
	case -1 * time.Nanosecond, -1 * time.Millisecond:
		return true
	default:
		return false
	}
}
<|endoftext|>