<commit_before>/*
Copyright 2017 The Nuclio Authors.

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

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

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

package rpc

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/nuclio/nuclio/pkg/common"
	"github.com/nuclio/nuclio/pkg/errors"
	"github.com/nuclio/nuclio/pkg/processor/runtime"
	"github.com/nuclio/nuclio/pkg/processor/status"

	"github.com/nuclio/logger"
	"github.com/nuclio/nuclio-sdk-go"
	"github.com/rs/xid"
)

// TODO: Find a better place (both on file system and configuration)
const (
	socketPathTemplate = "/tmp/nuclio-rpc-%s.sock"
	connectionTimeout  = 2 * time.Minute
)

type result struct {
	StatusCode   int                    `json:"status_code"`
	ContentType  string                 `json:"content_type"`
	Body         string                 `json:"body"`
	BodyEncoding string                 `json:"body_encoding"`
	Headers      map[string]interface{} `json:"headers"`

	DecodedBody []byte
	err         error
}

// Runtime is a runtime that communicates via unix domain socket
type AbstractRuntime struct {
	runtime.AbstractRuntime
	configuration  *runtime.Configuration
	eventEncoder   EventEncoder
	wrapperProcess *os.Process
	resultChan     chan *result
	functionLogger logger.Logger
	runtime        Runtime
	startChan      chan struct{}
	socketType     SocketType
}

type rpcLogRecord struct {
	DateTime string                 `json:"datetime"`
	Level    string                 `json:"level"`
	Message  string                 `json:"message"`
	With     map[string]interface{} `json:"with"`
}

// NewRPCRuntime returns a new RPC runtime
func NewAbstractRuntime(logger logger.Logger,
	configuration *runtime.Configuration,
	runtimeInstance Runtime) (*AbstractRuntime, error) {
	var err error

	abstractRuntime, err := runtime.NewAbstractRuntime(logger, configuration)
	if err != nil {
		return nil, errors.Wrap(err, "Can't create AbstractRuntime")
	}

	newRuntime := &AbstractRuntime{
		AbstractRuntime: *abstractRuntime,
		configuration:   configuration,
		runtime:         runtimeInstance,
		startChan:       make(chan struct{}, 1),
	}

	return newRuntime, nil
}

func (r *AbstractRuntime) Start() error {
	if err := r.startWrapper(); err != nil {
		r.SetStatus(status.Error)
		return errors.Wrap(err, "Failed to run wrapper")
	}

	r.SetStatus(status.Ready)
	return nil
}

// ProcessEvent processes an event
func (r *AbstractRuntime) ProcessEvent(event nuclio.Event, functionLogger logger.Logger) (interface{}, error) {

	if currentStatus := r.GetStatus(); currentStatus != status.Ready {
		return nil, errors.Errorf("Processor not ready (current status: %s)", currentStatus)
	}

	if currentStatus := r.GetStatus(); currentStatus != status.Ready {
		return nil, errors.Errorf("Processor not ready (current status: %s)", currentStatus)
	}

	r.functionLogger = functionLogger

	// We don't use defer to reset r.functionLogger since it decreases performance
	if err := r.eventEncoder.Encode(event); err != nil {
		r.functionLogger = nil
		return nil, errors.Wrapf(err, "Can't encode event: %+v", event)
	}

	result, ok := <-r.resultChan
	r.functionLogger = nil
	if !ok {
		msg := "Client disconnected"
		r.Logger.Error(msg)
		r.SetStatus(status.Error)
		r.functionLogger = nil
		return nil, errors.New(msg)
	}

	return nuclio.Response{
		Body:        result.DecodedBody,
		ContentType: result.ContentType,
		Headers:     result.Headers,
		StatusCode:  result.StatusCode,
	}, nil
}

// Stop stops the runtime
func (r *AbstractRuntime) Stop() error {
	if r.wrapperProcess != nil {
		err := r.wrapperProcess.Kill()
		if err != nil {
			r.SetStatus(status.Error)
			return errors.Wrap(err, "Can't kill wrapper process")
		}
	}

	r.SetStatus(status.Stopped)
	return nil
}

// Restart restarts the runtime
func (r *AbstractRuntime) Restart() error {
	if err := r.Stop(); err != nil {
		return err
	}

	// Send error for current event (non-blocking)
	select {
	case r.resultChan <- &result{
		StatusCode: http.StatusRequestTimeout,
		err:        errors.New("Runtime restarted"),
	}:

	default:
		r.Logger.Warn("Nothing waiting on result channel during restart. Continuing")
	}

	close(r.resultChan)
	if err := r.startWrapper(); err != nil {
		r.SetStatus(status.Error)
		return errors.Wrap(err, "Can't start wrapper process")
	}

	r.SetStatus(status.Ready)
	return nil
}

// GetSocketType returns the type of socket the runtime works with (unix/tcp)
func (r *AbstractRuntime) GetSocketType() SocketType {
	return UnixSocket
}

// WaitForStart returns whether the runtime supports sending an indication that it started
func (r *AbstractRuntime) WaitForStart() bool {
	return false
}

// SupportsRestart returns true if the runtime supports restart
func (r *AbstractRuntime) SupportsRestart() bool {
	return true
}

func (r *AbstractRuntime) startWrapper() error {
	var err error

	var listener net.Listener
	var address string

	if r.runtime.GetSocketType() == UnixSocket {
		listener, address, err = r.createUnixListener()
	} else {
		listener, address, err = r.createTCPListener()
	}

	if err != nil {
		return errors.Wrap(err, "Can't create listener")
	}

	wrapperProcess, err := r.runtime.RunWrapper(address)
	if err != nil {
		return errors.Wrap(err, "Can't run wrapper")
	}
	r.wrapperProcess = wrapperProcess
	go r.watchWrapperProcess()

	conn, err := listener.Accept()
	if err != nil {
		return errors.Wrap(err, "Can't get connection from wrapper")
	}

	r.Logger.Info("Wrapper connected")

	r.eventEncoder = r.runtime.GetEventEncoder(conn)
	r.resultChan = make(chan *result)
	go r.wrapperOutputHandler(conn, r.resultChan)

	// wait for start if required to
	if r.runtime.WaitForStart() {
		r.Logger.Debug("Waiting for start")

		<-r.startChan
	}

	r.Logger.Debug("Started")

	return nil
}

// Create a listener on unix domian docker, return listener, path to socket and error
func (r *AbstractRuntime) createUnixListener() (net.Listener, string, error) {
	socketPath := fmt.Sprintf(socketPathTemplate, xid.New().String())

	if common.FileExists(socketPath) {
		if err := os.Remove(socketPath); err != nil {
			return nil, "", errors.Wrapf(err, "Can't remove socket at %q", socketPath)
		}
	}

	r.Logger.DebugWith("Creating listener socket", "path", socketPath)

	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		return nil, "", errors.Wrapf(err, "Can't listen on %s", socketPath)
	}

	unixListener, ok := listener.(*net.UnixListener)
	if !ok {
		return nil, "", fmt.Errorf("Can't get underlying Unix listener")
	}

	if err = unixListener.SetDeadline(time.Now().Add(connectionTimeout)); err != nil {
		return nil, "", errors.Wrap(err, "Can't set deadline")
	}

	return listener, socketPath, nil
}

// Create a listener on TCP docker, return listener, port and error
func (r *AbstractRuntime) createTCPListener() (net.Listener, string, error) {
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, "", errors.Wrap(err, "Can't find free port")
	}

	tcpListener, ok := listener.(*net.TCPListener)
	if !ok {
		return nil, "", errors.Wrap(err, "Can't get underlying TCP listener")
	}
	if err = tcpListener.SetDeadline(time.Now().Add(connectionTimeout)); err != nil {
		return nil, "", errors.Wrap(err, "Can't set deadline")
	}

	port := listener.Addr().(*net.TCPAddr).Port

	return listener, fmt.Sprintf("%d", port), nil
}

func (r *AbstractRuntime) wrapperOutputHandler(conn io.Reader, resultChan chan *result) {

	// Reset might close outChan, which will cause panic when sending
	defer func() {
		if err := recover(); err != nil {
			r.Logger.WarnWith("panic handling wrapper output (Reset called?)")
		}
	}()

	outReader := bufio.NewReader(conn)

	// Read logs & output
	for {
		unmarshalledResult := &result{}
		var data []byte

		data, unmarshalledResult.err = outReader.ReadBytes('\n')

		if unmarshalledResult.err != nil {
			r.Logger.WarnWith("Failed to read from connection", "err", unmarshalledResult.err)
			resultChan <- unmarshalledResult
			continue
		}

		switch data[0] {
		case 'r':

			// try to unmarshall the result
			if unmarshalledResult.err = json.Unmarshal(data[1:], unmarshalledResult); unmarshalledResult.err != nil {
				r.resultChan <- unmarshalledResult
				continue
			}

			switch unmarshalledResult.BodyEncoding {
			case "text":
				unmarshalledResult.DecodedBody = []byte(unmarshalledResult.Body)
			case "base64":
				unmarshalledResult.DecodedBody, unmarshalledResult.err = base64.StdEncoding.DecodeString(unmarshalledResult.Body)
			default:
				unmarshalledResult.err = fmt.Errorf("Unknown body encoding - %q", unmarshalledResult.BodyEncoding)
			}

			// write back to result channel
			resultChan <- unmarshalledResult
		case 'm':
			r.handleReponseMetric(data[1:])
		case 'l':
			r.handleResponseLog(data[1:])
		case 's':
			r.handleStart()
		}
	}
}

func (r *AbstractRuntime) handleResponseLog(response []byte) {
	var logRecord rpcLogRecord

	if err := json.Unmarshal(response, &logRecord); err != nil {
		r.Logger.ErrorWith("Can't decode log", "error", err)
		return
	}

	logger := r.resolveFunctionLogger(r.functionLogger)
	logFunc := logger.DebugWith

	switch logRecord.Level {
	case "error", "critical", "fatal":
		logFunc = logger.ErrorWith
	case "warning":
		logFunc = logger.WarnWith
	case "info":
		logFunc = logger.InfoWith
	}

	vars := common.MapToSlice(logRecord.With)
	logFunc(logRecord.Message, vars...)
}

func (r *AbstractRuntime) handleReponseMetric(response []byte) {
	var metrics struct {
		DurationSec float64 `json:"duration"`
	}

	logger := r.resolveFunctionLogger(r.functionLogger)
	if err := json.Unmarshal(response, &metrics); err != nil {
		logger.ErrorWith("Can't decode metric", "error", err)
		return
	}

	if metrics.DurationSec == 0 {
		logger.ErrorWith("No duration in metrics", "metrics", metrics)
		return
	}

	r.Statistics.DurationMilliSecondsCount++
	r.Statistics.DurationMilliSecondsSum += uint64(metrics.DurationSec * 1000)
}

func (r *AbstractRuntime) handleStart() {
	r.startChan <- struct{}{}
}

// resolveFunctionLogger return either functionLogger if provided or root logger if not
func (r *AbstractRuntime) resolveFunctionLogger(functionLogger logger.Logger) logger.Logger {
	if functionLogger == nil {
		return r.Logger
	}
	return functionLogger
}

func (r *AbstractRuntime) newResultChan() {

	// We create the channel buffered so we won't block on sending
	r.resultChan = make(chan *result, 1)
}

func (r *AbstractRuntime) watchWrapperProcess() {
	procStatus, err := r.wrapperProcess.Wait()
	if r.GetStatus() == status.Ready && (err != nil || !procStatus.Success()) {
		r.Logger.ErrorWith("Unexpected termination of child process", "error", err, "status", procStatus.String())
	}
	r.SetStatus(status.Stopped)
	r.wrapperProcess = nil
	// TODO: Do we want to exit the processor here?
}
<commit_msg>Panic on runtime processWrapper crash (#1470)<commit_after>/*
Copyright 2017 The Nuclio Authors.

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

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

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

package rpc

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/nuclio/nuclio/pkg/common"
	"github.com/nuclio/nuclio/pkg/errors"
	"github.com/nuclio/nuclio/pkg/processor/runtime"
	"github.com/nuclio/nuclio/pkg/processor/status"

	"github.com/nuclio/logger"
	"github.com/nuclio/nuclio-sdk-go"
	"github.com/rs/xid"
)

// TODO: Find a better place (both on file system and configuration)
const (
	socketPathTemplate = "/tmp/nuclio-rpc-%s.sock"
	connectionTimeout  = 2 * time.Minute
)

type result struct {
	StatusCode   int                    `json:"status_code"`
	ContentType  string                 `json:"content_type"`
	Body         string                 `json:"body"`
	BodyEncoding string                 `json:"body_encoding"`
	Headers      map[string]interface{} `json:"headers"`

	DecodedBody []byte
	err         error
}

// Runtime is a runtime that communicates via unix domain socket
type AbstractRuntime struct {
	runtime.AbstractRuntime
	configuration  *runtime.Configuration
	eventEncoder   EventEncoder
	wrapperProcess *os.Process
	resultChan     chan *result
	functionLogger logger.Logger
	runtime        Runtime
	startChan      chan struct{}
	socketType     SocketType
}

type rpcLogRecord struct {
	DateTime string                 `json:"datetime"`
	Level    string                 `json:"level"`
	Message  string                 `json:"message"`
	With     map[string]interface{} `json:"with"`
}

// NewRPCRuntime returns a new RPC runtime
func NewAbstractRuntime(logger logger.Logger,
	configuration *runtime.Configuration,
	runtimeInstance Runtime) (*AbstractRuntime, error) {
	var err error

	abstractRuntime, err := runtime.NewAbstractRuntime(logger, configuration)
	if err != nil {
		return nil, errors.Wrap(err, "Can't create AbstractRuntime")
	}

	newRuntime := &AbstractRuntime{
		AbstractRuntime: *abstractRuntime,
		configuration:   configuration,
		runtime:         runtimeInstance,
		startChan:       make(chan struct{}, 1),
	}

	return newRuntime, nil
}

func (r *AbstractRuntime) Start() error {
	if err := r.startWrapper(); err != nil {
		r.SetStatus(status.Error)
		return errors.Wrap(err, "Failed to run wrapper")
	}

	r.SetStatus(status.Ready)
	return nil
}

// ProcessEvent processes an event
func (r *AbstractRuntime) ProcessEvent(event nuclio.Event, functionLogger logger.Logger) (interface{}, error) {

	if currentStatus := r.GetStatus(); currentStatus != status.Ready {
		return nil, errors.Errorf("Processor not ready (current status: %s)", currentStatus)
	}

	if currentStatus := r.GetStatus(); currentStatus != status.Ready {
		return nil, errors.Errorf("Processor not ready (current status: %s)", currentStatus)
	}

	r.functionLogger = functionLogger

	// We don't use defer to reset r.functionLogger since it decreases performance
	if err := r.eventEncoder.Encode(event); err != nil {
		r.functionLogger = nil
		return nil, errors.Wrapf(err, "Can't encode event: %+v", event)
	}

	result, ok := <-r.resultChan
	r.functionLogger = nil
	if !ok {
		msg := "Client disconnected"
		r.Logger.Error(msg)
		r.SetStatus(status.Error)
		r.functionLogger = nil
		return nil, errors.New(msg)
	}

	return nuclio.Response{
		Body:        result.DecodedBody,
		ContentType: result.ContentType,
		Headers:     result.Headers,
		StatusCode:  result.StatusCode,
	}, nil
}

// Stop stops the runtime
func (r *AbstractRuntime) Stop() error {
	if r.wrapperProcess != nil {
		err := r.wrapperProcess.Kill()
		if err != nil {
			r.SetStatus(status.Error)
			return errors.Wrap(err, "Can't kill wrapper process")
		}
	}

	r.SetStatus(status.Stopped)
	return nil
}

// Restart restarts the runtime
func (r *AbstractRuntime) Restart() error {
	if err := r.Stop(); err != nil {
		return err
	}

	// Send error for current event (non-blocking)
	select {
	case r.resultChan <- &result{
		StatusCode: http.StatusRequestTimeout,
		err:        errors.New("Runtime restarted"),
	}:

	default:
		r.Logger.Warn("Nothing waiting on result channel during restart. Continuing")
	}

	close(r.resultChan)
	if err := r.startWrapper(); err != nil {
		r.SetStatus(status.Error)
		return errors.Wrap(err, "Can't start wrapper process")
	}

	r.SetStatus(status.Ready)
	return nil
}

// GetSocketType returns the type of socket the runtime works with (unix/tcp)
func (r *AbstractRuntime) GetSocketType() SocketType {
	return UnixSocket
}

// WaitForStart returns whether the runtime supports sending an indication that it started
func (r *AbstractRuntime) WaitForStart() bool {
	return false
}

// SupportsRestart returns true if the runtime supports restart
func (r *AbstractRuntime) SupportsRestart() bool {
	return true
}

func (r *AbstractRuntime) startWrapper() error {
	var err error

	var listener net.Listener
	var address string

	if r.runtime.GetSocketType() == UnixSocket {
		listener, address, err = r.createUnixListener()
	} else {
		listener, address, err = r.createTCPListener()
	}

	if err != nil {
		return errors.Wrap(err, "Can't create listener")
	}

	wrapperProcess, err := r.runtime.RunWrapper(address)
	if err != nil {
		return errors.Wrap(err, "Can't run wrapper")
	}
	r.wrapperProcess = wrapperProcess
	go r.watchWrapperProcess()

	conn, err := listener.Accept()
	if err != nil {
		return errors.Wrap(err, "Can't get connection from wrapper")
	}

	r.Logger.Info("Wrapper connected")

	r.eventEncoder = r.runtime.GetEventEncoder(conn)
	r.resultChan = make(chan *result)
	go r.wrapperOutputHandler(conn, r.resultChan)

	// wait for start if required to
	if r.runtime.WaitForStart() {
		r.Logger.Debug("Waiting for start")

		<-r.startChan
	}

	r.Logger.Debug("Started")

	return nil
}

// Create a listener on unix domian docker, return listener, path to socket and error
func (r *AbstractRuntime) createUnixListener() (net.Listener, string, error) {
	socketPath := fmt.Sprintf(socketPathTemplate, xid.New().String())

	if common.FileExists(socketPath) {
		if err := os.Remove(socketPath); err != nil {
			return nil, "", errors.Wrapf(err, "Can't remove socket at %q", socketPath)
		}
	}

	r.Logger.DebugWith("Creating listener socket", "path", socketPath)

	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		return nil, "", errors.Wrapf(err, "Can't listen on %s", socketPath)
	}

	unixListener, ok := listener.(*net.UnixListener)
	if !ok {
		return nil, "", fmt.Errorf("Can't get underlying Unix listener")
	}

	if err = unixListener.SetDeadline(time.Now().Add(connectionTimeout)); err != nil {
		return nil, "", errors.Wrap(err, "Can't set deadline")
	}

	return listener, socketPath, nil
}

// Create a listener on TCP docker, return listener, port and error
func (r *AbstractRuntime) createTCPListener() (net.Listener, string, error) {
	listener, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, "", errors.Wrap(err, "Can't find free port")
	}

	tcpListener, ok := listener.(*net.TCPListener)
	if !ok {
		return nil, "", errors.Wrap(err, "Can't get underlying TCP listener")
	}
	if err = tcpListener.SetDeadline(time.Now().Add(connectionTimeout)); err != nil {
		return nil, "", errors.Wrap(err, "Can't set deadline")
	}

	port := listener.Addr().(*net.TCPAddr).Port

	return listener, fmt.Sprintf("%d", port), nil
}

func (r *AbstractRuntime) wrapperOutputHandler(conn io.Reader, resultChan chan *result) {

	// Reset might close outChan, which will cause panic when sending
	defer func() {
		if err := recover(); err != nil {
			r.Logger.WarnWith("panic handling wrapper output (Reset called?)")
		}
	}()

	outReader := bufio.NewReader(conn)

	// Read logs & output
	for {
		unmarshalledResult := &result{}
		var data []byte

		data, unmarshalledResult.err = outReader.ReadBytes('\n')

		if unmarshalledResult.err != nil {
			r.Logger.WarnWith("Failed to read from connection", "err", unmarshalledResult.err)
			resultChan <- unmarshalledResult
			continue
		}

		switch data[0] {
		case 'r':

			// try to unmarshall the result
			if unmarshalledResult.err = json.Unmarshal(data[1:], unmarshalledResult); unmarshalledResult.err != nil {
				r.resultChan <- unmarshalledResult
				continue
			}

			switch unmarshalledResult.BodyEncoding {
			case "text":
				unmarshalledResult.DecodedBody = []byte(unmarshalledResult.Body)
			case "base64":
				unmarshalledResult.DecodedBody, unmarshalledResult.err = base64.StdEncoding.DecodeString(unmarshalledResult.Body)
			default:
				unmarshalledResult.err = fmt.Errorf("Unknown body encoding - %q", unmarshalledResult.BodyEncoding)
			}

			// write back to result channel
			resultChan <- unmarshalledResult
		case 'm':
			r.handleReponseMetric(data[1:])
		case 'l':
			r.handleResponseLog(data[1:])
		case 's':
			r.handleStart()
		}
	}
}

func (r *AbstractRuntime) handleResponseLog(response []byte) {
	var logRecord rpcLogRecord

	if err := json.Unmarshal(response, &logRecord); err != nil {
		r.Logger.ErrorWith("Can't decode log", "error", err)
		return
	}

	logger := r.resolveFunctionLogger(r.functionLogger)
	logFunc := logger.DebugWith

	switch logRecord.Level {
	case "error", "critical", "fatal":
		logFunc = logger.ErrorWith
	case "warning":
		logFunc = logger.WarnWith
	case "info":
		logFunc = logger.InfoWith
	}

	vars := common.MapToSlice(logRecord.With)
	logFunc(logRecord.Message, vars...)
}

func (r *AbstractRuntime) handleReponseMetric(response []byte) {
	var metrics struct {
		DurationSec float64 `json:"duration"`
	}

	logger := r.resolveFunctionLogger(r.functionLogger)
	if err := json.Unmarshal(response, &metrics); err != nil {
		logger.ErrorWith("Can't decode metric", "error", err)
		return
	}

	if metrics.DurationSec == 0 {
		logger.ErrorWith("No duration in metrics", "metrics", metrics)
		return
	}

	r.Statistics.DurationMilliSecondsCount++
	r.Statistics.DurationMilliSecondsSum += uint64(metrics.DurationSec * 1000)
}

func (r *AbstractRuntime) handleStart() {
	r.startChan <- struct{}{}
}

// resolveFunctionLogger return either functionLogger if provided or root logger if not
func (r *AbstractRuntime) resolveFunctionLogger(functionLogger logger.Logger) logger.Logger {
	if functionLogger == nil {
		return r.Logger
	}
	return functionLogger
}

func (r *AbstractRuntime) newResultChan() {

	// We create the channel buffered so we won't block on sending
	r.resultChan = make(chan *result, 1)
}

func (r *AbstractRuntime) watchWrapperProcess() {
	procStatus, err := r.wrapperProcess.Wait()
	if r.GetStatus() == status.Ready && (err != nil || !procStatus.Success()) {
		r.Logger.ErrorWith("Unexpected termination of child process",
			"error", err,
			"status", procStatus.String())

		var panicMessage string
		if err != nil {
			panicMessage = err.Error()
		} else {
			panicMessage = procStatus.String()
		}

		panic(fmt.Sprintf("Wrapper process exited unexpectedly with: %s", panicMessage))
	}
	r.SetStatus(status.Stopped)
	r.wrapperProcess = nil
}
<|endoftext|>


<commit_before>package rkenodeconfigserver

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/pkg/errors"
	"github.com/rancher/norman/types/slice"
	"github.com/rancher/rancher/pkg/librke"
	"github.com/rancher/rancher/pkg/rkecerts"
	"github.com/rancher/rancher/pkg/rkeworker"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/tunnelserver"
	"github.com/rancher/rke/services"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
)

type RKENodeConfigServer struct {
	auth   *tunnelserver.Authorizer
	lookup *rkecerts.BundleLookup
}

func Handler(auth *tunnelserver.Authorizer, scaledContext *config.ScaledContext) http.Handler {
	return &RKENodeConfigServer{
		auth:   auth,
		lookup: rkecerts.NewLookup(scaledContext.Core.Namespaces(""), scaledContext.K8sClient.CoreV1()),
	}
}

func (n *RKENodeConfigServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// 404 tells the client to continue without plan
	// 5xx tells the client to try again later for plan

	client, ok, err := n.auth.Authorize(req)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	if !ok {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	if client.Node == nil {
		rw.WriteHeader(http.StatusNotFound)
		return
	}

	if client.Cluster.Status.Driver == "" {
		rw.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	if client.Cluster.Status.Driver != v3.ClusterDriverRKE {
		rw.WriteHeader(http.StatusNotFound)
		return
	}

	if client.Node.Status.NodeConfig == nil {
		rw.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	var nodeConfig *rkeworker.NodeConfig
	if isNonWorkerOnly(client.Node.Status.NodeConfig.Role) {
		nodeConfig, err = n.nonWorkerConfig(req.Context(), client.Cluster, client.Node)
	} else {
		if client.Cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
			rw.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		nodeConfig, err = n.nodeConfig(req.Context(), client.Cluster, client.Node)
	}

	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	rw.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(rw).Encode(nodeConfig); err != nil {
		logrus.Errorf("failed to write nodeConfig to agent: %v", err)
	}
}

func isNonWorkerOnly(role []string) bool {
	if slice.ContainsString(role, services.ETCDRole) ||
		slice.ContainsString(role, services.ControlRole) {
		return true
	}
	return false
}

func (n *RKENodeConfigServer) nonWorkerConfig(ctx context.Context, cluster *v3.Cluster, node *v3.Node) (*rkeworker.NodeConfig, error) {
	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig
	if rkeConfig == nil {
		rkeConfig = &v3.RancherKubernetesEngineConfig{}
	}

	rkeConfig = rkeConfig.DeepCopy()
	rkeConfig.Nodes = []v3.RKEConfigNode{
		*node.Status.NodeConfig,
	}
	rkeConfig.Nodes[0].Role = []string{services.WorkerRole, services.ETCDRole, services.ControlRole}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(ctx, rkeConfig, infos)
	if err != nil {
		return nil, err
	}

	nc := &rkeworker.NodeConfig{
		ClusterName: cluster.Name,
	}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == node.Status.NodeConfig.Address {
			nc.Processes = augmentProcesses(tempNode.Processes, false)
			return nc, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for non-worker %s", node.Status.NodeConfig.Address)
}

func (n *RKENodeConfigServer) nodeConfig(ctx context.Context, cluster *v3.Cluster, node *v3.Node) (*rkeworker.NodeConfig, error) {
	spec := cluster.Status.AppliedSpec

	bundle, err := n.lookup.Lookup(cluster)
	if err != nil {
		return nil, err
	}

	bundle = bundle.ForNode(spec.RancherKubernetesEngineConfig, node.Status.NodeConfig.Address)

	certString, err := bundle.Marshal()
	if err != nil {
		return nil, errors.Wrapf(err, "failed to marshall bundle")
	}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(ctx, spec.RancherKubernetesEngineConfig, infos)
	if err != nil {
		return nil, err
	}

	nc := &rkeworker.NodeConfig{
		Certs:       certString,
		ClusterName: cluster.Name,
	}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == node.Status.NodeConfig.Address {
			nc.Processes = augmentProcesses(tempNode.Processes, true)
			nc.Files = tempNode.Files
			return nc, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", node.Status.NodeConfig.Address)
}

func augmentProcesses(processes map[string]v3.Process, worker bool) map[string]v3.Process {
	var shared []string

	for _, process := range processes {
		for _, bind := range process.Binds {
			parts := strings.Split(bind, ":")
			if len(parts) > 2 && strings.Contains(parts[2], "shared") {
				shared = append(shared, parts[0])
			}
		}
	}

	if len(shared) > 0 {
		args := []string{"--", "share-root.sh"}
		args = append(args, shared...)

		processes["share-mnt"] = v3.Process{
			Name:          "share-mnt",
			Args:          args,
			Image:         settings.AgentImage.Get(),
			Binds:         []string{"/var/run:/var/run"},
			NetworkMode:   "host",
			RestartPolicy: "always",
			PidMode:       "host",
			Privileged:    true,
		}
	}

	if worker {
		// not sure if we really need this anymore
		delete(processes, "etcd")
	} else {
		if p, ok := processes["share-mnt"]; ok {
			processes = map[string]v3.Process{
				"share-mnt": p,
			}
		} else {
			processes = nil
		}
	}

	for _, p := range processes {
		for i, bind := range p.Binds {
			parts := strings.Split(bind, ":")
			if len(parts) > 1 && parts[1] == "/etc/kubernetes" {
				parts[0] = parts[1]
				p.Binds[i] = strings.Join(parts, ":")
			}
		}
	}

	return processes
}
<commit_msg>Look for slave mounts to setup with share-mnt<commit_after>package rkenodeconfigserver

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/pkg/errors"
	"github.com/rancher/norman/types/slice"
	"github.com/rancher/rancher/pkg/librke"
	"github.com/rancher/rancher/pkg/rkecerts"
	"github.com/rancher/rancher/pkg/rkeworker"
	"github.com/rancher/rancher/pkg/settings"
	"github.com/rancher/rancher/pkg/tunnelserver"
	"github.com/rancher/rke/services"
	"github.com/rancher/types/apis/management.cattle.io/v3"
	"github.com/rancher/types/config"
	"github.com/sirupsen/logrus"
)

type RKENodeConfigServer struct {
	auth   *tunnelserver.Authorizer
	lookup *rkecerts.BundleLookup
}

func Handler(auth *tunnelserver.Authorizer, scaledContext *config.ScaledContext) http.Handler {
	return &RKENodeConfigServer{
		auth:   auth,
		lookup: rkecerts.NewLookup(scaledContext.Core.Namespaces(""), scaledContext.K8sClient.CoreV1()),
	}
}

func (n *RKENodeConfigServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	// 404 tells the client to continue without plan
	// 5xx tells the client to try again later for plan

	client, ok, err := n.auth.Authorize(req)
	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	if !ok {
		rw.WriteHeader(http.StatusUnauthorized)
		return
	}

	if client.Node == nil {
		rw.WriteHeader(http.StatusNotFound)
		return
	}

	if client.Cluster.Status.Driver == "" {
		rw.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	if client.Cluster.Status.Driver != v3.ClusterDriverRKE {
		rw.WriteHeader(http.StatusNotFound)
		return
	}

	if client.Node.Status.NodeConfig == nil {
		rw.WriteHeader(http.StatusServiceUnavailable)
		return
	}

	var nodeConfig *rkeworker.NodeConfig
	if isNonWorkerOnly(client.Node.Status.NodeConfig.Role) {
		nodeConfig, err = n.nonWorkerConfig(req.Context(), client.Cluster, client.Node)
	} else {
		if client.Cluster.Status.AppliedSpec.RancherKubernetesEngineConfig == nil {
			rw.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		nodeConfig, err = n.nodeConfig(req.Context(), client.Cluster, client.Node)
	}

	if err != nil {
		rw.WriteHeader(http.StatusInternalServerError)
		rw.Write([]byte(err.Error()))
		return
	}

	rw.Header().Set("Content-Type", "application/json")

	if err := json.NewEncoder(rw).Encode(nodeConfig); err != nil {
		logrus.Errorf("failed to write nodeConfig to agent: %v", err)
	}
}

func isNonWorkerOnly(role []string) bool {
	if slice.ContainsString(role, services.ETCDRole) ||
		slice.ContainsString(role, services.ControlRole) {
		return true
	}
	return false
}

func (n *RKENodeConfigServer) nonWorkerConfig(ctx context.Context, cluster *v3.Cluster, node *v3.Node) (*rkeworker.NodeConfig, error) {
	rkeConfig := cluster.Status.AppliedSpec.RancherKubernetesEngineConfig
	if rkeConfig == nil {
		rkeConfig = &v3.RancherKubernetesEngineConfig{}
	}

	rkeConfig = rkeConfig.DeepCopy()
	rkeConfig.Nodes = []v3.RKEConfigNode{
		*node.Status.NodeConfig,
	}
	rkeConfig.Nodes[0].Role = []string{services.WorkerRole, services.ETCDRole, services.ControlRole}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(ctx, rkeConfig, infos)
	if err != nil {
		return nil, err
	}

	nc := &rkeworker.NodeConfig{
		ClusterName: cluster.Name,
	}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == node.Status.NodeConfig.Address {
			nc.Processes = augmentProcesses(tempNode.Processes, false)
			return nc, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for non-worker %s", node.Status.NodeConfig.Address)
}

func (n *RKENodeConfigServer) nodeConfig(ctx context.Context, cluster *v3.Cluster, node *v3.Node) (*rkeworker.NodeConfig, error) {
	spec := cluster.Status.AppliedSpec

	bundle, err := n.lookup.Lookup(cluster)
	if err != nil {
		return nil, err
	}

	bundle = bundle.ForNode(spec.RancherKubernetesEngineConfig, node.Status.NodeConfig.Address)

	certString, err := bundle.Marshal()
	if err != nil {
		return nil, errors.Wrapf(err, "failed to marshall bundle")
	}

	infos, err := librke.GetDockerInfo(node)
	if err != nil {
		return nil, err
	}

	plan, err := librke.New().GeneratePlan(ctx, spec.RancherKubernetesEngineConfig, infos)
	if err != nil {
		return nil, err
	}

	nc := &rkeworker.NodeConfig{
		Certs:       certString,
		ClusterName: cluster.Name,
	}

	for _, tempNode := range plan.Nodes {
		if tempNode.Address == node.Status.NodeConfig.Address {
			nc.Processes = augmentProcesses(tempNode.Processes, true)
			nc.Files = tempNode.Files
			return nc, nil
		}
	}

	return nil, fmt.Errorf("failed to find plan for %s", node.Status.NodeConfig.Address)
}

func augmentProcesses(processes map[string]v3.Process, worker bool) map[string]v3.Process {
	var shared []string

	for _, process := range processes {
		for _, bind := range process.Binds {
			parts := strings.Split(bind, ":")
			if len(parts) > 2 && (strings.Contains(parts[2], "shared") || strings.Contains(parts[2], "slave")) {
				shared = append(shared, parts[0])
			}
		}
	}

	if len(shared) > 0 {
		args := []string{"--", "share-root.sh"}
		args = append(args, shared...)

		processes["share-mnt"] = v3.Process{
			Name:          "share-mnt",
			Args:          args,
			Image:         settings.AgentImage.Get(),
			Binds:         []string{"/var/run:/var/run"},
			NetworkMode:   "host",
			RestartPolicy: "always",
			PidMode:       "host",
			Privileged:    true,
		}
	}

	if worker {
		// not sure if we really need this anymore
		delete(processes, "etcd")
	} else {
		if p, ok := processes["share-mnt"]; ok {
			processes = map[string]v3.Process{
				"share-mnt": p,
			}
		} else {
			processes = nil
		}
	}

	for _, p := range processes {
		for i, bind := range p.Binds {
			parts := strings.Split(bind, ":")
			if len(parts) > 1 && parts[1] == "/etc/kubernetes" {
				parts[0] = parts[1]
				p.Binds[i] = strings.Join(parts, ":")
			}
		}
	}

	return processes
}
<|endoftext|>


<commit_before>package build

import (
	"database/sql"
	"log"
	"os"

	_ "github.com/mattn/go-sqlite3" // SQLite driver used in database/sql
)

type sqlContainer struct {
	db       *sql.DB // TODO: when to db.Close
	filename string
}

func (c *sqlContainer) Init(purge bool) error {
	var err error
	if c.filename == "" {
		c.filename = "/tmp/builder.db"
	}
	if purge {
		os.Remove(c.filename)
	}

	c.db, err = sql.Open("sqlite3", c.filename)
	if err != nil {
		return err
	}
	sqlStmt := `
	create table if not exists builds (
		id text not null primary key,
		projectid text not null,
		script text not null,
		executortype text not null);
	create table if not exists stages (
		id integer not null primary key autoincrement,
		build integer not null,
		type text not null,
		timestamp int64 not null,
		name text not null,
		data text);
	create table if not exists output (
		id integer not null primary key autoincrement,
		build text not null,
		data text);
	`
	_, err = c.db.Exec(sqlStmt)
	if err != nil {
		log.Printf("%q: %s\n", err, sqlStmt)
		return err
	}
	return nil
}

func (c *sqlContainer) Builds() []string {
	builds := []string{}

	rows, err := c.db.Query("select id from builds")
	if err != nil {
		log.Println(err)
		return nil
	}
	defer rows.Close()
	for rows.Next() {
		var id string
		err = rows.Scan(&id)
		if err != nil {
			log.Println(err)
			return nil
		}
		builds = append(builds, id)
	}
	err = rows.Err()
	if err != nil {
		log.Println(err)
		return nil
	}
	return builds
}

func (c *sqlContainer) Build(ID string) (Build, error) {
	build := defaultBuild{BID: ID}
	stmt, err := c.db.Prepare("select projectid, script, executortype from builds where id = ?")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	err = stmt.QueryRow(ID).Scan(&build.BProjectID, &build.BScript, &build.BExecutorType)
	if err == sql.ErrNoRows {
		return nil, ErrNotFound
	} else if err != nil {
		return nil, err
	}

	stmt, err = c.db.Prepare("select type, timestamp, name, data from stages where build = ?")
	if err != nil {
		return nil, err
	}
	rows, err := stmt.Query(ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		stage := Stage{}
		err = rows.Scan((*string)(&stage.Type), &stage.Timestamp, &stage.Name, &stage.Data)
		if err != nil {
			return nil, err
		}
		err = build.AddStage(stage)
		if err != nil {
			return nil, err
		}
	}

	stmt, err = c.db.Prepare("select data from output where build = ? order by id")
	if err != nil {
		return nil, err
	}
	rows, err = stmt.Query(ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		data := []byte{}
		err = rows.Scan(&data)
		if err != nil {
			return nil, err
		}
		build.output = append(build.output, data...)
	}

	return &build, nil
}

func (c *sqlContainer) New(b Buildable) (Build, error) {
	build, err := New(b)
	if err != nil {
		return nil, err
	}
	stmt, err := c.db.Prepare("insert into builds(id, projectid, script, executortype) values(?, ?, ?, ?)")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	_, err = stmt.Exec(build.ID(), build.ProjectID(), build.Script(), build.ExecutorType())
	if err != nil {
		return nil, err
	}
	return build, nil
}

func (c *sqlContainer) AddStage(buildID string, stage Stage) error {
	buf := Stage{}
	previous := &buf
	tx, err := c.db.Begin()
	if err != nil {
		return err
	}

	stmt, err := tx.Prepare("select type, timestamp, name, data from stages where build = ? order by id desc limit 1")
	if err != nil {
		return err
	}
	t := ""
	err = stmt.QueryRow(buildID).Scan(&t, &buf.Timestamp, &buf.Name, &buf.Data)
	buf.Type = StageType(t)
	if err == sql.ErrNoRows {
		previous = nil
	} else if err != nil {
		return err
	}
	err = stage.ValidateWithPredecessor(previous)
	if err != nil {
		return err
	}
	stmt, err = tx.Prepare("insert into stages (build, type, timestamp, name, data) values (?, ?, ?, ?, ?)")
	if err != nil {
		return err
	}

	_, err = stmt.Exec(buildID, string(stage.Type), stage.Timestamp, stage.Name, stage.Data)
	if err != nil {
		return err
	}
	return tx.Commit()

}

func (c *sqlContainer) Output(buildID string, output []byte) error {
	stmt, err := c.db.Prepare("insert into output(build, data) values(?, ?)")
	if err != nil {
		return err
	}
	defer stmt.Close()
	_, err = stmt.Exec(buildID, output)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>List builds in descending id order<commit_after>package build

import (
	"database/sql"
	"log"
	"os"

	_ "github.com/mattn/go-sqlite3" // SQLite driver used in database/sql
)

type sqlContainer struct {
	db       *sql.DB // TODO: when to db.Close
	filename string
}

func (c *sqlContainer) Init(purge bool) error {
	var err error
	if c.filename == "" {
		c.filename = "/tmp/builder.db"
	}
	if purge {
		os.Remove(c.filename)
	}

	c.db, err = sql.Open("sqlite3", c.filename)
	if err != nil {
		return err
	}
	sqlStmt := `
	create table if not exists builds (
		id text not null primary key,
		projectid text not null,
		script text not null,
		executortype text not null);
	create table if not exists stages (
		id integer not null primary key autoincrement,
		build integer not null,
		type text not null,
		timestamp int64 not null,
		name text not null,
		data text);
	create table if not exists output (
		id integer not null primary key autoincrement,
		build text not null,
		data text);
	`
	_, err = c.db.Exec(sqlStmt)
	if err != nil {
		log.Printf("%q: %s\n", err, sqlStmt)
		return err
	}
	return nil
}

func (c *sqlContainer) Builds() []string {
	builds := []string{}

	rows, err := c.db.Query("select id from builds order by id desc")
	if err != nil {
		log.Println(err)
		return nil
	}
	defer rows.Close()
	for rows.Next() {
		var id string
		err = rows.Scan(&id)
		if err != nil {
			log.Println(err)
			return nil
		}
		builds = append(builds, id)
	}
	err = rows.Err()
	if err != nil {
		log.Println(err)
		return nil
	}
	return builds
}

func (c *sqlContainer) Build(ID string) (Build, error) {
	build := defaultBuild{BID: ID}
	stmt, err := c.db.Prepare("select projectid, script, executortype from builds where id = ?")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	err = stmt.QueryRow(ID).Scan(&build.BProjectID, &build.BScript, &build.BExecutorType)
	if err == sql.ErrNoRows {
		return nil, ErrNotFound
	} else if err != nil {
		return nil, err
	}

	stmt, err = c.db.Prepare("select type, timestamp, name, data from stages where build = ?")
	if err != nil {
		return nil, err
	}
	rows, err := stmt.Query(ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		stage := Stage{}
		err = rows.Scan((*string)(&stage.Type), &stage.Timestamp, &stage.Name, &stage.Data)
		if err != nil {
			return nil, err
		}
		err = build.AddStage(stage)
		if err != nil {
			return nil, err
		}
	}

	stmt, err = c.db.Prepare("select data from output where build = ? order by id")
	if err != nil {
		return nil, err
	}
	rows, err = stmt.Query(ID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		data := []byte{}
		err = rows.Scan(&data)
		if err != nil {
			return nil, err
		}
		build.output = append(build.output, data...)
	}

	return &build, nil
}

func (c *sqlContainer) New(b Buildable) (Build, error) {
	build, err := New(b)
	if err != nil {
		return nil, err
	}
	stmt, err := c.db.Prepare("insert into builds(id, projectid, script, executortype) values(?, ?, ?, ?)")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	_, err = stmt.Exec(build.ID(), build.ProjectID(), build.Script(), build.ExecutorType())
	if err != nil {
		return nil, err
	}
	return build, nil
}

func (c *sqlContainer) AddStage(buildID string, stage Stage) error {
	buf := Stage{}
	previous := &buf
	tx, err := c.db.Begin()
	if err != nil {
		return err
	}

	stmt, err := tx.Prepare("select type, timestamp, name, data from stages where build = ? order by id desc limit 1")
	if err != nil {
		return err
	}
	t := ""
	err = stmt.QueryRow(buildID).Scan(&t, &buf.Timestamp, &buf.Name, &buf.Data)
	buf.Type = StageType(t)
	if err == sql.ErrNoRows {
		previous = nil
	} else if err != nil {
		return err
	}
	err = stage.ValidateWithPredecessor(previous)
	if err != nil {
		return err
	}
	stmt, err = tx.Prepare("insert into stages (build, type, timestamp, name, data) values (?, ?, ?, ?, ?)")
	if err != nil {
		return err
	}

	_, err = stmt.Exec(buildID, string(stage.Type), stage.Timestamp, stage.Name, stage.Data)
	if err != nil {
		return err
	}
	return tx.Commit()

}

func (c *sqlContainer) Output(buildID string, output []byte) error {
	stmt, err := c.db.Prepare("insert into output(build, data) values(?, ?)")
	if err != nil {
		return err
	}
	defer stmt.Close()
	_, err = stmt.Exec(buildID, output)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Package signal2 is the supplement of the standard library of `signal`,
// such as `HandleSignal`.
package signal2

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/xgfone/go-tools/lifecycle"
)

// DefaultSignals is the default signals to be handled.
var DefaultSignals = []os.Signal{
	syscall.SIGTERM,
	syscall.SIGQUIT,
	syscall.SIGABRT,
	syscall.SIGINT,
}

// HandleSignal is the same as HandleSignalWithLifecycle, but using the global
// default lifecycle manager.
//
// It's equal to
//   HandleSignalWithLifecycle(lifecycle.GetDefaultManager(), signals...)
//
// Notice: If the signals are empty, it will be equal to
//   HandleSignal(DefaultSignals...)
func HandleSignal(signals ...os.Signal) {
	if len(signals) == 0 {
		signals = DefaultSignals
	}
	HandleSignalWithLifecycle(lifecycle.GetDefaultManager(), signals...)
}

// HandleSignalWithLifecycle wraps and handles the signals.
//
// The default wraps os.Interrupt. And you can pass the extra signals,
// syscall.SIGTERM, syscall.SIGQUIT, etc, such as
//   m := lifecycle.GetDefaultManager()
//   HandleSignalWithLifecycle(m, syscall.SIGTERM, syscall.SIGQUIT)
//
// For running it in a goroutine, use
//   go HandleSignalWithLifecycle(m, syscall.SIGTERM, syscall.SIGQUIT)
func HandleSignalWithLifecycle(m *lifecycle.Manager, signals ...os.Signal) {
	HandleSignalWithFunc(func() { m.Stop() }, os.Interrupt, signals...)
}

// HandleSignalWithFunc calls the function f when the signals are received.
func HandleSignalWithFunc(f func(), sig os.Signal, signals ...os.Signal) {
	ss := make(chan os.Signal, 1)
	signals = append(signals, sig)
	signal.Notify(ss, signals...)
	<-ss
	f()
}
<commit_msg>fix: loop to handle the signal<commit_after>// Package signal2 is the supplement of the standard library of `signal`,
// such as `HandleSignal`.
package signal2

import (
	"os"
	"os/signal"
	"syscall"

	"github.com/xgfone/go-tools/lifecycle"
)

// DefaultSignals is the default signals to be handled.
var DefaultSignals = []os.Signal{
	syscall.SIGTERM,
	syscall.SIGQUIT,
	syscall.SIGABRT,
	syscall.SIGINT,
}

// HandleSignal is the same as HandleSignalWithLifecycle, but using the global
// default lifecycle manager.
//
// It's equal to
//   HandleSignalWithLifecycle(lifecycle.GetDefaultManager(), signals...)
//
// Notice: If the signals are empty, it will be equal to
//   HandleSignal(DefaultSignals...)
func HandleSignal(signals ...os.Signal) {
	if len(signals) == 0 {
		signals = DefaultSignals
	}
	HandleSignalWithLifecycle(lifecycle.GetDefaultManager(), signals...)
}

// HandleSignalWithLifecycle wraps and handles the signals.
//
// The default wraps os.Interrupt. And you can pass the extra signals,
// syscall.SIGTERM, syscall.SIGQUIT, etc, such as
//   m := lifecycle.GetDefaultManager()
//   HandleSignalWithLifecycle(m, syscall.SIGTERM, syscall.SIGQUIT)
//
// For running it in a goroutine, use
//   go HandleSignalWithLifecycle(m, syscall.SIGTERM, syscall.SIGQUIT)
func HandleSignalWithLifecycle(m *lifecycle.Manager, signals ...os.Signal) {
	HandleSignalWithFunc(func() { m.Stop() }, os.Interrupt, signals...)
}

// HandleSignalWithFunc calls the function f when the signals are received.
func HandleSignalWithFunc(f func(), sig os.Signal, signals ...os.Signal) {
	ss := make(chan os.Signal, 1)
	signals = append(signals, sig)
	signal.Notify(ss, signals...)
	for {
		<-ss
		f()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

// Package webhook checks a webhook for configured operation admission
package webhook // import "k8s.io/kubernetes/plugin/pkg/admission/webhook"
<commit_msg>fix import warning<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

// Package webhook checks a webhook for configured operation admission
package webhook // import "k8s.io/apiserver/pkg/admission/plugin/webhook"
<|endoftext|>


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

// Package controller provides a Kubernetes controller for a TFJob resource.
package tensorflow

import (
	"fmt"
	"strconv"
	"strings"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"

	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
	"github.com/kubeflow/tf-operator/pkg/controller.v2/jobcontroller"
	tflogger "github.com/kubeflow/tf-operator/pkg/logger"
	train_util "github.com/kubeflow/tf-operator/pkg/util/train"
)

const (
	// tfConfig is the environment variable name of TensorFlow cluster spec.
	tfConfig = "TF_CONFIG"

	// podTemplateRestartPolicyReason is the warning reason when the restart
	// policy is setted in pod template.
	podTemplateRestartPolicyReason = "SettedPodTemplateRestartPolicy"
)

// reconcilePods checks and updates pods for each given TFReplicaSpec.
// It will requeue the tfjob in case of an error while creating/deleting pods.
func (tc *TFController) reconcilePods(
	tfjob *tfv1alpha2.TFJob,
	pods []*v1.Pod,
	rtype tfv1alpha2.TFReplicaType,
	spec *tfv1alpha2.TFReplicaSpec, rstatus map[string]v1.PodPhase) error {

	// Convert TFReplicaType to lower string.
	rt := strings.ToLower(string(rtype))
	logger := tflogger.LoggerForReplica(tfjob, rt)
	// Get all pods for the type rt.
	pods, err := tc.FilterPodsForReplicaType(pods, rt)
	if err != nil {
		return err
	}
	replicas := int(*spec.Replicas)
	restart := false
	worker0Completed := false

	initializeTFReplicaStatuses(tfjob, rtype)

	podSlices := tc.GetPodSlices(pods, replicas, logger)
	for index, podSlice := range podSlices {
		if len(podSlice) > 1 {
			logger.Warningf("We have too many pods for %s %d", rt, index)
			// TODO(gaocegege): Kill some pods.
		} else if len(podSlice) == 0 {
			logger.Infof("Need to create new pod: %s-%d", rt, index)
			err = tc.createNewPod(tfjob, rt, strconv.Itoa(index), spec)
			if err != nil {
				return err
			}
		} else {
			// Check the status of the current pod.
			pod := podSlice[0]
			// Get the exit code of the tensorflow container.
			var exitCode int32 = 0xbeef // magic number
			for _, status := range pod.Status.ContainerStatuses {
				state := status.State
				if status.Name == tfv1alpha2.DefaultContainerName && state.Terminated != nil {
					exitCode = state.Terminated.ExitCode
					logger.Infof("Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
					tc.Recorder.Eventf(tfjob, v1.EventTypeNormal, "Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
				}
			}
			// Check if the pod is retryable.
			if spec.RestartPolicy == tfv1alpha2.RestartPolicyExitCode {
				if pod.Status.Phase == v1.PodFailed && train_util.IsRetryableExitCode(exitCode) {
					logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
					if err := tc.PodControl.DeletePod(pod.Namespace, pod.Name, tfjob); err != nil {
						return err
					}
					restart = true
				}
			}

			// Check whether worker 0 is exited without error.
			if rtype == tfv1alpha2.TFReplicaTypeWorker && index == 0 && exitCode == 0 {
				worker0Completed = true
			}
			updateTFJobReplicaStatuses(tfjob, rtype, pod)
		}
	}

	return updateStatusSingle(tfjob, rtype, replicas, restart, worker0Completed)
}

// createNewPod creates a new pod for the given index and type.
func (tc *TFController) createNewPod(tfjob *tfv1alpha2.TFJob, rt, index string, spec *tfv1alpha2.TFReplicaSpec) error {
	tfjobKey, err := KeyFunc(tfjob)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("Couldn't get key for tfjob object %#v: %v", tfjob, err))
		return err
	}
	expectationPodsKey := jobcontroller.GenExpectationPodsKey(tfjobKey, rt)
	err = tc.Expectations.ExpectCreations(expectationPodsKey, 1)
	if err != nil {
		return err
	}
	logger := tflogger.LoggerForReplica(tfjob, rt)
	// Create OwnerReference.
	controllerRef := tc.GenOwnerReference(tfjob)

	// Set type and index for the worker.
	labels := tc.GenLabels(tfjob.Name)
	labels[tfReplicaTypeLabel] = rt
	labels[tfReplicaIndexLabel] = index

	podTemplate := spec.Template.DeepCopy()

	// Set name for the template.
	podTemplate.Name = jobcontroller.GenGeneralName(tfjob.Name, rt, index)

	if podTemplate.Labels == nil {
		podTemplate.Labels = make(map[string]string)
	}

	for key, value := range labels {
		podTemplate.Labels[key] = value
	}
	setSchedulerName(podTemplate, tfjob)
	if err := setClusterSpec(podTemplate, tfjob, rt, index); err != nil {
		return err
	}

	// Submit a warning event if the user specifies restart policy for
	// the pod template. We recommend to set it from the replica level.
	if podTemplate.Spec.RestartPolicy != v1.RestartPolicy("") {
		errMsg := "Restart policy in pod template will be overwritten by restart policy in replica spec"
		logger.Warning(errMsg)
		tc.Recorder.Event(tfjob, v1.EventTypeWarning, podTemplateRestartPolicyReason, errMsg)
	}
	setRestartPolicy(podTemplate, spec)

	err = tc.PodControl.CreatePodsWithControllerRef(tfjob.Namespace, podTemplate, tfjob, controllerRef)
	if err != nil && errors.IsTimeout(err) {
		// Pod is created but its initialization has timed out.
		// If the initialization is successful eventually, the
		// controller will observe the creation via the informer.
		// If the initialization fails, or if the pod keeps
		// uninitialized for a long time, the informer will not
		// receive any update, and the controller will create a new
		// pod when the expectation expires.
		return nil
	} else if err != nil {
		return err
	}
	return nil
}

func setSchedulerName(podTemplateSpec *v1.PodTemplateSpec, tfjob *tfv1alpha2.TFJob) {
	podTemplateSpec.Spec.SchedulerName = tfjob.Spec.SchedulerName
}

func setClusterSpec(podTemplateSpec *v1.PodTemplateSpec, tfjob *tfv1alpha2.TFJob, rt, index string) error {
	// Generate TF_CONFIG JSON string.
	tfConfigStr, err := genTFConfigJSONStr(tfjob, rt, index)
	if err != nil {
		return err
	}

	if tfConfigStr == "" {
		return nil
	}
	// Add TF_CONFIG environment variable.
	for i := range podTemplateSpec.Spec.Containers {
		if len(podTemplateSpec.Spec.Containers[i].Env) == 0 {
			podTemplateSpec.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
		}
		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, v1.EnvVar{
			Name:  tfConfig,
			Value: tfConfigStr,
		})
	}
	return nil
}

func setRestartPolicy(podTemplateSpec *v1.PodTemplateSpec, spec *tfv1alpha2.TFReplicaSpec) {
	if spec.RestartPolicy == tfv1alpha2.RestartPolicyExitCode {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicyNever
	} else {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicy(spec.RestartPolicy)
	}
}
<commit_msg>pod: Fix eventf (#819)<commit_after>// Copyright 2018 The Kubeflow Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package controller provides a Kubernetes controller for a TFJob resource.
package tensorflow

import (
	"fmt"
	"strconv"
	"strings"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"

	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
	"github.com/kubeflow/tf-operator/pkg/controller.v2/jobcontroller"
	tflogger "github.com/kubeflow/tf-operator/pkg/logger"
	train_util "github.com/kubeflow/tf-operator/pkg/util/train"
)

const (
	// tfConfig is the environment variable name of TensorFlow cluster spec.
	tfConfig = "TF_CONFIG"

	// podTemplateRestartPolicyReason is the warning reason when the restart
	// policy is setted in pod template.
	podTemplateRestartPolicyReason = "SettedPodTemplateRestartPolicy"
	// exitedWithCodeReason is the normal reason when the pod is exited because of the exit code.
	exitedWithCodeReason = "ExitedWithCode"
)

// reconcilePods checks and updates pods for each given TFReplicaSpec.
// It will requeue the tfjob in case of an error while creating/deleting pods.
func (tc *TFController) reconcilePods(
	tfjob *tfv1alpha2.TFJob,
	pods []*v1.Pod,
	rtype tfv1alpha2.TFReplicaType,
	spec *tfv1alpha2.TFReplicaSpec, rstatus map[string]v1.PodPhase) error {

	// Convert TFReplicaType to lower string.
	rt := strings.ToLower(string(rtype))
	logger := tflogger.LoggerForReplica(tfjob, rt)
	// Get all pods for the type rt.
	pods, err := tc.FilterPodsForReplicaType(pods, rt)
	if err != nil {
		return err
	}
	replicas := int(*spec.Replicas)
	restart := false
	worker0Completed := false

	initializeTFReplicaStatuses(tfjob, rtype)

	podSlices := tc.GetPodSlices(pods, replicas, logger)
	for index, podSlice := range podSlices {
		if len(podSlice) > 1 {
			logger.Warningf("We have too many pods for %s %d", rt, index)
			// TODO(gaocegege): Kill some pods.
		} else if len(podSlice) == 0 {
			logger.Infof("Need to create new pod: %s-%d", rt, index)
			err = tc.createNewPod(tfjob, rt, strconv.Itoa(index), spec)
			if err != nil {
				return err
			}
		} else {
			// Check the status of the current pod.
			pod := podSlice[0]
			// Get the exit code of the tensorflow container.
			var exitCode int32 = 0xbeef // magic number
			for _, status := range pod.Status.ContainerStatuses {
				state := status.State
				if status.Name == tfv1alpha2.DefaultContainerName && state.Terminated != nil {
					exitCode = state.Terminated.ExitCode
					logger.Infof("Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
					tc.Recorder.Eventf(tfjob, v1.EventTypeNormal, exitedWithCodeReason, "Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
				}
			}
			// Check if the pod is retryable.
			if spec.RestartPolicy == tfv1alpha2.RestartPolicyExitCode {
				if pod.Status.Phase == v1.PodFailed && train_util.IsRetryableExitCode(exitCode) {
					logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
					if err := tc.PodControl.DeletePod(pod.Namespace, pod.Name, tfjob); err != nil {
						return err
					}
					restart = true
				}
			}

			// Check whether worker 0 is exited without error.
			if rtype == tfv1alpha2.TFReplicaTypeWorker && index == 0 && exitCode == 0 {
				worker0Completed = true
			}
			updateTFJobReplicaStatuses(tfjob, rtype, pod)
		}
	}

	return updateStatusSingle(tfjob, rtype, replicas, restart, worker0Completed)
}

// createNewPod creates a new pod for the given index and type.
func (tc *TFController) createNewPod(tfjob *tfv1alpha2.TFJob, rt, index string, spec *tfv1alpha2.TFReplicaSpec) error {
	tfjobKey, err := KeyFunc(tfjob)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("Couldn't get key for tfjob object %#v: %v", tfjob, err))
		return err
	}
	expectationPodsKey := jobcontroller.GenExpectationPodsKey(tfjobKey, rt)
	err = tc.Expectations.ExpectCreations(expectationPodsKey, 1)
	if err != nil {
		return err
	}
	logger := tflogger.LoggerForReplica(tfjob, rt)
	// Create OwnerReference.
	controllerRef := tc.GenOwnerReference(tfjob)

	// Set type and index for the worker.
	labels := tc.GenLabels(tfjob.Name)
	labels[tfReplicaTypeLabel] = rt
	labels[tfReplicaIndexLabel] = index

	podTemplate := spec.Template.DeepCopy()

	// Set name for the template.
	podTemplate.Name = jobcontroller.GenGeneralName(tfjob.Name, rt, index)

	if podTemplate.Labels == nil {
		podTemplate.Labels = make(map[string]string)
	}

	for key, value := range labels {
		podTemplate.Labels[key] = value
	}
	setSchedulerName(podTemplate, tfjob)
	if err := setClusterSpec(podTemplate, tfjob, rt, index); err != nil {
		return err
	}

	// Submit a warning event if the user specifies restart policy for
	// the pod template. We recommend to set it from the replica level.
	if podTemplate.Spec.RestartPolicy != v1.RestartPolicy("") {
		errMsg := "Restart policy in pod template will be overwritten by restart policy in replica spec"
		logger.Warning(errMsg)
		tc.Recorder.Event(tfjob, v1.EventTypeWarning, podTemplateRestartPolicyReason, errMsg)
	}
	setRestartPolicy(podTemplate, spec)

	err = tc.PodControl.CreatePodsWithControllerRef(tfjob.Namespace, podTemplate, tfjob, controllerRef)
	if err != nil && errors.IsTimeout(err) {
		// Pod is created but its initialization has timed out.
		// If the initialization is successful eventually, the
		// controller will observe the creation via the informer.
		// If the initialization fails, or if the pod keeps
		// uninitialized for a long time, the informer will not
		// receive any update, and the controller will create a new
		// pod when the expectation expires.
		return nil
	} else if err != nil {
		return err
	}
	return nil
}

func setSchedulerName(podTemplateSpec *v1.PodTemplateSpec, tfjob *tfv1alpha2.TFJob) {
	podTemplateSpec.Spec.SchedulerName = tfjob.Spec.SchedulerName
}

func setClusterSpec(podTemplateSpec *v1.PodTemplateSpec, tfjob *tfv1alpha2.TFJob, rt, index string) error {
	// Generate TF_CONFIG JSON string.
	tfConfigStr, err := genTFConfigJSONStr(tfjob, rt, index)
	if err != nil {
		return err
	}

	if tfConfigStr == "" {
		return nil
	}
	// Add TF_CONFIG environment variable.
	for i := range podTemplateSpec.Spec.Containers {
		if len(podTemplateSpec.Spec.Containers[i].Env) == 0 {
			podTemplateSpec.Spec.Containers[i].Env = make([]v1.EnvVar, 0)
		}
		podTemplateSpec.Spec.Containers[i].Env = append(podTemplateSpec.Spec.Containers[i].Env, v1.EnvVar{
			Name:  tfConfig,
			Value: tfConfigStr,
		})
	}
	return nil
}

func setRestartPolicy(podTemplateSpec *v1.PodTemplateSpec, spec *tfv1alpha2.TFReplicaSpec) {
	if spec.RestartPolicy == tfv1alpha2.RestartPolicyExitCode {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicyNever
	} else {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicy(spec.RestartPolicy)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package stats

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"path"
	"time"

	restful "github.com/emicklei/go-restful"
	cadvisorapi "github.com/google/cadvisor/info/v1"
	"k8s.io/klog"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	statsapi "k8s.io/kubernetes/pkg/kubelet/apis/stats/v1alpha1"
	"k8s.io/kubernetes/pkg/kubelet/cm"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	"k8s.io/kubernetes/pkg/volume"
)

// Provider hosts methods required by stats handlers.
type Provider interface {
	// The following stats are provided by either CRI or cAdvisor.
	//
	// ListPodStats returns the stats of all the containers managed by pods.
	ListPodStats() ([]statsapi.PodStats, error)
	// ListPodStatsAndUpdateCPUNanoCoreUsage updates the cpu nano core usage for
	// the containers and returns the stats for all the pod-managed containers.
	ListPodCPUAndMemoryStats() ([]statsapi.PodStats, error)
	// ListPodStatsAndUpdateCPUNanoCoreUsage returns the stats of all the
	// containers managed by pods and force update the cpu usageNanoCores.
	// This is a workaround for CRI runtimes that do not integrate with
	// cadvisor. See https://github.com/kubernetes/kubernetes/issues/72788
	// for more details.
	ListPodStatsAndUpdateCPUNanoCoreUsage() ([]statsapi.PodStats, error)
	// ImageFsStats returns the stats of the image filesystem.
	ImageFsStats() (*statsapi.FsStats, error)

	// The following stats are provided by cAdvisor.
	//
	// GetCgroupStats returns the stats and the networking usage of the cgroup
	// with the specified cgroupName.
	GetCgroupStats(cgroupName string, updateStats bool) (*statsapi.ContainerStats, *statsapi.NetworkStats, error)
	// GetCgroupCPUAndMemoryStats returns the CPU and memory stats of the cgroup with the specified cgroupName.
	GetCgroupCPUAndMemoryStats(cgroupName string, updateStats bool) (*statsapi.ContainerStats, error)

	// RootFsStats returns the stats of the node root filesystem.
	RootFsStats() (*statsapi.FsStats, error)

	// The following stats are provided by cAdvisor for legacy usage.
	//
	// GetContainerInfo returns the information of the container with the
	// containerName managed by the pod with the uid.
	GetContainerInfo(podFullName string, uid types.UID, containerName string, req *cadvisorapi.ContainerInfoRequest) (*cadvisorapi.ContainerInfo, error)
	// GetRawContainerInfo returns the information of the container with the
	// containerName. If subcontainers is true, this function will return the
	// information of all the sub-containers as well.
	GetRawContainerInfo(containerName string, req *cadvisorapi.ContainerInfoRequest, subcontainers bool) (map[string]*cadvisorapi.ContainerInfo, error)

	// The following information is provided by Kubelet.
	//
	// GetPodByName returns the spec of the pod with the name in the specified
	// namespace.
	GetPodByName(namespace, name string) (*v1.Pod, bool)
	// GetNode returns the spec of the local node.
	GetNode() (*v1.Node, error)
	// GetNodeConfig returns the configuration of the local node.
	GetNodeConfig() cm.NodeConfig
	// ListVolumesForPod returns the stats of the volume used by the pod with
	// the podUID.
	ListVolumesForPod(podUID types.UID) (map[string]volume.Volume, bool)
	// GetPods returns the specs of all the pods running on this node.
	GetPods() []*v1.Pod

	// RlimitStats returns the rlimit stats of system.
	RlimitStats() (*statsapi.RlimitStats, error)

	// GetPodCgroupRoot returns the literal cgroupfs value for the cgroup containing all pods
	GetPodCgroupRoot() string

	// GetPodByCgroupfs provides the pod that maps to the specified cgroup literal, as well
	// as whether the pod was found.
	GetPodByCgroupfs(cgroupfs string) (*v1.Pod, bool)
}

type handler struct {
	provider        Provider
	summaryProvider SummaryProvider
}

// CreateHandlers creates the REST handlers for the stats.
func CreateHandlers(rootPath string, provider Provider, summaryProvider SummaryProvider, enableCAdvisorJSONEndpoints bool) *restful.WebService {
	h := &handler{provider, summaryProvider}

	ws := &restful.WebService{}
	ws.Path(rootPath).
		Produces(restful.MIME_JSON)

	type endpoint struct {
		path    string
		handler restful.RouteFunction
	}

	endpoints := []endpoint{
		{"/summary", h.handleSummary},
	}

	if enableCAdvisorJSONEndpoints {
		endpoints = append(endpoints,
			endpoint{"", h.handleStats},
			endpoint{"/container", h.handleSystemContainer},
			endpoint{"/{podName}/{containerName}", h.handlePodContainer},
			endpoint{"/{namespace}/{podName}/{uid}/{containerName}", h.handlePodContainer},
		)
	}

	for _, e := range endpoints {
		for _, method := range []string{"GET", "POST"} {
			ws.Route(ws.
				Method(method).
				Path(e.path).
				To(e.handler))
		}
	}

	return ws
}

type statsRequest struct {
	// The name of the container for which to request stats.
	// Default: /
	// +optional
	ContainerName string `json:"containerName,omitempty"`

	// Max number of stats to return.
	// If start and end time are specified this limit is ignored.
	// Default: 60
	// +optional
	NumStats int `json:"num_stats,omitempty"`

	// Start time for which to query information.
	// If omitted, the beginning of time is assumed.
	// +optional
	Start time.Time `json:"start,omitempty"`

	// End time for which to query information.
	// If omitted, current time is assumed.
	// +optional
	End time.Time `json:"end,omitempty"`

	// Whether to also include information from subcontainers.
	// Default: false.
	// +optional
	Subcontainers bool `json:"subcontainers,omitempty"`
}

func (r *statsRequest) cadvisorRequest() *cadvisorapi.ContainerInfoRequest {
	return &cadvisorapi.ContainerInfoRequest{
		NumStats: r.NumStats,
		Start:    r.Start,
		End:      r.End,
	}
}

func parseStatsRequest(request *restful.Request) (statsRequest, error) {
	// Default request.
	query := statsRequest{
		NumStats: 60,
	}

	err := json.NewDecoder(request.Request.Body).Decode(&query)
	if err != nil && err != io.EOF {
		return query, err
	}
	return query, nil
}

// Handles root container stats requests to /stats
func (h *handler) handleStats(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, "/stats", err)
		return
	}

	// Root container stats.
	statsMap, err := h.provider.GetRawContainerInfo("/", query.cadvisorRequest(), false)
	if err != nil {
		handleError(response, fmt.Sprintf("/stats %v", query), err)
		return
	}
	writeResponse(response, statsMap["/"])
}

// Handles stats summary requests to /stats/summary
// If "only_cpu_and_memory" GET param is true then only cpu and memory is returned in response.
func (h *handler) handleSummary(request *restful.Request, response *restful.Response) {
	onlyCPUAndMemory := false
	request.Request.ParseForm()
	if onlyCluAndMemoryParam, found := request.Request.Form["only_cpu_and_memory"]; found &&
		len(onlyCluAndMemoryParam) == 1 && onlyCluAndMemoryParam[0] == "true" {
		onlyCPUAndMemory = true
	}
	var summary *statsapi.Summary
	var err error
	if onlyCPUAndMemory {
		summary, err = h.summaryProvider.GetCPUAndMemoryStats()
	} else {
		// external calls to the summary API use cached stats
		forceStatsUpdate := false
		summary, err = h.summaryProvider.Get(forceStatsUpdate)
	}
	if err != nil {
		handleError(response, "/stats/summary", err)
	} else {
		writeResponse(response, summary)
	}
}

// Handles non-kubernetes container stats requests to /stats/container/
func (h *handler) handleSystemContainer(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, "/stats/container", err)
		return
	}

	// Non-Kubernetes container stats.
	containerName := path.Join("/", query.ContainerName)
	stats, err := h.provider.GetRawContainerInfo(
		containerName, query.cadvisorRequest(), query.Subcontainers)
	if err != nil {
		if _, ok := stats[containerName]; ok {
			// If the failure is partial, log it and return a best-effort response.
			klog.Errorf("Partial failure issuing GetRawContainerInfo(%v): %v", query, err)
		} else {
			handleError(response, fmt.Sprintf("/stats/container %v", query), err)
			return
		}
	}
	writeResponse(response, stats)
}

// Handles kubernetes pod/container stats requests to:
// /stats/<pod name>/<container name>
// /stats/<namespace>/<pod name>/<uid>/<container name>
func (h *handler) handlePodContainer(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, request.Request.URL.String(), err)
		return
	}

	// Default parameters.
	params := map[string]string{
		"namespace": metav1.NamespaceDefault,
		"uid":       "",
	}
	for k, v := range request.PathParameters() {
		params[k] = v
	}

	if params["podName"] == "" || params["containerName"] == "" {
		response.WriteErrorString(http.StatusBadRequest,
			fmt.Sprintf("Invalid pod container request: %v", params))
		return
	}

	pod, ok := h.provider.GetPodByName(params["namespace"], params["podName"])
	if !ok {
		klog.V(4).Infof("Container not found: %v", params)
		response.WriteError(http.StatusNotFound, kubecontainer.ErrContainerNotFound)
		return
	}
	stats, err := h.provider.GetContainerInfo(
		kubecontainer.GetPodFullName(pod),
		types.UID(params["uid"]),
		params["containerName"],
		query.cadvisorRequest())

	if err != nil {
		handleError(response, fmt.Sprintf("%s %v", request.Request.URL.String(), query), err)
		return
	}
	writeResponse(response, stats)
}

func writeResponse(response *restful.Response, stats interface{}) {
	if err := response.WriteAsJson(stats); err != nil {
		klog.Errorf("Error writing response: %v", err)
	}
}

// handleError serializes an error object into an HTTP response.
// request is provided for logging.
func handleError(response *restful.Response, request string, err error) {
	switch err {
	case kubecontainer.ErrContainerNotFound:
		response.WriteError(http.StatusNotFound, err)
	default:
		msg := fmt.Sprintf("Internal Error: %v", err)
		klog.Errorf("HTTP InternalServerError serving %s: %s", request, msg)
		response.WriteErrorString(http.StatusInternalServerError, msg)
	}
}
<commit_msg>Handling error returned by request.Request.ParseForm()<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package stats

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"path"
	"time"

	restful "github.com/emicklei/go-restful"
	cadvisorapi "github.com/google/cadvisor/info/v1"
	"k8s.io/klog"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	statsapi "k8s.io/kubernetes/pkg/kubelet/apis/stats/v1alpha1"
	"k8s.io/kubernetes/pkg/kubelet/cm"
	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
	"k8s.io/kubernetes/pkg/volume"
)

// Provider hosts methods required by stats handlers.
type Provider interface {
	// The following stats are provided by either CRI or cAdvisor.
	//
	// ListPodStats returns the stats of all the containers managed by pods.
	ListPodStats() ([]statsapi.PodStats, error)
	// ListPodStatsAndUpdateCPUNanoCoreUsage updates the cpu nano core usage for
	// the containers and returns the stats for all the pod-managed containers.
	ListPodCPUAndMemoryStats() ([]statsapi.PodStats, error)
	// ListPodStatsAndUpdateCPUNanoCoreUsage returns the stats of all the
	// containers managed by pods and force update the cpu usageNanoCores.
	// This is a workaround for CRI runtimes that do not integrate with
	// cadvisor. See https://github.com/kubernetes/kubernetes/issues/72788
	// for more details.
	ListPodStatsAndUpdateCPUNanoCoreUsage() ([]statsapi.PodStats, error)
	// ImageFsStats returns the stats of the image filesystem.
	ImageFsStats() (*statsapi.FsStats, error)

	// The following stats are provided by cAdvisor.
	//
	// GetCgroupStats returns the stats and the networking usage of the cgroup
	// with the specified cgroupName.
	GetCgroupStats(cgroupName string, updateStats bool) (*statsapi.ContainerStats, *statsapi.NetworkStats, error)
	// GetCgroupCPUAndMemoryStats returns the CPU and memory stats of the cgroup with the specified cgroupName.
	GetCgroupCPUAndMemoryStats(cgroupName string, updateStats bool) (*statsapi.ContainerStats, error)

	// RootFsStats returns the stats of the node root filesystem.
	RootFsStats() (*statsapi.FsStats, error)

	// The following stats are provided by cAdvisor for legacy usage.
	//
	// GetContainerInfo returns the information of the container with the
	// containerName managed by the pod with the uid.
	GetContainerInfo(podFullName string, uid types.UID, containerName string, req *cadvisorapi.ContainerInfoRequest) (*cadvisorapi.ContainerInfo, error)
	// GetRawContainerInfo returns the information of the container with the
	// containerName. If subcontainers is true, this function will return the
	// information of all the sub-containers as well.
	GetRawContainerInfo(containerName string, req *cadvisorapi.ContainerInfoRequest, subcontainers bool) (map[string]*cadvisorapi.ContainerInfo, error)

	// The following information is provided by Kubelet.
	//
	// GetPodByName returns the spec of the pod with the name in the specified
	// namespace.
	GetPodByName(namespace, name string) (*v1.Pod, bool)
	// GetNode returns the spec of the local node.
	GetNode() (*v1.Node, error)
	// GetNodeConfig returns the configuration of the local node.
	GetNodeConfig() cm.NodeConfig
	// ListVolumesForPod returns the stats of the volume used by the pod with
	// the podUID.
	ListVolumesForPod(podUID types.UID) (map[string]volume.Volume, bool)
	// GetPods returns the specs of all the pods running on this node.
	GetPods() []*v1.Pod

	// RlimitStats returns the rlimit stats of system.
	RlimitStats() (*statsapi.RlimitStats, error)

	// GetPodCgroupRoot returns the literal cgroupfs value for the cgroup containing all pods
	GetPodCgroupRoot() string

	// GetPodByCgroupfs provides the pod that maps to the specified cgroup literal, as well
	// as whether the pod was found.
	GetPodByCgroupfs(cgroupfs string) (*v1.Pod, bool)
}

type handler struct {
	provider        Provider
	summaryProvider SummaryProvider
}

// CreateHandlers creates the REST handlers for the stats.
func CreateHandlers(rootPath string, provider Provider, summaryProvider SummaryProvider, enableCAdvisorJSONEndpoints bool) *restful.WebService {
	h := &handler{provider, summaryProvider}

	ws := &restful.WebService{}
	ws.Path(rootPath).
		Produces(restful.MIME_JSON)

	type endpoint struct {
		path    string
		handler restful.RouteFunction
	}

	endpoints := []endpoint{
		{"/summary", h.handleSummary},
	}

	if enableCAdvisorJSONEndpoints {
		endpoints = append(endpoints,
			endpoint{"", h.handleStats},
			endpoint{"/container", h.handleSystemContainer},
			endpoint{"/{podName}/{containerName}", h.handlePodContainer},
			endpoint{"/{namespace}/{podName}/{uid}/{containerName}", h.handlePodContainer},
		)
	}

	for _, e := range endpoints {
		for _, method := range []string{"GET", "POST"} {
			ws.Route(ws.
				Method(method).
				Path(e.path).
				To(e.handler))
		}
	}

	return ws
}

type statsRequest struct {
	// The name of the container for which to request stats.
	// Default: /
	// +optional
	ContainerName string `json:"containerName,omitempty"`

	// Max number of stats to return.
	// If start and end time are specified this limit is ignored.
	// Default: 60
	// +optional
	NumStats int `json:"num_stats,omitempty"`

	// Start time for which to query information.
	// If omitted, the beginning of time is assumed.
	// +optional
	Start time.Time `json:"start,omitempty"`

	// End time for which to query information.
	// If omitted, current time is assumed.
	// +optional
	End time.Time `json:"end,omitempty"`

	// Whether to also include information from subcontainers.
	// Default: false.
	// +optional
	Subcontainers bool `json:"subcontainers,omitempty"`
}

func (r *statsRequest) cadvisorRequest() *cadvisorapi.ContainerInfoRequest {
	return &cadvisorapi.ContainerInfoRequest{
		NumStats: r.NumStats,
		Start:    r.Start,
		End:      r.End,
	}
}

func parseStatsRequest(request *restful.Request) (statsRequest, error) {
	// Default request.
	query := statsRequest{
		NumStats: 60,
	}

	err := json.NewDecoder(request.Request.Body).Decode(&query)
	if err != nil && err != io.EOF {
		return query, err
	}
	return query, nil
}

// Handles root container stats requests to /stats
func (h *handler) handleStats(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, "/stats", err)
		return
	}

	// Root container stats.
	statsMap, err := h.provider.GetRawContainerInfo("/", query.cadvisorRequest(), false)
	if err != nil {
		handleError(response, fmt.Sprintf("/stats %v", query), err)
		return
	}
	writeResponse(response, statsMap["/"])
}

// Handles stats summary requests to /stats/summary
// If "only_cpu_and_memory" GET param is true then only cpu and memory is returned in response.
func (h *handler) handleSummary(request *restful.Request, response *restful.Response) {
	onlyCPUAndMemory := false
	var err error
	err = request.Request.ParseForm()
	if err != nil {
		handleError(response, "/stats/summary", err)
		return
	}
	if onlyCluAndMemoryParam, found := request.Request.Form["only_cpu_and_memory"]; found &&
		len(onlyCluAndMemoryParam) == 1 && onlyCluAndMemoryParam[0] == "true" {
		onlyCPUAndMemory = true
	}
	var summary *statsapi.Summary
	if onlyCPUAndMemory {
		summary, err = h.summaryProvider.GetCPUAndMemoryStats()
	} else {
		// external calls to the summary API use cached stats
		forceStatsUpdate := false
		summary, err = h.summaryProvider.Get(forceStatsUpdate)
	}
	if err != nil {
		handleError(response, "/stats/summary", err)
	} else {
		writeResponse(response, summary)
	}
}

// Handles non-kubernetes container stats requests to /stats/container/
func (h *handler) handleSystemContainer(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, "/stats/container", err)
		return
	}

	// Non-Kubernetes container stats.
	containerName := path.Join("/", query.ContainerName)
	stats, err := h.provider.GetRawContainerInfo(
		containerName, query.cadvisorRequest(), query.Subcontainers)
	if err != nil {
		if _, ok := stats[containerName]; ok {
			// If the failure is partial, log it and return a best-effort response.
			klog.Errorf("Partial failure issuing GetRawContainerInfo(%v): %v", query, err)
		} else {
			handleError(response, fmt.Sprintf("/stats/container %v", query), err)
			return
		}
	}
	writeResponse(response, stats)
}

// Handles kubernetes pod/container stats requests to:
// /stats/<pod name>/<container name>
// /stats/<namespace>/<pod name>/<uid>/<container name>
func (h *handler) handlePodContainer(request *restful.Request, response *restful.Response) {
	query, err := parseStatsRequest(request)
	if err != nil {
		handleError(response, request.Request.URL.String(), err)
		return
	}

	// Default parameters.
	params := map[string]string{
		"namespace": metav1.NamespaceDefault,
		"uid":       "",
	}
	for k, v := range request.PathParameters() {
		params[k] = v
	}

	if params["podName"] == "" || params["containerName"] == "" {
		response.WriteErrorString(http.StatusBadRequest,
			fmt.Sprintf("Invalid pod container request: %v", params))
		return
	}

	pod, ok := h.provider.GetPodByName(params["namespace"], params["podName"])
	if !ok {
		klog.V(4).Infof("Container not found: %v", params)
		response.WriteError(http.StatusNotFound, kubecontainer.ErrContainerNotFound)
		return
	}
	stats, err := h.provider.GetContainerInfo(
		kubecontainer.GetPodFullName(pod),
		types.UID(params["uid"]),
		params["containerName"],
		query.cadvisorRequest())

	if err != nil {
		handleError(response, fmt.Sprintf("%s %v", request.Request.URL.String(), query), err)
		return
	}
	writeResponse(response, stats)
}

func writeResponse(response *restful.Response, stats interface{}) {
	if err := response.WriteAsJson(stats); err != nil {
		klog.Errorf("Error writing response: %v", err)
	}
}

// handleError serializes an error object into an HTTP response.
// request is provided for logging.
func handleError(response *restful.Response, request string, err error) {
	switch err {
	case kubecontainer.ErrContainerNotFound:
		response.WriteError(http.StatusNotFound, err)
	default:
		msg := fmt.Sprintf("Internal Error: %v", err)
		klog.Errorf("HTTP InternalServerError serving %s: %s", request, msg)
		response.WriteErrorString(http.StatusInternalServerError, msg)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package store

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

	utilfs "k8s.io/kubernetes/pkg/util/filesystem"
)

const (
	// Name prefix for the temporary files.
	tmpPrefix = "."
)

// FileStore is an implementation of the Store interface which stores data in files.
type FileStore struct {
	// Absolute path to the base directory for storing data files.
	directoryPath string

	// filesystem to use.
	filesystem utilfs.Filesystem
}

// NewFileStore returns an instance of FileStore.
func NewFileStore(path string, fs utilfs.Filesystem) (Store, error) {
	if err := ensureDirectory(fs, path); err != nil {
		return nil, err
	}
	return &FileStore{directoryPath: path, filesystem: fs}, nil
}

// Write writes the given data to a file named key.
func (f *FileStore) Write(key string, data []byte) error {
	if err := ValidateKey(key); err != nil {
		return err
	}
	if err := ensureDirectory(f.filesystem, f.directoryPath); err != nil {
		return err
	}

	return writeFile(f.filesystem, f.getPathByKey(key), data)
}

// Read reads the data from the file named key.
func (f *FileStore) Read(key string) ([]byte, error) {
	if err := ValidateKey(key); err != nil {
		return nil, err
	}
	bytes, err := f.filesystem.ReadFile(f.getPathByKey(key))
	if os.IsNotExist(err) {
		return bytes, ErrKeyNotFound
	}
	return bytes, err
}

// Delete deletes the key file.
func (f *FileStore) Delete(key string) error {
	if err := ValidateKey(key); err != nil {
		return err
	}
	return removePath(f.filesystem, f.getPathByKey(key))
}

// List returns all keys in the store.
func (f *FileStore) List() ([]string, error) {
	keys := make([]string, 0)
	files, err := f.filesystem.ReadDir(f.directoryPath)
	if err != nil {
		return keys, err
	}
	for _, f := range files {
		if !strings.HasPrefix(f.Name(), tmpPrefix) {
			keys = append(keys, f.Name())
		}
	}
	return keys, nil
}

// getPathByKey returns the full path of the file for the key.
func (f *FileStore) getPathByKey(key string) string {
	return filepath.Join(f.directoryPath, key)
}

// ensureDirectory creates the directory if it does not exist.
func ensureDirectory(fs utilfs.Filesystem, path string) error {
	if _, err := fs.Stat(path); err != nil {
		// MkdirAll returns nil if directory already exists.
		return fs.MkdirAll(path, 0755)
	}
	return nil
}

// writeFile writes data to path in a single transaction.
func writeFile(fs utilfs.Filesystem, path string, data []byte) (retErr error) {
	// Create a temporary file in the base directory of `path` with a prefix.
	tmpFile, err := fs.TempFile(filepath.Dir(path), tmpPrefix)
	if err != nil {
		return err
	}

	tmpPath := tmpFile.Name()
	shouldClose := true

	defer func() {
		// Close the file.
		if shouldClose {
			if err := tmpFile.Close(); err != nil {
				if retErr == nil {
					retErr = fmt.Errorf("close error: %v", err)
				} else {
					retErr = fmt.Errorf("failed to close temp file after error %v; close error: %v", retErr, err)
				}
			}
		}

		// Clean up the temp file on error.
		if retErr != nil && tmpPath != "" {
			if err := removePath(fs, tmpPath); err != nil {
				retErr = fmt.Errorf("failed to remove the temporary file (%q) after error %v; remove error: %v", tmpPath, retErr, err)
			}
		}
	}()

	// Write data.
	if _, err := tmpFile.Write(data); err != nil {
		return err
	}

	// Sync file.
	if err := tmpFile.Sync(); err != nil {
		return err
	}

	// Closing the file before renaming.
	err = tmpFile.Close()
	shouldClose = false
	if err != nil {
		return err
	}

	return fs.Rename(tmpPath, path)
}

func removePath(fs utilfs.Filesystem, path string) error {
	if err := fs.Remove(path); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}
<commit_msg>kubelet: remove superfluous function<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package store

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

	utilfs "k8s.io/kubernetes/pkg/util/filesystem"
)

const (
	// Name prefix for the temporary files.
	tmpPrefix = "."
)

// FileStore is an implementation of the Store interface which stores data in files.
type FileStore struct {
	// Absolute path to the base directory for storing data files.
	directoryPath string

	// filesystem to use.
	filesystem utilfs.Filesystem
}

// NewFileStore returns an instance of FileStore.
func NewFileStore(path string, fs utilfs.Filesystem) (Store, error) {
	if err := fs.MkdirAll(path, 0755); err != nil {
		return nil, err
	}
	return &FileStore{directoryPath: path, filesystem: fs}, nil
}

// Write writes the given data to a file named key.
func (f *FileStore) Write(key string, data []byte) error {
	if err := ValidateKey(key); err != nil {
		return err
	}
	if err := f.filesystem.MkdirAll(f.directoryPath, 0755); err != nil {
		return err
	}

	return writeFile(f.filesystem, f.getPathByKey(key), data)
}

// Read reads the data from the file named key.
func (f *FileStore) Read(key string) ([]byte, error) {
	if err := ValidateKey(key); err != nil {
		return nil, err
	}
	bytes, err := f.filesystem.ReadFile(f.getPathByKey(key))
	if os.IsNotExist(err) {
		return bytes, ErrKeyNotFound
	}
	return bytes, err
}

// Delete deletes the key file.
func (f *FileStore) Delete(key string) error {
	if err := ValidateKey(key); err != nil {
		return err
	}
	return removePath(f.filesystem, f.getPathByKey(key))
}

// List returns all keys in the store.
func (f *FileStore) List() ([]string, error) {
	keys := make([]string, 0)
	files, err := f.filesystem.ReadDir(f.directoryPath)
	if err != nil {
		return keys, err
	}
	for _, f := range files {
		if !strings.HasPrefix(f.Name(), tmpPrefix) {
			keys = append(keys, f.Name())
		}
	}
	return keys, nil
}

// getPathByKey returns the full path of the file for the key.
func (f *FileStore) getPathByKey(key string) string {
	return filepath.Join(f.directoryPath, key)
}

// writeFile writes data to path in a single transaction.
func writeFile(fs utilfs.Filesystem, path string, data []byte) (retErr error) {
	// Create a temporary file in the base directory of `path` with a prefix.
	tmpFile, err := fs.TempFile(filepath.Dir(path), tmpPrefix)
	if err != nil {
		return err
	}

	tmpPath := tmpFile.Name()
	shouldClose := true

	defer func() {
		// Close the file.
		if shouldClose {
			if err := tmpFile.Close(); err != nil {
				if retErr == nil {
					retErr = fmt.Errorf("close error: %v", err)
				} else {
					retErr = fmt.Errorf("failed to close temp file after error %v; close error: %v", retErr, err)
				}
			}
		}

		// Clean up the temp file on error.
		if retErr != nil && tmpPath != "" {
			if err := removePath(fs, tmpPath); err != nil {
				retErr = fmt.Errorf("failed to remove the temporary file (%q) after error %v; remove error: %v", tmpPath, retErr, err)
			}
		}
	}()

	// Write data.
	if _, err := tmpFile.Write(data); err != nil {
		return err
	}

	// Sync file.
	if err := tmpFile.Sync(); err != nil {
		return err
	}

	// Closing the file before renaming.
	err = tmpFile.Close()
	shouldClose = false
	if err != nil {
		return err
	}

	return fs.Rename(tmpPath, path)
}

func removePath(fs utilfs.Filesystem, path string) error {
	if err := fs.Remove(path); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package printers

import (
	"fmt"
	"reflect"
	"testing"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	metav1beta1 "k8s.io/apimachinery/pkg/apis/meta/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

type TestPrintType struct {
	Data string
}

func (obj *TestPrintType) GetObjectKind() schema.ObjectKind { return schema.EmptyObjectKind }
func (obj *TestPrintType) DeepCopyObject() runtime.Object {
	if obj == nil {
		return nil
	}
	clone := *obj
	return &clone
}

func PrintCustomType(obj *TestPrintType, options GenerateOptions) ([]metav1beta1.TableRow, error) {
	return []metav1beta1.TableRow{{Cells: []interface{}{obj.Data}}}, nil
}

func ErrorPrintHandler(obj *TestPrintType, options GenerateOptions) ([]metav1beta1.TableRow, error) {
	return nil, fmt.Errorf("ErrorPrintHandler error")
}

func TestCustomTypePrinting(t *testing.T) {
	columns := []metav1beta1.TableColumnDefinition{{Name: "Data"}}
	generator := NewTableGenerator()
	generator.TableHandler(columns, PrintCustomType)

	obj := TestPrintType{"test object"}
	table, err := generator.GenerateTable(&obj, GenerateOptions{})
	if err != nil {
		t.Fatalf("An error occurred generating the table for custom type: %#v", err)
	}

	expectedTable := &metav1.Table{
		ColumnDefinitions: []metav1.TableColumnDefinition{{Name: "Data"}},
		Rows:              []metav1.TableRow{{Cells: []interface{}{"test object"}}},
	}
	if !reflect.DeepEqual(expectedTable, table) {
		t.Errorf("Error generating table from custom type. Expected (%#v), got (%#v)", expectedTable, table)
	}
}

func TestPrintHandlerError(t *testing.T) {
	columns := []metav1beta1.TableColumnDefinition{{Name: "Data"}}
	generator := NewTableGenerator()
	generator.TableHandler(columns, ErrorPrintHandler)
	obj := TestPrintType{"test object"}
	_, err := generator.GenerateTable(&obj, GenerateOptions{})
	if err == nil || err.Error() != "ErrorPrintHandler error" {
		t.Errorf("Did not get the expected error: %#v", err)
	}
}
<commit_msg>Check for errors ahead of time<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package printers

import (
	"fmt"
	"reflect"
	"testing"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	metav1beta1 "k8s.io/apimachinery/pkg/apis/meta/v1beta1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

type TestPrintType struct {
	Data string
}

func (obj *TestPrintType) GetObjectKind() schema.ObjectKind { return schema.EmptyObjectKind }
func (obj *TestPrintType) DeepCopyObject() runtime.Object {
	if obj == nil {
		return nil
	}
	clone := *obj
	return &clone
}

func PrintCustomType(obj *TestPrintType, options GenerateOptions) ([]metav1beta1.TableRow, error) {
	return []metav1beta1.TableRow{{Cells: []interface{}{obj.Data}}}, nil
}

func ErrorPrintHandler(obj *TestPrintType, options GenerateOptions) ([]metav1beta1.TableRow, error) {
	return nil, fmt.Errorf("ErrorPrintHandler error")
}

func TestCustomTypePrinting(t *testing.T) {
	columns := []metav1beta1.TableColumnDefinition{{Name: "Data"}}
	generator := NewTableGenerator()
	err := generator.TableHandler(columns, PrintCustomType)
	if err != nil {
		t.Fatalf("An error occurred when adds a print handler with a given set of columns: %#v", err)
	}

	obj := TestPrintType{"test object"}
	table, err := generator.GenerateTable(&obj, GenerateOptions{})
	if err != nil {
		t.Fatalf("An error occurred generating the table for custom type: %#v", err)
	}

	expectedTable := &metav1.Table{
		ColumnDefinitions: []metav1.TableColumnDefinition{{Name: "Data"}},
		Rows:              []metav1.TableRow{{Cells: []interface{}{"test object"}}},
	}
	if !reflect.DeepEqual(expectedTable, table) {
		t.Errorf("Error generating table from custom type. Expected (%#v), got (%#v)", expectedTable, table)
	}
}

func TestPrintHandlerError(t *testing.T) {
	columns := []metav1beta1.TableColumnDefinition{{Name: "Data"}}
	generator := NewTableGenerator()
	err := generator.TableHandler(columns, ErrorPrintHandler)
	if err != nil {
		t.Fatalf("An error occurred when adds a print handler with a given set of columns: %#v", err)
	}

	obj := TestPrintType{"test object"}
	_, err = generator.GenerateTable(&obj, GenerateOptions{})
	if err == nil || err.Error() != "ErrorPrintHandler error" {
		t.Errorf("Did not get the expected error: %#v", err)
	}
}
<|endoftext|>