<commit_before>package forwarding

import (
	"sort"

	"github.com/pkg/errors"

	"github.com/mutagen-io/mutagen/pkg/filesystem"
	"github.com/mutagen-io/mutagen/pkg/logging"
	"github.com/mutagen-io/mutagen/pkg/identifier"
	"github.com/mutagen-io/mutagen/pkg/selection"
	"github.com/mutagen-io/mutagen/pkg/state"
	"github.com/mutagen-io/mutagen/pkg/url"
)

// Manager provides forwarding session management facilities. Its methods are
// safe for concurrent usage, so it can be easily exported via an RPC interface.
type Manager struct {
	// logger is the underlying logger.
	logger *logging.Logger
	// tracker tracks changes to session states.
	tracker *state.Tracker
	// sessionLock locks the sessions registry.
	sessionsLock *state.TrackingLock
	// sessions maps sessions to their respective controllers.
	sessions map[string]*controller
}

// NewManager creates a new Manager instance.
func NewManager(logger *logging.Logger) (*Manager, error) {
	// Create a tracker and corresponding lock to watch for state changes.
	tracker := state.NewTracker()
	sessionsLock := state.NewTrackingLock(tracker)

	// Create the session registry.
	sessions := make(map[string]*controller)

	// Load existing sessions.
	logger.Println("Looking for existing sessions")
	sessionsDirectory, err := pathForSession("")
	if err != nil {
		return nil, errors.Wrap(err, "unable to compute sessions directory")
	}
	sessionsDirectoryContents, err := filesystem.DirectoryContentsByPath(sessionsDirectory)
	if err != nil {
		return nil, errors.Wrap(err, "unable to read contents of sessions directory")
	}
	for _, c := range sessionsDirectoryContents {
		identifier := c.Name()
		logger.Println("Loading session", identifier)
		if controller, err := loadSession(logger.Sublogger(identifier), tracker, identifier); err != nil {
			continue
		} else {
			sessions[identifier] = controller
		}
	}

	// Success.
	logger.Println("Session manager initialized")
	return &Manager{
		logger:       logger,
		tracker:      tracker,
		sessionsLock: sessionsLock,
		sessions:     sessions,
	}, nil
}

// allControllers creates a list of all controllers managed by the manager.
func (m *Manager) allControllers() []*controller {
	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Generate a list of all controllers.
	controllers := make([]*controller, 0, len(m.sessions))
	for _, controller := range m.sessions {
		controllers = append(controllers, controller)
	}

	// Done.
	return controllers
}

// findControllersBySpecification generates a list of controllers matching the
// given specifications.
func (m *Manager) findControllersBySpecification(specifications []string) ([]*controller, error) {
	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Generate a list of controllers matching the specifications. We allow each
	// specification to match multiple controllers, so we store matches in a set
	// before converting them to a list. We do require that each specification
	// match at least one controller.
	controllerSet := make(map[*controller]bool)
	for _, specification := range specifications {
		var matched bool
		for _, controller := range m.sessions {
			if controller.session.Identifier == specification || controller.session.Name == specification {
				controllerSet[controller] = true
				matched = true
			}
		}
		if !matched {
			return nil, errors.Errorf("specification \"%s\" did not match any sessions", specification)
		}
	}

	// Convert the set to a list.
	controllers := make([]*controller, 0, len(controllerSet))
	for c := range controllerSet {
		controllers = append(controllers, c)
	}

	// Done.
	return controllers, nil
}

// findControllersByLabelSelector generates a list of controllers using the
// specified label selector.
func (m *Manager) findControllersByLabelSelector(labelSelector string) ([]*controller, error) {
	// Parse the label selector.
	selector, err := selection.ParseLabelSelector(labelSelector)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse label selector")
	}

	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Loop over controllers and look for matches.
	var controllers []*controller
	for _, controller := range m.sessions {
		if selector.Matches(controller.session.Labels) {
			controllers = append(controllers, controller)
		}
	}

	// Done.
	return controllers, nil
}

// selectControllers generates a list of controllers using the mechanism
// specified by the provided selection.
func (m *Manager) selectControllers(selection *selection.Selection) ([]*controller, error) {
	// Dispatch selection based on the requested mechanism.
	if selection.All {
		return m.allControllers(), nil
	} else if len(selection.Specifications) > 0 {
		return m.findControllersBySpecification(selection.Specifications)
	} else if selection.LabelSelector != "" {
		return m.findControllersByLabelSelector(selection.LabelSelector)
	} else {
		// TODO: Should we panic here instead?
		return nil, errors.New("invalid session selection")
	}
}

// Shutdown tells the manager to gracefully halt sessions.
func (m *Manager) Shutdown() {
	// Log the shutdown.
	m.logger.Println("Shutting down")

	// Poison state tracking to terminate monitoring.
	m.tracker.Poison()

	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Attempt to halt each session so that it can shutdown cleanly. Ignore but
	// log any that fail to halt.
	for _, controller := range m.sessions {
		m.logger.Println("Halting session", controller.session.Identifier)
		if err := controller.halt(controllerHaltModeShutdown, ""); err != nil {
			// TODO: Log this halt failure.
		}
	}
}

// Create tells the manager to create a new session.
func (m *Manager) Create(
	source, destination *url.URL,
	configuration, configurationSource, configurationDestination *Configuration,
	name string,
	labels map[string]string,
	paused bool,
	prompter string,
) (string, error) {
	// Create a unique session identifier.
	identifier, err := identifier.New(identifier.PrefixForwarding)
	if err != nil {
		return "", errors.Wrap(err, "unable to generate UUID for session")
	}

	// Attempt to create a session.
	controller, err := newSession(
		m.logger.Sublogger(identifier),
		m.tracker,
		identifier,
		source, destination,
		configuration, configurationSource, configurationDestination,
		name,
		labels,
		paused,
		prompter,
	)
	if err != nil {
		return "", err
	}

	// Register the controller.
	m.sessionsLock.Lock()
	m.sessions[controller.session.Identifier] = controller
	m.sessionsLock.Unlock()

	// Done.
	return controller.session.Identifier, nil
}

// List requests a state snapshot for the specified sessions.
func (m *Manager) List(selection *selection.Selection, previousStateIndex uint64) (uint64, []*State, error) {
	// Wait for a state change from the previous index.
	stateIndex, poisoned := m.tracker.WaitForChange(previousStateIndex)
	if poisoned {
		return 0, nil, errors.New("state tracking terminated")
	}

	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return 0, nil, errors.Wrap(err, "unable to locate requested sessions")
	}

	// Extract the state from each controller.
	states := make([]*State, len(controllers))
	for i, controller := range controllers {
		states[i] = controller.currentState()
	}

	// Sort session states by session creation time.
	sort.Slice(states, func(i, j int) bool {
		iTime := states[i].Session.CreationTime
		jTime := states[j].Session.CreationTime
		return iTime.Seconds < jTime.Seconds ||
			(iTime.Seconds == jTime.Seconds && iTime.Nanos < jTime.Nanos)
	})

	// Success.
	return stateIndex, states, nil
}

// Pause tells the manager to pause sessions matching the given specifications.
func (m *Manager) Pause(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to pause the sessions.
	for _, controller := range controllers {
		if err := controller.halt(controllerHaltModePause, prompter); err != nil {
			return errors.Wrap(err, "unable to pause session")
		}
	}

	// Success.
	return nil
}

// Resume tells the manager to resume sessions matching the given
// specifications.
func (m *Manager) Resume(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to resume.
	for _, controller := range controllers {
		if err := controller.resume(prompter); err != nil {
			return errors.Wrap(err, "unable to resume session")
		}
	}

	// Success.
	return nil
}

// Terminate tells the manager to terminate sessions matching the given
// specifications.
func (m *Manager) Terminate(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to terminate the sessions. Since we're terminating them, we're
	// responsible for removing them from the session map.
	for _, controller := range controllers {
		if err := controller.halt(controllerHaltModeTerminate, prompter); err != nil {
			return errors.Wrap(err, "unable to terminate session")
		}
		m.sessionsLock.Lock()
		delete(m.sessions, controller.session.Identifier)
		m.sessionsLock.Unlock()
	}

	// Success.
	return nil
}
<commit_msg>Fixed incorrect import formatting.<commit_after>package forwarding

import (
	"sort"

	"github.com/pkg/errors"

	"github.com/mutagen-io/mutagen/pkg/filesystem"
	"github.com/mutagen-io/mutagen/pkg/identifier"
	"github.com/mutagen-io/mutagen/pkg/logging"
	"github.com/mutagen-io/mutagen/pkg/selection"
	"github.com/mutagen-io/mutagen/pkg/state"
	"github.com/mutagen-io/mutagen/pkg/url"
)

// Manager provides forwarding session management facilities. Its methods are
// safe for concurrent usage, so it can be easily exported via an RPC interface.
type Manager struct {
	// logger is the underlying logger.
	logger *logging.Logger
	// tracker tracks changes to session states.
	tracker *state.Tracker
	// sessionLock locks the sessions registry.
	sessionsLock *state.TrackingLock
	// sessions maps sessions to their respective controllers.
	sessions map[string]*controller
}

// NewManager creates a new Manager instance.
func NewManager(logger *logging.Logger) (*Manager, error) {
	// Create a tracker and corresponding lock to watch for state changes.
	tracker := state.NewTracker()
	sessionsLock := state.NewTrackingLock(tracker)

	// Create the session registry.
	sessions := make(map[string]*controller)

	// Load existing sessions.
	logger.Println("Looking for existing sessions")
	sessionsDirectory, err := pathForSession("")
	if err != nil {
		return nil, errors.Wrap(err, "unable to compute sessions directory")
	}
	sessionsDirectoryContents, err := filesystem.DirectoryContentsByPath(sessionsDirectory)
	if err != nil {
		return nil, errors.Wrap(err, "unable to read contents of sessions directory")
	}
	for _, c := range sessionsDirectoryContents {
		identifier := c.Name()
		logger.Println("Loading session", identifier)
		if controller, err := loadSession(logger.Sublogger(identifier), tracker, identifier); err != nil {
			continue
		} else {
			sessions[identifier] = controller
		}
	}

	// Success.
	logger.Println("Session manager initialized")
	return &Manager{
		logger:       logger,
		tracker:      tracker,
		sessionsLock: sessionsLock,
		sessions:     sessions,
	}, nil
}

// allControllers creates a list of all controllers managed by the manager.
func (m *Manager) allControllers() []*controller {
	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Generate a list of all controllers.
	controllers := make([]*controller, 0, len(m.sessions))
	for _, controller := range m.sessions {
		controllers = append(controllers, controller)
	}

	// Done.
	return controllers
}

// findControllersBySpecification generates a list of controllers matching the
// given specifications.
func (m *Manager) findControllersBySpecification(specifications []string) ([]*controller, error) {
	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Generate a list of controllers matching the specifications. We allow each
	// specification to match multiple controllers, so we store matches in a set
	// before converting them to a list. We do require that each specification
	// match at least one controller.
	controllerSet := make(map[*controller]bool)
	for _, specification := range specifications {
		var matched bool
		for _, controller := range m.sessions {
			if controller.session.Identifier == specification || controller.session.Name == specification {
				controllerSet[controller] = true
				matched = true
			}
		}
		if !matched {
			return nil, errors.Errorf("specification \"%s\" did not match any sessions", specification)
		}
	}

	// Convert the set to a list.
	controllers := make([]*controller, 0, len(controllerSet))
	for c := range controllerSet {
		controllers = append(controllers, c)
	}

	// Done.
	return controllers, nil
}

// findControllersByLabelSelector generates a list of controllers using the
// specified label selector.
func (m *Manager) findControllersByLabelSelector(labelSelector string) ([]*controller, error) {
	// Parse the label selector.
	selector, err := selection.ParseLabelSelector(labelSelector)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse label selector")
	}

	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Loop over controllers and look for matches.
	var controllers []*controller
	for _, controller := range m.sessions {
		if selector.Matches(controller.session.Labels) {
			controllers = append(controllers, controller)
		}
	}

	// Done.
	return controllers, nil
}

// selectControllers generates a list of controllers using the mechanism
// specified by the provided selection.
func (m *Manager) selectControllers(selection *selection.Selection) ([]*controller, error) {
	// Dispatch selection based on the requested mechanism.
	if selection.All {
		return m.allControllers(), nil
	} else if len(selection.Specifications) > 0 {
		return m.findControllersBySpecification(selection.Specifications)
	} else if selection.LabelSelector != "" {
		return m.findControllersByLabelSelector(selection.LabelSelector)
	} else {
		// TODO: Should we panic here instead?
		return nil, errors.New("invalid session selection")
	}
}

// Shutdown tells the manager to gracefully halt sessions.
func (m *Manager) Shutdown() {
	// Log the shutdown.
	m.logger.Println("Shutting down")

	// Poison state tracking to terminate monitoring.
	m.tracker.Poison()

	// Grab the registry lock and defer its release.
	m.sessionsLock.Lock()
	defer m.sessionsLock.UnlockWithoutNotify()

	// Attempt to halt each session so that it can shutdown cleanly. Ignore but
	// log any that fail to halt.
	for _, controller := range m.sessions {
		m.logger.Println("Halting session", controller.session.Identifier)
		if err := controller.halt(controllerHaltModeShutdown, ""); err != nil {
			// TODO: Log this halt failure.
		}
	}
}

// Create tells the manager to create a new session.
func (m *Manager) Create(
	source, destination *url.URL,
	configuration, configurationSource, configurationDestination *Configuration,
	name string,
	labels map[string]string,
	paused bool,
	prompter string,
) (string, error) {
	// Create a unique session identifier.
	identifier, err := identifier.New(identifier.PrefixForwarding)
	if err != nil {
		return "", errors.Wrap(err, "unable to generate UUID for session")
	}

	// Attempt to create a session.
	controller, err := newSession(
		m.logger.Sublogger(identifier),
		m.tracker,
		identifier,
		source, destination,
		configuration, configurationSource, configurationDestination,
		name,
		labels,
		paused,
		prompter,
	)
	if err != nil {
		return "", err
	}

	// Register the controller.
	m.sessionsLock.Lock()
	m.sessions[controller.session.Identifier] = controller
	m.sessionsLock.Unlock()

	// Done.
	return controller.session.Identifier, nil
}

// List requests a state snapshot for the specified sessions.
func (m *Manager) List(selection *selection.Selection, previousStateIndex uint64) (uint64, []*State, error) {
	// Wait for a state change from the previous index.
	stateIndex, poisoned := m.tracker.WaitForChange(previousStateIndex)
	if poisoned {
		return 0, nil, errors.New("state tracking terminated")
	}

	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return 0, nil, errors.Wrap(err, "unable to locate requested sessions")
	}

	// Extract the state from each controller.
	states := make([]*State, len(controllers))
	for i, controller := range controllers {
		states[i] = controller.currentState()
	}

	// Sort session states by session creation time.
	sort.Slice(states, func(i, j int) bool {
		iTime := states[i].Session.CreationTime
		jTime := states[j].Session.CreationTime
		return iTime.Seconds < jTime.Seconds ||
			(iTime.Seconds == jTime.Seconds && iTime.Nanos < jTime.Nanos)
	})

	// Success.
	return stateIndex, states, nil
}

// Pause tells the manager to pause sessions matching the given specifications.
func (m *Manager) Pause(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to pause the sessions.
	for _, controller := range controllers {
		if err := controller.halt(controllerHaltModePause, prompter); err != nil {
			return errors.Wrap(err, "unable to pause session")
		}
	}

	// Success.
	return nil
}

// Resume tells the manager to resume sessions matching the given
// specifications.
func (m *Manager) Resume(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to resume.
	for _, controller := range controllers {
		if err := controller.resume(prompter); err != nil {
			return errors.Wrap(err, "unable to resume session")
		}
	}

	// Success.
	return nil
}

// Terminate tells the manager to terminate sessions matching the given
// specifications.
func (m *Manager) Terminate(selection *selection.Selection, prompter string) error {
	// Extract the controllers for the sessions of interest.
	controllers, err := m.selectControllers(selection)
	if err != nil {
		return errors.Wrap(err, "unable to locate requested sessions")
	}

	// Attempt to terminate the sessions. Since we're terminating them, we're
	// responsible for removing them from the session map.
	for _, controller := range controllers {
		if err := controller.halt(controllerHaltModeTerminate, prompter); err != nil {
			return errors.Wrap(err, "unable to terminate session")
		}
		m.sessionsLock.Lock()
		delete(m.sessions, controller.session.Identifier)
		m.sessionsLock.Unlock()
	}

	// Success.
	return nil
}
<|endoftext|>


<commit_before>package isolation

import "os/exec"
import "strconv"
import "strings"
import "errors"

// CPUInfo defines data needed for CPU topology
type CPUInfo struct {
	sockets        int
	physicalCores  int
	threadsPerCore int
	cacheL1i       int
	cacheL1d       int
	cacheL2        int
	cacheL3        int
}

// const for human readable byte sizes
const (
	BYTE     = 1.0
	KILOBYTE = 1024 * BYTE
	MEGABYTE = 1024 * KILOBYTE
	GIGABYTE = 1024 * MEGABYTE
	TERABYTE = 1024 * GIGABYTE
)

// NewCPUInfo instance creation.
func NewCPUInfo(cores int, threads int, l1i int, l1d int, l2 int, l3 int) *CPUInfo {
	return &CPUInfo{physicalCores: cores, threadsPerCore: threads, cacheL1i: l1i, cacheL1d: l1d, cacheL2: l2, cacheL3: l3}
}

// UnitsOfBytes returns KILOBYTES, MEGHABYTES, etc as detected in the input string
func UnitsOfBytes(s string) (int, error) {
	if strings.Contains(s, "K") {
		return KILOBYTE, nil
	} else if strings.Contains(s, "M") {
		return MEGABYTE, nil
	} else if strings.Contains(s, "G") {
		return GIGABYTE, nil
	} else if strings.Contains(s, "T") {
		return TERABYTE, nil
	} else {
		err := errors.New("Unexpected input error")
		return BYTE, err
	}
}

// DiscoverCPU removes the specified cgroup
func (cputopo *CPUInfo) Discover() error {

	out, err := exec.Command("lscpu").Output()
	outstring := strings.TrimSpace(string(out))
	lines := strings.Split(outstring, "\n")

	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "Socket(s)":
			t, _ := strconv.Atoi(value)
			cputopo.sockets = int(t)
		case "Core(s) per socket":
			t, _ := strconv.Atoi(value)
			cputopo.physicalCores = int(t)
		case "Thread(s) per core":
			t, _ := strconv.Atoi(value)
			cputopo.threadsPerCore = int(t)
		case "L1d cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL1d = int(t) * multiplier
			}
		case "L1i cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL1i = int(t) * multiplier
			}
		case "L2 cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL2 = int(t) * multiplier
			}
		case "L3 cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL3 = int(t) * multiplier
			}

		}

	}

	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Created CPU topology and discovered L1 instruction and data, L2 & L3 Caches sizes<commit_after>package isolation

import "os/exec"
import "strconv"
import "strings"
import "errors"

// CPUInfo defines data needed for CPU topology
type CPUInfo struct {
	sockets        int
	physicalCores  int
	threadsPerCore int
	cacheL1i       int
	cacheL1d       int
	cacheL2        int
	cacheL3        int
}

// const for human readable byte sizes
const (
	BYTE     = 1.0
	KILOBYTE = 1024 * BYTE
	MEGABYTE = 1024 * KILOBYTE
	GIGABYTE = 1024 * MEGABYTE
	TERABYTE = 1024 * GIGABYTE
)

// NewCPUInfo instance creation.
func NewCPUInfo(cores int, threads int, l1i int, l1d int, l2 int, l3 int) *CPUInfo {
	return &CPUInfo{physicalCores: cores, threadsPerCore: threads, cacheL1i: l1i, cacheL1d: l1d, cacheL2: l2, cacheL3: l3}
}

// UnitsOfBytes returns KILOBYTES, MEGHABYTES, etc as detected in the input string
func UnitsOfBytes(s string) (int, error) {
	if strings.Contains(s, "K") {
		return KILOBYTE, nil
	} else if strings.Contains(s, "M") {
		return MEGABYTE, nil
	} else if strings.Contains(s, "G") {
		return GIGABYTE, nil
	} else if strings.Contains(s, "T") {
		return TERABYTE, nil
	} else {
		err := errors.New("Unexpected input error")
		return BYTE, err
	}
}

// Discover CPU removes the specified cgroup
func (cputopo *CPUInfo) Discover() error {

	out, err := exec.Command("lscpu").Output()
	outstring := strings.TrimSpace(string(out))
	lines := strings.Split(outstring, "\n")

	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "Socket(s)":
			t, _ := strconv.Atoi(value)
			cputopo.sockets = int(t)
		case "Core(s) per socket":
			t, _ := strconv.Atoi(value)
			cputopo.physicalCores = int(t)
		case "Thread(s) per core":
			t, _ := strconv.Atoi(value)
			cputopo.threadsPerCore = int(t)
		case "L1d cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL1d = int(t) * multiplier
			}
		case "L1i cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL1i = int(t) * multiplier
			}
		case "L2 cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL2 = int(t) * multiplier
			}
		case "L3 cache":
			tFmt := value[:len(value)-1]
			t, _ := strconv.Atoi(tFmt)

			multiplier, err := UnitsOfBytes(value)
			if err == nil {
				cputopo.cacheL3 = int(t) * multiplier
			}

		}

	}

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 cmd

import (
	"fmt"
	"io"
	"net/url"
	"os"
	"os/signal"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/remotecommand"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	utilerrors "k8s.io/kubernetes/pkg/util/errors"
)

const (
	attach_example = `# Get output from running pod 123456-7890, using the first container by default
$ kubectl attach 123456-7890

# Get output from ruby-container from pod 123456-7890
$ kubectl attach 123456-7890 -c ruby-container

# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
# and sends stdout/stderr from 'bash' back to the client
$ kubectl attach 123456-7890 -c ruby-container -i -t`
)

func NewCmdAttach(f *cmdutil.Factory, cmdIn io.Reader, cmdOut, cmdErr io.Writer) *cobra.Command {
	options := &AttachOptions{
		In:  cmdIn,
		Out: cmdOut,
		Err: cmdErr,

		Attach: &DefaultRemoteAttach{},
	}
	cmd := &cobra.Command{
		Use:     "attach POD -c CONTAINER",
		Short:   "Attach to a running container.",
		Long:    "Attach to a process that is already running inside an existing container.",
		Example: attach_example,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.Run())
		},
	}
	// TODO support UID
	cmd.Flags().StringVarP(&options.ContainerName, "container", "c", "", "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&options.Stdin, "stdin", "i", false, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&options.TTY, "tty", "t", false, "Stdin is a TTY")
	return cmd
}

// RemoteAttach defines the interface accepted by the Attach command - provided for test stubbing
type RemoteAttach interface {
	Attach(method string, url *url.URL, config *client.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool) error
}

// DefaultRemoteAttach is the standard implementation of attaching
type DefaultRemoteAttach struct{}

func (*DefaultRemoteAttach) Attach(method string, url *url.URL, config *client.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool) error {
	exec, err := remotecommand.NewExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(stdin, stdout, stderr, tty)
}

// AttachOptions declare the arguments accepted by the Exec command
type AttachOptions struct {
	Namespace     string
	PodName       string
	ContainerName string
	Stdin         bool
	TTY           bool

	In  io.Reader
	Out io.Writer
	Err io.Writer

	Attach RemoteAttach
	Client *client.Client
	Config *client.Config
}

// Complete verifies command line arguments and loads data from the command environment
func (p *AttachOptions) Complete(f *cmdutil.Factory, cmd *cobra.Command, argsIn []string) error {
	if len(argsIn) == 0 {
		return cmdutil.UsageError(cmd, "POD is required for attach")
	}
	if len(argsIn) > 1 {
		return cmdutil.UsageError(cmd, fmt.Sprintf("expected a single argument: POD, saw %d: %s", len(argsIn), argsIn))
	}
	p.PodName = argsIn[0]

	namespace, _, err := f.DefaultNamespace()
	if err != nil {
		return err
	}
	p.Namespace = namespace

	config, err := f.ClientConfig()
	if err != nil {
		return err
	}
	p.Config = config

	client, err := f.Client()
	if err != nil {
		return err
	}
	p.Client = client

	return nil
}

// Validate checks that the provided attach options are specified.
func (p *AttachOptions) Validate() error {
	allErrs := []error{}
	if len(p.PodName) == 0 {
		allErrs = append(allErrs, fmt.Errorf("pod name must be specified"))
	}
	if p.Out == nil || p.Err == nil {
		allErrs = append(allErrs, fmt.Errorf("both output and error output must be provided"))
	}
	if p.Attach == nil || p.Client == nil || p.Config == nil {
		allErrs = append(allErrs, fmt.Errorf("client, client config, and attach must be provided"))
	}
	return utilerrors.NewAggregate(allErrs)
}

// Run executes a validated remote execution against a pod.
func (p *AttachOptions) Run() error {
	pod, err := p.Client.Pods(p.Namespace).Get(p.PodName)
	if err != nil {
		return err
	}

	if pod.Status.Phase != api.PodRunning {
		return fmt.Errorf("pod %s is not running and cannot be attached to; current phase is %s", p.PodName, pod.Status.Phase)
	}

	var stdin io.Reader
	tty := p.TTY

	containerToAttach := p.GetContainer(pod)
	if tty && !containerToAttach.TTY {
		tty = false
		fmt.Fprintf(p.Err, "Unable to use a TTY - container %s doesn't allocate one\n", containerToAttach.Name)
	}

	// TODO: refactor with terminal helpers from the edit utility once that is merged
	if p.Stdin {
		stdin = p.In
		if tty {
			if file, ok := stdin.(*os.File); ok {
				inFd := file.Fd()
				if term.IsTerminal(inFd) {
					oldState, err := term.SetRawTerminal(inFd)
					if err != nil {
						glog.Fatal(err)
					}
					fmt.Fprintln(p.Out, "\nHit enter for command prompt")
					// this handles a clean exit, where the command finished
					defer term.RestoreTerminal(inFd, oldState)

					// SIGINT is handled by term.SetRawTerminal (it runs a goroutine that listens
					// for SIGINT and restores the terminal before exiting)

					// this handles SIGTERM
					sigChan := make(chan os.Signal, 1)
					signal.Notify(sigChan, syscall.SIGTERM)
					go func() {
						<-sigChan
						term.RestoreTerminal(inFd, oldState)
						os.Exit(0)
					}()
				} else {
					fmt.Fprintln(p.Err, "STDIN is not a terminal")
				}
			} else {
				tty = false
				fmt.Fprintln(p.Err, "Unable to use a TTY - input is not the right kind of file")
			}
		}
	}

	// TODO: consider abstracting into a client invocation or client helper
	req := p.Client.RESTClient.Post().
		Resource("pods").
		Name(pod.Name).
		Namespace(pod.Namespace).
		SubResource("attach")
	req.VersionedParams(&api.PodAttachOptions{
		Container: containerToAttach.Name,
		Stdin:     stdin != nil,
		Stdout:    p.Out != nil,
		Stderr:    p.Err != nil,
		TTY:       tty,
	}, api.Scheme)

	return p.Attach.Attach("POST", req.URL(), p.Config, stdin, p.Out, p.Err, tty)
}

// GetContainer returns the container to attach to, with a fallback.
func (p *AttachOptions) GetContainer(pod *api.Pod) api.Container {
	if len(p.ContainerName) > 0 {
		for _, container := range pod.Spec.Containers {
			if container.Name == p.ContainerName {
				return container
			}
		}
	}

	glog.V(4).Infof("defaulting container name to %s", pod.Spec.Containers[0].Name)
	return pod.Spec.Containers[0]
}

// GetContainerName returns the name of the container to attach to, with a fallback.
func (p *AttachOptions) GetContainerName(pod *api.Pod) string {
	return p.GetContainer(pod).Name
}
<commit_msg>Print how to reattch to seesion when session ends<commit_after>/*
Copyright 2014 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 cmd

import (
	"fmt"
	"io"
	"net/url"
	"os"
	"os/signal"
	"syscall"

	"github.com/docker/docker/pkg/term"
	"github.com/golang/glog"
	"github.com/spf13/cobra"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/remotecommand"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	utilerrors "k8s.io/kubernetes/pkg/util/errors"
)

const (
	attach_example = `# Get output from running pod 123456-7890, using the first container by default
$ kubectl attach 123456-7890

# Get output from ruby-container from pod 123456-7890
$ kubectl attach 123456-7890 -c ruby-container

# Switch to raw terminal mode, sends stdin to 'bash' in ruby-container from pod 123456-7890
# and sends stdout/stderr from 'bash' back to the client
$ kubectl attach 123456-7890 -c ruby-container -i -t`
)

func NewCmdAttach(f *cmdutil.Factory, cmdIn io.Reader, cmdOut, cmdErr io.Writer) *cobra.Command {
	options := &AttachOptions{
		In:  cmdIn,
		Out: cmdOut,
		Err: cmdErr,

		Attach: &DefaultRemoteAttach{},
	}
	cmd := &cobra.Command{
		Use:     "attach POD -c CONTAINER",
		Short:   "Attach to a running container.",
		Long:    "Attach to a process that is already running inside an existing container.",
		Example: attach_example,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(options.Complete(f, cmd, args))
			cmdutil.CheckErr(options.Validate())
			cmdutil.CheckErr(options.Run())
		},
	}
	// TODO support UID
	cmd.Flags().StringVarP(&options.ContainerName, "container", "c", "", "Container name. If omitted, the first container in the pod will be chosen")
	cmd.Flags().BoolVarP(&options.Stdin, "stdin", "i", false, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&options.TTY, "tty", "t", false, "Stdin is a TTY")
	return cmd
}

// RemoteAttach defines the interface accepted by the Attach command - provided for test stubbing
type RemoteAttach interface {
	Attach(method string, url *url.URL, config *client.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool) error
}

// DefaultRemoteAttach is the standard implementation of attaching
type DefaultRemoteAttach struct{}

func (*DefaultRemoteAttach) Attach(method string, url *url.URL, config *client.Config, stdin io.Reader, stdout, stderr io.Writer, tty bool) error {
	exec, err := remotecommand.NewExecutor(config, method, url)
	if err != nil {
		return err
	}
	return exec.Stream(stdin, stdout, stderr, tty)
}

// AttachOptions declare the arguments accepted by the Exec command
type AttachOptions struct {
	Namespace     string
	PodName       string
	ContainerName string
	Stdin         bool
	TTY           bool

	In  io.Reader
	Out io.Writer
	Err io.Writer

	Attach RemoteAttach
	Client *client.Client
	Config *client.Config
}

// Complete verifies command line arguments and loads data from the command environment
func (p *AttachOptions) Complete(f *cmdutil.Factory, cmd *cobra.Command, argsIn []string) error {
	if len(argsIn) == 0 {
		return cmdutil.UsageError(cmd, "POD is required for attach")
	}
	if len(argsIn) > 1 {
		return cmdutil.UsageError(cmd, fmt.Sprintf("expected a single argument: POD, saw %d: %s", len(argsIn), argsIn))
	}
	p.PodName = argsIn[0]

	namespace, _, err := f.DefaultNamespace()
	if err != nil {
		return err
	}
	p.Namespace = namespace

	config, err := f.ClientConfig()
	if err != nil {
		return err
	}
	p.Config = config

	client, err := f.Client()
	if err != nil {
		return err
	}
	p.Client = client

	return nil
}

// Validate checks that the provided attach options are specified.
func (p *AttachOptions) Validate() error {
	allErrs := []error{}
	if len(p.PodName) == 0 {
		allErrs = append(allErrs, fmt.Errorf("pod name must be specified"))
	}
	if p.Out == nil || p.Err == nil {
		allErrs = append(allErrs, fmt.Errorf("both output and error output must be provided"))
	}
	if p.Attach == nil || p.Client == nil || p.Config == nil {
		allErrs = append(allErrs, fmt.Errorf("client, client config, and attach must be provided"))
	}
	return utilerrors.NewAggregate(allErrs)
}

// Run executes a validated remote execution against a pod.
func (p *AttachOptions) Run() error {
	pod, err := p.Client.Pods(p.Namespace).Get(p.PodName)
	if err != nil {
		return err
	}

	if pod.Status.Phase != api.PodRunning {
		return fmt.Errorf("pod %s is not running and cannot be attached to; current phase is %s", p.PodName, pod.Status.Phase)
	}

	var stdin io.Reader
	tty := p.TTY

	containerToAttach := p.GetContainer(pod)
	if tty && !containerToAttach.TTY {
		tty = false
		fmt.Fprintf(p.Err, "Unable to use a TTY - container %s doesn't allocate one\n", containerToAttach.Name)
	}

	// TODO: refactor with terminal helpers from the edit utility once that is merged
	if p.Stdin {
		stdin = p.In
		if tty {
			if file, ok := stdin.(*os.File); ok {
				inFd := file.Fd()
				if term.IsTerminal(inFd) {
					oldState, err := term.SetRawTerminal(inFd)
					if err != nil {
						glog.Fatal(err)
					}
					fmt.Fprintln(p.Out, "\nHit enter for command prompt")
					// this handles a clean exit, where the command finished
					defer term.RestoreTerminal(inFd, oldState)

					// SIGINT is handled by term.SetRawTerminal (it runs a goroutine that listens
					// for SIGINT and restores the terminal before exiting)

					// this handles SIGTERM
					sigChan := make(chan os.Signal, 1)
					signal.Notify(sigChan, syscall.SIGTERM)
					go func() {
						<-sigChan
						term.RestoreTerminal(inFd, oldState)
						os.Exit(0)
					}()
				} else {
					fmt.Fprintln(p.Err, "STDIN is not a terminal")
				}
			} else {
				tty = false
				fmt.Fprintln(p.Err, "Unable to use a TTY - input is not the right kind of file")
			}
		}
	}

	// TODO: consider abstracting into a client invocation or client helper
	req := p.Client.RESTClient.Post().
		Resource("pods").
		Name(pod.Name).
		Namespace(pod.Namespace).
		SubResource("attach")
	req.VersionedParams(&api.PodAttachOptions{
		Container: containerToAttach.Name,
		Stdin:     stdin != nil,
		Stdout:    p.Out != nil,
		Stderr:    p.Err != nil,
		TTY:       tty,
	}, api.Scheme)

	err = p.Attach.Attach("POST", req.URL(), p.Config, stdin, p.Out, p.Err, tty)
	if err != nil {
		return err
	}
	if p.Stdin && tty && pod.Spec.RestartPolicy == api.RestartPolicyAlways {
		fmt.Fprintf(p.Out, "Session ended, resume using 'kubectl attach %s -c %s -i -t' command when the pod is running\n", pod.Name, containerToAttach.Name)
	}
	return nil
}

// GetContainer returns the container to attach to, with a fallback.
func (p *AttachOptions) GetContainer(pod *api.Pod) api.Container {
	if len(p.ContainerName) > 0 {
		for _, container := range pod.Spec.Containers {
			if container.Name == p.ContainerName {
				return container
			}
		}
	}

	glog.V(4).Infof("defaulting container name to %s", pod.Spec.Containers[0].Name)
	return pod.Spec.Containers[0]
}

// GetContainerName returns the name of the container to attach to, with a fallback.
func (p *AttachOptions) GetContainerName(pod *api.Pod) string {
	return p.GetContainer(pod).Name
}
<|endoftext|>


<commit_before>package proxystorage

import (
	"context"
	"fmt"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/sirupsen/logrus"

	"github.com/jacksontj/promxy/pkg/promhttputil"
	"github.com/jacksontj/promxy/pkg/remote"

	proxyconfig "github.com/jacksontj/promxy/pkg/config"
	"github.com/jacksontj/promxy/pkg/promclient"
	"github.com/jacksontj/promxy/pkg/proxyquerier"
	"github.com/jacksontj/promxy/pkg/servergroup"
)

type proxyStorageState struct {
	sgs            []*servergroup.ServerGroup
	client         promclient.API
	cfg            *proxyconfig.PromxyConfig
	remoteStorage  *remote.Storage
	appender       storage.Appender
	appenderCloser func() error
}

// Ready blocks until all servergroups are ready
func (p *proxyStorageState) Ready() {
	for _, sg := range p.sgs {
		<-sg.Ready
	}
}

// Cancel this state
func (p *proxyStorageState) Cancel(n *proxyStorageState) {
	if p.sgs != nil {
		for _, sg := range p.sgs {
			sg.Cancel()
		}
	}
	// We call close if the new one is nil, or if the appanders don't match
	if n == nil || p.appender != n.appender {
		if p.appenderCloser != nil {
			p.appenderCloser()
		}
	}
}

// NewProxyStorage creates a new ProxyStorage
func NewProxyStorage() (*ProxyStorage, error) {
	return &ProxyStorage{}, nil
}

// ProxyStorage implements prometheus' Storage interface
type ProxyStorage struct {
	state atomic.Value
}

// GetState returns the current state of the ProxyStorage
func (p *ProxyStorage) GetState() *proxyStorageState {
	tmp := p.state.Load()
	if sg, ok := tmp.(*proxyStorageState); ok {
		return sg
	}
	return &proxyStorageState{}
}

// ApplyConfig updates the current state of this ProxyStorage
func (p *ProxyStorage) ApplyConfig(c *proxyconfig.Config) error {
	oldState := p.GetState() // Fetch the old state

	failed := false

	apis := make([]promclient.API, len(c.ServerGroups))
	newState := &proxyStorageState{
		sgs: make([]*servergroup.ServerGroup, len(c.ServerGroups)),
		cfg: &c.PromxyConfig,
	}
	for i, sgCfg := range c.ServerGroups {
		tmp := servergroup.New()
		if err := tmp.ApplyConfig(sgCfg); err != nil {
			failed = true
			logrus.Errorf("Error applying config to server group: %s", err)
		}
		newState.sgs[i] = tmp
		apis[i] = tmp
	}
	newState.client = promclient.NewMultiAPI(apis, model.TimeFromUnix(0), nil, len(apis))

	if failed {
		newState.Cancel(nil)
		return fmt.Errorf("Error Applying Config to one or more server group(s)")
	}

	// Check for remote_write (for appender)
	if c.PromConfig.RemoteWriteConfigs != nil {
		if oldState.remoteStorage != nil {
			if err := oldState.remoteStorage.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			// if it was an appenderstub we just need to replace
		} else {
			// TODO: configure path?
			remote := remote.NewStorage(nil, func() (int64, error) { return 0, nil }, 1*time.Second)
			if err := remote.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.remoteStorage = remote
			var err error
			newState.appender, err = remote.Appender()
			if err != nil {
				return errors.Wrap(err, "unable to create remote_write appender")
			}
			newState.appenderCloser = remote.Close
		}
	} else {
		newState.appender = &appenderStub{}
	}

	newState.Ready()        // Wait for the newstate to be ready
	p.state.Store(newState) // Store the new state
	if oldState != nil && oldState.appender != newState.appender {
		oldState.Cancel(newState) // Cancel the old one
	}

	return nil
}

// Querier returns a new Querier on the storage.
func (p *ProxyStorage) Querier(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
	state := p.GetState()
	return &proxyquerier.ProxyQuerier{
		ctx,
		timestamp.Time(mint),
		timestamp.Time(maxt),
		state.client,

		state.cfg,
	}, nil
}

// StartTime returns the oldest timestamp stored in the storage.
func (p *ProxyStorage) StartTime() (int64, error) {
	return 0, nil
}

// Appender returns a new appender against the storage.
func (p *ProxyStorage) Appender() (storage.Appender, error) {
	state := p.GetState()
	return state.appender, nil
}

// Close releases the resources of the Querier.
func (p *ProxyStorage) Close() error { return nil }

// NodeReplacer replaces promql Nodes with more efficient-to-fetch ones. This works by taking lower-layer
// chunks of the query, farming them out to prometheus hosts, then stitching the results back together.
// An example would be a sum, we can sum multiple sums and come up with the same result -- so we do.
// There are a few ground rules for this:
//      - Children cannot be AggregateExpr: aggregates have their own combining logic, so its not safe to send a subquery with additional aggregations
//      - offsets within the subtree must match: if they don't then we'll get mismatched data, so we wait until we are far enough down the tree that they converge
//      - Don't reduce accuracy/granularity: the intention of this is to get the correct data faster, meaning correctness overrules speed.
func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *promql.EvalStmt, node promql.Node) (promql.Node, error) {

	isAgg := func(node promql.Node) bool {
		_, ok := node.(*promql.AggregateExpr)
		return ok
	}

	// If there is a child that is an aggregator we cannot do anything (as they have their own
	// rules around combining). We'll skip this node and let a lower layer take this on
	aggFinder := &BooleanFinder{Func: isAgg}
	offsetFinder := &OffsetFinder{}

	visitor := &MultiVisitor{[]promql.Visitor{aggFinder, offsetFinder}}

	if _, err := promql.Walk(ctx, visitor, s, node, nil, nil); err != nil {
		return nil, err
	}

	if aggFinder.Found > 0 {
		// If there was a single agg and that was us, then we're okay
		if !(isAgg(node) && aggFinder.Found == 1) {
			return nil, nil
		}
	}

	// If the tree below us is not all the same offset, then we can't do anything below -- we'll need
	// to wait until further in execution where they all match
	var offset time.Duration

	// If we couldn't find an offset, then something is wrong-- lets skip
	// Also if there was an error, skip
	if !offsetFinder.Found || offsetFinder.Error != nil {
		return nil, nil
	}
	offset = offsetFinder.Offset

	// Function to recursivelt remove offset. This is needed as we're using
	// the node API to String() the query to downstreams. Promql's iterators require
	// that the time be the absolute time, whereas the API returns them based on the
	// range you ask for (with the offset being implicit)
	// TODO: rename
	removeOffset := func() error {
		_, err := promql.Walk(ctx, &OffsetRemover{}, s, node, nil, nil)
		return err
	}

	state := p.GetState()
	switch n := node.(type) {
	// Some AggregateExprs can be composed (meaning they are "reentrant". If the aggregation op
	// is reentrant/composable then we'll do so, otherwise we let it fall through to normal query mechanisms
	case *promql.AggregateExpr:
		logrus.Debugf("AggregateExpr %v", n)

		var result model.Value
		var warnings api.Warnings
		var err error

		// Not all Aggregation functions are composable, so we'll do what we can
		switch n.Op {
		// All "reentrant" cases (meaning they can be done repeatedly and the outcome doesn't change)
		case promql.ItemSum, promql.ItemMin, promql.ItemMax, promql.ItemTopK, promql.ItemBottomK:
			removeOffset()

			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

		// Convert avg into sum() / count()
		case promql.ItemAvg:
			// Replace with sum() / count()
			return &promql.BinaryExpr{
				Op: promql.ItemDIV,
				LHS: &promql.AggregateExpr{
					Op:       promql.ItemSum,
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},

				RHS: &promql.AggregateExpr{
					Op:       promql.ItemCount,
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},
				VectorMatching: &promql.VectorMatching{Card: promql.CardOneToOne},
			}, nil

		// For count we simply need to change this to a sum over the data we get back
		case promql.ItemCount:
			removeOffset()

			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}
			n.Op = promql.ItemSum

			// To aggregate count_values we simply sum(count_values(key, metric)) by (key)
		case promql.ItemCountValues:

			// First we must fetch the data into a vectorselector
			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series, promhttputil.WarningsConvert(warnings))

			// Replace with sum(count_values()) BY (label)
			return &promql.AggregateExpr{
				Op:       promql.ItemSum,
				Expr:     ret,
				Grouping: []string{n.Param.(*promql.StringLiteral).Val},
			}, nil

		case promql.ItemQuantile:
			// DO NOTHING
			// this caltulates an actual quantile over the resulting data
			// as such there is no way to reduce the load necessary here. If
			// the query is something like quantile(sum(foo)) then the inner aggregation
			// will reduce the required data

		// Both of these cases require some mechanism of knowing what labels to do the aggregation on.
		// WIthout that knowledge we require pulling all of the data in, so we do nothing
		case promql.ItemStddev:
			// DO NOTHING
		case promql.ItemStdvar:
			// DO NOTHING

		}

		if result != nil {
			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series, promhttputil.WarningsConvert(warnings))
			n.Expr = ret
			return n, nil
		}

	// Call is for things such as rate() etc. This can be sent directly to the
	// prometheus node to answer
	case *promql.Call:
		logrus.Debugf("call %v %v", n, n.Type())
		removeOffset()

		var result model.Value
		var warnings api.Warnings
		var err error
		if s.Interval > 0 {
			result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
				Start: s.Start.Add(-offset - promql.LookbackDelta),
				End:   s.End.Add(-offset),
				Step:  s.Interval,
			})
		} else {
			result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
		}

		if err != nil {
			return nil, errors.Cause(err)
		}
		iterators := promclient.IteratorsForValue(result)
		series := make([]storage.Series, len(iterators))
		for i, iterator := range iterators {
			series[i] = &proxyquerier.Series{iterator}
		}

		ret := &promql.VectorSelector{Offset: offset}
		ret.SetSeries(series, promhttputil.WarningsConvert(warnings))
		return ret, nil

	// If we are simply fetching a Vector then we can fetch the data using the same step that
	// the query came in as (reducing the amount of data we need to fetch)
	// If we are simply fetching data, we skip here to let it fall through to the normal
	// storage API
	case *promql.VectorSelector:
		// Do Nothing
		return nil, nil

	// If we hit this someone is asking for a matrix directly, if so then we don't
	// have anyway to ask for less-- since this is exactly what they are asking for
	case *promql.MatrixSelector:
		// DO NOTHING

	default:
		logrus.Debugf("default %v %s", n, reflect.TypeOf(n))

	}
	return nil, nil
}
<commit_msg>Add NodeReplacer for VectorSelector<commit_after>package proxystorage

import (
	"context"
	"fmt"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/api"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/timestamp"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/storage"
	"github.com/sirupsen/logrus"

	"github.com/jacksontj/promxy/pkg/promhttputil"
	"github.com/jacksontj/promxy/pkg/remote"

	proxyconfig "github.com/jacksontj/promxy/pkg/config"
	"github.com/jacksontj/promxy/pkg/promclient"
	"github.com/jacksontj/promxy/pkg/proxyquerier"
	"github.com/jacksontj/promxy/pkg/servergroup"
)

type proxyStorageState struct {
	sgs            []*servergroup.ServerGroup
	client         promclient.API
	cfg            *proxyconfig.PromxyConfig
	remoteStorage  *remote.Storage
	appender       storage.Appender
	appenderCloser func() error
}

// Ready blocks until all servergroups are ready
func (p *proxyStorageState) Ready() {
	for _, sg := range p.sgs {
		<-sg.Ready
	}
}

// Cancel this state
func (p *proxyStorageState) Cancel(n *proxyStorageState) {
	if p.sgs != nil {
		for _, sg := range p.sgs {
			sg.Cancel()
		}
	}
	// We call close if the new one is nil, or if the appanders don't match
	if n == nil || p.appender != n.appender {
		if p.appenderCloser != nil {
			p.appenderCloser()
		}
	}
}

// NewProxyStorage creates a new ProxyStorage
func NewProxyStorage() (*ProxyStorage, error) {
	return &ProxyStorage{}, nil
}

// ProxyStorage implements prometheus' Storage interface
type ProxyStorage struct {
	state atomic.Value
}

// GetState returns the current state of the ProxyStorage
func (p *ProxyStorage) GetState() *proxyStorageState {
	tmp := p.state.Load()
	if sg, ok := tmp.(*proxyStorageState); ok {
		return sg
	}
	return &proxyStorageState{}
}

// ApplyConfig updates the current state of this ProxyStorage
func (p *ProxyStorage) ApplyConfig(c *proxyconfig.Config) error {
	oldState := p.GetState() // Fetch the old state

	failed := false

	apis := make([]promclient.API, len(c.ServerGroups))
	newState := &proxyStorageState{
		sgs: make([]*servergroup.ServerGroup, len(c.ServerGroups)),
		cfg: &c.PromxyConfig,
	}
	for i, sgCfg := range c.ServerGroups {
		tmp := servergroup.New()
		if err := tmp.ApplyConfig(sgCfg); err != nil {
			failed = true
			logrus.Errorf("Error applying config to server group: %s", err)
		}
		newState.sgs[i] = tmp
		apis[i] = tmp
	}
	newState.client = promclient.NewMultiAPI(apis, model.TimeFromUnix(0), nil, len(apis))

	if failed {
		newState.Cancel(nil)
		return fmt.Errorf("Error Applying Config to one or more server group(s)")
	}

	// Check for remote_write (for appender)
	if c.PromConfig.RemoteWriteConfigs != nil {
		if oldState.remoteStorage != nil {
			if err := oldState.remoteStorage.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			// if it was an appenderstub we just need to replace
		} else {
			// TODO: configure path?
			remote := remote.NewStorage(nil, func() (int64, error) { return 0, nil }, 1*time.Second)
			if err := remote.ApplyConfig(&c.PromConfig); err != nil {
				return err
			}
			newState.remoteStorage = remote
			var err error
			newState.appender, err = remote.Appender()
			if err != nil {
				return errors.Wrap(err, "unable to create remote_write appender")
			}
			newState.appenderCloser = remote.Close
		}
	} else {
		newState.appender = &appenderStub{}
	}

	newState.Ready()        // Wait for the newstate to be ready
	p.state.Store(newState) // Store the new state
	if oldState != nil && oldState.appender != newState.appender {
		oldState.Cancel(newState) // Cancel the old one
	}

	return nil
}

// Querier returns a new Querier on the storage.
func (p *ProxyStorage) Querier(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
	state := p.GetState()
	return &proxyquerier.ProxyQuerier{
		ctx,
		timestamp.Time(mint),
		timestamp.Time(maxt),
		state.client,

		state.cfg,
	}, nil
}

// StartTime returns the oldest timestamp stored in the storage.
func (p *ProxyStorage) StartTime() (int64, error) {
	return 0, nil
}

// Appender returns a new appender against the storage.
func (p *ProxyStorage) Appender() (storage.Appender, error) {
	state := p.GetState()
	return state.appender, nil
}

// Close releases the resources of the Querier.
func (p *ProxyStorage) Close() error { return nil }

// NodeReplacer replaces promql Nodes with more efficient-to-fetch ones. This works by taking lower-layer
// chunks of the query, farming them out to prometheus hosts, then stitching the results back together.
// An example would be a sum, we can sum multiple sums and come up with the same result -- so we do.
// There are a few ground rules for this:
//      - Children cannot be AggregateExpr: aggregates have their own combining logic, so its not safe to send a subquery with additional aggregations
//      - offsets within the subtree must match: if they don't then we'll get mismatched data, so we wait until we are far enough down the tree that they converge
//      - Don't reduce accuracy/granularity: the intention of this is to get the correct data faster, meaning correctness overrules speed.
func (p *ProxyStorage) NodeReplacer(ctx context.Context, s *promql.EvalStmt, node promql.Node) (promql.Node, error) {

	isAgg := func(node promql.Node) bool {
		_, ok := node.(*promql.AggregateExpr)
		return ok
	}

	// If there is a child that is an aggregator we cannot do anything (as they have their own
	// rules around combining). We'll skip this node and let a lower layer take this on
	aggFinder := &BooleanFinder{Func: isAgg}
	offsetFinder := &OffsetFinder{}

	visitor := &MultiVisitor{[]promql.Visitor{aggFinder, offsetFinder}}

	if _, err := promql.Walk(ctx, visitor, s, node, nil, nil); err != nil {
		return nil, err
	}

	if aggFinder.Found > 0 {
		// If there was a single agg and that was us, then we're okay
		if !(isAgg(node) && aggFinder.Found == 1) {
			return nil, nil
		}
	}

	// If the tree below us is not all the same offset, then we can't do anything below -- we'll need
	// to wait until further in execution where they all match
	var offset time.Duration

	// If we couldn't find an offset, then something is wrong-- lets skip
	// Also if there was an error, skip
	if !offsetFinder.Found || offsetFinder.Error != nil {
		return nil, nil
	}
	offset = offsetFinder.Offset

	// Function to recursivelt remove offset. This is needed as we're using
	// the node API to String() the query to downstreams. Promql's iterators require
	// that the time be the absolute time, whereas the API returns them based on the
	// range you ask for (with the offset being implicit)
	// TODO: rename
	removeOffset := func() error {
		_, err := promql.Walk(ctx, &OffsetRemover{}, s, node, nil, nil)
		return err
	}

	state := p.GetState()
	switch n := node.(type) {
	// Some AggregateExprs can be composed (meaning they are "reentrant". If the aggregation op
	// is reentrant/composable then we'll do so, otherwise we let it fall through to normal query mechanisms
	case *promql.AggregateExpr:
		logrus.Debugf("AggregateExpr %v", n)

		var result model.Value
		var warnings api.Warnings
		var err error

		// Not all Aggregation functions are composable, so we'll do what we can
		switch n.Op {
		// All "reentrant" cases (meaning they can be done repeatedly and the outcome doesn't change)
		case promql.ItemSum, promql.ItemMin, promql.ItemMax, promql.ItemTopK, promql.ItemBottomK:
			removeOffset()

			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

		// Convert avg into sum() / count()
		case promql.ItemAvg:
			// Replace with sum() / count()
			return &promql.BinaryExpr{
				Op: promql.ItemDIV,
				LHS: &promql.AggregateExpr{
					Op:       promql.ItemSum,
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},

				RHS: &promql.AggregateExpr{
					Op:       promql.ItemCount,
					Expr:     n.Expr,
					Param:    n.Param,
					Grouping: n.Grouping,
					Without:  n.Without,
				},
				VectorMatching: &promql.VectorMatching{Card: promql.CardOneToOne},
			}, nil

		// For count we simply need to change this to a sum over the data we get back
		case promql.ItemCount:
			removeOffset()

			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}
			n.Op = promql.ItemSum

			// To aggregate count_values we simply sum(count_values(key, metric)) by (key)
		case promql.ItemCountValues:

			// First we must fetch the data into a vectorselector
			if s.Interval > 0 {
				result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
					Start: s.Start.Add(-offset - promql.LookbackDelta),
					End:   s.End.Add(-offset),
					Step:  s.Interval,
				})
			} else {
				result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
			}

			if err != nil {
				return nil, errors.Cause(err)
			}

			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series, promhttputil.WarningsConvert(warnings))

			// Replace with sum(count_values()) BY (label)
			return &promql.AggregateExpr{
				Op:       promql.ItemSum,
				Expr:     ret,
				Grouping: []string{n.Param.(*promql.StringLiteral).Val},
			}, nil

		case promql.ItemQuantile:
			// DO NOTHING
			// this caltulates an actual quantile over the resulting data
			// as such there is no way to reduce the load necessary here. If
			// the query is something like quantile(sum(foo)) then the inner aggregation
			// will reduce the required data

		// Both of these cases require some mechanism of knowing what labels to do the aggregation on.
		// WIthout that knowledge we require pulling all of the data in, so we do nothing
		case promql.ItemStddev:
			// DO NOTHING
		case promql.ItemStdvar:
			// DO NOTHING

		}

		if result != nil {
			iterators := promclient.IteratorsForValue(result)

			series := make([]storage.Series, len(iterators))
			for i, iterator := range iterators {
				series[i] = &proxyquerier.Series{iterator}
			}

			ret := &promql.VectorSelector{Offset: offset}
			ret.SetSeries(series, promhttputil.WarningsConvert(warnings))
			n.Expr = ret
			return n, nil
		}

	// Call is for things such as rate() etc. This can be sent directly to the
	// prometheus node to answer
	case *promql.Call:
		logrus.Debugf("call %v %v", n, n.Type())
		removeOffset()

		var result model.Value
		var warnings api.Warnings
		var err error
		if s.Interval > 0 {
			result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
				Start: s.Start.Add(-offset - promql.LookbackDelta),
				End:   s.End.Add(-offset),
				Step:  s.Interval,
			})
		} else {
			result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
		}

		if err != nil {
			return nil, errors.Cause(err)
		}
		iterators := promclient.IteratorsForValue(result)
		series := make([]storage.Series, len(iterators))
		for i, iterator := range iterators {
			series[i] = &proxyquerier.Series{iterator}
		}

		ret := &promql.VectorSelector{Offset: offset}
		ret.SetSeries(series, promhttputil.WarningsConvert(warnings))
		return ret, nil

	// If we are simply fetching a Vector then we can fetch the data using the same step that
	// the query came in as (reducing the amount of data we need to fetch)
	case *promql.VectorSelector:
		// If the vector selector already has the data we can skip
		if n.HasSeries() {
			return nil, nil
		}

		var result model.Value
		var warnings api.Warnings
		var err error
		if s.Interval > 0 {
			result, warnings, err = state.client.QueryRange(ctx, n.String(), v1.Range{
				Start: s.Start.Add(-offset - promql.LookbackDelta),
				End:   s.End.Add(-offset),
				Step:  s.Interval,
			})
		} else {
			result, warnings, err = state.client.Query(ctx, n.String(), s.Start.Add(-offset))
		}

		if err != nil {
			return nil, errors.Cause(err)
		}

		iterators := promclient.IteratorsForValue(result)
		series := make([]storage.Series, len(iterators))
		for i, iterator := range iterators {
			series[i] = &proxyquerier.Series{iterator}
		}
		n.SetSeries(series, promhttputil.WarningsConvert(warnings))

	// If we hit this someone is asking for a matrix directly, if so then we don't
	// have anyway to ask for less-- since this is exactly what they are asking for
	case *promql.MatrixSelector:
		// DO NOTHING

	default:
		logrus.Debugf("default %v %s", n, reflect.TypeOf(n))

	}
	return nil, nil
}
<|endoftext|>


<commit_before>package signal // import "github.com/docker/docker/pkg/signal"

import (
	"syscall"
	"testing"

	"gotest.tools/assert"
	is "gotest.tools/assert/cmp"
)

func TestParseSignal(t *testing.T) {
	_, checkAtoiError := ParseSignal("0")
	assert.Check(t, is.Error(checkAtoiError, "Invalid signal: 0"))

	_, error := ParseSignal("SIG")
	assert.Check(t, is.Error(error, "Invalid signal: SIG"))

	for sigStr := range SignalMap {
		responseSignal, error := ParseSignal(sigStr)
		assert.Check(t, error)
		signal := SignalMap[sigStr]
		assert.Check(t, is.DeepEqual(signal, responseSignal))
	}
}

func TestValidSignalForPlatform(t *testing.T) {
	isValidSignal := ValidSignalForPlatform(syscall.Signal(0))
	assert.Check(t, is.Equal(false, isValidSignal))

	for _, sigN := range SignalMap {
		isValidSignal = ValidSignalForPlatform(syscall.Signal(sigN))
		assert.Check(t, is.Equal(true, isValidSignal))
	}
}
<commit_msg>pkg/signal: remove unnescessary conversion (unconvert)<commit_after>package signal // import "github.com/docker/docker/pkg/signal"

import (
	"syscall"
	"testing"

	"gotest.tools/assert"
	is "gotest.tools/assert/cmp"
)

func TestParseSignal(t *testing.T) {
	_, checkAtoiError := ParseSignal("0")
	assert.Check(t, is.Error(checkAtoiError, "Invalid signal: 0"))

	_, error := ParseSignal("SIG")
	assert.Check(t, is.Error(error, "Invalid signal: SIG"))

	for sigStr := range SignalMap {
		responseSignal, error := ParseSignal(sigStr)
		assert.Check(t, error)
		signal := SignalMap[sigStr]
		assert.Check(t, is.DeepEqual(signal, responseSignal))
	}
}

func TestValidSignalForPlatform(t *testing.T) {
	isValidSignal := ValidSignalForPlatform(syscall.Signal(0))
	assert.Check(t, is.Equal(false, isValidSignal))

	for _, sigN := range SignalMap {
		isValidSignal = ValidSignalForPlatform(sigN)
		assert.Check(t, is.Equal(true, isValidSignal))
	}
}
<|endoftext|>


<commit_before>package command

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

var (
	copy CopyOptions
)

type CopyOptions struct {
	filerGrpcPort  *int
	master         *string
	include        *string
	replication    *string
	collection     *string
	ttl            *string
	maxMB          *int
	grpcDialOption grpc.DialOption
	masterClient   *wdclient.MasterClient
}

func init() {
	cmdCopy.Run = runCopy // break init cycle
	cmdCopy.IsDebug = cmdCopy.Flag.Bool("debug", false, "verbose debug information")
	copy.master = cmdCopy.Flag.String("master", "localhost:9333", "SeaweedFS master location")
	copy.include = cmdCopy.Flag.String("include", "", "pattens of files to copy, e.g., *.pdf, *.html, ab?d.txt, works together with -dir")
	copy.replication = cmdCopy.Flag.String("replication", "", "replication type")
	copy.collection = cmdCopy.Flag.String("collection", "", "optional collection name")
	copy.ttl = cmdCopy.Flag.String("ttl", "", "time to live, e.g.: 1m, 1h, 1d, 1M, 1y")
	copy.maxMB = cmdCopy.Flag.Int("maxMB", 32, "split files larger than the limit")
	copy.filerGrpcPort = cmdCopy.Flag.Int("filer.port.grpc", 0, "filer grpc server listen port, default to filer port + 10000")
}

var cmdCopy = &Command{
	UsageLine: "filer.copy file_or_dir1 [file_or_dir2 file_or_dir3] http://localhost:8888/path/to/a/folder/",
	Short:     "copy one or a list of files to a filer folder",
	Long: `copy one or a list of files, or batch copy one whole folder recursively, to a filer folder

  It can copy one or a list of files or folders.

  If copying a whole folder recursively:
  All files under the folder and subfolders will be copyed.
  Optional parameter "-include" allows you to specify the file name patterns.

  If "maxMB" is set to a positive number, files larger than it would be split into chunks.

  `,
}

func runCopy(cmd *Command, args []string) bool {

	weed_server.LoadConfiguration("security", false)

	if len(args) <= 1 {
		return false
	}
	filerDestination := args[len(args)-1]
	fileOrDirs := args[0 : len(args)-1]

	filerUrl, err := url.Parse(filerDestination)
	if err != nil {
		fmt.Printf("The last argument should be a URL on filer: %v\n", err)
		return false
	}
	urlPath := filerUrl.Path
	if !strings.HasSuffix(urlPath, "/") {
		fmt.Printf("The last argument should be a folder and end with \"/\": %v\n", err)
		return false
	}

	if filerUrl.Port() == "" {
		fmt.Printf("The filer port should be specified.\n")
		return false
	}

	filerPort, parseErr := strconv.ParseUint(filerUrl.Port(), 10, 64)
	if parseErr != nil {
		fmt.Printf("The filer port parse error: %v\n", parseErr)
		return false
	}

	filerGrpcPort := filerPort + 10000
	if *copy.filerGrpcPort != 0 {
		filerGrpcPort = uint64(*copy.filerGrpcPort)
	}

	filerGrpcAddress := fmt.Sprintf("%s:%d", filerUrl.Hostname(), filerGrpcPort)
	copy.grpcDialOption = security.LoadClientTLS(viper.Sub("grpc"), "client")

	copy.masterClient = wdclient.NewMasterClient(context.Background(), copy.grpcDialOption, "client", strings.Split(*copy.master, ","))
	go copy.masterClient.KeepConnectedToMaster()
	copy.masterClient.WaitUntilConnected()

	for _, fileOrDir := range fileOrDirs {
		if !doEachCopy(context.Background(), fileOrDir, filerUrl.Host, filerGrpcAddress, copy.grpcDialOption, urlPath) {
			return false
		}
	}
	return true
}

func doEachCopy(ctx context.Context, fileOrDir string, filerAddress, filerGrpcAddress string, grpcDialOption grpc.DialOption, path string) bool {
	f, err := os.Open(fileOrDir)
	if err != nil {
		fmt.Printf("Failed to open file %s: %v\n", fileOrDir, err)
		if _, ok := err.(*os.PathError); ok {
			fmt.Printf("skipping %s\n", fileOrDir)
			return true
		}
		return false
	}
	defer f.Close()

	fi, err := f.Stat()
	if err != nil {
		fmt.Printf("Failed to get stat for file %s: %v\n", fileOrDir, err)
		return false
	}

	mode := fi.Mode()
	if mode.IsDir() {
		files, _ := ioutil.ReadDir(fileOrDir)
		for _, subFileOrDir := range files {
			if !doEachCopy(ctx, fileOrDir+"/"+subFileOrDir.Name(), filerAddress, filerGrpcAddress, grpcDialOption, path+fi.Name()+"/") {
				return false
			}
		}
		return true
	}

	// this is a regular file
	if *copy.include != "" {
		if ok, _ := filepath.Match(*copy.include, filepath.Base(fileOrDir)); !ok {
			return true
		}
	}

	// find the chunk count
	chunkSize := int64(*copy.maxMB * 1024 * 1024)
	chunkCount := 1
	if chunkSize > 0 && fi.Size() > chunkSize {
		chunkCount = int(fi.Size()/chunkSize) + 1
	}

	if chunkCount == 1 {
		return uploadFileAsOne(ctx, filerAddress, filerGrpcAddress, grpcDialOption, path, f, fi)
	}

	return uploadFileInChunks(ctx, filerAddress, filerGrpcAddress, grpcDialOption, path, f, fi, chunkCount, chunkSize)
}

func uploadFileAsOne(ctx context.Context, filerAddress, filerGrpcAddress string, grpcDialOption grpc.DialOption, urlFolder string, f *os.File, fi os.FileInfo) bool {

	// upload the file content
	fileName := filepath.Base(f.Name())
	mimeType := detectMimeType(f)

	var chunks []*filer_pb.FileChunk

	if fi.Size() > 0 {

		// assign a volume
		assignResult, err := operation.Assign(copy.masterClient.GetMaster(), grpcDialOption, &operation.VolumeAssignRequest{
			Count:       1,
			Replication: *copy.replication,
			Collection:  *copy.collection,
			Ttl:         *copy.ttl,
		})
		if err != nil {
			fmt.Printf("Failed to assign from %s: %v\n", *copy.master, err)
		}

		targetUrl := "http://" + assignResult.Url + "/" + assignResult.Fid

		uploadResult, err := operation.Upload(targetUrl, fileName, f, false, mimeType, nil, assignResult.Auth)
		if err != nil {
			fmt.Printf("upload data %v to %s: %v\n", fileName, targetUrl, err)
			return false
		}
		if uploadResult.Error != "" {
			fmt.Printf("upload %v to %s result: %v\n", fileName, targetUrl, uploadResult.Error)
			return false
		}
		fmt.Printf("uploaded %s to %s\n", fileName, targetUrl)

		chunks = append(chunks, &filer_pb.FileChunk{
			FileId: assignResult.Fid,
			Offset: 0,
			Size:   uint64(uploadResult.Size),
			Mtime:  time.Now().UnixNano(),
			ETag:   uploadResult.ETag,
		})

		fmt.Printf("copied %s => http://%s%s%s\n", fileName, filerAddress, urlFolder, fileName)
	}

	if err := withFilerClient(ctx, filerGrpcAddress, grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		request := &filer_pb.CreateEntryRequest{
			Directory: urlFolder,
			Entry: &filer_pb.Entry{
				Name: fileName,
				Attributes: &filer_pb.FuseAttributes{
					Crtime:      time.Now().Unix(),
					Mtime:       time.Now().Unix(),
					Gid:         uint32(os.Getgid()),
					Uid:         uint32(os.Getuid()),
					FileSize:    uint64(fi.Size()),
					FileMode:    uint32(fi.Mode()),
					Mime:        mimeType,
					Replication: *copy.replication,
					Collection:  *copy.collection,
					TtlSec:      int32(util.ParseInt(*copy.ttl, 0)),
				},
				Chunks: chunks,
			},
		}

		if _, err := client.CreateEntry(ctx, request); err != nil {
			return fmt.Errorf("update fh: %v", err)
		}
		return nil
	}); err != nil {
		fmt.Printf("upload data %v to http://%s%s%s: %v\n", fileName, filerAddress, urlFolder, fileName, err)
		return false
	}

	return true
}

func uploadFileInChunks(ctx context.Context, filerAddress, filerGrpcAddress string, grpcDialOption grpc.DialOption, urlFolder string, f *os.File, fi os.FileInfo, chunkCount int, chunkSize int64) bool {

	fileName := filepath.Base(f.Name())
	mimeType := detectMimeType(f)

	var chunks []*filer_pb.FileChunk

	for i := int64(0); i < int64(chunkCount); i++ {

		// assign a volume
		assignResult, err := operation.Assign(copy.masterClient.GetMaster(), grpcDialOption, &operation.VolumeAssignRequest{
			Count:       1,
			Replication: *copy.replication,
			Collection:  *copy.collection,
			Ttl:         *copy.ttl,
		})
		if err != nil {
			fmt.Printf("Failed to assign from %s: %v\n", *copy.master, err)
		}

		targetUrl := "http://" + assignResult.Url + "/" + assignResult.Fid

		uploadResult, err := operation.Upload(targetUrl,
			fileName+"-"+strconv.FormatInt(i+1, 10),
			io.LimitReader(f, chunkSize),
			false, "application/octet-stream", nil, assignResult.Auth)
		if err != nil {
			fmt.Printf("upload data %v to %s: %v\n", fileName, targetUrl, err)
			return false
		}
		if uploadResult.Error != "" {
			fmt.Printf("upload %v to %s result: %v\n", fileName, targetUrl, uploadResult.Error)
			return false
		}
		chunks = append(chunks, &filer_pb.FileChunk{
			FileId: assignResult.Fid,
			Offset: i * chunkSize,
			Size:   uint64(uploadResult.Size),
			Mtime:  time.Now().UnixNano(),
			ETag:   uploadResult.ETag,
		})
		fmt.Printf("uploaded %s-%d to %s [%d,%d)\n", fileName, i+1, targetUrl, i*chunkSize, i*chunkSize+int64(uploadResult.Size))
	}

	if err := withFilerClient(ctx, filerGrpcAddress, grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		request := &filer_pb.CreateEntryRequest{
			Directory: urlFolder,
			Entry: &filer_pb.Entry{
				Name: fileName,
				Attributes: &filer_pb.FuseAttributes{
					Crtime:      time.Now().Unix(),
					Mtime:       time.Now().Unix(),
					Gid:         uint32(os.Getgid()),
					Uid:         uint32(os.Getuid()),
					FileSize:    uint64(fi.Size()),
					FileMode:    uint32(fi.Mode()),
					Mime:        mimeType,
					Replication: *copy.replication,
					Collection:  *copy.collection,
					TtlSec:      int32(util.ParseInt(*copy.ttl, 0)),
				},
				Chunks: chunks,
			},
		}

		if _, err := client.CreateEntry(ctx, request); err != nil {
			return fmt.Errorf("update fh: %v", err)
		}
		return nil
	}); err != nil {
		fmt.Printf("upload data %v to http://%s%s%s: %v\n", fileName, filerAddress, urlFolder, fileName, err)
		return false
	}

	fmt.Printf("copied %s => http://%s%s%s\n", fileName, filerAddress, urlFolder, fileName)

	return true
}

func detectMimeType(f *os.File) string {
	head := make([]byte, 512)
	f.Seek(0, io.SeekStart)
	n, err := f.Read(head)
	if err == io.EOF {
		return ""
	}
	if err != nil {
		fmt.Printf("read head of %v: %v\n", f.Name(), err)
		return "application/octet-stream"
	}
	f.Seek(0, io.SeekStart)
	mimeType := http.DetectContentType(head[:n])
	return mimeType
}

func withFilerClient(ctx context.Context, filerAddress string, grpcDialOption grpc.DialOption, fn func(filer_pb.SeaweedFilerClient) error) error {

	return util.WithCachedGrpcClient(ctx, func(clientConn *grpc.ClientConn) error {
		client := filer_pb.NewSeaweedFilerClient(clientConn)
		return fn(client)
	}, filerAddress, grpcDialOption)

}
<commit_msg>weed filer.copy: parallelize the copying to increase throughput<commit_after>package command

import (
	"context"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/security"
	"github.com/chrislusf/seaweedfs/weed/server"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/chrislusf/seaweedfs/weed/wdclient"
	"github.com/spf13/viper"
	"google.golang.org/grpc"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	copy      CopyOptions
	waitGroup sync.WaitGroup
)

type CopyOptions struct {
	filerGrpcPort  *int
	master         *string
	include        *string
	replication    *string
	collection     *string
	ttl            *string
	maxMB          *int
	grpcDialOption grpc.DialOption
	masterClient   *wdclient.MasterClient
	concurrency    *int
}

func init() {
	cmdCopy.Run = runCopy // break init cycle
	cmdCopy.IsDebug = cmdCopy.Flag.Bool("debug", false, "verbose debug information")
	copy.master = cmdCopy.Flag.String("master", "localhost:9333", "SeaweedFS master location")
	copy.include = cmdCopy.Flag.String("include", "", "pattens of files to copy, e.g., *.pdf, *.html, ab?d.txt, works together with -dir")
	copy.replication = cmdCopy.Flag.String("replication", "", "replication type")
	copy.collection = cmdCopy.Flag.String("collection", "", "optional collection name")
	copy.ttl = cmdCopy.Flag.String("ttl", "", "time to live, e.g.: 1m, 1h, 1d, 1M, 1y")
	copy.maxMB = cmdCopy.Flag.Int("maxMB", 32, "split files larger than the limit")
	copy.filerGrpcPort = cmdCopy.Flag.Int("filer.port.grpc", 0, "filer grpc server listen port, default to filer port + 10000")
	copy.concurrency = cmdCopy.Flag.Int("c", 8, "concurrent file copy goroutines")
}

var cmdCopy = &Command{
	UsageLine: "filer.copy file_or_dir1 [file_or_dir2 file_or_dir3] http://localhost:8888/path/to/a/folder/",
	Short:     "copy one or a list of files to a filer folder",
	Long: `copy one or a list of files, or batch copy one whole folder recursively, to a filer folder

  It can copy one or a list of files or folders.

  If copying a whole folder recursively:
  All files under the folder and subfolders will be copyed.
  Optional parameter "-include" allows you to specify the file name patterns.

  If "maxMB" is set to a positive number, files larger than it would be split into chunks.

  `,
}

func runCopy(cmd *Command, args []string) bool {

	weed_server.LoadConfiguration("security", false)

	if len(args) <= 1 {
		return false
	}
	filerDestination := args[len(args)-1]
	fileOrDirs := args[0 : len(args)-1]

	filerUrl, err := url.Parse(filerDestination)
	if err != nil {
		fmt.Printf("The last argument should be a URL on filer: %v\n", err)
		return false
	}
	urlPath := filerUrl.Path
	if !strings.HasSuffix(urlPath, "/") {
		fmt.Printf("The last argument should be a folder and end with \"/\": %v\n", err)
		return false
	}

	if filerUrl.Port() == "" {
		fmt.Printf("The filer port should be specified.\n")
		return false
	}

	filerPort, parseErr := strconv.ParseUint(filerUrl.Port(), 10, 64)
	if parseErr != nil {
		fmt.Printf("The filer port parse error: %v\n", parseErr)
		return false
	}

	filerGrpcPort := filerPort + 10000
	if *copy.filerGrpcPort != 0 {
		filerGrpcPort = uint64(*copy.filerGrpcPort)
	}

	filerGrpcAddress := fmt.Sprintf("%s:%d", filerUrl.Hostname(), filerGrpcPort)
	copy.grpcDialOption = security.LoadClientTLS(viper.Sub("grpc"), "client")

	copy.masterClient = wdclient.NewMasterClient(context.Background(), copy.grpcDialOption, "client", strings.Split(*copy.master, ","))
	go copy.masterClient.KeepConnectedToMaster()
	copy.masterClient.WaitUntilConnected()

	fileCopyTaskChan := make(chan FileCopyTask, *copy.concurrency)

	ctx := context.Background()

	go func() {
		defer close(fileCopyTaskChan)
		for _, fileOrDir := range fileOrDirs {
			if err := genFileCopyTask(fileOrDir, urlPath, fileCopyTaskChan); err != nil {
				fmt.Fprintf(os.Stderr, "gen file list error: %v\n", err)
				break
			}
		}
	}()
	for i := 0; i < *copy.concurrency; i++ {
		waitGroup.Add(1)
		go func() {
			defer waitGroup.Done()
			worker := FileCopyWorker{
				options:          &copy,
				filerHost:        filerUrl.Host,
				filerGrpcAddress: filerGrpcAddress,
			}
			if err := worker.copyFiles(ctx, fileCopyTaskChan); err != nil {
				fmt.Fprintf(os.Stderr, "copy file error: %v\n", err)
				return
			}
		}()
	}
	waitGroup.Wait()

	return true
}

func genFileCopyTask(fileOrDir string, destPath string, fileCopyTaskChan chan FileCopyTask) error {

	fi, err := os.Stat(fileOrDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to get stat for file %s: %v\n", fileOrDir, err)
		return nil
	}

	mode := fi.Mode()
	if mode.IsDir() {
		files, _ := ioutil.ReadDir(fileOrDir)
		for _, subFileOrDir := range files {
			if err = genFileCopyTask(fileOrDir+"/"+subFileOrDir.Name(), destPath+fi.Name()+"/", fileCopyTaskChan); err != nil {
				return err
			}
		}
		return nil
	}

	fileCopyTaskChan <- FileCopyTask{
		sourceLocation:     fileOrDir,
		destinationUrlPath: destPath,
		fileSize:           fi.Size(),
		fileMode:           fi.Mode(),
	}

	return nil
}

type FileCopyWorker struct {
	options          *CopyOptions
	filerHost        string
	filerGrpcAddress string
}

func (worker *FileCopyWorker) copyFiles(ctx context.Context, fileCopyTaskChan chan FileCopyTask) error {
	for task := range fileCopyTaskChan {
		if err := worker.doEachCopy(ctx, task); err != nil {
			return err
		}
	}
	return nil
}

type FileCopyTask struct {
	sourceLocation     string
	destinationUrlPath string
	fileSize           int64
	fileMode           os.FileMode
}

func (worker *FileCopyWorker) doEachCopy(ctx context.Context, task FileCopyTask) error {

	f, err := os.Open(task.sourceLocation)
	if err != nil {
		fmt.Printf("Failed to open file %s: %v\n", task.sourceLocation, err)
		if _, ok := err.(*os.PathError); ok {
			fmt.Printf("skipping %s\n", task.sourceLocation)
			return nil
		}
		return err
	}
	defer f.Close()

	// this is a regular file
	if *worker.options.include != "" {
		if ok, _ := filepath.Match(*worker.options.include, filepath.Base(task.sourceLocation)); !ok {
			return nil
		}
	}

	// find the chunk count
	chunkSize := int64(*worker.options.maxMB * 1024 * 1024)
	chunkCount := 1
	if chunkSize > 0 && task.fileSize > chunkSize {
		chunkCount = int(task.fileSize/chunkSize) + 1
	}

	if chunkCount == 1 {
		return worker.uploadFileAsOne(ctx, task, f)
	}

	return worker.uploadFileInChunks(ctx, task, f, chunkCount, chunkSize)
}

func (worker *FileCopyWorker) uploadFileAsOne(ctx context.Context, task FileCopyTask, f *os.File) error {

	// upload the file content
	fileName := filepath.Base(f.Name())
	mimeType := detectMimeType(f)

	var chunks []*filer_pb.FileChunk

	if task.fileSize > 0 {

		// assign a volume
		assignResult, err := operation.Assign(worker.options.masterClient.GetMaster(), worker.options.grpcDialOption, &operation.VolumeAssignRequest{
			Count:       1,
			Replication: *worker.options.replication,
			Collection:  *worker.options.collection,
			Ttl:         *worker.options.ttl,
		})
		if err != nil {
			fmt.Printf("Failed to assign from %s: %v\n", *worker.options.master, err)
		}

		targetUrl := "http://" + assignResult.Url + "/" + assignResult.Fid

		uploadResult, err := operation.Upload(targetUrl, fileName, f, false, mimeType, nil, assignResult.Auth)
		if err != nil {
			return fmt.Errorf("upload data %v to %s: %v\n", fileName, targetUrl, err)
		}
		if uploadResult.Error != "" {
			return fmt.Errorf("upload %v to %s result: %v\n", fileName, targetUrl, uploadResult.Error)
		}
		fmt.Printf("uploaded %s to %s\n", fileName, targetUrl)

		chunks = append(chunks, &filer_pb.FileChunk{
			FileId: assignResult.Fid,
			Offset: 0,
			Size:   uint64(uploadResult.Size),
			Mtime:  time.Now().UnixNano(),
			ETag:   uploadResult.ETag,
		})

		fmt.Printf("copied %s => http://%s%s%s\n", fileName, worker.filerHost, task.destinationUrlPath, fileName)
	}

	if err := withFilerClient(ctx, worker.filerGrpcAddress, worker.options.grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		request := &filer_pb.CreateEntryRequest{
			Directory: task.destinationUrlPath,
			Entry: &filer_pb.Entry{
				Name: fileName,
				Attributes: &filer_pb.FuseAttributes{
					Crtime:      time.Now().Unix(),
					Mtime:       time.Now().Unix(),
					Gid:         uint32(os.Getgid()),
					Uid:         uint32(os.Getuid()),
					FileSize:    uint64(task.fileSize),
					FileMode:    uint32(task.fileMode),
					Mime:        mimeType,
					Replication: *worker.options.replication,
					Collection:  *worker.options.collection,
					TtlSec:      int32(util.ParseInt(*worker.options.ttl, 0)),
				},
				Chunks: chunks,
			},
		}

		if _, err := client.CreateEntry(ctx, request); err != nil {
			return fmt.Errorf("update fh: %v", err)
		}
		return nil
	}); err != nil {
		return fmt.Errorf("upload data %v to http://%s%s%s: %v\n", fileName, worker.filerHost, task.destinationUrlPath, fileName, err)
	}

	return nil
}

func (worker *FileCopyWorker) uploadFileInChunks(ctx context.Context, task FileCopyTask, f *os.File, chunkCount int, chunkSize int64) error {

	fileName := filepath.Base(f.Name())
	mimeType := detectMimeType(f)

	var chunks []*filer_pb.FileChunk

	for i := int64(0); i < int64(chunkCount); i++ {

		// assign a volume
		assignResult, err := operation.Assign(worker.options.masterClient.GetMaster(), worker.options.grpcDialOption, &operation.VolumeAssignRequest{
			Count:       1,
			Replication: *worker.options.replication,
			Collection:  *worker.options.collection,
			Ttl:         *worker.options.ttl,
		})
		if err != nil {
			fmt.Printf("Failed to assign from %s: %v\n", *worker.options.master, err)
		}

		targetUrl := "http://" + assignResult.Url + "/" + assignResult.Fid

		uploadResult, err := operation.Upload(targetUrl,
			fileName+"-"+strconv.FormatInt(i+1, 10),
			io.LimitReader(f, chunkSize),
			false, "application/octet-stream", nil, assignResult.Auth)
		if err != nil {
			return fmt.Errorf("upload data %v to %s: %v\n", fileName, targetUrl, err)
		}
		if uploadResult.Error != "" {
			return fmt.Errorf("upload %v to %s result: %v\n", fileName, targetUrl, uploadResult.Error)
		}
		chunks = append(chunks, &filer_pb.FileChunk{
			FileId: assignResult.Fid,
			Offset: i * chunkSize,
			Size:   uint64(uploadResult.Size),
			Mtime:  time.Now().UnixNano(),
			ETag:   uploadResult.ETag,
		})
		fmt.Printf("uploaded %s-%d to %s [%d,%d)\n", fileName, i+1, targetUrl, i*chunkSize, i*chunkSize+int64(uploadResult.Size))
	}

	if err := withFilerClient(ctx, worker.filerGrpcAddress, worker.options.grpcDialOption, func(client filer_pb.SeaweedFilerClient) error {
		request := &filer_pb.CreateEntryRequest{
			Directory: task.destinationUrlPath,
			Entry: &filer_pb.Entry{
				Name: fileName,
				Attributes: &filer_pb.FuseAttributes{
					Crtime:      time.Now().Unix(),
					Mtime:       time.Now().Unix(),
					Gid:         uint32(os.Getgid()),
					Uid:         uint32(os.Getuid()),
					FileSize:    uint64(task.fileSize),
					FileMode:    uint32(task.fileMode),
					Mime:        mimeType,
					Replication: *worker.options.replication,
					Collection:  *worker.options.collection,
					TtlSec:      int32(util.ParseInt(*worker.options.ttl, 0)),
				},
				Chunks: chunks,
			},
		}

		if _, err := client.CreateEntry(ctx, request); err != nil {
			return fmt.Errorf("update fh: %v", err)
		}
		return nil
	}); err != nil {
		return fmt.Errorf("upload data %v to http://%s%s%s: %v\n", fileName, worker.filerHost, task.destinationUrlPath, fileName, err)
	}

	fmt.Printf("copied %s => http://%s%s%s\n", fileName, worker.filerHost, task.destinationUrlPath, fileName)

	return nil
}

func detectMimeType(f *os.File) string {
	head := make([]byte, 512)
	f.Seek(0, io.SeekStart)
	n, err := f.Read(head)
	if err == io.EOF {
		return ""
	}
	if err != nil {
		fmt.Printf("read head of %v: %v\n", f.Name(), err)
		return "application/octet-stream"
	}
	f.Seek(0, io.SeekStart)
	mimeType := http.DetectContentType(head[:n])
	return mimeType
}

func withFilerClient(ctx context.Context, filerAddress string, grpcDialOption grpc.DialOption, fn func(filer_pb.SeaweedFilerClient) error) error {

	return util.WithCachedGrpcClient(ctx, func(clientConn *grpc.ClientConn) error {
		client := filer_pb.NewSeaweedFilerClient(clientConn)
		return fn(client)
	}, filerAddress, grpcDialOption)

}
<|endoftext|>


<commit_before>package weed_server

import (
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/chrislusf/raft"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"github.com/gorilla/mux"
)

type RaftServer struct {
	peers      []string // initial peers to join with
	raftServer raft.Server
	dataDir    string
	httpAddr   string
	router     *mux.Router
	topo       *topology.Topology
}

func NewRaftServer(r *mux.Router, peers []string, httpAddr string, dataDir string, topo *topology.Topology, pulseSeconds int) *RaftServer {
	s := &RaftServer{
		peers:    peers,
		httpAddr: httpAddr,
		dataDir:  dataDir,
		router:   r,
		topo:     topo,
	}

	if glog.V(4) {
		raft.SetLogLevel(2)
	}

	raft.RegisterCommand(&topology.MaxVolumeIdCommand{})

	var err error
	transporter := raft.NewHTTPTransporter("/cluster", time.Second)
	transporter.Transport.MaxIdleConnsPerHost = 1024
	transporter.Transport.IdleConnTimeout = time.Second
	transporter.Transport.ResponseHeaderTimeout = time.Second
	glog.V(0).Infof("Starting RaftServer with %v", httpAddr)

	// Clear old cluster configurations if peers are changed
	if oldPeers, changed := isPeersChanged(s.dataDir, httpAddr, s.peers); changed {
		glog.V(0).Infof("Peers Change: %v => %v", oldPeers, s.peers)
		os.RemoveAll(path.Join(s.dataDir, "conf"))
		os.RemoveAll(path.Join(s.dataDir, "log"))
		os.RemoveAll(path.Join(s.dataDir, "snapshot"))
	}

	s.raftServer, err = raft.NewServer(s.httpAddr, s.dataDir, transporter, nil, topo, "")
	if err != nil {
		glog.V(0).Infoln(err)
		return nil
	}
	transporter.Install(s.raftServer, s)
	s.raftServer.SetHeartbeatInterval(500 * time.Millisecond)
	s.raftServer.SetElectionTimeout(time.Duration(pulseSeconds) * 500 * time.Millisecond)
	s.raftServer.Start()

	s.router.HandleFunc("/cluster/status", s.statusHandler).Methods("GET")

	for _, peer := range s.peers {
		s.raftServer.AddPeer(peer, "http://"+peer)
	}
	time.Sleep(time.Duration(1000+rand.Int31n(3000)) * time.Millisecond)
	if s.raftServer.IsLogEmpty() {
		// Initialize the server by joining itself.
		glog.V(0).Infoln("Initializing new cluster")

		_, err := s.raftServer.Do(&raft.DefaultJoinCommand{
			Name:             s.raftServer.Name(),
			ConnectionString: "http://" + s.httpAddr,
		})

		if err != nil {
			glog.V(0).Infoln(err)
			return nil
		}
	}

	glog.V(0).Infof("current cluster leader: %v", s.raftServer.Leader())

	return s
}

func (s *RaftServer) Peers() (members []string) {
	peers := s.raftServer.Peers()

	for _, p := range peers {
		members = append(members, strings.TrimPrefix(p.ConnectionString, "http://"))
	}

	return
}

func isPeersChanged(dir string, self string, peers []string) (oldPeers []string, changed bool) {
	confPath := path.Join(dir, "conf")
	// open conf file
	b, err := ioutil.ReadFile(confPath)
	if err != nil {
		return oldPeers, true
	}
	conf := &raft.Config{}
	if err = json.Unmarshal(b, conf); err != nil {
		return oldPeers, true
	}

	for _, p := range conf.Peers {
		oldPeers = append(oldPeers, strings.TrimPrefix(p.ConnectionString, "http://"))
	}
	oldPeers = append(oldPeers, self)

	if len(peers) == 0 && len(oldPeers) <= 1 {
		return oldPeers, false
	}

	sort.Strings(peers)
	sort.Strings(oldPeers)

	return oldPeers, !reflect.DeepEqual(peers, oldPeers)

}
<commit_msg>raft: use the first master to bootstrap the election<commit_after>package weed_server

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/chrislusf/raft"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/topology"
	"github.com/gorilla/mux"
)

type RaftServer struct {
	peers      []string // initial peers to join with
	raftServer raft.Server
	dataDir    string
	httpAddr   string
	router     *mux.Router
	topo       *topology.Topology
}

func NewRaftServer(r *mux.Router, peers []string, httpAddr string, dataDir string, topo *topology.Topology, pulseSeconds int) *RaftServer {
	s := &RaftServer{
		peers:    peers,
		httpAddr: httpAddr,
		dataDir:  dataDir,
		router:   r,
		topo:     topo,
	}

	if glog.V(4) {
		raft.SetLogLevel(2)
	}

	raft.RegisterCommand(&topology.MaxVolumeIdCommand{})

	var err error
	transporter := raft.NewHTTPTransporter("/cluster", time.Second)
	transporter.Transport.MaxIdleConnsPerHost = 1024
	transporter.Transport.IdleConnTimeout = time.Second
	transporter.Transport.ResponseHeaderTimeout = time.Second
	glog.V(0).Infof("Starting RaftServer with %v", httpAddr)

	// Clear old cluster configurations if peers are changed
	if oldPeers, changed := isPeersChanged(s.dataDir, httpAddr, s.peers); changed {
		glog.V(0).Infof("Peers Change: %v => %v", oldPeers, s.peers)
		os.RemoveAll(path.Join(s.dataDir, "conf"))
		os.RemoveAll(path.Join(s.dataDir, "log"))
		os.RemoveAll(path.Join(s.dataDir, "snapshot"))
	}

	s.raftServer, err = raft.NewServer(s.httpAddr, s.dataDir, transporter, nil, topo, "")
	if err != nil {
		glog.V(0).Infoln(err)
		return nil
	}
	transporter.Install(s.raftServer, s)
	s.raftServer.SetHeartbeatInterval(500 * time.Millisecond)
	s.raftServer.SetElectionTimeout(time.Duration(pulseSeconds) * 500 * time.Millisecond)
	s.raftServer.Start()

	s.router.HandleFunc("/cluster/status", s.statusHandler).Methods("GET")

	for _, peer := range s.peers {
		s.raftServer.AddPeer(peer, "http://"+peer)
	}

	if s.raftServer.IsLogEmpty() && isTheFirstOne(httpAddr, s.peers) {
		// Initialize the server by joining itself.
		glog.V(0).Infoln("Initializing new cluster")

		_, err := s.raftServer.Do(&raft.DefaultJoinCommand{
			Name:             s.raftServer.Name(),
			ConnectionString: "http://" + s.httpAddr,
		})

		if err != nil {
			glog.V(0).Infoln(err)
			return nil
		}
	}

	glog.V(0).Infof("current cluster leader: %v", s.raftServer.Leader())

	return s
}

func (s *RaftServer) Peers() (members []string) {
	peers := s.raftServer.Peers()

	for _, p := range peers {
		members = append(members, strings.TrimPrefix(p.ConnectionString, "http://"))
	}

	return
}

func isPeersChanged(dir string, self string, peers []string) (oldPeers []string, changed bool) {
	confPath := path.Join(dir, "conf")
	// open conf file
	b, err := ioutil.ReadFile(confPath)
	if err != nil {
		return oldPeers, true
	}
	conf := &raft.Config{}
	if err = json.Unmarshal(b, conf); err != nil {
		return oldPeers, true
	}

	for _, p := range conf.Peers {
		oldPeers = append(oldPeers, strings.TrimPrefix(p.ConnectionString, "http://"))
	}
	oldPeers = append(oldPeers, self)

	if len(peers) == 0 && len(oldPeers) <= 1 {
		return oldPeers, false
	}

	sort.Strings(peers)
	sort.Strings(oldPeers)

	return oldPeers, !reflect.DeepEqual(peers, oldPeers)

}

func isTheFirstOne(self string, peers []string) bool {
	sort.Strings(peers)
	if len(peers)<=0{
		return true
	}
	return self == peers[0]
}
<|endoftext|>