<commit_before>package stdcopy

import (
	"encoding/binary"
	"errors"
	"io"

	"github.com/Sirupsen/logrus"
)

const (
	StdWriterPrefixLen = 8
	StdWriterFdIndex   = 0
	StdWriterSizeIndex = 4
)

type StdType [StdWriterPrefixLen]byte

var (
	Stdin  StdType = StdType{0: 0}
	Stdout StdType = StdType{0: 1}
	Stderr StdType = StdType{0: 2}
)

type StdWriter struct {
	io.Writer
	prefix  StdType
	sizeBuf []byte
}

func (w *StdWriter) Write(buf []byte) (n int, err error) {
	var n1, n2 int
	if w == nil || w.Writer == nil {
		return 0, errors.New("Writer not instanciated")
	}
	binary.BigEndian.PutUint32(w.prefix[4:], uint32(len(buf)))
	n1, err = w.Writer.Write(w.prefix[:])
	if err != nil {
		n = n1 - StdWriterPrefixLen
	} else {
		n2, err = w.Writer.Write(buf)
		n = n1 + n2 - StdWriterPrefixLen
	}
	if n < 0 {
		n = 0
	}
	return
}

// NewStdWriter instanciates a new Writer.
// Everything written to it will be encapsulated using a custom format,
// and written to the underlying `w` stream.
// This allows multiple write streams (e.g. stdout and stderr) to be muxed into a single connection.
// `t` indicates the id of the stream to encapsulate.
// It can be utils.Stdin, utils.Stdout, utils.Stderr.
func NewStdWriter(w io.Writer, t StdType) *StdWriter {
	if len(t) != StdWriterPrefixLen {
		return nil
	}

	return &StdWriter{
		Writer:  w,
		prefix:  t,
		sizeBuf: make([]byte, 4),
	}
}

var ErrInvalidStdHeader = errors.New("Unrecognized input header")

// StdCopy is a modified version of io.Copy.
//
// StdCopy will demultiplex `src`, assuming that it contains two streams,
// previously multiplexed together using a StdWriter instance.
// As it reads from `src`, StdCopy will write to `dstout` and `dsterr`.
//
// StdCopy will read until it hits EOF on `src`. It will then return a nil error.
// In other words: if `err` is non nil, it indicates a real underlying error.
//
// `written` will hold the total number of bytes written to `dstout` and `dsterr`.
func StdCopy(dstout, dsterr io.Writer, src io.Reader) (written int64, err error) {
	var (
		buf       = make([]byte, 32*1024+StdWriterPrefixLen+1)
		bufLen    = len(buf)
		nr, nw    int
		er, ew    error
		out       io.Writer
		frameSize int
	)

	for {
		// Make sure we have at least a full header
		for nr < StdWriterPrefixLen {
			var nr2 int
			nr2, er = src.Read(buf[nr:])
			nr += nr2
			if er == io.EOF {
				if nr < StdWriterPrefixLen {
					logrus.Debugf("Corrupted prefix: %v", buf[:nr])
					return written, nil
				}
				break
			}
			if er != nil {
				logrus.Debugf("Error reading header: %s", er)
				return 0, er
			}
		}

		// Check the first byte to know where to write
		switch buf[StdWriterFdIndex] {
		case 0:
			fallthrough
		case 1:
			// Write on stdout
			out = dstout
		case 2:
			// Write on stderr
			out = dsterr
		default:
			logrus.Debugf("Error selecting output fd: (%d)", buf[StdWriterFdIndex])
			return 0, ErrInvalidStdHeader
		}

		// Retrieve the size of the frame
		frameSize = int(binary.BigEndian.Uint32(buf[StdWriterSizeIndex : StdWriterSizeIndex+4]))
		logrus.Debugf("framesize: %d", frameSize)

		// Check if the buffer is big enough to read the frame.
		// Extend it if necessary.
		if frameSize+StdWriterPrefixLen > bufLen {
			logrus.Debugf("Extending buffer cap by %d (was %d)", frameSize+StdWriterPrefixLen-bufLen+1, len(buf))
			buf = append(buf, make([]byte, frameSize+StdWriterPrefixLen-bufLen+1)...)
			bufLen = len(buf)
		}

		// While the amount of bytes read is less than the size of the frame + header, we keep reading
		for nr < frameSize+StdWriterPrefixLen {
			var nr2 int
			nr2, er = src.Read(buf[nr:])
			nr += nr2
			if er == io.EOF {
				if nr < frameSize+StdWriterPrefixLen {
					logrus.Debugf("Corrupted frame: %v", buf[StdWriterPrefixLen:nr])
					return written, nil
				}
				break
			}
			if er != nil {
				logrus.Debugf("Error reading frame: %s", er)
				return 0, er
			}
		}

		// Write the retrieved frame (without header)
		nw, ew = out.Write(buf[StdWriterPrefixLen : frameSize+StdWriterPrefixLen])
		if ew != nil {
			logrus.Debugf("Error writing frame: %s", ew)
			return 0, ew
		}
		// If the frame has not been fully written: error
		if nw != frameSize {
			logrus.Debugf("Error Short Write: (%d on %d)", nw, frameSize)
			return 0, io.ErrShortWrite
		}
		written += int64(nw)

		// Move the rest of the buffer to the beginning
		copy(buf, buf[frameSize+StdWriterPrefixLen:])
		// Move the index
		nr -= frameSize + StdWriterPrefixLen
	}
}
<commit_msg>correct pkg/stdcopy NewStdWriter function comments<commit_after>package stdcopy

import (
	"encoding/binary"
	"errors"
	"io"

	"github.com/Sirupsen/logrus"
)

const (
	StdWriterPrefixLen = 8
	StdWriterFdIndex   = 0
	StdWriterSizeIndex = 4
)

type StdType [StdWriterPrefixLen]byte

var (
	Stdin  StdType = StdType{0: 0}
	Stdout StdType = StdType{0: 1}
	Stderr StdType = StdType{0: 2}
)

type StdWriter struct {
	io.Writer
	prefix  StdType
	sizeBuf []byte
}

func (w *StdWriter) Write(buf []byte) (n int, err error) {
	var n1, n2 int
	if w == nil || w.Writer == nil {
		return 0, errors.New("Writer not instanciated")
	}
	binary.BigEndian.PutUint32(w.prefix[4:], uint32(len(buf)))
	n1, err = w.Writer.Write(w.prefix[:])
	if err != nil {
		n = n1 - StdWriterPrefixLen
	} else {
		n2, err = w.Writer.Write(buf)
		n = n1 + n2 - StdWriterPrefixLen
	}
	if n < 0 {
		n = 0
	}
	return
}

// NewStdWriter instanciates a new Writer.
// Everything written to it will be encapsulated using a custom format,
// and written to the underlying `w` stream.
// This allows multiple write streams (e.g. stdout and stderr) to be muxed into a single connection.
// `t` indicates the id of the stream to encapsulate.
// It can be stdcopy.Stdin, stdcopy.Stdout, stdcopy.Stderr.
func NewStdWriter(w io.Writer, t StdType) *StdWriter {
	if len(t) != StdWriterPrefixLen {
		return nil
	}

	return &StdWriter{
		Writer:  w,
		prefix:  t,
		sizeBuf: make([]byte, 4),
	}
}

var ErrInvalidStdHeader = errors.New("Unrecognized input header")

// StdCopy is a modified version of io.Copy.
//
// StdCopy will demultiplex `src`, assuming that it contains two streams,
// previously multiplexed together using a StdWriter instance.
// As it reads from `src`, StdCopy will write to `dstout` and `dsterr`.
//
// StdCopy will read until it hits EOF on `src`. It will then return a nil error.
// In other words: if `err` is non nil, it indicates a real underlying error.
//
// `written` will hold the total number of bytes written to `dstout` and `dsterr`.
func StdCopy(dstout, dsterr io.Writer, src io.Reader) (written int64, err error) {
	var (
		buf       = make([]byte, 32*1024+StdWriterPrefixLen+1)
		bufLen    = len(buf)
		nr, nw    int
		er, ew    error
		out       io.Writer
		frameSize int
	)

	for {
		// Make sure we have at least a full header
		for nr < StdWriterPrefixLen {
			var nr2 int
			nr2, er = src.Read(buf[nr:])
			nr += nr2
			if er == io.EOF {
				if nr < StdWriterPrefixLen {
					logrus.Debugf("Corrupted prefix: %v", buf[:nr])
					return written, nil
				}
				break
			}
			if er != nil {
				logrus.Debugf("Error reading header: %s", er)
				return 0, er
			}
		}

		// Check the first byte to know where to write
		switch buf[StdWriterFdIndex] {
		case 0:
			fallthrough
		case 1:
			// Write on stdout
			out = dstout
		case 2:
			// Write on stderr
			out = dsterr
		default:
			logrus.Debugf("Error selecting output fd: (%d)", buf[StdWriterFdIndex])
			return 0, ErrInvalidStdHeader
		}

		// Retrieve the size of the frame
		frameSize = int(binary.BigEndian.Uint32(buf[StdWriterSizeIndex : StdWriterSizeIndex+4]))
		logrus.Debugf("framesize: %d", frameSize)

		// Check if the buffer is big enough to read the frame.
		// Extend it if necessary.
		if frameSize+StdWriterPrefixLen > bufLen {
			logrus.Debugf("Extending buffer cap by %d (was %d)", frameSize+StdWriterPrefixLen-bufLen+1, len(buf))
			buf = append(buf, make([]byte, frameSize+StdWriterPrefixLen-bufLen+1)...)
			bufLen = len(buf)
		}

		// While the amount of bytes read is less than the size of the frame + header, we keep reading
		for nr < frameSize+StdWriterPrefixLen {
			var nr2 int
			nr2, er = src.Read(buf[nr:])
			nr += nr2
			if er == io.EOF {
				if nr < frameSize+StdWriterPrefixLen {
					logrus.Debugf("Corrupted frame: %v", buf[StdWriterPrefixLen:nr])
					return written, nil
				}
				break
			}
			if er != nil {
				logrus.Debugf("Error reading frame: %s", er)
				return 0, er
			}
		}

		// Write the retrieved frame (without header)
		nw, ew = out.Write(buf[StdWriterPrefixLen : frameSize+StdWriterPrefixLen])
		if ew != nil {
			logrus.Debugf("Error writing frame: %s", ew)
			return 0, ew
		}
		// If the frame has not been fully written: error
		if nw != frameSize {
			logrus.Debugf("Error Short Write: (%d on %d)", nw, frameSize)
			return 0, io.ErrShortWrite
		}
		written += int64(nw)

		// Move the rest of the buffer to the beginning
		copy(buf, buf[frameSize+StdWriterPrefixLen:])
		// Move the index
		nr -= frameSize + StdWriterPrefixLen
	}
}
<|endoftext|>


<commit_before>package initialize

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/coreos/coreos-cloudinit/system"
)

const (
	locksmithUnit = "locksmithd.service"
)

// updateOption represents a configurable update option, which, if set, will be
// written into update.conf, replacing any existing value for the option
type updateOption struct {
	key    string   // key used to configure this option in cloud-config
	valid  []string // valid values for the option
	prefix string   // prefix for the option in the update.conf file
	value  string   // used to store the new value in update.conf (including prefix)
	seen   bool     // whether the option has been seen in any existing update.conf
}

// updateOptions defines the update options understood by cloud-config.
// The keys represent the string used in cloud-config to configure the option.
var updateOptions = []*updateOption{
	&updateOption{
		key:    "reboot-strategy",
		prefix: "REBOOT_STRATEGY=",
		valid:  []string{"best-effort", "etcd-lock", "reboot", "off"},
	},
	&updateOption{
		key:    "group",
		prefix: "GROUP=",
		valid:  []string{"master", "beta", "alpha", "stable"},
	},
	&updateOption{
		key:    "server",
		prefix: "SERVER=",
	},
}

// isValid checks whether a supplied value is valid for this option
func (uo updateOption) isValid(val string) bool {
	if len(uo.valid) == 0 {
		return true
	}
	for _, v := range uo.valid {
		if val == v {
			return true
		}
	}
	return false
}

type UpdateConfig map[string]string

// File generates an `/etc/coreos/update.conf` file (if any update
// configuration options are set in cloud-config) by either rewriting the
// existing file on disk, or starting from `/usr/share/coreos/update.conf`
func (uc UpdateConfig) File(root string) (*system.File, error) {
	if len(uc) < 1 {
		return nil, nil
	}

	var out string

	// Generate the list of possible substitutions to be performed based on the options that are configured
	subs := make([]*updateOption, 0)
	for _, uo := range updateOptions {
		val, ok := uc[uo.key]
		if !ok {
			continue
		}
		if !uo.isValid(val) {
			return nil, errors.New(fmt.Sprintf("invalid value %v for option %v (valid options: %v)", val, uo.key, uo.valid))
		}
		uo.value = uo.prefix + val
		subs = append(subs, uo)
	}

	etcUpdate := path.Join(root, "etc", "coreos", "update.conf")
	usrUpdate := path.Join(root, "usr", "share", "coreos", "update.conf")

	conf, err := os.Open(etcUpdate)
	if os.IsNotExist(err) {
		conf, err = os.Open(usrUpdate)
	}
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(conf)

	for scanner.Scan() {
		line := scanner.Text()
		for _, s := range subs {
			if strings.HasPrefix(line, s.prefix) {
				line = s.value
				s.seen = true
				break
			}
		}
		out += line
		out += "\n"
		if err := scanner.Err(); err != nil {
			return nil, err
		}
	}

	for _, s := range subs {
		if !s.seen {
			out += s.value
			out += "\n"
		}
	}

	return &system.File{
		Path:               path.Join("etc", "coreos", "update.conf"),
		RawFilePermissions: "0644",
		Content:            out,
	}, nil
}

// Units generates units for the cloud-init initializer to act on:
// - a locksmith system.Unit, if "reboot-strategy" was set in cloud-config
// - an update_engine system.Unit, if "group" was set in cloud-config
func (uc UpdateConfig) Units(root string) ([]system.Unit, error) {
	var units []system.Unit
	if strategy, ok := uc["reboot-strategy"]; ok {
		ls := &system.Unit{
			Name:    locksmithUnit,
			Command: "restart",
			Mask:    false,
			Runtime: true,
		}

		if strategy == "off" {
			ls.Command = "stop"
			ls.Mask = true
		}
		units = append(units, *ls)
	}

	rue := false
	if _, ok := uc["group"]; ok {
		rue = true
	}
	if _, ok := uc["server"]; ok {
		rue = true
	}
	if rue {
		ue := system.Unit{
			Name:    "update-engine",
			Command: "restart",
		}
		units = append(units, ue)
	}

	return units, nil
}
<commit_msg>fix(update): Relax requirements of update group value<commit_after>package initialize

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/coreos/coreos-cloudinit/system"
)

const (
	locksmithUnit = "locksmithd.service"
)

// updateOption represents a configurable update option, which, if set, will be
// written into update.conf, replacing any existing value for the option
type updateOption struct {
	key    string   // key used to configure this option in cloud-config
	valid  []string // valid values for the option
	prefix string   // prefix for the option in the update.conf file
	value  string   // used to store the new value in update.conf (including prefix)
	seen   bool     // whether the option has been seen in any existing update.conf
}

// updateOptions defines the update options understood by cloud-config.
// The keys represent the string used in cloud-config to configure the option.
var updateOptions = []*updateOption{
	&updateOption{
		key:    "reboot-strategy",
		prefix: "REBOOT_STRATEGY=",
		valid:  []string{"best-effort", "etcd-lock", "reboot", "off"},
	},
	&updateOption{
		key:    "group",
		prefix: "GROUP=",
	},
	&updateOption{
		key:    "server",
		prefix: "SERVER=",
	},
}

// isValid checks whether a supplied value is valid for this option
func (uo updateOption) isValid(val string) bool {
	if len(uo.valid) == 0 {
		return true
	}
	for _, v := range uo.valid {
		if val == v {
			return true
		}
	}
	return false
}

type UpdateConfig map[string]string

// File generates an `/etc/coreos/update.conf` file (if any update
// configuration options are set in cloud-config) by either rewriting the
// existing file on disk, or starting from `/usr/share/coreos/update.conf`
func (uc UpdateConfig) File(root string) (*system.File, error) {
	if len(uc) < 1 {
		return nil, nil
	}

	var out string

	// Generate the list of possible substitutions to be performed based on the options that are configured
	subs := make([]*updateOption, 0)
	for _, uo := range updateOptions {
		val, ok := uc[uo.key]
		if !ok {
			continue
		}
		if !uo.isValid(val) {
			return nil, errors.New(fmt.Sprintf("invalid value %v for option %v (valid options: %v)", val, uo.key, uo.valid))
		}
		uo.value = uo.prefix + val
		subs = append(subs, uo)
	}

	etcUpdate := path.Join(root, "etc", "coreos", "update.conf")
	usrUpdate := path.Join(root, "usr", "share", "coreos", "update.conf")

	conf, err := os.Open(etcUpdate)
	if os.IsNotExist(err) {
		conf, err = os.Open(usrUpdate)
	}
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(conf)

	for scanner.Scan() {
		line := scanner.Text()
		for _, s := range subs {
			if strings.HasPrefix(line, s.prefix) {
				line = s.value
				s.seen = true
				break
			}
		}
		out += line
		out += "\n"
		if err := scanner.Err(); err != nil {
			return nil, err
		}
	}

	for _, s := range subs {
		if !s.seen {
			out += s.value
			out += "\n"
		}
	}

	return &system.File{
		Path:               path.Join("etc", "coreos", "update.conf"),
		RawFilePermissions: "0644",
		Content:            out,
	}, nil
}

// Units generates units for the cloud-init initializer to act on:
// - a locksmith system.Unit, if "reboot-strategy" was set in cloud-config
// - an update_engine system.Unit, if "group" was set in cloud-config
func (uc UpdateConfig) Units(root string) ([]system.Unit, error) {
	var units []system.Unit
	if strategy, ok := uc["reboot-strategy"]; ok {
		ls := &system.Unit{
			Name:    locksmithUnit,
			Command: "restart",
			Mask:    false,
			Runtime: true,
		}

		if strategy == "off" {
			ls.Command = "stop"
			ls.Mask = true
		}
		units = append(units, *ls)
	}

	rue := false
	if _, ok := uc["group"]; ok {
		rue = true
	}
	if _, ok := uc["server"]; ok {
		rue = true
	}
	if rue {
		ue := system.Unit{
			Name:    "update-engine",
			Command: "restart",
		}
		units = append(units, ue)
	}

	return units, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 util

import (
	"strings"
	"testing"
)

func TestGetBinaryDownloadURL(t *testing.T) {
	testData := []struct {
		version     string
		platform    string
		expectedURL string
	}{
		{"v0.0.1", "linux", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-linux-amd64"},
		{"v0.0.1", "darwin", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-darwin-amd64"},
		{"v0.0.1", "windows", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-windows-amd64.exe"},
	}

	for _, tt := range testData {
		url := GetBinaryDownloadURL(tt.version, tt.platform)
		if url != tt.expectedURL {
			t.Fatalf("Expected '%s' but got '%s'", tt.expectedURL, url)
		}
	}

}

func TestReplaceChars(t *testing.T) {
	testData := []struct {
		src         []string
		replacer    *strings.Replacer
		expectedRes []string
	}{
		{[]string{"abc%def", "%Y%"}, strings.NewReplacer("%", "X"), []string{"abcXdef", "XYX"}},
	}

	for _, tt := range testData {
		res := ReplaceChars(tt.src, tt.replacer)
		for i, val := range res {
			if val != tt.expectedRes[i] {
				t.Fatalf("Expected '%s' but got '%s'", tt.expectedRes, res)
			}
		}
	}
}

func TestConcatStrings(t *testing.T) {
	testData := []struct {
		src         []string
		prefix      string
		postfix     string
		expectedRes []string
	}{
		{[]string{"abc", ""}, "xx", "yy", []string{"xxabcyy", "xxyy"}},
		{[]string{"abc", ""}, "", "", []string{"abc", ""}},
	}

	for _, tt := range testData {
		res := ConcatStrings(tt.src, tt.prefix, tt.postfix)
		for i, val := range res {
			if val != tt.expectedRes[i] {
				t.Fatalf("Expected '%s' but got '%s'", tt.expectedRes, res)
			}
		}
	}
}
<commit_msg>pkg/util: add CalculateSizeInMB test<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 util

import (
	"strings"
	"testing"
)

func TestGetBinaryDownloadURL(t *testing.T) {
	testData := []struct {
		version     string
		platform    string
		expectedURL string
	}{
		{"v0.0.1", "linux", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-linux-amd64"},
		{"v0.0.1", "darwin", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-darwin-amd64"},
		{"v0.0.1", "windows", "https://storage.googleapis.com/minikube/releases/v0.0.1/minikube-windows-amd64.exe"},
	}

	for _, tt := range testData {
		url := GetBinaryDownloadURL(tt.version, tt.platform)
		if url != tt.expectedURL {
			t.Fatalf("Expected '%s' but got '%s'", tt.expectedURL, url)
		}
	}

}

func TestReplaceChars(t *testing.T) {
	testData := []struct {
		src         []string
		replacer    *strings.Replacer
		expectedRes []string
	}{
		{[]string{"abc%def", "%Y%"}, strings.NewReplacer("%", "X"), []string{"abcXdef", "XYX"}},
	}

	for _, tt := range testData {
		res := ReplaceChars(tt.src, tt.replacer)
		for i, val := range res {
			if val != tt.expectedRes[i] {
				t.Fatalf("Expected '%s' but got '%s'", tt.expectedRes, res)
			}
		}
	}
}

func TestConcatStrings(t *testing.T) {
	testData := []struct {
		src         []string
		prefix      string
		postfix     string
		expectedRes []string
	}{
		{[]string{"abc", ""}, "xx", "yy", []string{"xxabcyy", "xxyy"}},
		{[]string{"abc", ""}, "", "", []string{"abc", ""}},
	}

	for _, tt := range testData {
		res := ConcatStrings(tt.src, tt.prefix, tt.postfix)
		for i, val := range res {
			if val != tt.expectedRes[i] {
				t.Fatalf("Expected '%s' but got '%s'", tt.expectedRes, res)
			}
		}
	}
}

func TestCalculateSizeInMB(t *testing.T) {
	testData := []struct {
		size           string
		expectedNumber int
	}{
		{"1024kb", 1},
		{"1024KB", 1},
		{"1024mb", 1024},
		{"1024b", 0},
	}

	for _, tt := range testData {
		number := CalculateSizeInMB(tt.size)
		if number != tt.expectedNumber {
			t.Fatalf("Expected '%d'' but got '%d'", tt.expectedNumber, number)
		}
	}
}
<|endoftext|>


<commit_before>package initiate

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/markelog/eclectica/console"
	"github.com/markelog/eclectica/io"
	"github.com/markelog/eclectica/rc"
	"github.com/markelog/eclectica/variables"
)

type Init struct {
	language           string
	plugins            []string
	shouldRestartShell bool
}

func New(language string, plugins []string) *Init {
	init := &Init{
		language:           language,
		plugins:            plugins,
		shouldRestartShell: false,
	}

	return init
}

func (init *Init) CheckShell() {
	init.shouldRestartShell = init.needRestartShell()
}

func (init *Init) Initiate() (err error) {
	_, err = io.CreateDir(variables.DefaultInstall)
	if err != nil {
		return
	}

	command := init.composeCommand()

	err = rc.New(command).Add()
	if err != nil {
		return
	}

	return
}

func (init *Init) RestartShell() {
	if init.shouldRestartShell {
		console.Shell()
	}
}

func (init *Init) needRestartShell() bool {
	if strings.Contains(os.Getenv("PATH"), variables.DefaultInstall) == false {
		return true
	}

	ecPath := fmt.Sprintf("/%s/%s", ".eclectica/bin", init.language)
	output, _ := exec.Command("type", init.language).Output()
	out := string(output)

	if strings.Contains(out, ecPath) == false {
		return true
	}

	return false
}

func (init *Init) composeCommand() string {
	result := "# Eclectic stuff\n"

	// For golang
	result += "export PATH=" + variables.Path("go") + "/bin:$PATH\n"

	// For shared modules
	shared := filepath.Join(variables.Home(), "shared")
	result += "export PATH=" + shared + "/bin:$PATH\n"

	for _, language := range init.plugins {
		result += "export PATH=" +
			filepath.Join(variables.Home(), language, "current/bin") + ":$PATH\n"
	}

	result += "export PATH=" + variables.DefaultInstall + ":$PATH\n"

	return result
}
<commit_msg>Improve rc insertion<commit_after>package initiate

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/markelog/eclectica/console"
	"github.com/markelog/eclectica/io"
	"github.com/markelog/eclectica/rc"
	"github.com/markelog/eclectica/variables"
)

type Init struct {
	language           string
	plugins            []string
	shouldRestartShell bool
}

func New(language string, plugins []string) *Init {
	init := &Init{
		language:           language,
		plugins:            plugins,
		shouldRestartShell: false,
	}

	return init
}

func (init *Init) CheckShell() {
	init.shouldRestartShell = init.needRestartShell()
}

func (init *Init) Initiate() (err error) {
	_, err = io.CreateDir(variables.DefaultInstall)
	if err != nil {
		return
	}

	command := init.composeCommand()

	err = rc.New(command).Add()
	if err != nil {
		return
	}

	return
}

func (init *Init) RestartShell() {
	if init.shouldRestartShell {
		console.Shell()
	}
}

func (init *Init) needRestartShell() bool {
	if strings.Contains(os.Getenv("PATH"), variables.DefaultInstall) == false {
		return true
	}

	ecPath := fmt.Sprintf("/%s/%s", ".eclectica/bin", init.language)
	output, _ := exec.Command("type", init.language).Output()
	out := string(output)

	if strings.Contains(out, ecPath) == false {
		return true
	}

	return false
}

func (init *Init) composeCommand() string {
	result := "# Eclectic stuff\n"

	for _, language := range init.plugins {
		result += "export PATH=" +
			filepath.Join(variables.Home(), language, "current/bin") + ":$PATH\n"
	}

	// Eclectica binaries
	result += "export PATH=" + variables.DefaultInstall + ":$PATH\n"

	// For shared modules
	shared := filepath.Join(variables.Base(), "shared")
	result += "export PATH=" + shared + "/bin:$PATH\n"

	return result
}
<|endoftext|>


<commit_before>package version

import (
	"fmt"
	"runtime"
)

// Version represents the current version of Delve.
type Version struct {
	Major    string
	Minor    string
	Patch    string
	Metadata string
	Build    string
}

var (
	// DelveVersion is the current version of Delve.
	DelveVersion = Version{
		Major: "1", Minor: "7", Patch: "2", Metadata: "",
		//TODO(aarzilli): before updating this to 1.8.0 re-enable staticcheck test
		Build: "$Id$",
	}
)

func (v Version) String() string {
	ver := fmt.Sprintf("Version: %s.%s.%s", v.Major, v.Minor, v.Patch)
	if v.Metadata != "" {
		ver += "-" + v.Metadata
	}
	return fmt.Sprintf("%s\nBuild: %s", ver, v.Build)
}

var buildInfo = func() string {
	return ""
}

func BuildInfo() string {
	return fmt.Sprintf("%s\n%s", runtime.Version(), buildInfo())
}
<commit_msg>v1.7.3 (#2796)<commit_after>package version

import (
	"fmt"
	"runtime"
)

// Version represents the current version of Delve.
type Version struct {
	Major    string
	Minor    string
	Patch    string
	Metadata string
	Build    string
}

var (
	// DelveVersion is the current version of Delve.
	DelveVersion = Version{
		Major: "1", Minor: "7", Patch: "3", Metadata: "",
		//TODO(aarzilli): before updating this to 1.8.0 re-enable staticcheck test
		Build: "$Id$",
	}
)

func (v Version) String() string {
	ver := fmt.Sprintf("Version: %s.%s.%s", v.Major, v.Minor, v.Patch)
	if v.Metadata != "" {
		ver += "-" + v.Metadata
	}
	return fmt.Sprintf("%s\nBuild: %s", ver, v.Build)
}

var buildInfo = func() string {
	return ""
}

func BuildInfo() string {
	return fmt.Sprintf("%s\n%s", runtime.Version(), buildInfo())
}
<|endoftext|>


<commit_before>package logtap

import (
	"fmt"
	"github.com/boltdb/bolt"
	"github.com/nanobox-core/hatchet"
	"net/http"
	"strconv"
)

// HistoricalDrain matches the drain interface
type HistoricalDrain struct {
	port int
	max  int
	log  hatchet.Logger
	db   *bolt.DB
}

// NewHistoricalDrain returns a new instance of a HistoricalDrain
func NewHistoricalDrain(port int, file string, max int) *HistoricalDrain {
	db, err := bolt.Open(file, 0644, nil)
	if err != nil {
		db, err = bolt.Open("./bolt.db", 0644, nil)
	}
	return &HistoricalDrain{
		port: port,
		max:  max,
		db:   db,
	}
}

// Start starts the http listener.
// The listener on every request returns a json hash of logs of some arbitrary size
// default size is 100
func (h *HistoricalDrain) Start() {
	go func() {
		http.HandleFunc("/", h.handler)
		http.ListenAndServe(":"+strconv.Itoa(h.port), nil)
	}()
}

// handler handles any web request with any path and returns logs
// this makes it so a client that talks to pagodabox's logvac
// can communicate with this system
func (h *HistoricalDrain) handler(w http.ResponseWriter, r *http.Request) {
	var limit int64
	if i, err := strconv.ParseInt(r.FormValue("limit"), 10, 64); err == nil {
		limit = i
	} else {
		limit = 10000
	}
	h.log.Info("limit: %d", limit)
	h.db.View(func(tx *bolt.Tx) error {
		// Create a new bucket.
		b := tx.Bucket([]byte("log"))
		c := b.Cursor()

		// move the curser along so we can start dropping logs
		// in the right order at the right place
		if int64(b.Stats().KeyN) > limit {
			c.First()
			move_forward := int64(b.Stats().KeyN) - limit
			for i := int64(1); i < move_forward; i++ {
				c.Next()
			}
		} else {
			c.First()
		}

		for k, v := c.Next(); k != nil; k, v = c.Next() {
			fmt.Fprintf(w, "%s - %s", k, v)
		}

		return nil
	})

}

// SetLogger really allows the logtap main struct
// to assign its own logger to the historical drain
func (h *HistoricalDrain) SetLogger(l hatchet.Logger) {
	h.log = l
}

// write drops data into a capped collection of logs
// if we hit the limit the last log item will be removed from the beginning
func (h *HistoricalDrain) Write(msg Message) {
	h.log.Info("[Historical][write] message: (%s)%s", msg.Time.String(), msg.Content)
	h.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("log"))
		if err != nil {
			h.log.Error("[Historical][write] ERROR:" + err.Error())
			return err
		}
		err = bucket.Put([]byte(msg.Time.String()), []byte(msg.Content))
		if err != nil {
			h.log.Error("[Historical][write] ERROR:" + err.Error())
			return err
		}

		if bucket.Stats().KeyN > h.max {
			delete_count := bucket.Stats().KeyN - h.max
			c := bucket.Cursor()
			for i := 0; i < delete_count; i++ {
				c.First()
				c.Delete()
			}
		}

		return nil
	})

}
<commit_msg>log error messaging on historical drain<commit_after>package logtap

import (
	"fmt"
	"github.com/boltdb/bolt"
	"github.com/nanobox-core/hatchet"
	"net/http"
	"strconv"
)

// HistoricalDrain matches the drain interface
type HistoricalDrain struct {
	port int
	max  int
	log  hatchet.Logger
	db   *bolt.DB
}

// NewHistoricalDrain returns a new instance of a HistoricalDrain
func NewHistoricalDrain(port int, file string, max int) *HistoricalDrain {
	db, err := bolt.Open(file, 0644, nil)
	if err != nil {
		db, err = bolt.Open("./bolt.db", 0644, nil)
	}
	return &HistoricalDrain{
		port: port,
		max:  max,
		db:   db,
	}
}

// Start starts the http listener.
// The listener on every request returns a json hash of logs of some arbitrary size
// default size is 100
func (h *HistoricalDrain) Start() {
	go func() {
		http.HandleFunc("/", h.handler)
		err := http.ListenAndServe(":"+strconv.Itoa(h.port), nil)
		if err != nil {
			h.log.Error(err.Error())
		}
	}()
}

// handler handles any web request with any path and returns logs
// this makes it so a client that talks to pagodabox's logvac
// can communicate with this system
func (h *HistoricalDrain) handler(w http.ResponseWriter, r *http.Request) {
	var limit int64
	if i, err := strconv.ParseInt(r.FormValue("limit"), 10, 64); err == nil {
		limit = i
	} else {
		limit = 10000
	}
	h.log.Info("limit: %d", limit)
	h.db.View(func(tx *bolt.Tx) error {
		// Create a new bucket.
		b := tx.Bucket([]byte("log"))
		c := b.Cursor()

		// move the curser along so we can start dropping logs
		// in the right order at the right place
		if int64(b.Stats().KeyN) > limit {
			c.First()
			move_forward := int64(b.Stats().KeyN) - limit
			for i := int64(1); i < move_forward; i++ {
				c.Next()
			}
		} else {
			c.First()
		}

		for k, v := c.Next(); k != nil; k, v = c.Next() {
			fmt.Fprintf(w, "%s - %s", k, v)
		}

		return nil
	})

}

// SetLogger really allows the logtap main struct
// to assign its own logger to the historical drain
func (h *HistoricalDrain) SetLogger(l hatchet.Logger) {
	h.log = l
}

// write drops data into a capped collection of logs
// if we hit the limit the last log item will be removed from the beginning
func (h *HistoricalDrain) Write(msg Message) {
	h.log.Info("[Historical][write] message: (%s)%s", msg.Time.String(), msg.Content)
	h.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte("log"))
		if err != nil {
			h.log.Error("[Historical][write] ERROR:" + err.Error())
			return err
		}
		err = bucket.Put([]byte(msg.Time.String()), []byte(msg.Content))
		if err != nil {
			h.log.Error("[Historical][write] ERROR:" + err.Error())
			return err
		}

		if bucket.Stats().KeyN > h.max {
			delete_count := bucket.Stats().KeyN - h.max
			c := bucket.Cursor()
			for i := 0; i < delete_count; i++ {
				c.First()
				c.Delete()
			}
		}

		return nil
	})

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"runtime"
	"sort"
	"strings"
)

const LastAltered = "Dec 24, 2020"

/*
REVISION HISTORY
================
 2 Dec 19 -- Started development of make directory bookmark, using dirbkmk.go as a start.  Will focus on Windows, as bash already has dirb.
			   This routine will write the map of bookmarks as a gob file, that makedirbkmk.go will read.  This is copying code from
			   primes2.go and makeprimeslice.go.
			   It will need these commands:
				 s -- save bookmark
				 a -- about
				 d -- delete bookmark
				 p -- print bookmarks
				 h -- help
 4 Dec 19 -- Will make both directoryAliasesMap and bookmark global.  I'm not doing concurrency anyway.
			 os.Args = makedirbkmk.exe cmd bkmk-name target-directory
		   subscript =      0           1      2           3
			  len    =      1           2      3           4
28 Dec 19 -- changed cd to cdd, so it will change drive and directory at same time.
14 Jan 20 -- my dirsave() also needed to change cd to cdd.  Just done.
16 Sep 20 -- Added sl cmd, as in dirb, as synynom for p.  And will prompt if trying to overwrite a bookmark.
24 Dec 20 -- Will now sort output from Print command.
*/

const bookmarkfilename = "bookmarkfile.gob"

var HomeDir string             // global because it is also needed in my dirsave rtn.
var bookmark map[string]string // used in all routines.

type bkmkslicetype struct {
	key, value string
}

func main() {
	fmt.Println(" makedirbkmk written in Go, last altered", LastAltered)
	sep := string(os.PathSeparator)
	if runtime.GOOS == "linux" {
		HomeDir = os.Getenv("HOME")
	} else if runtime.GOOS == "windows" {
		HomeDir = os.Getenv("userprofile")
	} else {
		fmt.Println(" not running on expected platform.  Will exit.  In fact, probably won't even compile.")
		os.Exit(1)
	}
	target := "cdd" + " " + HomeDir + sep
	fullbookmarkfilename := HomeDir + sep + bookmarkfilename

	// read or init directory bookmark file
	_, err := os.Stat(fullbookmarkfilename)
	if err == nil { // need to read in bookmarkfile
		bookmarkfile, err := os.Open(fullbookmarkfilename)
		if err != nil {
			log.Fatalln(" cannot open", fullbookmarkfilename, " as input bookmark file, because of", err)
		}
		defer bookmarkfile.Close()
		decoder := gob.NewDecoder(bookmarkfile)
		err = decoder.Decode(&bookmark)
		if err != nil {
			log.Fatalln(" cannot decode", fullbookmarkfilename, ", error is", err, ".  Aborting")
		}
		bookmarkfile.Close()
		fmt.Println(" Bookmarks read in from", fullbookmarkfilename)
		fmt.Println()

	} else { // need to init bookmarkfile
		bookmark = make(map[string]string, 10)

		bookmark["docs"] = target + "Documents"
		bookmark["doc"] = target + "Documents"
		bookmark["inet"] = target + "Downloads"
		bookmark["inetdnld"] = target + "Downloads"
		bookmark["vid"] = target + "Videos"
		bookmark["go"] = target + "go"
		bookmark["src"] = target + "go" + sep + "src"
		bookmark["pic"] = target + "Pictures"
		bookmark["pics"] = target + "Pictures"
		bookmark["winx"] = target + "Videos" + sep + "winxvideos"
		bookmark["bin"] = target + "go" + sep + "bin"

		fmt.Println("Bookmark's initialized.")
		fmt.Println()
	}

	execname, _ := os.Executable()
	ExecFI, _ := os.Stat(execname)
	ExecTimeStamp := ExecFI.ModTime().Format("Mon Jan-2-2006_15:04:05 MST")
	help := func() {
		fmt.Println(" dirbkmk, a Directory Bookmark program written in Go.  Last altered", LastAltered)
		fmt.Println(" HomeDir is", HomeDir, ", ", ExecFI.Name(), "timestamp is", ExecTimeStamp, ". ")
		fmt.Println(" Full exec is", execname, ".  Fullbookmark is", fullbookmarkfilename)
		fmt.Println()
		fmt.Println(" s -- save current directory or entered directory name")
		fmt.Println(" a -- about.")
		fmt.Println(" d -- delete entry given on same line.")
		fmt.Println(" p, sl -- print out bookmark list.")
		fmt.Println(" h -- help.")
		fmt.Println()
		fmt.Println()
	}

	ch := ""
	if len(os.Args) < 2 {
		fmt.Println(" usage: makedirbkmk [s|a|d|p|h].  Note that there is no dash preceeding the commands, as these are not options.")
		os.Exit(0)
	} else {
		ch = strings.ToLower(os.Args[1])
	}
	fmt.Println(" os.Args is", os.Args)
	fmt.Println()

	switch ch {
	case "s": // save current directory or entered directory name
		dirsave()
		fmt.Println()
		fmt.Println()
		for idx, valu := range bookmark {
			fmt.Printf(" bookmark[%s] = %s \n", idx, valu)
		}
		fmt.Println()
		fmt.Println()

	case "a": // about.
		help()

	case "d": // delete entry given on same line
		direntrydel()

	case "p", "sl": // print out bookmark list
		bkmkslice := make([]bkmkslicetype, 0, len(bookmark))
		for idx, valu := range bookmark {
			bkmk := bkmkslicetype{idx, valu}  // structured literal syntax
			bkmkslice = append(bkmkslice, bkmk)
		}
		sortless := func (i,j int) bool {
			return bkmkslice[i].key < bkmkslice[j].key
		}
		sort.Slice(bkmkslice, sortless)
		for _, bkmk := range bkmkslice {
			fmt.Printf(" bookmark[%s] = %s\n", bkmk.key, bkmk.value)
		}
		fmt.Println()
		fmt.Println()

	case "h": // help
		help()

	default:
		fmt.Println(" command not recognized.", ch, "was entered.")
		os.Exit(0)
	}

	// write out bookmarkfile
	bookmarkfile, er := os.Create(fullbookmarkfilename)
	if er != nil {
		log.Fatalln(" could not create bookmarkfile upon exiting, because", er)
	}
	defer bookmarkfile.Close()

	encoder := gob.NewEncoder(bookmarkfile)
	e := encoder.Encode(&bookmark)
	if e != nil {
		log.Println(" could not encode bookmarkfile upon exiting, because", e)
	}
	bookmarkfile.Close()
} // end main

// ------------------------------------------ dirsave -----------------------------------------------
func dirsave() { // implement s (save) command
	//    os.Args = makedirbkmk.exe cmd bkmk-name target-directory
	//    len    =         1         2      3           4
	//    subscript =      0         1      2           3

	workingdir, err := os.Getwd()
	if err != nil {
		log.Fatalln(" could not get working directory because", err)
	}

	if len(os.Args) <= 2 { // only have cmd without a bookmark name to save
		log.Fatalln(" need bookmark name on command line to use s command.")
	} else if len(os.Args) == 3 { // no directory target name on command line.  Use current directory
		bkmkname := os.Args[2]
		_, ok := bookmark[bkmkname]
		if ok { // ok will be true if name already exists in the map.
			fmt.Print(" bookmark named ", bkmkname, " already exists.  Overwrite? [y/N] ")
			ans := ""
			fmt.Scanln(&ans)
			ans = strings.ToLower(ans)
			if strings.HasPrefix(ans, "y") {
				bookmark[bkmkname] = "cdd " + workingdir
				fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
			} else {
				fmt.Println(" save bookmark command ignored.")
			}
		} else { // bookmark name is not already in the map.
			bookmark[bkmkname] = "cdd " + workingdir
			fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
		}
	} else if len(os.Args) == 4 { // have potential directory target on command line
		target := os.Args[3]
		if strings.ContainsRune(target, ':') {
			directoryAliasesMap := GetDirectoryAliases()
			target = ProcessDirectoryAliases(directoryAliasesMap, target)
		} else if strings.Contains(target, "~") {
			target = strings.Replace(target, "~", HomeDir, 1)
		}
		// verify that target is a valid directory or symlink name.
		_, err := os.Lstat(target)
		if err == nil {
			bkmkname := os.Args[2]
			_, ok := bookmark[bkmkname]
			if ok { // ok will be true if name already exists in the map
				fmt.Print(" bookmark named ", bkmkname, " already exists.  Overwrite? [y/N] ")
				ans := ""
				fmt.Scanln(&ans)
				ans = strings.ToLower(ans)
				if strings.HasPrefix(ans, "y") {
					bookmark[bkmkname] = "cdd " + target
					fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
				} else {
					fmt.Println(" save bookmark command ignored.")
				}
			} else { // bookmark name is not already in the map
				bookmark[bkmkname] = "cdd " + target
				fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
			}
		} else {
			log.Fatalln(" Lstat call for", target, "failed with error of", err)
		}
	}

} // end dirsave

// ------------------------------------------ direntrydel ----------------------------------
func direntrydel() { // implement d (delete) command.
	//    os.Args = makedirbkmk.exe cmd bkmk-name target-directory
	//    len    =         1         2      3           4
	//    subscript =      0         1      2           3

	if len(os.Args) <= 2 {
		log.Fatalln(" need bookmark name on command line to be deleted.  os.Args is ", os.Args, ", len(os.Args) is", len(os.Args))
	}

	targt, ok := bookmark[os.Args[2]]
	if ok {
		delete(bookmark, os.Args[2])
	} else {
		log.Println(os.Args[2], " not in bookmark map.")
	}
	fmt.Printf(" deleted bookmark[%s] which referenced %s", os.Args[2], targt)
} // end direntrydel

// --------------------------- MakeSubst -------------------------------------------
func MakeSubst(instr string, r1, r2 rune) string {

	inRune := make([]rune, len(instr))
	if !strings.ContainsRune(instr, r1) {
		return instr
	}

	for i, s := range instr {
		if s == r1 {
			s = r2
		}
		inRune[i] = s // was byte(s) before I made this a slice of runes.
	}
	return string(inRune)
} // makesubst

//------------------------------ GetDirectoryAliases ----------------------------------------
func GetDirectoryAliases() map[string]string { // Env variable is diraliases.

	s := os.Getenv("diraliases")
	if len(s) == 0 {
		return nil
	}

	s = MakeSubst(s, '_', ' ') // substitute the underscore, _, for a space
	directoryAliasesMap := make(map[string]string, 10)
	//anAliasMap := make(dirAliasMapType,1)

	dirAliasSlice := strings.Fields(s)

	for _, aliasPair := range dirAliasSlice {
		if string(aliasPair[len(aliasPair)-1]) != "\\" {
			aliasPair = aliasPair + "\\"
		}
		aliasPair = MakeSubst(aliasPair, '-', ' ') // substitute a dash,-, for a space
		splitAlias := strings.Fields(aliasPair)
		directoryAliasesMap[splitAlias[0]] = splitAlias[1]
	}
	return directoryAliasesMap
} // end GetDirectoryAliases

// ------------------------------ ProcessDirectoryAliases ---------------------------
func ProcessDirectoryAliases(aliasesMap map[string]string, cmdline string) string {

	idx := strings.IndexRune(cmdline, ':')
	if idx < 2 { // note that if rune is not found, function returns -1.
		return cmdline
	}
	aliasesMap = GetDirectoryAliases()
	aliasName := cmdline[:idx] // substring of directory alias not including the colon, :
	aliasValue, ok := aliasesMap[aliasName]
	if !ok {
		return cmdline
	}
	PathnFile := cmdline[idx+1:]
	completeValue := aliasValue + PathnFile
	fmt.Println("in ProcessDirectoryAliases and complete value is", completeValue)
	return completeValue
} // ProcessDirectoryAliases
<commit_msg>05/10/2021 05:13:12 PM makedirbkmk/makedirbkmk.go<commit_after>package main

import (
	"encoding/gob"
	"fmt"
	"log"
	"os"
	"runtime"
	"sort"
	"strings"
)

const LastAltered = "May 10, 2021"

/*
REVISION HISTORY
================
 2 Dec 19 -- Started development of make directory bookmark, using dirbkmk.go as a start.  Will focus on Windows, as bash already has dirb.
			   This routine will write the map of bookmarks as a gob file, that makedirbkmk.go will read.  This is copying code from
			   primes2.go and makeprimeslice.go.
			   It will need these commands:
				 s -- save bookmark
				 a -- about
				 d -- delete bookmark
				 p -- print bookmarks
				 h -- help
 4 Dec 19 -- Will make both directoryAliasesMap and bookmark global.  I'm not doing concurrency anyway.
			 os.Args = makedirbkmk.exe cmd bkmk-name target-directory
		   subscript =      0           1      2           3
			  len    =      1           2      3           4
28 Dec 19 -- changed cd to cdd, so it will change drive and directory at same time.
14 Jan 20 -- my dirsave() also needed to change cd to cdd.  Just done.
16 Sep 20 -- Added sl cmd, as in dirb, as synynom for p.  And will prompt if trying to overwrite a bookmark.
24 Dec 20 -- Will now sort output from Print command.
10 May 21 -- Wrote dirPrint() so will sort output when other commands display the map, esp the save cmd.
*/

const bookmarkfilename = "bookmarkfile.gob"

var HomeDir string             // global because it is also needed in my dirsave rtn.
var bookmark map[string]string // used in all routines.

type bkmkslicetype struct {
	key, value string
}

func main() {
	fmt.Println(" makedirbkmk written in Go, last altered", LastAltered, "and compiled w/", runtime.Version())
	sep := string(os.PathSeparator)
	HomeDir, err := os.UserHomeDir() // this routine became available in Go 1.12
	if err != nil {
		fmt.Println(err,"Exiting")
		os.Exit(1)
	}
	target := "cdd" + " " + HomeDir + sep
	fullbookmarkfilename := HomeDir + sep + bookmarkfilename

	// read or init directory bookmark file
	_, err = os.Stat(fullbookmarkfilename)
	if err == nil { // need to read in bookmarkfile
		bookmarkfile, err := os.Open(fullbookmarkfilename)
		if err != nil {
			log.Fatalln(" cannot open", fullbookmarkfilename, " as input bookmark file, because of", err)
		}
		defer bookmarkfile.Close()
		decoder := gob.NewDecoder(bookmarkfile)
		err = decoder.Decode(&bookmark)
		if err != nil {
			log.Fatalln(" cannot decode", fullbookmarkfilename, ", error is", err, ".  Aborting")
		}
		bookmarkfile.Close()
		fmt.Println(" Bookmarks read in from", fullbookmarkfilename)
		fmt.Println()

	} else { // need to init bookmarkfile
		bookmark = make(map[string]string, 10)

		bookmark["docs"] = target + "Documents"
		bookmark["doc"] = target + "Documents"
		bookmark["inet"] = target + "Downloads"
		bookmark["inetdnld"] = target + "Downloads"
		bookmark["vid"] = target + "Videos"
		bookmark["go"] = target + "go"
		bookmark["src"] = target + "go" + sep + "src"
		bookmark["pic"] = target + "Pictures"
		bookmark["pics"] = target + "Pictures"
		bookmark["winx"] = target + "Videos" + sep + "winxvideos"
		bookmark["bin"] = target + "go" + sep + "bin"

		fmt.Println("Bookmark's initialized.")
		fmt.Println()
	}

	execname, _ := os.Executable()
	ExecFI, _ := os.Stat(execname)
	ExecTimeStamp := ExecFI.ModTime().Format("Mon Jan-2-2006_15:04:05 MST")
	help := func() {
		fmt.Println(" makedirbkmk, a Directory Bookmark program written in Go.  Last altered", LastAltered)
		fmt.Println(" HomeDir is", HomeDir, ", ", ExecFI.Name(), "timestamp is", ExecTimeStamp, ". ")
		fmt.Println(" Full exec is", execname, ".  Fullbookmark is", fullbookmarkfilename)
		fmt.Println()
		fmt.Println(" s -- save current directory or entered directory name")
		fmt.Println(" a -- about.")
		fmt.Println(" d -- delete entry given on same line.")
		fmt.Println(" p, sl -- print out bookmark list.")
		fmt.Println(" h -- help.")
		fmt.Println()
		fmt.Println()
	}

	ch := ""
	if len(os.Args) < 2 {
		fmt.Println(" usage: makedirbkmk [s|a|d|p|h].  Note that there is no dash preceeding the commands, as these are not options.")
		os.Exit(0)
	} else {
		ch = strings.ToLower(os.Args[1])
	}
	fmt.Println(" os.Args is", os.Args)
	fmt.Println()

	switch ch {
	case "s": // save current directory or entered directory name
		dirsave()
		fmt.Println()
		fmt.Println()
//		for idx, valu := range bookmark {
//			fmt.Printf(" bookmark[%s] = %s \n", idx, valu)
//		}
		dirPrint()
		fmt.Println()
		fmt.Println()

	case "a": // about.
		help()

	case "d": // delete entry given on same line
		direntrydel()

	case "p", "sl": // print out bookmark list
		/*
		bkmkslice := make([]bkmkslicetype, 0, len(bookmark))
		for idx, valu := range bookmark {
			bkmk := bkmkslicetype{idx, valu}  // structured literal syntax
			bkmkslice = append(bkmkslice, bkmk)
		}
		sortless := func (i,j int) bool {
			return bkmkslice[i].key < bkmkslice[j].key
		}
		sort.Slice(bkmkslice, sortless)
		for _, bkmk := range bkmkslice {
			fmt.Printf(" bookmark[%s] = %s\n", bkmk.key, bkmk.value)
		}
		 */
		dirPrint()
		fmt.Println()
		fmt.Println()

	case "h": // help
		help()

	default:
		fmt.Println(" command not recognized.", ch, "was entered.")
		os.Exit(0)
	}

	// write out bookmarkfile
	bookmarkfile, er := os.Create(fullbookmarkfilename)
	if er != nil {
		log.Fatalln(" could not create bookmarkfile upon exiting, because", er)
	}
	defer bookmarkfile.Close()

	encoder := gob.NewEncoder(bookmarkfile)
	e := encoder.Encode(&bookmark)
	if e != nil {
		log.Println(" could not encode bookmarkfile upon exiting, because", e)
	}
	bookmarkfile.Close()
} // end main

// ------------------------------------------ dirsave -----------------------------------------------
func dirsave() { // implement s (save) command
	//    os.Args = makedirbkmk.exe cmd bkmk-name target-directory
	//    len    =         1         2      3           4
	//    subscript =      0         1      2           3

	workingdir, err := os.Getwd()
	if err != nil {
		log.Fatalln(" could not get working directory because", err)
	}

	if len(os.Args) <= 2 { // only have cmd without a bookmark name to save
		log.Fatalln(" need bookmark name on command line to use s command.")
	} else if len(os.Args) == 3 { // no directory target name on command line.  Use current directory
		bkmkname := os.Args[2]
		_, ok := bookmark[bkmkname]
		if ok { // ok will be true if name already exists in the map.
			fmt.Print(" bookmark named ", bkmkname, " already exists.  Overwrite? [y/N] ")
			ans := ""
			fmt.Scanln(&ans)
			ans = strings.ToLower(ans)
			if strings.HasPrefix(ans, "y") {
				bookmark[bkmkname] = "cdd " + workingdir
				fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
			} else {
				fmt.Println(" save bookmark command ignored.")
			}
		} else { // bookmark name is not already in the map.
			bookmark[bkmkname] = "cdd " + workingdir
			fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
		}
	} else if len(os.Args) == 4 { // have potential directory target on command line
		target := os.Args[3]
		if strings.ContainsRune(target, ':') {
			directoryAliasesMap := GetDirectoryAliases()
			target = ProcessDirectoryAliases(directoryAliasesMap, target)
		} else if strings.Contains(target, "~") {
			target = strings.Replace(target, "~", HomeDir, 1)
		}
		// verify that target is a valid directory or symlink name.
		_, err := os.Lstat(target)
		if err == nil {
			bkmkname := os.Args[2]
			_, ok := bookmark[bkmkname]
			if ok { // ok will be true if name already exists in the map
				fmt.Print(" bookmark named ", bkmkname, " already exists.  Overwrite? [y/N] ")
				ans := ""
				fmt.Scanln(&ans)
				ans = strings.ToLower(ans)
				if strings.HasPrefix(ans, "y") {
					bookmark[bkmkname] = "cdd " + target
					fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
				} else {
					fmt.Println(" save bookmark command ignored.")
				}
			} else { // bookmark name is not already in the map
				bookmark[bkmkname] = "cdd " + target
				fmt.Printf(" created bookmark[%s] = %s \n", os.Args[2], bookmark[os.Args[2]])
			}
		} else {
			log.Fatalln(" Lstat call for", target, "failed with error of", err)
		}
	}

} // end dirsave

// ------------------------------------------ direntrydel ----------------------------------
func direntrydel() { // implement d (delete) command.
	//    os.Args = makedirbkmk.exe cmd bkmk-name target-directory
	//    len    =         1         2      3           4
	//    subscript =      0         1      2           3

	if len(os.Args) <= 2 {
		log.Fatalln(" need bookmark name on command line to be deleted.  os.Args is ", os.Args, ", len(os.Args) is", len(os.Args))
	}

	targt, ok := bookmark[os.Args[2]]
	if ok {
		delete(bookmark, os.Args[2])
	} else {
		log.Println(os.Args[2], " not in bookmark map.")
	}
	fmt.Printf(" deleted bookmark[%s] which referenced %s", os.Args[2], targt)
} // end direntrydel

// --------------------------- MakeSubst -------------------------------------------
func MakeSubst(instr string, r1, r2 rune) string {

	inRune := make([]rune, len(instr))
	if !strings.ContainsRune(instr, r1) {
		return instr
	}

	for i, s := range instr {
		if s == r1 {
			s = r2
		}
		inRune[i] = s // was byte(s) before I made this a slice of runes.
	}
	return string(inRune)
} // makesubst

//------------------------------ GetDirectoryAliases ----------------------------------------
func GetDirectoryAliases() map[string]string { // Env variable is diraliases.

	s := os.Getenv("diraliases")
	if len(s) == 0 {
		return nil
	}

	s = MakeSubst(s, '_', ' ') // substitute the underscore, _, for a space
	directoryAliasesMap := make(map[string]string, 10)
	//anAliasMap := make(dirAliasMapType,1)

	dirAliasSlice := strings.Fields(s)

	for _, aliasPair := range dirAliasSlice {
		if string(aliasPair[len(aliasPair)-1]) != "\\" {
			aliasPair = aliasPair + "\\"
		}
		aliasPair = MakeSubst(aliasPair, '-', ' ') // substitute a dash,-, for a space
		splitAlias := strings.Fields(aliasPair)
		directoryAliasesMap[splitAlias[0]] = splitAlias[1]
	}
	return directoryAliasesMap
} // end GetDirectoryAliases

// ------------------------------ ProcessDirectoryAliases ---------------------------
func ProcessDirectoryAliases(aliasesMap map[string]string, cmdline string) string {

	idx := strings.IndexRune(cmdline, ':')
	if idx < 2 { // note that if rune is not found, function returns -1.
		return cmdline
	}
	aliasesMap = GetDirectoryAliases()
	aliasName := cmdline[:idx] // substring of directory alias not including the colon, :
	aliasValue, ok := aliasesMap[aliasName]
	if !ok {
		return cmdline
	}
	PathnFile := cmdline[idx+1:]
	completeValue := aliasValue + PathnFile
	fmt.Println("in ProcessDirectoryAliases and complete value is", completeValue)
	return completeValue
} // ProcessDirectoryAliases

func dirPrint() {
	bkmkslice := make([]bkmkslicetype, 0, len(bookmark))
	for idx, valu := range bookmark {
		bkmk := bkmkslicetype{idx, valu}  // structured literal syntax
		bkmkslice = append(bkmkslice, bkmk)
	}
	sortless := func (i,j int) bool {
		return bkmkslice[i].key < bkmkslice[j].key
	}
	sort.Slice(bkmkslice, sortless)
	for _, bkmk := range bkmkslice {
		fmt.Printf(" bookmark[%s] = %s\n", bkmk.key, bkmk.value)
	}
}<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestGetAlive(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(alive))
	defer ts.Close()

	res, err := http.Get(ts.URL)
	if err != nil {
		t.Fatalf(err.Error())
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		t.Fatalf(err.Error())
	}
	found := string(body[:])
	expected := "I'm alive"
	if found != expected {
		t.Errorf("Expected\n---\n%s\n---\nbut found\n---\n%s\n---\n", expected, found)
	}
}
<commit_msg>from unit tests to integration tests<commit_after>package main

import (
	"io/ioutil"
	"net/http"
	"net/url"
	"testing"
)

const HOST = "http://localhost:8765/"

func TestGetAlive(t *testing.T) {
	response, err := http.Get(HOST)
	if err != nil {
		t.Fatalf("Unexpected error found - %s", err.Error())
	}
	defer response.Body.Close()
	body, err := ioutil.ReadAll(response.Body)
	found := string(body[:])

	expected := "I'm alive"
	if found != expected {
		t.Errorf("Expected\n---\n%s\n---\nbut found\n---\n%s\n---\n", expected, found)
	}
}

func TestPlay(t *testing.T) {
	filename := "test_sounds/beep9.mp3"
	url := HOST + "play/" + url.QueryEscape(filename)

	client := &http.Client{}
	request, err := http.NewRequest("PUT", url, nil)
	res, err := client.Do(request)

	if err != nil {
		t.Fatalf("Unexpected error found - %s", err.Error())
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		t.Fatalf(err.Error())
	}
	found := string(body[:])
	expected := "Playing, test_sounds/beep9.mp3!"
	if found != expected {
		t.Errorf("Expected\n---\n%s\n---\nbut found\n---\n%s\n---\n", expected, found)
	}
}
<|endoftext|>


<commit_before>package filesystem

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"

	"github.com/colinmarc/hdfs"
)

// Requirement:
//   Hadoop/HDFS version: 2
//   WebHDFS (HDFS REST)

type hdfsConfig struct {
	namenodeAddr string
	webHdfsAddr  string
	user         string
}

type HdfsClient struct {
	client *hdfs.Client
	hdfsConfig
}

func NewHdfsClient(namenodeAddr, webHdfsAddr, user string) (Client, error) {
	client, err := hdfs.NewForUser(namenodeAddr, user)
	if err != nil {
		return nil, err
	}
	return &HdfsClient{
		client: client,
		hdfsConfig: hdfsConfig{
			namenodeAddr: namenodeAddr,
			webHdfsAddr:  webHdfsAddr,
			user:         user,
		},
	}, nil
}

func (c *HdfsClient) Remove(name string) error {
	return c.client.Remove(name)
}

func (c *HdfsClient) OpenReadCloser(name string) (io.ReadCloser, error) {
	return c.client.Open(name)
}

func (c *HdfsClient) OpenWriteCloser(name string) (io.WriteCloser, error) {
	exist, err := c.Exists(name)
	if err != nil {
		return nil, err
	}
	if !exist {
		err := c.client.CreateEmptyFile(name)
		if err != nil {
			return nil, err
		}
	}
	return &HdfsFile{
		path:       name,
		logger:     log.New(os.Stdout, "", log.Lshortfile|log.LstdFlags),
		hdfsConfig: c.hdfsConfig,
	}, nil
}

func (c *HdfsClient) Exists(name string) (bool, error) {
	var (
		err error
		ret bool
	)
	for retry := 0; retry < 3; retry++ {
		_, err = c.client.Stat(name)
		ret, err = existCommon(err)
		if err == nil {
			return ret, err
		}
		log.Printf("Recovering HDFS client with namenode %s for user %s. ret: %v", c.hdfsConfig.namenodeAddr, c.hdfsConfig.user, c.Recover())
	}
	return false, err
}

func (c *HdfsClient) Rename(oldpath, newpath string) error {
	return c.client.Rename(oldpath, newpath)
}

func (c *HdfsClient) Recover() error {
	var err error
	c.client, err = hdfs.NewForUser(c.hdfsConfig.namenodeAddr, c.hdfsConfig.user)
	return err
}

// only supports '*', '?'
// Syntax:
//    /user/hdfs/etl*/part.*
func (c *HdfsClient) Glob(pattern string) (matches []string, err error) {
	if pattern == "" {
		return nil, fmt.Errorf("Glob pattern shouldn't be empty")
	}
	if pattern[len(pattern)-1] == '/' {
		return nil, fmt.Errorf("Glob pattern shouldn't be a directory")
	}
	// names will have all the pathnames of the pattern.
	// e.g. "/a/b/c" => [a, b, c]
	var names []string
	for path.Dir(pattern) != "/" {
		names = append(names, path.Base(pattern))
		pattern = path.Dir(pattern)
	}
	names = append(names, pattern[1:len(pattern)])
	for i, j := 0, len(names)-1; i < j; i, j = i+1, j-1 {
		names[i], names[j] = names[j], names[i]
	}
	return c.glob("/", names)
}

func (c *HdfsClient) glob(dir string, names []string) (m []string, err error) {
	name := names[0]
	var dirs []string
	if hasMeta(name) {
		fileInfos, err := c.client.ReadDir(dir)
		if err != nil {
			return nil, err
		}
		for _, fi := range fileInfos {
			matched, err := path.Match(name, fi.Name())
			if err != nil {
				return nil, err
			}
			if matched {
				dirs = append(dirs, path.Join(dir, fi.Name()))
			}
		}
	} else {
		dirs = append(dirs, path.Join(dir, name))
	}
	for _, pathname := range dirs {
		if len(names) == 1 {
			exist, err := c.Exists(pathname)
			if err != nil {
				return nil, err
			}
			if exist {
				m = append(m, pathname)
			}
		} else {
			return c.glob(pathname, names[1:len(names)])

		}
	}
	return
}

func hasMeta(name string) bool {
	return strings.ContainsAny(name, "*?")
}

type HdfsFile struct {
	path   string
	logger *log.Logger
	hdfsConfig
	// buffer
}

// REST docs:
// http://hadoop.apache.org/docs/stable/hadoop-project-dist/hadoop-hdfs/WebHDFS.html#Append_to_a_File
func (f *HdfsFile) Write(b []byte) (int, error) {
	tr := &http.Transport{}
	urlStr := buildNamenodeURL(f.webHdfsAddr, f.path, f.user)

	// POST request to namenode. We shouldn't follow redirect and should get
	// the datanode URL in response.
	req, err := http.NewRequest("POST", urlStr, nil)
	if err != nil {
		return 0, fmt.Errorf("Write: NewRequest failed: %v", err)
	}
	// no redirect

	var resp *http.Response
	var loc string
	for retry := 0; retry < 3 && (err != nil || loc == ""); retry++ {
		resp, err := tr.RoundTrip(req)
		if err != nil {
			continue
		}
		defer resp.Body.Close()
		loc = resp.Header.Get("Location")
	}

	if loc == "" {
		return 0, fmt.Errorf("Write: Failed retrieving datanode location.")
	}

	u, err := url.ParseRequestURI(loc)
	if err != nil {
		return 0, fmt.Errorf("Write: ParseRequestURI failed: %v", err)
	}
	// POST request to datanode.
	resp, err = http.Post(u.String(), "application/octet-stream", bytes.NewBuffer(b))
	for retry := 0; retry < 3 && err != nil; retry++ {
		resp, err = http.Post(u.String(), "application/octet-stream", bytes.NewBuffer(b))
	}
	if err != nil {
		return 0, fmt.Errorf("Write: POST to datanode (%s) failed: %v", u.String(), err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		// Reason will show the java stack trace for the error.
		return 0, fmt.Errorf("Status code isn't OK. Response: %v\nReason: %v", resp, explain(resp.Body))
	}
	return len(b), nil
}

func (f *HdfsFile) Close() error {
	return nil
}

func buildNamenodeURL(webHdfsAddr, name, user string) string {
	u := &url.URL{
		Scheme: "http",
		Host:   webHdfsAddr,
		Path:   path.Join("webhdfs", "v1", name),
	}
	q := u.Query()
	q.Set("op", "APPEND")
	q.Set("user.name", user)
	u.RawQuery = q.Encode()
	return u.String()
}

func explain(r io.Reader) interface{} {
	body, _ := ioutil.ReadAll(r)
	var reason interface{}
	json.Unmarshal(body, &reason)
	return reason
}
<commit_msg>Added sleep interval in between retrials.<commit_after>package filesystem

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"time"

	"github.com/colinmarc/hdfs"
)

// Requirement:
//   Hadoop/HDFS version: 2
//   WebHDFS (HDFS REST)

type hdfsConfig struct {
	namenodeAddr string
	webHdfsAddr  string
	user         string
}

type HdfsClient struct {
	client *hdfs.Client
	hdfsConfig
}

func NewHdfsClient(namenodeAddr, webHdfsAddr, user string) (Client, error) {
	client, err := hdfs.NewForUser(namenodeAddr, user)
	if err != nil {
		return nil, err
	}
	return &HdfsClient{
		client: client,
		hdfsConfig: hdfsConfig{
			namenodeAddr: namenodeAddr,
			webHdfsAddr:  webHdfsAddr,
			user:         user,
		},
	}, nil
}

func (c *HdfsClient) Remove(name string) error {
	return c.client.Remove(name)
}

func (c *HdfsClient) OpenReadCloser(name string) (io.ReadCloser, error) {
	return c.client.Open(name)
}

func (c *HdfsClient) OpenWriteCloser(name string) (io.WriteCloser, error) {
	exist, err := c.Exists(name)
	if err != nil {
		return nil, err
	}
	if !exist {
		err := c.client.CreateEmptyFile(name)
		if err != nil {
			return nil, err
		}
	}
	return &HdfsFile{
		path:       name,
		logger:     log.New(os.Stdout, "", log.Lshortfile|log.LstdFlags),
		hdfsConfig: c.hdfsConfig,
	}, nil
}

func (c *HdfsClient) Exists(name string) (bool, error) {
	var (
		err error
		ret bool
	)
	for retry := 0; retry < 3; retry++ {
		_, err = c.client.Stat(name)
		ret, err = existCommon(err)
		if err == nil {
			return ret, err
		}
		time.Sleep(1 * time.Second)
		log.Printf("Recovering HDFS client with namenode %s for user %s. ret: %v", c.hdfsConfig.namenodeAddr, c.hdfsConfig.user, c.Recover())
	}
	return false, err
}

func (c *HdfsClient) Rename(oldpath, newpath string) error {
	return c.client.Rename(oldpath, newpath)
}

func (c *HdfsClient) Recover() error {
	var err error
	c.client, err = hdfs.NewForUser(c.hdfsConfig.namenodeAddr, c.hdfsConfig.user)
	return err
}

// only supports '*', '?'
// Syntax:
//    /user/hdfs/etl*/part.*
func (c *HdfsClient) Glob(pattern string) (matches []string, err error) {
	if pattern == "" {
		return nil, fmt.Errorf("Glob pattern shouldn't be empty")
	}
	if pattern[len(pattern)-1] == '/' {
		return nil, fmt.Errorf("Glob pattern shouldn't be a directory")
	}
	// names will have all the pathnames of the pattern.
	// e.g. "/a/b/c" => [a, b, c]
	var names []string
	for path.Dir(pattern) != "/" {
		names = append(names, path.Base(pattern))
		pattern = path.Dir(pattern)
	}
	names = append(names, pattern[1:len(pattern)])
	for i, j := 0, len(names)-1; i < j; i, j = i+1, j-1 {
		names[i], names[j] = names[j], names[i]
	}
	return c.glob("/", names)
}

func (c *HdfsClient) glob(dir string, names []string) (m []string, err error) {
	name := names[0]
	var dirs []string
	if hasMeta(name) {
		fileInfos, err := c.client.ReadDir(dir)
		if err != nil {
			return nil, err
		}
		for _, fi := range fileInfos {
			matched, err := path.Match(name, fi.Name())
			if err != nil {
				return nil, err
			}
			if matched {
				dirs = append(dirs, path.Join(dir, fi.Name()))
			}
		}
	} else {
		dirs = append(dirs, path.Join(dir, name))
	}
	for _, pathname := range dirs {
		if len(names) == 1 {
			exist, err := c.Exists(pathname)
			if err != nil {
				return nil, err
			}
			if exist {
				m = append(m, pathname)
			}
		} else {
			return c.glob(pathname, names[1:len(names)])

		}
	}
	return
}

func hasMeta(name string) bool {
	return strings.ContainsAny(name, "*?")
}

type HdfsFile struct {
	path   string
	logger *log.Logger
	hdfsConfig
	// buffer
}

// REST docs:
// http://hadoop.apache.org/docs/stable/hadoop-project-dist/hadoop-hdfs/WebHDFS.html#Append_to_a_File
func (f *HdfsFile) Write(b []byte) (int, error) {
	tr := &http.Transport{}
	urlStr := buildNamenodeURL(f.webHdfsAddr, f.path, f.user)

	// POST request to namenode. We shouldn't follow redirect and should get
	// the datanode URL in response.
	req, err := http.NewRequest("POST", urlStr, nil)
	if err != nil {
		return 0, fmt.Errorf("Write: NewRequest failed: %v", err)
	}
	// no redirect

	var resp *http.Response
	var loc string
	for retry := 0; retry < 3 && (err != nil || loc == ""); retry++ {
		time.Sleep(300 * time.Millisecond)
		resp, err := tr.RoundTrip(req)
		if err != nil {
			continue
		}
		defer resp.Body.Close()
		loc = resp.Header.Get("Location")
	}

	if loc == "" {
		return 0, fmt.Errorf("Write: Failed retrieving datanode location.")
	}

	u, err := url.ParseRequestURI(loc)
	if err != nil {
		return 0, fmt.Errorf("Write: ParseRequestURI failed: %v", err)
	}
	// POST request to datanode.
	resp, err = http.Post(u.String(), "application/octet-stream", bytes.NewBuffer(b))
	for retry := 0; retry < 3 && err != nil; retry++ {
		time.Sleep(1 * time.Second)
		resp, err = http.Post(u.String(), "application/octet-stream", bytes.NewBuffer(b))
	}
	if err != nil {
		return 0, fmt.Errorf("Write: POST to datanode (%s) failed: %v", u.String(), err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		// Reason will show the java stack trace for the error.
		return 0, fmt.Errorf("Status code isn't OK. Response: %v\nReason: %v", resp, explain(resp.Body))
	}
	return len(b), nil
}

func (f *HdfsFile) Close() error {
	return nil
}

func buildNamenodeURL(webHdfsAddr, name, user string) string {
	u := &url.URL{
		Scheme: "http",
		Host:   webHdfsAddr,
		Path:   path.Join("webhdfs", "v1", name),
	}
	q := u.Query()
	q.Set("op", "APPEND")
	q.Set("user.name", user)
	u.RawQuery = q.Encode()
	return u.String()
}

func explain(r io.Reader) interface{} {
	body, _ := ioutil.ReadAll(r)
	var reason interface{}
	json.Unmarshal(body, &reason)
	return reason
}
<|endoftext|>


<commit_before>package redis

import (
	"encoding/json"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/adjust/redismq"
	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/yaml.v2"
)

const (
	recvBuffer = 100
)

type Config struct {
	Host       string `yaml:"host"`
	Port       int    `yaml:"port"`
	Db         int64  `yaml:"db"`
	Password   string `yaml:"password"`
	QueueName  string `yaml:"queue_name"`
	JsonDecode bool   `yaml:"json_decode"`
}

type RedisInputServer struct {
	config   Config
	receiver input.Receiver
	term     chan bool
}

func init() {
	input.Register("redis", &RedisInputServer{
		term: make(chan bool, 1),
	})
}

func redisGet(redisServer *RedisInputServer, consumer *redismq.Consumer) error {
	consumer.ResetWorking()
	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	for {
		unacked := consumer.GetUnackedLength()

		if unacked > 0 {
			log.Printf("Requeued %d messages\n", unacked)
			consumer.RequeueWorking()
		}

		packages, err := consumer.MultiGet(recvBuffer)

		if err == nil {
			numPackages := len(packages)

			if numPackages > 0 {
				rateCounter.Incr(int64(numPackages))
				err = packages[numPackages-1].MultiAck()

				if err != nil {
					log.Println("Failed to ack", err)
				}
			}

			for i := range packages {
				var ev buffer.Event
				payload := string(packages[i].Payload)
				ev.Text = &payload

				if redisServer.config.JsonDecode {
					decoder := json.NewDecoder(strings.NewReader(payload))
					decoder.UseNumber()

					err = decoder.Decode(&ev.Fields)

					if err != nil {
						continue
					}
				}

				redisServer.receiver.Send(&ev)
			}
		} else {
			log.Printf("Error reading from Redis: %s, sleeping", err)
			time.Sleep(2 * time.Second)
		}
	}

	return nil
}

func (redisServer *RedisInputServer) Init(config yaml.MapSlice, receiver input.Receiver) error {
	var redisConfig *Config

	// go-yaml doesn't have a great way to partially unmarshal YAML data
	// See https://github.com/go-yaml/yaml/issues/13
	yamlConfig, _ := yaml.Marshal(config)

	if err := yaml.Unmarshal(yamlConfig, &redisConfig); err != nil {
		return fmt.Errorf("Error parsing Redis config: %v", err)
	}

	redisServer.config = *redisConfig
	redisServer.receiver = receiver

	return nil
}

func (redisServer *RedisInputServer) Start() error {
	log.Println("Starting Redis input")
	port := strconv.Itoa(redisServer.config.Port)

	// Create Redis queue
	queue := redismq.CreateQueue(redisServer.config.Host,
		port,
		redisServer.config.Password,
		redisServer.config.Db,
		redisServer.config.QueueName)

	consumer, err := queue.AddConsumer("redisInput")

	if err != nil {
		log.Println("Error opening Redis input")
		return err
	}

	go redisGet(redisServer, consumer)

	for {
		select {
		case <-redisServer.term:
			log.Println("Redis input server received term signal")
			return nil
		}
	}

	return nil
}

func (redisServer *RedisInputServer) Stop() error {
	redisServer.term <- true
	return nil
}
<commit_msg>Display queue name in Redis input<commit_after>package redis

import (
	"encoding/json"
	"fmt"
	"log"
	"strconv"
	"strings"
	"time"

	"github.com/adjust/redismq"
	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/yaml.v2"
)

const (
	recvBuffer = 100
)

type Config struct {
	Host       string `yaml:"host"`
	Port       int    `yaml:"port"`
	Db         int64  `yaml:"db"`
	Password   string `yaml:"password"`
	QueueName  string `yaml:"queue_name"`
	JsonDecode bool   `yaml:"json_decode"`
}

type RedisInputServer struct {
	config   Config
	receiver input.Receiver
	term     chan bool
}

func init() {
	input.Register("redis", &RedisInputServer{
		term: make(chan bool, 1),
	})
}

func redisGet(redisServer *RedisInputServer, consumer *redismq.Consumer) error {
	consumer.ResetWorking()
	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	for {
		unacked := consumer.GetUnackedLength()

		if unacked > 0 {
			log.Printf("Requeued %d messages\n", unacked)
			consumer.RequeueWorking()
		}

		packages, err := consumer.MultiGet(recvBuffer)

		if err == nil {
			numPackages := len(packages)

			if numPackages > 0 {
				rateCounter.Incr(int64(numPackages))
				err = packages[numPackages-1].MultiAck()

				if err != nil {
					log.Println("Failed to ack", err)
				}
			}

			for i := range packages {
				var ev buffer.Event
				payload := string(packages[i].Payload)
				ev.Text = &payload

				if redisServer.config.JsonDecode {
					decoder := json.NewDecoder(strings.NewReader(payload))
					decoder.UseNumber()

					err = decoder.Decode(&ev.Fields)

					if err != nil {
						continue
					}
				}

				redisServer.receiver.Send(&ev)
			}
		} else {
			log.Printf("Error reading from Redis: %s, sleeping", err)
			time.Sleep(2 * time.Second)
		}
	}

	return nil
}

func (redisServer *RedisInputServer) Init(config yaml.MapSlice, receiver input.Receiver) error {
	var redisConfig *Config

	// go-yaml doesn't have a great way to partially unmarshal YAML data
	// See https://github.com/go-yaml/yaml/issues/13
	yamlConfig, _ := yaml.Marshal(config)

	if err := yaml.Unmarshal(yamlConfig, &redisConfig); err != nil {
		return fmt.Errorf("Error parsing Redis config: %v", err)
	}

	redisServer.config = *redisConfig
	redisServer.receiver = receiver

	return nil
}

func (redisServer *RedisInputServer) Start() error {
	log.Println("Starting Redis input on queue", redisServer.config.QueueName)
	port := strconv.Itoa(redisServer.config.Port)

	// Create Redis queue
	queue := redismq.CreateQueue(redisServer.config.Host,
		port,
		redisServer.config.Password,
		redisServer.config.Db,
		redisServer.config.QueueName)

	consumer, err := queue.AddConsumer("redisInput")

	if err != nil {
		log.Println("Error opening Redis input")
		return err
	}

	go redisGet(redisServer, consumer)

	for {
		select {
		case <-redisServer.term:
			log.Println("Redis input server received term signal")
			return nil
		}
	}

	return nil
}

func (redisServer *RedisInputServer) Stop() error {
	redisServer.term <- true
	return nil
}
<|endoftext|>


<commit_before>package jwt

import (
	"fmt"
	"log"

	"github.com/dgrijalva/jwt-go"
	"github.com/neptulon/neptulon"
)

type token struct {
	Token string `json:"token"`
}

// HMAC is JSON Web Token authentication using HMAC.
// If successful, token context will be store with the key "userid" in session.
// If unsuccessful, connection will be closed right away.
func HMAC(password string) func(ctx *neptulon.ReqCtx) error {
	pass := []byte(password)

	return func(ctx *neptulon.ReqCtx) error {
		if _, ok := ctx.Conn.Session.GetOk("userid"); ok {
			return nil
		}

		var t token
		if err := ctx.Params(&t); err != nil {
			ctx.Conn.Close()
			return err
		}

		jt, err := jwt.Parse(t.Token, func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("middleware: jwt: unexpected signing method: %v", token.Header["alg"])
			}
			return pass, nil
		})

		if err != nil || !jt.Valid {
			ctx.Conn.Close()
			return fmt.Errorf("middleware: jwt: invalid JWT authentication attempt: %v: %v: %v", err, ctx.Conn.RemoteAddr(), t.Token)
		}

		userID := jt.Claims["userid"].(string)
		ctx.Conn.Session.Set("userid", userID)
		log.Printf("middleware: jwt: client authenticated. user: %v, conn: %v, ip: %v", userID, ctx.Conn.ID, ctx.Conn.RemoteAddr())
		return ctx.Next()
	}
}
<commit_msg>jwt: bring back next()<commit_after>package jwt

import (
	"fmt"
	"log"

	"github.com/dgrijalva/jwt-go"
	"github.com/neptulon/neptulon"
)

type token struct {
	Token string `json:"token"`
}

// HMAC is JSON Web Token authentication using HMAC.
// If successful, token context will be store with the key "userid" in session.
// If unsuccessful, connection will be closed right away.
func HMAC(password string) func(ctx *neptulon.ReqCtx) error {
	pass := []byte(password)

	return func(ctx *neptulon.ReqCtx) error {
		if _, ok := ctx.Conn.Session.GetOk("userid"); ok {
			return ctx.Next()
		}

		var t token
		if err := ctx.Params(&t); err != nil {
			ctx.Conn.Close()
			return err
		}

		jt, err := jwt.Parse(t.Token, func(token *jwt.Token) (interface{}, error) {
			if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
				return nil, fmt.Errorf("middleware: jwt: unexpected signing method: %v", token.Header["alg"])
			}
			return pass, nil
		})

		if err != nil || !jt.Valid {
			ctx.Conn.Close()
			return fmt.Errorf("middleware: jwt: invalid JWT authentication attempt: %v: %v: %v", err, ctx.Conn.RemoteAddr(), t.Token)
		}

		userID := jt.Claims["userid"].(string)
		ctx.Conn.Session.Set("userid", userID)
		log.Printf("middleware: jwt: client authenticated. user: %v, conn: %v, ip: %v", userID, ctx.Conn.ID, ctx.Conn.RemoteAddr())
		return ctx.Next()
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package pd

import (
	"bufio"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/deadline"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/msgpb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/kvproto/pkg/util"
	"github.com/pingcap/pd/pkg/metrics"
	"github.com/twinj/uuid"
)

const (
	pdRPCPrefix      = "/pd/rpc"
	connectPDTimeout = time.Second * 3
	netIOTimeout     = time.Second
)

const (
	readBufferSize  = 8 * 1024
	writeBufferSize = 8 * 1024
)

const maxPipelineRequest = 10000

// errInvalidResponse represents response message is invalid.
var errInvalidResponse = errors.New("invalid response")

type tsoRequest struct {
	done     chan error
	physical int64
	logical  int64
}

type storeRequest struct {
	pbReq  *pdpb.GetStoreRequest
	done   chan error
	pbResp *pdpb.GetStoreResponse
}

type regionRequest struct {
	pbReq  *pdpb.GetRegionRequest
	done   chan error
	pbResp *pdpb.GetRegionResponse
}

type clusterConfigRequest struct {
	pbReq  *pdpb.GetClusterConfigRequest
	done   chan error
	pbResp *pdpb.GetClusterConfigResponse
}

type rpcWorker struct {
	addr      string
	clusterID uint64
	requests  chan interface{}
	wg        sync.WaitGroup
	quit      chan struct{}
}

func newRPCWorker(addr string, clusterID uint64) *rpcWorker {
	w := &rpcWorker{
		addr:      addr,
		clusterID: clusterID,
		requests:  make(chan interface{}, maxPipelineRequest),
		quit:      make(chan struct{}),
	}
	w.wg.Add(1)
	go w.work()
	return w
}

func (w *rpcWorker) stop(err error) {
	close(w.quit)
	w.wg.Wait()

	for i := 0; i < len(w.requests); i++ {
		req := <-w.requests
		switch r := req.(type) {
		case *tsoRequest:
			r.done <- err
		case *clusterConfigRequest:
			r.done <- err
		case *storeRequest:
			r.done <- err
		case *regionRequest:
			r.done <- err
		}
	}
}

func (w *rpcWorker) work() {
	defer w.wg.Done()

RECONNECT:
	log.Infof("[pd] connect to pd server %v", w.addr)
	conn, err := rpcConnect(w.addr)
	if err != nil {
		log.Warnf("[pd] failed connect pd server: %v, will retry later", err)

		select {
		case <-time.After(netIOTimeout):
			goto RECONNECT
		case <-w.quit:
			return
		}
	}

	reader := bufio.NewReaderSize(deadline.NewDeadlineReader(conn, netIOTimeout), readBufferSize)
	writer := bufio.NewWriterSize(deadline.NewDeadlineWriter(conn, netIOTimeout), writeBufferSize)
	readwriter := bufio.NewReadWriter(reader, writer)

	for {
		var pending []interface{}
		select {
		case req := <-w.requests:
			pending = append(pending, req)
		POP_ALL:
			for {
				select {
				case req := <-w.requests:
					pending = append(pending, req)
				default:
					break POP_ALL
				}
			}
			if ok := w.handleRequests(pending, readwriter); !ok {
				conn.Close()
				goto RECONNECT
			}
		case <-w.quit:
			conn.Close()
			return
		}
	}
}

func (w *rpcWorker) handleRequests(requests []interface{}, conn *bufio.ReadWriter) bool {
	var tsoRequests []*tsoRequest
	ok := true
	for _, req := range requests {
		switch r := req.(type) {
		case *tsoRequest:
			tsoRequests = append(tsoRequests, r)
		case *storeRequest:
			storeResp, err := w.getStoreFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = storeResp
				r.done <- nil
			}
		case *regionRequest:
			regionResp, err := w.getRegionFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = regionResp
				r.done <- nil
			}
		case *clusterConfigRequest:
			clusterConfigResp, err := w.getClusterConfigFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = clusterConfigResp
				r.done <- nil
			}
		default:
			log.Errorf("[pd] invalid request %v", r)
		}
	}

	// Check whether we have tso requests to process.
	if len(tsoRequests) > 0 {
		ts, err := w.getTSFromRemote(conn, len(tsoRequests))
		if err != nil {
			ok = false
			log.Error(err)
		}
		logicalHigh := ts.GetLogical()
		physical := ts.GetPhysical()
		for _, req := range tsoRequests {
			if err != nil {
				req.done <- err
			} else {
				req.physical = physical
				req.logical = logicalHigh
				req.done <- nil
				logicalHigh--
			}
		}
	}

	return ok
}

var msgID uint64

func newMsgID() uint64 {
	return atomic.AddUint64(&msgID, 1)
}

func (w *rpcWorker) getTSFromRemote(conn *bufio.ReadWriter, n int) (pdpb.Timestamp, error) {
	var timestampHigh = pdpb.Timestamp{}
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType: pdpb.CommandType_Tso,
		Tso: &pdpb.TsoRequest{
			Count: uint32(n),
		},
	}
	resp, err := w.callRPC(conn, req)
	if err != nil {
		return timestampHigh, errors.Trace(err)
	}
	if resp.GetTso() == nil {
		return timestampHigh, errors.New("[pd] tso filed in rpc response not set")
	}
	timestampHigh = resp.GetTso().GetTimestamp()
	if resp.GetTso().GetCount() != uint32(n) {
		return timestampHigh, errors.New("[pd] tso length in rpc response is incorrect")
	}
	return timestampHigh, nil
}

func (w *rpcWorker) getStoreFromRemote(conn *bufio.ReadWriter, storeReq *pdpb.GetStoreRequest) (*pdpb.GetStoreResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:  pdpb.CommandType_GetStore,
		GetStore: storeReq,
	}
	resp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if resp.GetGetStore() == nil {
		return nil, errors.New("[pd] GetStore filed in rpc response not set")
	}
	return resp.GetGetStore(), nil
}

func (w *rpcWorker) getRegionFromRemote(conn *bufio.ReadWriter, regionReq *pdpb.GetRegionRequest) (*pdpb.GetRegionResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:   pdpb.CommandType_GetRegion,
		GetRegion: regionReq,
	}
	rsp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if rsp.GetGetRegion() == nil {
		return nil, errors.New("[pd] GetRegion field in rpc response not set")
	}
	return rsp.GetGetRegion(), nil
}

func (w *rpcWorker) getClusterConfigFromRemote(conn *bufio.ReadWriter, clusterConfigReq *pdpb.GetClusterConfigRequest) (*pdpb.GetClusterConfigResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:          pdpb.CommandType_GetClusterConfig,
		GetClusterConfig: clusterConfigReq,
	}
	rsp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if rsp.GetGetClusterConfig() == nil {
		return nil, errors.New("[pd] GetClusterConfig field in rpc response not set")
	}
	return rsp.GetGetClusterConfig(), nil
}

func (w *rpcWorker) callRPC(conn *bufio.ReadWriter, req *pdpb.Request) (resp *pdpb.Response, err error) {
	// Record some metrics.
	start := time.Now()
	label := metrics.GetCmdLabel(req)
	defer func() {
		if err == nil {
			cmdCounter.WithLabelValues(label).Inc()
			cmdDuration.WithLabelValues(label).Observe(time.Since(start).Seconds())
		} else {
			cmdFailedCounter.WithLabelValues(label).Inc()
			cmdFailedDuration.WithLabelValues(label).Observe(time.Since(start).Seconds())
		}
	}()

	msg := &msgpb.Message{
		MsgType: msgpb.MessageType_PdReq,
		PdReq:   req,
	}
	if err = util.WriteMessage(conn, newMsgID(), msg); err != nil {
		return nil, errors.Errorf("[pd] rpc failed: %v", err)
	}
	conn.Flush()
	if _, err = util.ReadMessage(conn, msg); err != nil {
		return nil, errors.Errorf("[pd] rpc failed: %v", err)
	}
	if msg.GetMsgType() != msgpb.MessageType_PdResp {
		return nil, errors.Trace(errInvalidResponse)
	}
	resp = msg.GetPdResp()
	if err = w.checkResponse(resp); err != nil {
		return nil, errors.Trace(err)
	}

	return resp, nil
}

func (w *rpcWorker) checkResponse(resp *pdpb.Response) error {
	header := resp.GetHeader()
	if header == nil {
		return errors.New("[pd] rpc response header not set")
	}
	if err := header.GetError(); err != nil {
		return errors.Errorf("[pd] rpc response with error: %v", err)
	}
	return nil
}

func parseUrls(s string) ([]url.URL, error) {
	items := strings.Split(s, ",")
	urls := make([]url.URL, 0, len(items))
	for _, item := range items {
		u, err := url.Parse(item)
		if err != nil {
			return nil, errors.Trace(err)
		}

		urls = append(urls, *u)
	}

	return urls, nil
}

func rpcConnect(addr string) (net.Conn, error) {
	req, err := http.NewRequest("GET", pdRPCPrefix, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}

	urls, err := parseUrls(addr)
	if err != nil {
		return nil, errors.Trace(err)
	}

	for _, url := range urls {
		var conn net.Conn
		switch url.Scheme {
		// used in tests
		case "unix", "unixs":
			conn, err = net.DialTimeout("unix", url.Host, connectPDTimeout)
		default:
			conn, err = net.DialTimeout("tcp", url.Host, connectPDTimeout)
		}

		if err != nil {
			continue
		}
		err = req.Write(conn)
		if err != nil {
			conn.Close()
			continue
		}
		return conn, nil
	}

	return nil, errors.Errorf("connect to %s failed", addr)
}
<commit_msg>pd-client: fix a bug that hangs requests forever. (#323)<commit_after>// Copyright 2016 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package pd

import (
	"bufio"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/deadline"
	"github.com/ngaut/log"
	"github.com/pingcap/kvproto/pkg/msgpb"
	"github.com/pingcap/kvproto/pkg/pdpb"
	"github.com/pingcap/kvproto/pkg/util"
	"github.com/pingcap/pd/pkg/metrics"
	"github.com/twinj/uuid"
)

const (
	pdRPCPrefix      = "/pd/rpc"
	connectPDTimeout = time.Second * 3
	netIOTimeout     = time.Second
)

const (
	readBufferSize  = 8 * 1024
	writeBufferSize = 8 * 1024
)

const maxPipelineRequest = 10000

// errInvalidResponse represents response message is invalid.
var errInvalidResponse = errors.New("invalid response")

type tsoRequest struct {
	done     chan error
	physical int64
	logical  int64
}

type storeRequest struct {
	pbReq  *pdpb.GetStoreRequest
	done   chan error
	pbResp *pdpb.GetStoreResponse
}

type regionRequest struct {
	pbReq  *pdpb.GetRegionRequest
	done   chan error
	pbResp *pdpb.GetRegionResponse
}

type clusterConfigRequest struct {
	pbReq  *pdpb.GetClusterConfigRequest
	done   chan error
	pbResp *pdpb.GetClusterConfigResponse
}

type rpcWorker struct {
	addr      string
	clusterID uint64
	requests  chan interface{}
	wg        sync.WaitGroup
	quit      chan struct{}
}

func newRPCWorker(addr string, clusterID uint64) *rpcWorker {
	w := &rpcWorker{
		addr:      addr,
		clusterID: clusterID,
		requests:  make(chan interface{}, maxPipelineRequest),
		quit:      make(chan struct{}),
	}
	w.wg.Add(1)
	go w.work()
	return w
}

func (w *rpcWorker) stop(err error) {
	close(w.quit)
	w.wg.Wait()

	n := len(w.requests)
	for i := 0; i < n; i++ {
		req := <-w.requests
		switch r := req.(type) {
		case *tsoRequest:
			r.done <- err
		case *clusterConfigRequest:
			r.done <- err
		case *storeRequest:
			r.done <- err
		case *regionRequest:
			r.done <- err
		}
	}
}

func (w *rpcWorker) work() {
	defer w.wg.Done()

RECONNECT:
	log.Infof("[pd] connect to pd server %v", w.addr)
	conn, err := rpcConnect(w.addr)
	if err != nil {
		log.Warnf("[pd] failed connect pd server: %v, will retry later", err)

		select {
		case <-time.After(netIOTimeout):
			goto RECONNECT
		case <-w.quit:
			return
		}
	}

	reader := bufio.NewReaderSize(deadline.NewDeadlineReader(conn, netIOTimeout), readBufferSize)
	writer := bufio.NewWriterSize(deadline.NewDeadlineWriter(conn, netIOTimeout), writeBufferSize)
	readwriter := bufio.NewReadWriter(reader, writer)

	for {
		var pending []interface{}
		select {
		case req := <-w.requests:
			pending = append(pending, req)
		POP_ALL:
			for {
				select {
				case req := <-w.requests:
					pending = append(pending, req)
				default:
					break POP_ALL
				}
			}
			if ok := w.handleRequests(pending, readwriter); !ok {
				conn.Close()
				goto RECONNECT
			}
		case <-w.quit:
			conn.Close()
			return
		}
	}
}

func (w *rpcWorker) handleRequests(requests []interface{}, conn *bufio.ReadWriter) bool {
	var tsoRequests []*tsoRequest
	ok := true
	for _, req := range requests {
		switch r := req.(type) {
		case *tsoRequest:
			tsoRequests = append(tsoRequests, r)
		case *storeRequest:
			storeResp, err := w.getStoreFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = storeResp
				r.done <- nil
			}
		case *regionRequest:
			regionResp, err := w.getRegionFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = regionResp
				r.done <- nil
			}
		case *clusterConfigRequest:
			clusterConfigResp, err := w.getClusterConfigFromRemote(conn, r.pbReq)
			if err != nil {
				ok = false
				log.Error(err)
				r.done <- err
			} else {
				r.pbResp = clusterConfigResp
				r.done <- nil
			}
		default:
			log.Errorf("[pd] invalid request %v", r)
		}
	}

	// Check whether we have tso requests to process.
	if len(tsoRequests) > 0 {
		ts, err := w.getTSFromRemote(conn, len(tsoRequests))
		if err != nil {
			ok = false
			log.Error(err)
		}
		logicalHigh := ts.GetLogical()
		physical := ts.GetPhysical()
		for _, req := range tsoRequests {
			if err != nil {
				req.done <- err
			} else {
				req.physical = physical
				req.logical = logicalHigh
				req.done <- nil
				logicalHigh--
			}
		}
	}

	return ok
}

var msgID uint64

func newMsgID() uint64 {
	return atomic.AddUint64(&msgID, 1)
}

func (w *rpcWorker) getTSFromRemote(conn *bufio.ReadWriter, n int) (pdpb.Timestamp, error) {
	var timestampHigh = pdpb.Timestamp{}
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType: pdpb.CommandType_Tso,
		Tso: &pdpb.TsoRequest{
			Count: uint32(n),
		},
	}
	resp, err := w.callRPC(conn, req)
	if err != nil {
		return timestampHigh, errors.Trace(err)
	}
	if resp.GetTso() == nil {
		return timestampHigh, errors.New("[pd] tso filed in rpc response not set")
	}
	timestampHigh = resp.GetTso().GetTimestamp()
	if resp.GetTso().GetCount() != uint32(n) {
		return timestampHigh, errors.New("[pd] tso length in rpc response is incorrect")
	}
	return timestampHigh, nil
}

func (w *rpcWorker) getStoreFromRemote(conn *bufio.ReadWriter, storeReq *pdpb.GetStoreRequest) (*pdpb.GetStoreResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:  pdpb.CommandType_GetStore,
		GetStore: storeReq,
	}
	resp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if resp.GetGetStore() == nil {
		return nil, errors.New("[pd] GetStore filed in rpc response not set")
	}
	return resp.GetGetStore(), nil
}

func (w *rpcWorker) getRegionFromRemote(conn *bufio.ReadWriter, regionReq *pdpb.GetRegionRequest) (*pdpb.GetRegionResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:   pdpb.CommandType_GetRegion,
		GetRegion: regionReq,
	}
	rsp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if rsp.GetGetRegion() == nil {
		return nil, errors.New("[pd] GetRegion field in rpc response not set")
	}
	return rsp.GetGetRegion(), nil
}

func (w *rpcWorker) getClusterConfigFromRemote(conn *bufio.ReadWriter, clusterConfigReq *pdpb.GetClusterConfigRequest) (*pdpb.GetClusterConfigResponse, error) {
	req := &pdpb.Request{
		Header: &pdpb.RequestHeader{
			Uuid:      uuid.NewV4().Bytes(),
			ClusterId: w.clusterID,
		},
		CmdType:          pdpb.CommandType_GetClusterConfig,
		GetClusterConfig: clusterConfigReq,
	}
	rsp, err := w.callRPC(conn, req)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if rsp.GetGetClusterConfig() == nil {
		return nil, errors.New("[pd] GetClusterConfig field in rpc response not set")
	}
	return rsp.GetGetClusterConfig(), nil
}

func (w *rpcWorker) callRPC(conn *bufio.ReadWriter, req *pdpb.Request) (resp *pdpb.Response, err error) {
	// Record some metrics.
	start := time.Now()
	label := metrics.GetCmdLabel(req)
	defer func() {
		if err == nil {
			cmdCounter.WithLabelValues(label).Inc()
			cmdDuration.WithLabelValues(label).Observe(time.Since(start).Seconds())
		} else {
			cmdFailedCounter.WithLabelValues(label).Inc()
			cmdFailedDuration.WithLabelValues(label).Observe(time.Since(start).Seconds())
		}
	}()

	msg := &msgpb.Message{
		MsgType: msgpb.MessageType_PdReq,
		PdReq:   req,
	}
	if err = util.WriteMessage(conn, newMsgID(), msg); err != nil {
		return nil, errors.Errorf("[pd] rpc failed: %v", err)
	}
	conn.Flush()
	if _, err = util.ReadMessage(conn, msg); err != nil {
		return nil, errors.Errorf("[pd] rpc failed: %v", err)
	}
	if msg.GetMsgType() != msgpb.MessageType_PdResp {
		return nil, errors.Trace(errInvalidResponse)
	}
	resp = msg.GetPdResp()
	if err = w.checkResponse(resp); err != nil {
		return nil, errors.Trace(err)
	}

	return resp, nil
}

func (w *rpcWorker) checkResponse(resp *pdpb.Response) error {
	header := resp.GetHeader()
	if header == nil {
		return errors.New("[pd] rpc response header not set")
	}
	if err := header.GetError(); err != nil {
		return errors.Errorf("[pd] rpc response with error: %v", err)
	}
	return nil
}

func parseUrls(s string) ([]url.URL, error) {
	items := strings.Split(s, ",")
	urls := make([]url.URL, 0, len(items))
	for _, item := range items {
		u, err := url.Parse(item)
		if err != nil {
			return nil, errors.Trace(err)
		}

		urls = append(urls, *u)
	}

	return urls, nil
}

func rpcConnect(addr string) (net.Conn, error) {
	req, err := http.NewRequest("GET", pdRPCPrefix, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}

	urls, err := parseUrls(addr)
	if err != nil {
		return nil, errors.Trace(err)
	}

	for _, url := range urls {
		var conn net.Conn
		switch url.Scheme {
		// used in tests
		case "unix", "unixs":
			conn, err = net.DialTimeout("unix", url.Host, connectPDTimeout)
		default:
			conn, err = net.DialTimeout("tcp", url.Host, connectPDTimeout)
		}

		if err != nil {
			continue
		}
		err = req.Write(conn)
		if err != nil {
			conn.Close()
			continue
		}
		return conn, nil
	}

	return nil, errors.Errorf("connect to %s failed", addr)
}
<|endoftext|>