<commit_before>package restapi

import (
	"crypto/tls"
	"log"
	"net/http"
	"strings"

	errors "github.com/go-openapi/errors"
	runtime "github.com/go-openapi/runtime"
	middleware "github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/swag"
	"golang.org/x/net/context"

	"github.com/influxdata/mrfusion"
	"github.com/influxdata/mrfusion/dist"
	"github.com/influxdata/mrfusion/handlers"
	"github.com/influxdata/mrfusion/influx"
	"github.com/influxdata/mrfusion/mock"
	"github.com/influxdata/mrfusion/restapi/operations"
)

// This file is safe to edit. Once it exists it will not be overwritten

//go:generate swagger generate server --target .. --name  --spec ../swagger.yaml --with-context

var devFlags = struct {
	Develop bool `short:"d" long:"develop" description:"Run server in develop mode."`
}{}

var influxFlags = struct {
	Server string `short:"s" long:"server" description:"Full URL of InfluxDB server (http://localhost:8086)"`
}{}

func configureFlags(api *operations.MrFusionAPI) {
	api.CommandLineOptionsGroups = []swag.CommandLineOptionsGroup{
		swag.CommandLineOptionsGroup{
			ShortDescription: "Develop Mode server",
			LongDescription:  "Server will use the ui/build directory directly.",
			Options:          &devFlags,
		},
		swag.CommandLineOptionsGroup{
			ShortDescription: "Default Time Series Backend",
			LongDescription:  "Specify the url of an InfxluDB server",
			Options:          &influxFlags,
		},
	}
}

func assets() mrfusion.Assets {
	if devFlags.Develop {
		return &dist.DebugAssets{
			Dir:     "ui/build",
			Default: "ui/build/index.html",
		}
	}
	return &dist.BindataAssets{
		Prefix:  "ui/build",
		Default: "index.html",
	}
}

func configureAPI(api *operations.MrFusionAPI) http.Handler {
	// configure the api here
	api.ServeError = errors.ServeError

	// Set your custom logger if needed. Default one is log.Printf
	// Expected interface func(string, ...interface{})
	//
	// Example:
	// s.api.Logger = log.Printf

	api.JSONConsumer = runtime.JSONConsumer()

	api.JSONProducer = runtime.JSONProducer()

	mockHandler := mock.NewHandler()

	api.DeleteDashboardsIDHandler = operations.DeleteDashboardsIDHandlerFunc(func(ctx context.Context, params operations.DeleteDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteDashboardsID has not yet been implemented")
	})
	api.DeleteSourcesIDHandler = operations.DeleteSourcesIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesID has not yet been implemented")
	})
	api.DeleteSourcesIDRolesRoleIDHandler = operations.DeleteSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDRolesRoleID has not yet been implemented")
	})
	api.DeleteSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.DeleteSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDUsersUserIDExplorationsExplorationIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDUsersUserIDExplorationsExplorationID has not yet been implemented")
	})
	api.DeleteSourcesIDUsersUserIDHandler = operations.DeleteSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDUsersUserID has not yet been implemented")
	})
	api.GetHandler = operations.GetHandlerFunc(func(ctx context.Context, params operations.GetParams) middleware.Responder {
		return middleware.NotImplemented("operation .Get has not yet been implemented")
	})
	api.GetDashboardsHandler = operations.GetDashboardsHandlerFunc(func(ctx context.Context, params operations.GetDashboardsParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetDashboards has not yet been implemented")
	})
	api.GetDashboardsIDHandler = operations.GetDashboardsIDHandlerFunc(func(ctx context.Context, params operations.GetDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetDashboardsID has not yet been implemented")
	})

	api.GetSourcesHandler = operations.GetSourcesHandlerFunc(mockHandler.Sources)
	api.GetSourcesIDHandler = operations.GetSourcesIDHandlerFunc(mockHandler.SourcesID)

	api.GetSourcesIDPermissionsHandler = operations.GetSourcesIDPermissionsHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDPermissionsParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDPermissions has not yet been implemented")
	})
	api.GetSourcesIDRolesHandler = operations.GetSourcesIDRolesHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDRolesParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDRoles has not yet been implemented")
	})
	api.GetSourcesIDRolesRoleIDHandler = operations.GetSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDRolesRoleID has not yet been implemented")
	})

	api.GetSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.GetSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(mockHandler.Exploration)

	api.GetSourcesIDUsersHandler = operations.GetSourcesIDUsersHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDUsersParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDUsers has not yet been implemented")
	})
	api.GetSourcesIDUsersUserIDHandler = operations.GetSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDUsersUserID has not yet been implemented")
	})

	api.GetSourcesIDUsersUserIDExplorationsHandler = operations.GetSourcesIDUsersUserIDExplorationsHandlerFunc(mockHandler.Explorations)
	api.PatchSourcesIDHandler = operations.PatchSourcesIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesID has not yet been implemented")
	})
	api.PatchSourcesIDRolesRoleIDHandler = operations.PatchSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDRolesRoleID has not yet been implemented")
	})
	api.PatchSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.PatchSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDUsersUserIDExplorationsExplorationIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDUsersUserIDExplorationsExplorationID has not yet been implemented")
	})
	api.PatchSourcesIDUsersUserIDHandler = operations.PatchSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDUsersUserID has not yet been implemented")
	})
	api.PostDashboardsHandler = operations.PostDashboardsHandlerFunc(func(ctx context.Context, params operations.PostDashboardsParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostDashboards has not yet been implemented")
	})
	api.PostSourcesHandler = operations.PostSourcesHandlerFunc(func(ctx context.Context, params operations.PostSourcesParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSources has not yet been implemented")
	})

	if len(influxFlags.Server) > 0 {
		c, err := influx.NewClient(influxFlags.Server)
		if err != nil {
			panic(err)
		}
		h := handlers.InfluxProxy{
			TimeSeries: c,
		}
		api.PostSourcesIDProxyHandler = operations.PostSourcesIDProxyHandlerFunc(h.Proxy)
	} else {
		api.PostSourcesIDProxyHandler = operations.PostSourcesIDProxyHandlerFunc(mockHandler.Proxy)
	}

	api.PostSourcesIDRolesHandler = operations.PostSourcesIDRolesHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDRolesParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDRoles has not yet been implemented")
	})
	api.PostSourcesIDUsersHandler = operations.PostSourcesIDUsersHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDUsersParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDUsers has not yet been implemented")
	})
	api.PostSourcesIDUsersUserIDExplorationsHandler = operations.PostSourcesIDUsersUserIDExplorationsHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDUsersUserIDExplorationsParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDUsersUserIDExplorations has not yet been implemented")
	})
	api.PutDashboardsIDHandler = operations.PutDashboardsIDHandlerFunc(func(ctx context.Context, params operations.PutDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PutDashboardsID has not yet been implemented")
	})

	api.GetSourcesIDMonitoredHandler = operations.GetSourcesIDMonitoredHandlerFunc(mockHandler.MonitoredServices)

	api.ServerShutdown = func() {}

	handler := setupGlobalMiddleware(api.Serve(setupMiddlewares))
	return handler
}

// The TLS configuration before HTTPS server starts.
func configureTLS(tlsConfig *tls.Config) {
	// Make all necessary changes to the TLS configuration here.
}

// The middleware configuration is for the handler executors. These do not apply to the swagger.json document.
// The middleware executes after routing but before authentication, binding and validation
func setupMiddlewares(handler http.Handler) http.Handler {
	return handler
}

// The middleware configuration happens before anything, this middleware also applies to serving the swagger.json document.
// So this is a good place to plug in a panic handling middleware, logging and metrics
func setupGlobalMiddleware(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
		if strings.Contains(r.URL.Path, "/chronograf/v1") {
			handler.ServeHTTP(w, r)
			return
		} else if r.URL.Path == "//" {
			http.Redirect(w, r, "/index.html", http.StatusFound)
		} else {
			assets().Handler().ServeHTTP(w, r)
			return
		}
	})
}
<commit_msg>Include env variable for influx server<commit_after>package restapi

import (
	"crypto/tls"
	"log"
	"net/http"
	"strings"

	errors "github.com/go-openapi/errors"
	runtime "github.com/go-openapi/runtime"
	middleware "github.com/go-openapi/runtime/middleware"
	"github.com/go-openapi/swag"
	"golang.org/x/net/context"

	"github.com/influxdata/mrfusion"
	"github.com/influxdata/mrfusion/dist"
	"github.com/influxdata/mrfusion/handlers"
	"github.com/influxdata/mrfusion/influx"
	"github.com/influxdata/mrfusion/mock"
	"github.com/influxdata/mrfusion/restapi/operations"
)

// This file is safe to edit. Once it exists it will not be overwritten

//go:generate swagger generate server --target .. --name  --spec ../swagger.yaml --with-context

var devFlags = struct {
	Develop bool `short:"d" long:"develop" description:"Run server in develop mode."`
}{}

var influxFlags = struct {
	Server string `short:"s" long:"server" description:"Full URL of InfluxDB server (http://localhost:8086)" env:"INFLUX_HOST"`
}{}

func configureFlags(api *operations.MrFusionAPI) {
	api.CommandLineOptionsGroups = []swag.CommandLineOptionsGroup{
		swag.CommandLineOptionsGroup{
			ShortDescription: "Develop Mode server",
			LongDescription:  "Server will use the ui/build directory directly.",
			Options:          &devFlags,
		},
		swag.CommandLineOptionsGroup{
			ShortDescription: "Default Time Series Backend",
			LongDescription:  "Specify the url of an InfxluDB server",
			Options:          &influxFlags,
		},
	}
}

func assets() mrfusion.Assets {
	if devFlags.Develop {
		return &dist.DebugAssets{
			Dir:     "ui/build",
			Default: "ui/build/index.html",
		}
	}
	return &dist.BindataAssets{
		Prefix:  "ui/build",
		Default: "index.html",
	}
}

func configureAPI(api *operations.MrFusionAPI) http.Handler {
	// configure the api here
	api.ServeError = errors.ServeError

	// Set your custom logger if needed. Default one is log.Printf
	// Expected interface func(string, ...interface{})
	//
	// Example:
	// s.api.Logger = log.Printf

	api.JSONConsumer = runtime.JSONConsumer()

	api.JSONProducer = runtime.JSONProducer()

	mockHandler := mock.NewHandler()

	api.DeleteDashboardsIDHandler = operations.DeleteDashboardsIDHandlerFunc(func(ctx context.Context, params operations.DeleteDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteDashboardsID has not yet been implemented")
	})
	api.DeleteSourcesIDHandler = operations.DeleteSourcesIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesID has not yet been implemented")
	})
	api.DeleteSourcesIDRolesRoleIDHandler = operations.DeleteSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDRolesRoleID has not yet been implemented")
	})
	api.DeleteSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.DeleteSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDUsersUserIDExplorationsExplorationIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDUsersUserIDExplorationsExplorationID has not yet been implemented")
	})
	api.DeleteSourcesIDUsersUserIDHandler = operations.DeleteSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.DeleteSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .DeleteSourcesIDUsersUserID has not yet been implemented")
	})
	api.GetHandler = operations.GetHandlerFunc(func(ctx context.Context, params operations.GetParams) middleware.Responder {
		return middleware.NotImplemented("operation .Get has not yet been implemented")
	})
	api.GetDashboardsHandler = operations.GetDashboardsHandlerFunc(func(ctx context.Context, params operations.GetDashboardsParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetDashboards has not yet been implemented")
	})
	api.GetDashboardsIDHandler = operations.GetDashboardsIDHandlerFunc(func(ctx context.Context, params operations.GetDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetDashboardsID has not yet been implemented")
	})

	api.GetSourcesHandler = operations.GetSourcesHandlerFunc(mockHandler.Sources)
	api.GetSourcesIDHandler = operations.GetSourcesIDHandlerFunc(mockHandler.SourcesID)

	api.GetSourcesIDPermissionsHandler = operations.GetSourcesIDPermissionsHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDPermissionsParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDPermissions has not yet been implemented")
	})
	api.GetSourcesIDRolesHandler = operations.GetSourcesIDRolesHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDRolesParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDRoles has not yet been implemented")
	})
	api.GetSourcesIDRolesRoleIDHandler = operations.GetSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDRolesRoleID has not yet been implemented")
	})

	api.GetSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.GetSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(mockHandler.Exploration)

	api.GetSourcesIDUsersHandler = operations.GetSourcesIDUsersHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDUsersParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDUsers has not yet been implemented")
	})
	api.GetSourcesIDUsersUserIDHandler = operations.GetSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.GetSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .GetSourcesIDUsersUserID has not yet been implemented")
	})

	api.GetSourcesIDUsersUserIDExplorationsHandler = operations.GetSourcesIDUsersUserIDExplorationsHandlerFunc(mockHandler.Explorations)
	api.PatchSourcesIDHandler = operations.PatchSourcesIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesID has not yet been implemented")
	})
	api.PatchSourcesIDRolesRoleIDHandler = operations.PatchSourcesIDRolesRoleIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDRolesRoleIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDRolesRoleID has not yet been implemented")
	})
	api.PatchSourcesIDUsersUserIDExplorationsExplorationIDHandler = operations.PatchSourcesIDUsersUserIDExplorationsExplorationIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDUsersUserIDExplorationsExplorationIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDUsersUserIDExplorationsExplorationID has not yet been implemented")
	})
	api.PatchSourcesIDUsersUserIDHandler = operations.PatchSourcesIDUsersUserIDHandlerFunc(func(ctx context.Context, params operations.PatchSourcesIDUsersUserIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PatchSourcesIDUsersUserID has not yet been implemented")
	})
	api.PostDashboardsHandler = operations.PostDashboardsHandlerFunc(func(ctx context.Context, params operations.PostDashboardsParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostDashboards has not yet been implemented")
	})
	api.PostSourcesHandler = operations.PostSourcesHandlerFunc(func(ctx context.Context, params operations.PostSourcesParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSources has not yet been implemented")
	})

	if len(influxFlags.Server) > 0 {
		c, err := influx.NewClient(influxFlags.Server)
		if err != nil {
			panic(err)
		}
		h := handlers.InfluxProxy{
			TimeSeries: c,
		}
		api.PostSourcesIDProxyHandler = operations.PostSourcesIDProxyHandlerFunc(h.Proxy)
	} else {
		api.PostSourcesIDProxyHandler = operations.PostSourcesIDProxyHandlerFunc(mockHandler.Proxy)
	}

	api.PostSourcesIDRolesHandler = operations.PostSourcesIDRolesHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDRolesParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDRoles has not yet been implemented")
	})
	api.PostSourcesIDUsersHandler = operations.PostSourcesIDUsersHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDUsersParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDUsers has not yet been implemented")
	})
	api.PostSourcesIDUsersUserIDExplorationsHandler = operations.PostSourcesIDUsersUserIDExplorationsHandlerFunc(func(ctx context.Context, params operations.PostSourcesIDUsersUserIDExplorationsParams) middleware.Responder {
		return middleware.NotImplemented("operation .PostSourcesIDUsersUserIDExplorations has not yet been implemented")
	})
	api.PutDashboardsIDHandler = operations.PutDashboardsIDHandlerFunc(func(ctx context.Context, params operations.PutDashboardsIDParams) middleware.Responder {
		return middleware.NotImplemented("operation .PutDashboardsID has not yet been implemented")
	})

	api.GetSourcesIDMonitoredHandler = operations.GetSourcesIDMonitoredHandlerFunc(mockHandler.MonitoredServices)

	api.ServerShutdown = func() {}

	handler := setupGlobalMiddleware(api.Serve(setupMiddlewares))
	return handler
}

// The TLS configuration before HTTPS server starts.
func configureTLS(tlsConfig *tls.Config) {
	// Make all necessary changes to the TLS configuration here.
}

// The middleware configuration is for the handler executors. These do not apply to the swagger.json document.
// The middleware executes after routing but before authentication, binding and validation
func setupMiddlewares(handler http.Handler) http.Handler {
	return handler
}

// The middleware configuration happens before anything, this middleware also applies to serving the swagger.json document.
// So this is a good place to plug in a panic handling middleware, logging and metrics
func setupGlobalMiddleware(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
		if strings.Contains(r.URL.Path, "/chronograf/v1") {
			handler.ServeHTTP(w, r)
			return
		} else if r.URL.Path == "//" {
			http.Redirect(w, r, "/index.html", http.StatusFound)
		} else {
			assets().Handler().ServeHTTP(w, r)
			return
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build !plan9
// +build !plan9

package cpio

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/ls"
	"github.com/u-root/u-root/pkg/uio"
	"github.com/u-root/u-root/pkg/upath"
	"golang.org/x/sys/unix"
)

var modeMap = map[uint64]os.FileMode{
	modeSocket:  os.ModeSocket,
	modeSymlink: os.ModeSymlink,
	modeFile:    0,
	modeBlock:   os.ModeDevice,
	modeDir:     os.ModeDir,
	modeChar:    os.ModeCharDevice,
	modeFIFO:    os.ModeNamedPipe,
}

// setModes sets the modes, changing the easy ones first and the harder ones last.
// In this way, we set as much as we can before bailing out.
// N.B.: if you set something with S_ISUID, then change the owner,
// the kernel (Linux, OSX, etc.) clears S_ISUID (a good idea). So, the simple thing:
// Do the chmod operations in order of difficulty, and give up as soon as we fail.
// Set the basic permissions -- not including SUID, GUID, etc.
// Set the times
// Set the owner
// Set ALL the mode bits, in case we need to do SUID, etc. If we could not
// set the owner, we won't even try this operation of course, so we won't
// have SUID incorrectly set for the wrong user.
func setModes(r Record) error {
	if err := os.Chmod(r.Name, toFileMode(r)&os.ModePerm); err != nil {
		return err
	}
	/*if err := os.Chtimes(r.Name, time.Time{}, time.Unix(int64(r.MTime), 0)); err != nil {
		return err
	}*/
	if err := os.Chown(r.Name, int(r.UID), int(r.GID)); err != nil {
		return err
	}
	if err := os.Chmod(r.Name, toFileMode(r)); err != nil {
		return err
	}
	return nil
}

func toFileMode(r Record) os.FileMode {
	m := os.FileMode(perm(r))
	if r.Mode&unix.S_ISUID != 0 {
		m |= os.ModeSetuid
	}
	if r.Mode&unix.S_ISGID != 0 {
		m |= os.ModeSetgid
	}
	if r.Mode&unix.S_ISVTX != 0 {
		m |= os.ModeSticky
	}
	return m
}

func perm(r Record) uint32 {
	return uint32(r.Mode) & modePermissions
}

func dev(r Record) int {
	return int(r.Rmajor<<8 | r.Rminor)
}

func linuxModeToFileType(m uint64) (os.FileMode, error) {
	if t, ok := modeMap[m&modeTypeMask]; ok {
		return t, nil
	}
	return 0, fmt.Errorf("invalid file type %#o", m&modeTypeMask)
}

// CreateFile creates a local file for f relative to the current working
// directory.
//
// CreateFile will attempt to set all metadata for the file, including
// ownership, times, and permissions.
func CreateFile(f Record) error {
	return CreateFileInRoot(f, ".", true)
}

// CreateFileInRoot creates a local file for f relative to rootDir.
//
// It will attempt to set all metadata for the file, including ownership,
// times, and permissions. If these fail, it only returns an error if
// forcePriv is true.
//
// Block and char device creation will only return error if forcePriv is true.
func CreateFileInRoot(f Record, rootDir string, forcePriv bool) error {
	m, err := linuxModeToFileType(f.Mode)
	if err != nil {
		return err
	}

	f.Name, err = upath.SafeFilepathJoin(rootDir, f.Name)
	if err != nil {
		// The behavior is to skip files which are unsafe due to
		// zipslip, but continue extracting everything else.
		log.Printf("Warning: Skipping file %q due to: %v", f.Name, err)
		return nil
	}
	dir := filepath.Dir(f.Name)
	// The problem: many cpio archives do not specify the directories and
	// hence the permissions. They just specify the whole path.  In order
	// to create files in these directories, we have to make them at least
	// mode 755.
	if _, err := os.Stat(dir); os.IsNotExist(err) && len(dir) > 0 {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return fmt.Errorf("CreateFileInRoot %q: %v", f.Name, err)
		}
	}

	switch m {
	case os.ModeSocket, os.ModeNamedPipe:
		return fmt.Errorf("%q: type %v: cannot create IPC endpoints", f.Name, m)

	case os.ModeSymlink:
		content, err := io.ReadAll(uio.Reader(f))
		if err != nil {
			return err
		}
		return os.Symlink(string(content), f.Name)

	case os.FileMode(0):
		nf, err := os.Create(f.Name)
		if err != nil {
			return err
		}
		defer nf.Close()
		if _, err := io.Copy(nf, uio.Reader(f)); err != nil {
			return err
		}

	case os.ModeDir:
		if err := os.MkdirAll(f.Name, toFileMode(f)); err != nil {
			return err
		}

	case os.ModeDevice:
		if err := mknod(f.Name, perm(f)|syscall.S_IFBLK, dev(f)); err != nil && forcePriv {
			return err
		}

	case os.ModeCharDevice:
		if err := mknod(f.Name, perm(f)|syscall.S_IFCHR, dev(f)); err != nil && forcePriv {
			return err
		}

	default:
		return fmt.Errorf("%v: Unknown type %#o", f.Name, m)
	}

	if err := setModes(f); err != nil && forcePriv {
		return err
	}
	return nil
}

// Inumber and devnumbers are unique to Unix-like
// operating systems. You can not uniquely disambiguate a file in a
// Unix system with just an inumber, you need a device number too.
// To handle hard links (unique to Unix) we need to figure out if a
// given file has been seen before. To do this we see if a file has the
// same [dev,ino] tuple as one we have seen. If so, we won't bother
// reading it in.

type devInode struct {
	dev uint64
	ino uint64
}

// A Recorder is a structure that contains variables used to calculate
// file parameters such as inode numbers for a CPIO file. The life-time
// of a Record structure is meant to be the same as the construction of a
// single CPIO archive. Do not reuse between CPIOs if you don't know what
// you're doing.
type Recorder struct {
	inodeMap map[devInode]Info
	inumber  uint64
}

// Certain elements of the file can not be set by cpio:
// the Inode #
// the Dev
// maintaining these elements leaves us with a non-reproducible
// output stream. In this function, we figure out what inumber
// we need to use, and clear out anything we can.
// We always zero the Dev.
// We try to find the matching inode. If found, we use its inumber.
// If not, we get a new inumber for it and save the inode away.
// This eliminates two of the messier parts of creating reproducible
// output streams.
func (r *Recorder) inode(i Info) (Info, bool) {
	d := devInode{dev: i.Dev, ino: i.Ino}
	i.Dev = 0

	if d, ok := r.inodeMap[d]; ok {
		i.Ino = d.Ino
		return i, true
	}

	i.Ino = r.inumber
	r.inumber++
	r.inodeMap[d] = i

	return i, false
}

// GetRecord returns a cpio Record for the given path on the local file system.
//
// GetRecord does not follow symlinks. If path is a symlink, the record
// returned will reflect that symlink.
func (r *Recorder) GetRecord(path string) (Record, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return Record{}, err
	}

	sys := fi.Sys().(*syscall.Stat_t)
	info, done := r.inode(sysInfo(path, sys))

	switch fi.Mode() & os.ModeType {
	case 0: // Regular file.
		if done {
			return Record{Info: info}, nil
		}
		return Record{Info: info, ReaderAt: uio.NewLazyFile(path)}, nil

	case os.ModeSymlink:
		linkname, err := os.Readlink(path)
		if err != nil {
			return Record{}, err
		}
		return StaticRecord([]byte(linkname), info), nil

	default:
		return StaticRecord(nil, info), nil
	}
}

// NewRecorder creates a new Recorder.
//
// A recorder is a structure that contains variables used to calculate
// file parameters such as inode numbers for a CPIO file. The life-time
// of a Record structure is meant to be the same as the construction of a
// single CPIO archive. Do not reuse between CPIOs if you don't know what
// you're doing.
func NewRecorder() *Recorder {
	return &Recorder{make(map[devInode]Info), 2}
}

// LSInfoFromRecord converts a Record to be usable with the ls package for
// listing files.
func LSInfoFromRecord(rec Record) ls.FileInfo {
	var target string

	mode := modeFromLinux(rec.Mode)
	if mode&os.ModeType == os.ModeSymlink {
		if l, err := uio.ReadAll(rec); err != nil {
			target = err.Error()
		} else {
			target = string(l)
		}
	}

	return ls.FileInfo{
		Name:          rec.Name,
		Mode:          mode,
		Rdev:          unix.Mkdev(uint32(rec.Rmajor), uint32(rec.Rminor)),
		UID:           uint32(rec.UID),
		GID:           uint32(rec.GID),
		Size:          int64(rec.FileSize),
		MTime:         time.Unix(int64(rec.MTime), 0).UTC(),
		SymlinkTarget: target,
	}
}
<commit_msg>cpio: fix to allow a same file as source for multiple extra files<commit_after>// Copyright 2013-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build !plan9
// +build !plan9

package cpio

import (
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/ls"
	"github.com/u-root/u-root/pkg/uio"
	"github.com/u-root/u-root/pkg/upath"
	"golang.org/x/sys/unix"
)

var modeMap = map[uint64]os.FileMode{
	modeSocket:  os.ModeSocket,
	modeSymlink: os.ModeSymlink,
	modeFile:    0,
	modeBlock:   os.ModeDevice,
	modeDir:     os.ModeDir,
	modeChar:    os.ModeCharDevice,
	modeFIFO:    os.ModeNamedPipe,
}

// setModes sets the modes, changing the easy ones first and the harder ones last.
// In this way, we set as much as we can before bailing out.
// N.B.: if you set something with S_ISUID, then change the owner,
// the kernel (Linux, OSX, etc.) clears S_ISUID (a good idea). So, the simple thing:
// Do the chmod operations in order of difficulty, and give up as soon as we fail.
// Set the basic permissions -- not including SUID, GUID, etc.
// Set the times
// Set the owner
// Set ALL the mode bits, in case we need to do SUID, etc. If we could not
// set the owner, we won't even try this operation of course, so we won't
// have SUID incorrectly set for the wrong user.
func setModes(r Record) error {
	if err := os.Chmod(r.Name, toFileMode(r)&os.ModePerm); err != nil {
		return err
	}
	/*if err := os.Chtimes(r.Name, time.Time{}, time.Unix(int64(r.MTime), 0)); err != nil {
		return err
	}*/
	if err := os.Chown(r.Name, int(r.UID), int(r.GID)); err != nil {
		return err
	}
	if err := os.Chmod(r.Name, toFileMode(r)); err != nil {
		return err
	}
	return nil
}

func toFileMode(r Record) os.FileMode {
	m := os.FileMode(perm(r))
	if r.Mode&unix.S_ISUID != 0 {
		m |= os.ModeSetuid
	}
	if r.Mode&unix.S_ISGID != 0 {
		m |= os.ModeSetgid
	}
	if r.Mode&unix.S_ISVTX != 0 {
		m |= os.ModeSticky
	}
	return m
}

func perm(r Record) uint32 {
	return uint32(r.Mode) & modePermissions
}

func dev(r Record) int {
	return int(r.Rmajor<<8 | r.Rminor)
}

func linuxModeToFileType(m uint64) (os.FileMode, error) {
	if t, ok := modeMap[m&modeTypeMask]; ok {
		return t, nil
	}
	return 0, fmt.Errorf("invalid file type %#o", m&modeTypeMask)
}

// CreateFile creates a local file for f relative to the current working
// directory.
//
// CreateFile will attempt to set all metadata for the file, including
// ownership, times, and permissions.
func CreateFile(f Record) error {
	return CreateFileInRoot(f, ".", true)
}

// CreateFileInRoot creates a local file for f relative to rootDir.
//
// It will attempt to set all metadata for the file, including ownership,
// times, and permissions. If these fail, it only returns an error if
// forcePriv is true.
//
// Block and char device creation will only return error if forcePriv is true.
func CreateFileInRoot(f Record, rootDir string, forcePriv bool) error {
	m, err := linuxModeToFileType(f.Mode)
	if err != nil {
		return err
	}

	f.Name, err = upath.SafeFilepathJoin(rootDir, f.Name)
	if err != nil {
		// The behavior is to skip files which are unsafe due to
		// zipslip, but continue extracting everything else.
		log.Printf("Warning: Skipping file %q due to: %v", f.Name, err)
		return nil
	}
	dir := filepath.Dir(f.Name)
	// The problem: many cpio archives do not specify the directories and
	// hence the permissions. They just specify the whole path.  In order
	// to create files in these directories, we have to make them at least
	// mode 755.
	if _, err := os.Stat(dir); os.IsNotExist(err) && len(dir) > 0 {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return fmt.Errorf("CreateFileInRoot %q: %v", f.Name, err)
		}
	}

	switch m {
	case os.ModeSocket, os.ModeNamedPipe:
		return fmt.Errorf("%q: type %v: cannot create IPC endpoints", f.Name, m)

	case os.ModeSymlink:
		content, err := io.ReadAll(uio.Reader(f))
		if err != nil {
			return err
		}
		return os.Symlink(string(content), f.Name)

	case os.FileMode(0):
		nf, err := os.Create(f.Name)
		if err != nil {
			return err
		}
		defer nf.Close()
		if _, err := io.Copy(nf, uio.Reader(f)); err != nil {
			return err
		}

	case os.ModeDir:
		if err := os.MkdirAll(f.Name, toFileMode(f)); err != nil {
			return err
		}

	case os.ModeDevice:
		if err := mknod(f.Name, perm(f)|syscall.S_IFBLK, dev(f)); err != nil && forcePriv {
			return err
		}

	case os.ModeCharDevice:
		if err := mknod(f.Name, perm(f)|syscall.S_IFCHR, dev(f)); err != nil && forcePriv {
			return err
		}

	default:
		return fmt.Errorf("%v: Unknown type %#o", f.Name, m)
	}

	if err := setModes(f); err != nil && forcePriv {
		return err
	}
	return nil
}

// Inumber and devnumbers are unique to Unix-like
// operating systems. You can not uniquely disambiguate a file in a
// Unix system with just an inumber, you need a device number too.
// To handle hard links (unique to Unix) we need to figure out if a
// given file has been seen before. To do this we see if a file has the
// same [dev,ino] tuple as one we have seen. If so, we won't bother
// reading it in.

type devInode struct {
	dev uint64
	ino uint64
}

// A Recorder is a structure that contains variables used to calculate
// file parameters such as inode numbers for a CPIO file. The life-time
// of a Record structure is meant to be the same as the construction of a
// single CPIO archive. Do not reuse between CPIOs if you don't know what
// you're doing.
type Recorder struct {
	inodeMap map[devInode]Info
	inumber  uint64
}

// Certain elements of the file can not be set by cpio:
// the Inode #
// the Dev
// maintaining these elements leaves us with a non-reproducible
// output stream. In this function, we figure out what inumber
// we need to use, and clear out anything we can.
// We always zero the Dev.
// We try to find the matching inode. If found, we use its inumber.
// If not, we get a new inumber for it and save the inode away.
// This eliminates two of the messier parts of creating reproducible
// output streams.
// The second return value indicates whether it is a hardlink or not.
func (r *Recorder) inode(i Info) (Info, bool) {
	d := devInode{dev: i.Dev, ino: i.Ino}
	i.Dev = 0

	if d, ok := r.inodeMap[d]; ok {
		i.Ino = d.Ino
		return i, d.Name != i.Name
	}

	i.Ino = r.inumber
	r.inumber++
	r.inodeMap[d] = i

	return i, false
}

// GetRecord returns a cpio Record for the given path on the local file system.
//
// GetRecord does not follow symlinks. If path is a symlink, the record
// returned will reflect that symlink.
func (r *Recorder) GetRecord(path string) (Record, error) {
	fi, err := os.Lstat(path)
	if err != nil {
		return Record{}, err
	}

	sys := fi.Sys().(*syscall.Stat_t)
	info, done := r.inode(sysInfo(path, sys))

	switch fi.Mode() & os.ModeType {
	case 0: // Regular file.
		if done {
			return Record{Info: info}, nil
		}
		return Record{Info: info, ReaderAt: uio.NewLazyFile(path)}, nil

	case os.ModeSymlink:
		linkname, err := os.Readlink(path)
		if err != nil {
			return Record{}, err
		}
		return StaticRecord([]byte(linkname), info), nil

	default:
		return StaticRecord(nil, info), nil
	}
}

// NewRecorder creates a new Recorder.
//
// A recorder is a structure that contains variables used to calculate
// file parameters such as inode numbers for a CPIO file. The life-time
// of a Record structure is meant to be the same as the construction of a
// single CPIO archive. Do not reuse between CPIOs if you don't know what
// you're doing.
func NewRecorder() *Recorder {
	return &Recorder{make(map[devInode]Info), 2}
}

// LSInfoFromRecord converts a Record to be usable with the ls package for
// listing files.
func LSInfoFromRecord(rec Record) ls.FileInfo {
	var target string

	mode := modeFromLinux(rec.Mode)
	if mode&os.ModeType == os.ModeSymlink {
		if l, err := uio.ReadAll(rec); err != nil {
			target = err.Error()
		} else {
			target = string(l)
		}
	}

	return ls.FileInfo{
		Name:          rec.Name,
		Mode:          mode,
		Rdev:          unix.Mkdev(uint32(rec.Rmajor), uint32(rec.Rminor)),
		UID:           uint32(rec.UID),
		GID:           uint32(rec.GID),
		Size:          int64(rec.FileSize),
		MTime:         time.Unix(int64(rec.MTime), 0).UTC(),
		SymlinkTarget: target,
	}
}
<|endoftext|>


<commit_before>package files

import (
	"errors"
	"os"
	"path/filepath"
	"sync"
	"github.com/materials-commons/gohandy/file"
)

// A TreeEntry is passed along the channel to the process function.
type TreeEntry struct {
	Path  string
	Finfo os.FileInfo
}

// ProcessFunc is the signature for a function that processes each file entry.
type ProcessFunc func(done <-chan struct{}, f <-chan TreeEntry, result chan<- string)

// IgnoreFunc is a function definition that tells walkFiles when it should
// ignore files and directory paths.
type IgnoreFunc func(path string, finfo os.FileInfo) bool

// PWalker implements a parallel walker. It holds the options and processing
// methods to use during the walk.
type PWalker struct {
	// NumParallel is the number of go routines to start for processing.
	NumParallel int

	// ProcessFn is the function to use to process entries.
	ProcessFn ProcessFunc

	// IgnoreFn is the method to test if an entry should be ignored. If
	// IgnoreFn is nil then all entries will be process.
	IgnoreFn IgnoreFunc

	// ProcessDirs tells the walker whether directory entries should also
	// be passed to ProcessFn. It defaults to false, so by default directories
	// are not passed to ProcessFn.
	ProcessDirs bool
}

// walkFiles walks a directory tree. It ignores files and paths with permission errors. All other errors
// will cause it to stop walking the tree. For every path segment it will call ignorePathFn and skip that
// file or sub directory if it returns true. Otherwise all regular files are sent along the filesChan
// for further processing.
func (p *PWalker) walkFiles(done <-chan struct{}, root string) (<-chan TreeEntry, <-chan error) {
	// filesChan channel is where we send regular files that aren't ignored.
	filesChan := make(chan TreeEntry)

	// errChan sends errors found while walking the tree.
	errChan := make(chan error, 1)

	// Walk the tree in separate go routine. Send back results and errors in the filesChan and errChan channels.
	go func() {
		defer close(filesChan)
		errChan <- filepath.Walk(root, func(path string, finfo os.FileInfo, err error) error {
			switch {
			case err != nil && os.IsPermission(err):
				// Permission errors are ignored. Just continue walking the tree
				// without processing the file or directory.
				return nil

			case err != nil:
				// All other errors cause walking to stop.
				return err

			case p.IgnoreFn(path, finfo):
				// if ignorePathFn returns true then skip processing this entry.
				if finfo.IsDir() {
					// If entry is a directory, then skip processing that
					// entry sub tree.
					return filepath.SkipDir
				}
				return nil

			case !finfo.Mode().IsRegular() && !finfo.IsDir():
				// Only process regular files. This means the following types
				// are not processed: Symbol Link, Named Pip, Socket,
				// and Device files.
				//
				// Directories are a special case and are handled in the default
				// clause, where regulars files are also handled.
				return nil

			default:
				// Directories aren't being processed, so skip this entry.
				if finfo.IsDir() && !p.ProcessDirs {
					return nil
				}

				// This is an entry we can process, so send it along the fileChan
				exfi, _ := file.ExStatFromFileInfo(finfo, path)
				entry := TreeEntry{
					Path:  path,
					Finfo: exfi,
				}
				select {
				case filesChan <- entry:
				case <-done:
					// if we receive a message on the done channel, then stop
					// walking the tree.
					return errors.New("walk canceled")
				}
				return nil
			}

		})
	}()
	return filesChan, errChan
}

// PWalk will walk a tree and process files in it in parallel. fn is the process function. "n" go routines will
// be started with fn for processing. The ignorePathFn determines whether a file or directory path should be
// ignored. ignorePathFn can be nil. If it is nil then no files or paths are ignored.
func (p *PWalker) PWalk(root string) (<-chan string, <-chan error) {
	done := make(chan struct{})
	defer close(done)

	if p.IgnoreFn == nil {
		p.IgnoreFn = neverIgnore
	}

	if p.NumParallel < 1 {
		p.NumParallel = 3
	}

	filesChan, errChan := p.walkFiles(done, root)

	// results holds the results of processing each entry. Its value is
	// ignored and only exists so that PWalk will block until all
	// processing has completed.
	results := make(chan string)

	// Setup a WaitGroup equal to the number of routines processing
	// file entries.
	var wg sync.WaitGroup
	wg.Add(p.NumParallel)

	// Start "n" fn routines to process files. When each one completes
	// processing all entries along its channel signal done.
	for i := 0; i < p.NumParallel; i++ {
		go func() {
			p.ProcessFn(done, filesChan, results)
			wg.Done()
		}()
	}

	go func() {
		wg.Wait()
		close(results)
	}()
	for r := range results {
		var _ = r
	}
	return results, errChan
}

// neverIgnore is the default method for checking if an entry should be ignored.
// It never ignores an entry.
func neverIgnore(_ string, _ os.FileInfo) bool {
	return false
}
<commit_msg>Run goimports.<commit_after>package files

import (
	"errors"
	"os"
	"path/filepath"
	"sync"

	"github.com/materials-commons/gohandy/file"
)

// A TreeEntry is passed along the channel to the process function.
type TreeEntry struct {
	Path  string
	Finfo os.FileInfo
}

// ProcessFunc is the signature for a function that processes each file entry.
type ProcessFunc func(done <-chan struct{}, f <-chan TreeEntry, result chan<- string)

// IgnoreFunc is a function definition that tells walkFiles when it should
// ignore files and directory paths.
type IgnoreFunc func(path string, finfo os.FileInfo) bool

// PWalker implements a parallel walker. It holds the options and processing
// methods to use during the walk.
type PWalker struct {
	// NumParallel is the number of go routines to start for processing.
	NumParallel int

	// ProcessFn is the function to use to process entries.
	ProcessFn ProcessFunc

	// IgnoreFn is the method to test if an entry should be ignored. If
	// IgnoreFn is nil then all entries will be process.
	IgnoreFn IgnoreFunc

	// ProcessDirs tells the walker whether directory entries should also
	// be passed to ProcessFn. It defaults to false, so by default directories
	// are not passed to ProcessFn.
	ProcessDirs bool
}

// walkFiles walks a directory tree. It ignores files and paths with permission errors. All other errors
// will cause it to stop walking the tree. For every path segment it will call ignorePathFn and skip that
// file or sub directory if it returns true. Otherwise all regular files are sent along the filesChan
// for further processing.
func (p *PWalker) walkFiles(done <-chan struct{}, root string) (<-chan TreeEntry, <-chan error) {
	// filesChan channel is where we send regular files that aren't ignored.
	filesChan := make(chan TreeEntry)

	// errChan sends errors found while walking the tree.
	errChan := make(chan error, 1)

	// Walk the tree in separate go routine. Send back results and errors in the filesChan and errChan channels.
	go func() {
		defer close(filesChan)
		errChan <- filepath.Walk(root, func(path string, finfo os.FileInfo, err error) error {
			switch {
			case err != nil && os.IsPermission(err):
				// Permission errors are ignored. Just continue walking the tree
				// without processing the file or directory.
				return nil

			case err != nil:
				// All other errors cause walking to stop.
				return err

			case p.IgnoreFn(path, finfo):
				// if ignorePathFn returns true then skip processing this entry.
				if finfo.IsDir() {
					// If entry is a directory, then skip processing that
					// entry sub tree.
					return filepath.SkipDir
				}
				return nil

			case !finfo.Mode().IsRegular() && !finfo.IsDir():
				// Only process regular files. This means the following types
				// are not processed: Symbol Link, Named Pip, Socket,
				// and Device files.
				//
				// Directories are a special case and are handled in the default
				// clause, where regulars files are also handled.
				return nil

			default:
				// Directories aren't being processed, so skip this entry.
				if finfo.IsDir() && !p.ProcessDirs {
					return nil
				}

				// This is an entry we can process, so send it along the fileChan
				exfi, _ := file.ExStatFromFileInfo(finfo, path)
				entry := TreeEntry{
					Path:  path,
					Finfo: exfi,
				}
				select {
				case filesChan <- entry:
				case <-done:
					// if we receive a message on the done channel, then stop
					// walking the tree.
					return errors.New("walk canceled")
				}
				return nil
			}

		})
	}()
	return filesChan, errChan
}

// PWalk will walk a tree and process files in it in parallel. fn is the process function. "n" go routines will
// be started with fn for processing. The ignorePathFn determines whether a file or directory path should be
// ignored. ignorePathFn can be nil. If it is nil then no files or paths are ignored.
func (p *PWalker) PWalk(root string) (<-chan string, <-chan error) {
	done := make(chan struct{})
	defer close(done)

	if p.IgnoreFn == nil {
		p.IgnoreFn = neverIgnore
	}

	if p.NumParallel < 1 {
		p.NumParallel = 3
	}

	filesChan, errChan := p.walkFiles(done, root)

	// results holds the results of processing each entry. Its value is
	// ignored and only exists so that PWalk will block until all
	// processing has completed.
	results := make(chan string)

	// Setup a WaitGroup equal to the number of routines processing
	// file entries.
	var wg sync.WaitGroup
	wg.Add(p.NumParallel)

	// Start "n" fn routines to process files. When each one completes
	// processing all entries along its channel signal done.
	for i := 0; i < p.NumParallel; i++ {
		go func() {
			p.ProcessFn(done, filesChan, results)
			wg.Done()
		}()
	}

	go func() {
		wg.Wait()
		close(results)
	}()
	for r := range results {
		var _ = r
	}
	return results, errChan
}

// neverIgnore is the default method for checking if an entry should be ignored.
// It never ignores an entry.
func neverIgnore(_ string, _ os.FileInfo) bool {
	return false
}
<|endoftext|>


<commit_before>// +build unittest

package log

import (
	"testing"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestLog(t *testing.T) {
	t.Run("podLogTest", podLogTest)
}

func podLogTest(t *testing.T) {
	pod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "testpod",
			Namespace: "testnamespace",
		},
	}
	PodLog(pod).Infof("valid pod info")
	pod = nil
	PodLog(pod).Infof("nil pod info")

	pod = &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:            "testpod",
			Namespace:       "testnamespace",
			OwnerReferences: []metav1.OwnerReference{},
		},
	}
	PodLog(pod).Infof("empty pod owner")
	pod.OwnerReferences = append(pod.OwnerReferences,
		metav1.OwnerReference{
			Kind:       "testkind",
			Name:       "testname",
			Controller: nil,
		})
	PodLog(pod).Infof("pod owner nil controller")
	controller := true
	pod.OwnerReferences[0].Controller = &controller
	PodLog(pod).Infof("pod owner controller")
}
<commit_msg>Add UTs for all the log APIs<commit_after>// +build unittest

package log

import (
	"testing"

	crdv1 "github.com/kubernetes-incubator/external-storage/snapshot/pkg/apis/crd/v1"
	storkv1 "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	appv1 "k8s.io/api/apps/v1"
	appv1beta1 "k8s.io/api/apps/v1beta1"
	appv1beta2 "k8s.io/api/apps/v1beta2"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestLog(t *testing.T) {
	t.Run("podLogTest", podLogTest)
	t.Run("deploymentLogTest", deploymentLogTest)
	t.Run("statefulsetLogTest", statefulsetLogTest)
	t.Run("snapshotLogTest", snapshotLogTest)
	t.Run("migrationLogTest", migrationLogTest)
	t.Run("ruleLogTest", ruleLogTest)
}

func podLogTest(t *testing.T) {
	pod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "testpod",
			Namespace: "testnamespace",
		},
	}
	PodLog(pod).Infof("valid pod info")
	pod = nil
	PodLog(pod).Infof("nil pod info")

	pod = &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:            "testpod",
			Namespace:       "testnamespace",
			OwnerReferences: []metav1.OwnerReference{},
		},
	}
	PodLog(pod).Infof("empty pod owner")
	pod.OwnerReferences = append(pod.OwnerReferences,
		metav1.OwnerReference{
			Kind:       "testkind",
			Name:       "testname",
			Controller: nil,
		})
	PodLog(pod).Infof("pod owner nil controller")
	controller := true
	pod.OwnerReferences[0].Controller = &controller
	PodLog(pod).Infof("pod owner controller")
}

func deploymentLogTest(t *testing.T) {
	metadata := metav1.ObjectMeta{
		Name:      "testdeployment",
		Namespace: "testnamespace",
	}
	deploymentv1 := &appv1.Deployment{
		ObjectMeta: metadata}
	DeploymentV1Log(deploymentv1).Infof("deploymentv1 log")
	DeploymentV1Log(nil).Infof("deploymentv1 nil log")
	deploymentv1beta1 := &appv1beta1.Deployment{
		ObjectMeta: metadata}
	DeploymentV1Beta1Log(deploymentv1beta1).Infof("deploymentv1beta1 log")
	DeploymentV1Beta1Log(nil).Infof("deploymentv1beta1 nil log")
	deploymentv1beta2 := &appv1beta2.Deployment{
		ObjectMeta: metadata}
	DeploymentV1Beta2Log(deploymentv1beta2).Infof("deploymentv1beta2 log")
	DeploymentV1Beta2Log(nil).Infof("deploymentv1beta2 nil log")
}
func statefulsetLogTest(t *testing.T) {
	metadata := metav1.ObjectMeta{
		Name:      "teststatefulset",
		Namespace: "testnamespace",
	}
	statefulsetv1 := &appv1.StatefulSet{
		ObjectMeta: metadata}
	StatefulSetV1Log(statefulsetv1).Infof("statefulsetv1 log")
	StatefulSetV1Log(nil).Infof("statefulsetv1 nil log")
	statefulsetv1beta1 := &appv1beta1.StatefulSet{
		ObjectMeta: metadata}
	StatefulSetV1Beta1Log(statefulsetv1beta1).Infof("statefulsetv1beta1 log")
	StatefulSetV1Beta1Log(nil).Infof("statefulsetv1beta1 nil log")
	statefulsetv1beta2 := &appv1beta2.StatefulSet{
		ObjectMeta: metadata}
	StatefulSetV1Beta2Log(statefulsetv1beta2).Infof("statefulsetv1beta2 log")
	StatefulSetV1Beta2Log(nil).Infof("statefulsetv1beta2 nil log")
}

func snapshotLogTest(t *testing.T) {
	metadata := metav1.ObjectMeta{
		Name:      "testsnapshot",
		Namespace: "testnamespace",
	}
	snapshot := &crdv1.VolumeSnapshot{
		Metadata: metadata,
	}
	SnapshotLog(snapshot).Infof("snapshot log")
	SnapshotLog(nil).Infof("snapshot nil log")
}

func migrationLogTest(t *testing.T) {
	metadata := metav1.ObjectMeta{
		Name:      "testmigration",
		Namespace: "testnamespace",
	}
	migration := &storkv1.Migration{
		ObjectMeta: metadata,
	}
	MigrationLog(migration).Infof("migration log")
	MigrationLog(nil).Infof("migration nil log")
}

func ruleLogTest(t *testing.T) {
	metadata := metav1.ObjectMeta{
		Name:      "testrule",
		Namespace: "testnamespace",
	}
	rule := &storkv1.Rule{
		ObjectMeta: metadata,
	}
	migration := &storkv1.Migration{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "testmigration",
			Namespace: "testnamespace",
		},
	}
	snapshot := &crdv1.VolumeSnapshot{
		Metadata: metav1.ObjectMeta{
			Name:      "testsnapshot",
			Namespace: "testnamespace",
		},
	}

	RuleLog(rule, snapshot).Infof("rule with snapshot log")
	RuleLog(rule, migration).Infof("rule with migration log")
	RuleLog(rule, nil).Infof("rule with nil object log")
}
<|endoftext|>


<commit_before>package project

import (
	"errors"
	"fmt"

	"github.com/golang/protobuf/proto"
	git "gopkg.in/libgit2/git2go.v23"

	pb "rsprd.com/spread/pkg/spreadproto"
)

func (p *Project) AddObjectToIndex(obj *pb.Object) error {
	info := obj.GetInfo()
	if info == nil {
		return ErrNilObjectInfo
	}

	data, err := proto.Marshal(obj)
	if err != nil {
		return fmt.Errorf("could not encode object: %v", err)
	}

	oid, err := p.repo.CreateBlobFromBuffer(data)
	if err != nil {
		return fmt.Errorf("could not write Object as blob in Git repo: %v", err)
	}

	entry := &git.IndexEntry{
		Mode: git.FilemodeBlob,
		Size: uint32(len(data)),
		Id:   oid,
		Path: info.Path,
	}

	index, err := p.repo.Index()
	if err != nil {
		return fmt.Errorf("could not retreive index: %v", err)
	}

	return index.Add(entry)
}

var (
	ErrNilObjectInfo = errors.New("an object's Info field cannot be nil")
)
<commit_msg>index works<commit_after>package project

import (
	"errors"
	"fmt"

	"github.com/golang/protobuf/proto"
	git "gopkg.in/libgit2/git2go.v23"

	pb "rsprd.com/spread/pkg/spreadproto"
)

func (p *Project) AddObjectToIndex(obj *pb.Object) error {
	info := obj.GetInfo()
	if info == nil {
		return ErrNilObjectInfo
	}

	data, err := proto.Marshal(obj)
	if err != nil {
		return fmt.Errorf("could not encode object: %v", err)
	}

	oid, err := p.repo.CreateBlobFromBuffer(data)
	if err != nil {
		return fmt.Errorf("could not write Object as blob in Git repo: %v", err)
	}

	entry := &git.IndexEntry{
		Mode: git.FilemodeBlob,
		Size: uint32(len(data)),
		Id:   oid,
		Path: info.Path,
	}

	index, err := p.repo.Index()
	if err != nil {
		return fmt.Errorf("could not retreive index: %v", err)
	}

	err = index.Add(entry)
	if err != nil {
		return err
	}

	return index.Write()
}

var (
	ErrNilObjectInfo = errors.New("an object's Info field cannot be nil")
)
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"bytes"

	"github.com/BurntSushi/toml"

	"github.com/CodisLabs/codis/pkg/utils/bytesize"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/timesize"
)

const DefaultConfig = `
##################################################
#                                                #
#                  Codis-Proxy                   #
#                                                #
##################################################

# Set Codis Product Name/Auth.
product_name = "codis-demo"
product_auth = ""

# Set bind address for admin(rpc), tcp only.
admin_addr = "0.0.0.0:11080"

# Set bind address for proxy, proto_type can be "tcp", "tcp4", "tcp6", "unix" or "unixpacket".
proto_type = "tcp4"
proxy_addr = "0.0.0.0:19000"

# Set jodis address & session timeout, only accept "zookeeper" & "etcd".
jodis_name = ""
jodis_addr = ""
jodis_timeout = "20s"
jodis_compatible = false

# Set datacenter of proxy.
proxy_datacenter = ""

# Set max number of alive sessions.
proxy_max_clients = 1000

# Set max offheap memory size. (0 to disable)
proxy_max_offheap_size = "1024mb"

# Set heap placeholder to reduce GC frequency.
proxy_heap_placeholder = "256mb"

# Proxy will ping backend redis in a predefined interval. (0 to disable)
backend_ping_period = "5s"

# Set backend recv buffer size & timeout.
backend_recv_bufsize = "128kb"
backend_recv_timeout = "50s"

# Set backend send buffer & timeout.
backend_send_bufsize = "128kb"
backend_send_timeout = "50s"

# Set backend pipeline buffer size.
backend_max_pipeline = 1024

# Set backend parallel connections per server
backend_primary_parallel = 1
backend_replica_parallel = 1

# Set backend tcp keepalive period. (0 to disable)
backend_keepalive_period = "75s"

# If there is no request from client for a long time, the connection will be closed. (0 to disable)
# Set session recv buffer size & timeout.
session_recv_bufsize = "128kb"
session_recv_timeout = "30m"

# Set session send buffer size & timeout.
session_send_bufsize = "64kb"
session_send_timeout = "30s"

# Make sure this is higher than the max number of requests for each pipeline request, or your client may be blocked.
# Set session pipeline buffer size.
session_max_pipeline = 512

# Set session tcp keepalive period. (0 to disable)
session_keepalive_period = "75s"

# Set metrics server (such as http://localhost:28000), proxy will report json formatted metrics to specified server in a predefined period.
metrics_report_server = ""
metrics_report_period = "1s"

# Set influxdb server (such as http://localhost:8086), proxy will report metrics to influxdb.
metrics_report_influxdb_server = ""
metrics_report_influxdb_period = "1s"
metrics_report_influxdb_username = ""
metrics_report_influxdb_password = ""
metrics_report_influxdb_database = ""
`

type Config struct {
	ProtoType string `toml:"proto_type" json:"proto_type"`
	ProxyAddr string `toml:"proxy_addr" json:"proxy_addr"`
	AdminAddr string `toml:"admin_addr" json:"admin_addr"`

	HostProxy string `toml:"-" json:"-"`
	HostAdmin string `toml:"-" json:"-"`

	JodisName       string            `toml:"jodis_name" json:"jodis_name"`
	JodisAddr       string            `toml:"jodis_addr" json:"jodis_addr"`
	JodisTimeout    timesize.Duration `toml:"jodis_timeout" json:"jodis_timeout"`
	JodisCompatible bool              `toml:"jodis_compatible" json:"jodis_compatible"`

	ProductName string `toml:"product_name" json:"product_name"`
	ProductAuth string `toml:"product_auth" json:"-"`

	ProxyDataCenter      string         `toml:"proxy_datacenter" json:"proxy_datacenter"`
	ProxyMaxClients      int            `toml:"proxy_max_clients" json:"proxy_max_clients"`
	ProxyMaxOffheapBytes bytesize.Int64 `toml:"proxy_max_offheap_size" json:"proxy_max_offheap_size"`
	ProxyHeapPlaceholder bytesize.Int64 `toml:"proxy_heap_placeholder" json:"proxy_heap_placeholder"`

	BackendPingPeriod      timesize.Duration `toml:"backend_ping_period" json:"backend_ping_period"`
	BackendRecvBufsize     bytesize.Int64    `toml:"backend_recv_bufsize" json:"backend_recv_bufsize"`
	BackendRecvTimeout     timesize.Duration `toml:"backend_recv_timeout" json:"backend_recv_timeout"`
	BackendSendBufsize     bytesize.Int64    `toml:"backend_send_bufsize" json:"backend_send_bufsize"`
	BackendSendTimeout     timesize.Duration `toml:"backend_send_timeout" json:"backend_send_timeout"`
	BackendMaxPipeline     int               `toml:"backend_max_pipeline" json:"backend_max_pipeline"`
	BackendPrimaryParallel int               `toml:"backend_primary_parallel" json:"backend_primary_parallel"`
	BackendReplicaParallel int               `toml:"backend_replica_parallel" json:"backend_replica_parallel"`
	BackendKeepAlivePeriod timesize.Duration `toml:"backend_keepalive_period" json:"backend_keepalive_period"`

	SessionRecvBufsize     bytesize.Int64    `toml:"session_recv_bufsize" json:"session_recv_bufsize"`
	SessionRecvTimeout     timesize.Duration `toml:"session_recv_timeout" json:"session_recv_timeout"`
	SessionSendBufsize     bytesize.Int64    `toml:"session_send_bufsize" json:"session_send_bufsize"`
	SessionSendTimeout     timesize.Duration `toml:"session_send_timeout" json:"session_send_timeout"`
	SessionMaxPipeline     int               `toml:"session_max_pipeline" json:"session_max_pipeline"`
	SessionKeepAlivePeriod timesize.Duration `toml:"session_keepalive_period" json:"session_keepalive_period"`

	MetricsReportServer           string            `toml:"metrics_report_server" json:"metrics_report_server"`
	MetricsReportPeriod           timesize.Duration `toml:"metrics_report_period" json:"metrics_report_period"`
	MetricsReportInfluxdbServer   string            `toml:"metrics_report_influxdb_server" json:"metrics_report_influxdb_server"`
	MetricsReportInfluxdbPeriod   timesize.Duration `toml:"metrics_report_influxdb_period" json:"metrics_report_influxdb_period"`
	MetricsReportInfluxdbUsername string            `toml:"metrics_report_influxdb_username" json:"metrics_report_influxdb_username"`
	MetricsReportInfluxdbPassword string            `toml:"metrics_report_influxdb_password" json:"-"`
	MetricsReportInfluxdbDatabase string            `toml:"metrics_report_influxdb_database" json:"metrics_report_influxdb_database"`
}

func NewDefaultConfig() *Config {
	c := &Config{}
	if _, err := toml.Decode(DefaultConfig, c); err != nil {
		log.PanicErrorf(err, "decode toml failed")
	}
	if err := c.Validate(); err != nil {
		log.PanicErrorf(err, "validate config failed")
	}
	return c
}

func (c *Config) LoadFromFile(path string) error {
	_, err := toml.DecodeFile(path, c)
	if err != nil {
		return errors.Trace(err)
	}
	return c.Validate()
}

func (c *Config) String() string {
	var b bytes.Buffer
	e := toml.NewEncoder(&b)
	e.Indent = "    "
	e.Encode(c)
	return b.String()
}

func (c *Config) Validate() error {
	if c.ProtoType == "" {
		return errors.New("invalid proto_type")
	}
	if c.ProxyAddr == "" {
		return errors.New("invalid proxy_addr")
	}
	if c.AdminAddr == "" {
		return errors.New("invalid admin_addr")
	}
	if c.JodisName != "" {
		if c.JodisAddr == "" {
			return errors.New("invalid jodis_addr")
		}
		if c.JodisTimeout < 0 {
			return errors.New("invalid jodis_timeout")
		}
	}
	if c.ProductName == "" {
		return errors.New("invalid product_name")
	}
	if c.ProxyMaxClients < 0 {
		return errors.New("invalid proxy_max_clients")
	}

	const MaxInt = bytesize.Int64(^uint(0) >> 1)

	if d := c.ProxyMaxOffheapBytes; d < 0 || d > MaxInt {
		return errors.New("invalid proxy_max_offheap_size")
	}
	if d := c.ProxyHeapPlaceholder; d < 0 || d > MaxInt {
		return errors.New("invalid proxy_heap_placeholder")
	}
	if c.BackendPingPeriod < 0 {
		return errors.New("invalid backend_ping_period")
	}

	if d := c.BackendRecvBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid backend_recv_bufsize")
	}
	if c.BackendRecvTimeout < 0 {
		return errors.New("invalid backend_recv_timeout")
	}
	if d := c.BackendSendBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid backend_send_bufsize")
	}
	if c.BackendSendTimeout < 0 {
		return errors.New("invalid backend_send_timeout")
	}
	if c.BackendMaxPipeline < 0 {
		return errors.New("invalid backend_max_pipeline")
	}
	if c.BackendPrimaryParallel < 0 {
		return errors.New("invalid backend_primary_parallel")
	}
	if c.BackendReplicaParallel < 0 {
		return errors.New("invalid backend_replica_parallel")
	}
	if c.BackendKeepAlivePeriod < 0 {
		return errors.New("invalid backend_keepalive_period")
	}

	if d := c.SessionRecvBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid session_recv_bufsize")
	}
	if c.SessionRecvTimeout < 0 {
		return errors.New("invalid session_recv_timeout")
	}
	if d := c.SessionSendBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid session_send_bufsize")
	}
	if c.SessionSendTimeout < 0 {
		return errors.New("invalid session_send_timeout")
	}
	if c.SessionMaxPipeline < 0 {
		return errors.New("invalid session_max_pipeline")
	}
	if c.SessionKeepAlivePeriod < 0 {
		return errors.New("invalid session_keepalive_period")
	}

	if c.MetricsReportPeriod < 0 {
		return errors.New("invalid metrics_report_period")
	}
	if c.MetricsReportInfluxdbPeriod < 0 {
		return errors.New("invalid metrics_report_influxdb_period")
	}
	return nil
}
<commit_msg>proxy: set default value of backend_{send/recv}_timeout to 30s<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"bytes"

	"github.com/BurntSushi/toml"

	"github.com/CodisLabs/codis/pkg/utils/bytesize"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
	"github.com/CodisLabs/codis/pkg/utils/timesize"
)

const DefaultConfig = `
##################################################
#                                                #
#                  Codis-Proxy                   #
#                                                #
##################################################

# Set Codis Product Name/Auth.
product_name = "codis-demo"
product_auth = ""

# Set bind address for admin(rpc), tcp only.
admin_addr = "0.0.0.0:11080"

# Set bind address for proxy, proto_type can be "tcp", "tcp4", "tcp6", "unix" or "unixpacket".
proto_type = "tcp4"
proxy_addr = "0.0.0.0:19000"

# Set jodis address & session timeout, only accept "zookeeper" & "etcd".
jodis_name = ""
jodis_addr = ""
jodis_timeout = "20s"
jodis_compatible = false

# Set datacenter of proxy.
proxy_datacenter = ""

# Set max number of alive sessions.
proxy_max_clients = 1000

# Set max offheap memory size. (0 to disable)
proxy_max_offheap_size = "1024mb"

# Set heap placeholder to reduce GC frequency.
proxy_heap_placeholder = "256mb"

# Proxy will ping backend redis in a predefined interval. (0 to disable)
backend_ping_period = "5s"

# Set backend recv buffer size & timeout.
backend_recv_bufsize = "128kb"
backend_recv_timeout = "30s"

# Set backend send buffer & timeout.
backend_send_bufsize = "128kb"
backend_send_timeout = "30s"

# Set backend pipeline buffer size.
backend_max_pipeline = 1024

# Set backend parallel connections per server
backend_primary_parallel = 1
backend_replica_parallel = 1

# Set backend tcp keepalive period. (0 to disable)
backend_keepalive_period = "75s"

# If there is no request from client for a long time, the connection will be closed. (0 to disable)
# Set session recv buffer size & timeout.
session_recv_bufsize = "128kb"
session_recv_timeout = "30m"

# Set session send buffer size & timeout.
session_send_bufsize = "64kb"
session_send_timeout = "30s"

# Make sure this is higher than the max number of requests for each pipeline request, or your client may be blocked.
# Set session pipeline buffer size.
session_max_pipeline = 512

# Set session tcp keepalive period. (0 to disable)
session_keepalive_period = "75s"

# Set metrics server (such as http://localhost:28000), proxy will report json formatted metrics to specified server in a predefined period.
metrics_report_server = ""
metrics_report_period = "1s"

# Set influxdb server (such as http://localhost:8086), proxy will report metrics to influxdb.
metrics_report_influxdb_server = ""
metrics_report_influxdb_period = "1s"
metrics_report_influxdb_username = ""
metrics_report_influxdb_password = ""
metrics_report_influxdb_database = ""
`

type Config struct {
	ProtoType string `toml:"proto_type" json:"proto_type"`
	ProxyAddr string `toml:"proxy_addr" json:"proxy_addr"`
	AdminAddr string `toml:"admin_addr" json:"admin_addr"`

	HostProxy string `toml:"-" json:"-"`
	HostAdmin string `toml:"-" json:"-"`

	JodisName       string            `toml:"jodis_name" json:"jodis_name"`
	JodisAddr       string            `toml:"jodis_addr" json:"jodis_addr"`
	JodisTimeout    timesize.Duration `toml:"jodis_timeout" json:"jodis_timeout"`
	JodisCompatible bool              `toml:"jodis_compatible" json:"jodis_compatible"`

	ProductName string `toml:"product_name" json:"product_name"`
	ProductAuth string `toml:"product_auth" json:"-"`

	ProxyDataCenter      string         `toml:"proxy_datacenter" json:"proxy_datacenter"`
	ProxyMaxClients      int            `toml:"proxy_max_clients" json:"proxy_max_clients"`
	ProxyMaxOffheapBytes bytesize.Int64 `toml:"proxy_max_offheap_size" json:"proxy_max_offheap_size"`
	ProxyHeapPlaceholder bytesize.Int64 `toml:"proxy_heap_placeholder" json:"proxy_heap_placeholder"`

	BackendPingPeriod      timesize.Duration `toml:"backend_ping_period" json:"backend_ping_period"`
	BackendRecvBufsize     bytesize.Int64    `toml:"backend_recv_bufsize" json:"backend_recv_bufsize"`
	BackendRecvTimeout     timesize.Duration `toml:"backend_recv_timeout" json:"backend_recv_timeout"`
	BackendSendBufsize     bytesize.Int64    `toml:"backend_send_bufsize" json:"backend_send_bufsize"`
	BackendSendTimeout     timesize.Duration `toml:"backend_send_timeout" json:"backend_send_timeout"`
	BackendMaxPipeline     int               `toml:"backend_max_pipeline" json:"backend_max_pipeline"`
	BackendPrimaryParallel int               `toml:"backend_primary_parallel" json:"backend_primary_parallel"`
	BackendReplicaParallel int               `toml:"backend_replica_parallel" json:"backend_replica_parallel"`
	BackendKeepAlivePeriod timesize.Duration `toml:"backend_keepalive_period" json:"backend_keepalive_period"`

	SessionRecvBufsize     bytesize.Int64    `toml:"session_recv_bufsize" json:"session_recv_bufsize"`
	SessionRecvTimeout     timesize.Duration `toml:"session_recv_timeout" json:"session_recv_timeout"`
	SessionSendBufsize     bytesize.Int64    `toml:"session_send_bufsize" json:"session_send_bufsize"`
	SessionSendTimeout     timesize.Duration `toml:"session_send_timeout" json:"session_send_timeout"`
	SessionMaxPipeline     int               `toml:"session_max_pipeline" json:"session_max_pipeline"`
	SessionKeepAlivePeriod timesize.Duration `toml:"session_keepalive_period" json:"session_keepalive_period"`

	MetricsReportServer           string            `toml:"metrics_report_server" json:"metrics_report_server"`
	MetricsReportPeriod           timesize.Duration `toml:"metrics_report_period" json:"metrics_report_period"`
	MetricsReportInfluxdbServer   string            `toml:"metrics_report_influxdb_server" json:"metrics_report_influxdb_server"`
	MetricsReportInfluxdbPeriod   timesize.Duration `toml:"metrics_report_influxdb_period" json:"metrics_report_influxdb_period"`
	MetricsReportInfluxdbUsername string            `toml:"metrics_report_influxdb_username" json:"metrics_report_influxdb_username"`
	MetricsReportInfluxdbPassword string            `toml:"metrics_report_influxdb_password" json:"-"`
	MetricsReportInfluxdbDatabase string            `toml:"metrics_report_influxdb_database" json:"metrics_report_influxdb_database"`
}

func NewDefaultConfig() *Config {
	c := &Config{}
	if _, err := toml.Decode(DefaultConfig, c); err != nil {
		log.PanicErrorf(err, "decode toml failed")
	}
	if err := c.Validate(); err != nil {
		log.PanicErrorf(err, "validate config failed")
	}
	return c
}

func (c *Config) LoadFromFile(path string) error {
	_, err := toml.DecodeFile(path, c)
	if err != nil {
		return errors.Trace(err)
	}
	return c.Validate()
}

func (c *Config) String() string {
	var b bytes.Buffer
	e := toml.NewEncoder(&b)
	e.Indent = "    "
	e.Encode(c)
	return b.String()
}

func (c *Config) Validate() error {
	if c.ProtoType == "" {
		return errors.New("invalid proto_type")
	}
	if c.ProxyAddr == "" {
		return errors.New("invalid proxy_addr")
	}
	if c.AdminAddr == "" {
		return errors.New("invalid admin_addr")
	}
	if c.JodisName != "" {
		if c.JodisAddr == "" {
			return errors.New("invalid jodis_addr")
		}
		if c.JodisTimeout < 0 {
			return errors.New("invalid jodis_timeout")
		}
	}
	if c.ProductName == "" {
		return errors.New("invalid product_name")
	}
	if c.ProxyMaxClients < 0 {
		return errors.New("invalid proxy_max_clients")
	}

	const MaxInt = bytesize.Int64(^uint(0) >> 1)

	if d := c.ProxyMaxOffheapBytes; d < 0 || d > MaxInt {
		return errors.New("invalid proxy_max_offheap_size")
	}
	if d := c.ProxyHeapPlaceholder; d < 0 || d > MaxInt {
		return errors.New("invalid proxy_heap_placeholder")
	}
	if c.BackendPingPeriod < 0 {
		return errors.New("invalid backend_ping_period")
	}

	if d := c.BackendRecvBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid backend_recv_bufsize")
	}
	if c.BackendRecvTimeout < 0 {
		return errors.New("invalid backend_recv_timeout")
	}
	if d := c.BackendSendBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid backend_send_bufsize")
	}
	if c.BackendSendTimeout < 0 {
		return errors.New("invalid backend_send_timeout")
	}
	if c.BackendMaxPipeline < 0 {
		return errors.New("invalid backend_max_pipeline")
	}
	if c.BackendPrimaryParallel < 0 {
		return errors.New("invalid backend_primary_parallel")
	}
	if c.BackendReplicaParallel < 0 {
		return errors.New("invalid backend_replica_parallel")
	}
	if c.BackendKeepAlivePeriod < 0 {
		return errors.New("invalid backend_keepalive_period")
	}

	if d := c.SessionRecvBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid session_recv_bufsize")
	}
	if c.SessionRecvTimeout < 0 {
		return errors.New("invalid session_recv_timeout")
	}
	if d := c.SessionSendBufsize; d < 0 || d > MaxInt {
		return errors.New("invalid session_send_bufsize")
	}
	if c.SessionSendTimeout < 0 {
		return errors.New("invalid session_send_timeout")
	}
	if c.SessionMaxPipeline < 0 {
		return errors.New("invalid session_max_pipeline")
	}
	if c.SessionKeepAlivePeriod < 0 {
		return errors.New("invalid session_keepalive_period")
	}

	if c.MetricsReportPeriod < 0 {
		return errors.New("invalid metrics_report_period")
	}
	if c.MetricsReportInfluxdbPeriod < 0 {
		return errors.New("invalid metrics_report_influxdb_period")
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.RWMutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := range s.slots {
		s.slots[i].id = i
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.closed = true

	for i := range s.slots {
		s.fillSlot(&models.Slot{Id: i})
	}
}

func (s *Router) GetGroupIds() map[int]bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	var groups = make(map[int]bool)
	for i := range s.slots {
		if gid := s.slots[i].backend.id; gid != 0 {
			groups[gid] = true
		}
		if gid := s.slots[i].migrate.id; gid != 0 {
			groups[gid] = true
		}
	}
	return groups
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	slots := make([]*models.Slot, len(s.slots))
	for i := range s.slots {
		slots[i] = s.slots[i].snapshot(true)
	}
	return slots
}

func (s *Router) GetSlot(id int) *models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if id < 0 || id >= len(s.slots) {
		return nil
	}
	slot := &s.slots[id]
	return slot.snapshot(true)
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(m *models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	if m.Id < 0 || m.Id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	s.fillSlot(m)
	return nil
}

func (s *Router) KeepAlive() error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	var id = Hash(hkey) % uint32(len(s.slots))
	slot := &s.slots[id]
	return slot.forward(r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(r, nil)
}

func (s *Router) dispatchAddr(r *Request, addr string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	bc := s.getBackendConn(addr, false)
	if bc == nil {
		return false
	} else {
		bc.PushBack(r)
		s.putBackendConn(bc)
		return true
	}
}

func (s *Router) getBackendConn(addr string, create bool) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else if create {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	} else {
		return nil
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) fillSlot(m *models.Slot) {
	slot := &s.slots[m.Id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend.bc)
	slot.backend.bc = nil
	slot.backend.id = 0
	s.putBackendConn(slot.migrate.bc)
	slot.migrate.bc = nil
	slot.migrate.id = 0
	for i := range slot.replicaGroups {
		for _, bc := range slot.replicaGroups[i] {
			s.putBackendConn(bc)
		}
	}
	slot.replicaGroups = nil

	if addr := m.BackendAddr; len(addr) != 0 {
		slot.backend.bc = s.getBackendConn(addr, true)
		slot.backend.id = m.BackendAddrId
	}
	if from := m.MigrateFrom; len(from) != 0 {
		slot.migrate.bc = s.getBackendConn(from, true)
		slot.migrate.id = m.MigrateFromId
	}
	for i := range m.ReplicaGroups {
		var group []*SharedBackendConn
		for _, addr := range m.ReplicaGroups[i] {
			group = append(group, s.getBackendConn(addr, true))
		}
		slot.replicaGroups = append(slot.replicaGroups, group)
	}

	if !m.Locked {
		slot.unblock()
	}
	if !s.closed {
		if slot.migrate.bc != nil {
			log.Warnf("fill slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
				slot.id, slot.backend.bc.Addr(), slot.migrate.bc.Addr(), slot.lock.hold)
		} else {
			log.Warnf("fill slot %04d, backend.addr = %s, locked = %t",
				slot.id, slot.backend.bc.Addr(), slot.lock.hold)
		}
	}
}

func (s *Router) SwitchMasters(masters map[int]string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	for i := range s.slots {
		s.trySwitchMaster(i, masters)
	}
	return nil
}

func (s *Router) trySwitchMaster(id int, masters map[int]string) {
	var refill = false
	var m = s.slots[id].snapshot(false)

	if addr := masters[m.BackendAddrId]; len(addr) != 0 {
		if addr != m.BackendAddr {
			m.BackendAddr = addr
			refill = true
		}
	}
	if from := masters[m.MigrateFromId]; len(from) != 0 {
		if from != m.MigrateFrom {
			m.MigrateFrom = from
			refill = true
		}
	}
	if !refill {
		return
	}
	log.Warnf("refill slot %04d +switch-master", id)

	s.fillSlot(m)
}
<commit_msg>proxy: adjust log format<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.RWMutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := range s.slots {
		s.slots[i].id = i
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.closed = true

	for i := range s.slots {
		s.fillSlot(&models.Slot{Id: i})
	}
}

func (s *Router) GetGroupIds() map[int]bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	var groups = make(map[int]bool)
	for i := range s.slots {
		if gid := s.slots[i].backend.id; gid != 0 {
			groups[gid] = true
		}
		if gid := s.slots[i].migrate.id; gid != 0 {
			groups[gid] = true
		}
	}
	return groups
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	slots := make([]*models.Slot, len(s.slots))
	for i := range s.slots {
		slots[i] = s.slots[i].snapshot(true)
	}
	return slots
}

func (s *Router) GetSlot(id int) *models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if id < 0 || id >= len(s.slots) {
		return nil
	}
	slot := &s.slots[id]
	return slot.snapshot(true)
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(m *models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	if m.Id < 0 || m.Id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	s.fillSlot(m)
	return nil
}

func (s *Router) KeepAlive() error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	var id = Hash(hkey) % uint32(len(s.slots))
	slot := &s.slots[id]
	return slot.forward(r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(r, nil)
}

func (s *Router) dispatchAddr(r *Request, addr string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	bc := s.getBackendConn(addr, false)
	if bc == nil {
		return false
	} else {
		bc.PushBack(r)
		s.putBackendConn(bc)
		return true
	}
}

func (s *Router) getBackendConn(addr string, create bool) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else if create {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	} else {
		return nil
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) fillSlot(m *models.Slot) {
	slot := &s.slots[m.Id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend.bc)
	slot.backend.bc = nil
	slot.backend.id = 0
	s.putBackendConn(slot.migrate.bc)
	slot.migrate.bc = nil
	slot.migrate.id = 0
	for i := range slot.replicaGroups {
		for _, bc := range slot.replicaGroups[i] {
			s.putBackendConn(bc)
		}
	}
	slot.replicaGroups = nil

	if addr := m.BackendAddr; len(addr) != 0 {
		slot.backend.bc = s.getBackendConn(addr, true)
		slot.backend.id = m.BackendAddrId
	}
	if from := m.MigrateFrom; len(from) != 0 {
		slot.migrate.bc = s.getBackendConn(from, true)
		slot.migrate.id = m.MigrateFromId
	}
	for i := range m.ReplicaGroups {
		var group []*SharedBackendConn
		for _, addr := range m.ReplicaGroups[i] {
			group = append(group, s.getBackendConn(addr, true))
		}
		slot.replicaGroups = append(slot.replicaGroups, group)
	}

	if !m.Locked {
		slot.unblock()
	}
	if !s.closed {
		if slot.migrate.bc != nil {
			log.Warnf("fill   slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
				slot.id, slot.backend.bc.Addr(), slot.migrate.bc.Addr(), slot.lock.hold)
		} else {
			log.Warnf("fill   slot %04d, backend.addr = %s, locked = %t",
				slot.id, slot.backend.bc.Addr(), slot.lock.hold)
		}
	}
}

func (s *Router) SwitchMasters(masters map[int]string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	for i := range s.slots {
		s.trySwitchMaster(i, masters)
	}
	return nil
}

func (s *Router) trySwitchMaster(id int, masters map[int]string) {
	var refill = false
	var m = s.slots[id].snapshot(false)

	if addr := masters[m.BackendAddrId]; len(addr) != 0 {
		if addr != m.BackendAddr {
			m.BackendAddr = addr
			refill = true
		}
	}
	if from := masters[m.MigrateFromId]; len(from) != 0 {
		if from != m.MigrateFrom {
			m.MigrateFrom = from
			refill = true
		}
	}
	if !refill {
		return
	}
	log.Warnf("refill slot %04d +switch-master", id)

	s.fillSlot(m)
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.Mutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := 0; i < len(s.slots); i++ {
		s.slots[i].id = uint32(i)
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	for i := 0; i < len(s.slots); i++ {
		s.resetSlot(i)
	}
	s.closed = true
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.Lock()
	defer s.mu.Unlock()
	slots := make([]*models.Slot, 0, len(s.slots))
	for i := 0; i < len(s.slots); i++ {
		slot := &s.slots[i]
		slots = append(slots, &models.Slot{
			Id:          i,
			BackendAddr: slot.backend.Addr(),
			MigrateFrom: slot.migrate.Addr(),
			Locked:      slot.lock.hold,
		})
	}
	return slots
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(id int, addr, from string, locked bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	return s.fillSlot(id, addr, from, locked)
}

func (s *Router) KeepAlive() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	slot := &s.slots[hashSlot(hkey)]
	return slot.forward(r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(r, nil)
}

func (s *Router) getBackendConn(addr string) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) resetSlot(id int) {
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	slot.reset()

	slot.unblock()
}

func (s *Router) fillSlot(id int, addr, from string, locked bool) error {
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	slot.reset()

	if len(addr) != 0 {
		slot.backend = s.getBackendConn(addr)
	}
	if len(from) != 0 {
		slot.migrate = s.getBackendConn(from)
	}

	if !locked {
		slot.unblock()
	}

	if slot.migrate != nil {
		log.Warnf("fill slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
			id, slot.backend.Addr(), slot.migrate.Addr(), locked)
	} else {
		log.Warnf("fill slot %04d, backend.addr = %s, locked = %t",
			id, slot.backend.Addr(), locked)
	}
	return nil
}
<commit_msg>router: use RWLock in Router<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.RWMutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := 0; i < len(s.slots); i++ {
		s.slots[i].id = uint32(i)
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	for i := 0; i < len(s.slots); i++ {
		s.resetSlot(i)
	}
	s.closed = true
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	slots := make([]*models.Slot, 0, len(s.slots))
	for i := 0; i < len(s.slots); i++ {
		slot := &s.slots[i]
		slots = append(slots, &models.Slot{
			Id:          i,
			BackendAddr: slot.backend.Addr(),
			MigrateFrom: slot.migrate.Addr(),
			Locked:      slot.lock.hold,
		})
	}
	return slots
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(id int, addr, from string, locked bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	return s.fillSlot(id, addr, from, locked)
}

func (s *Router) KeepAlive() error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	slot := &s.slots[hashSlot(hkey)]
	return slot.forward(r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(r, nil)
}

func (s *Router) getBackendConn(addr string, create bool) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else if create {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	} else {
		return nil
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) resetSlot(id int) {
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	slot.reset()

	slot.unblock()
}

func (s *Router) fillSlot(id int, addr, from string, locked bool) error {
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	slot.reset()

	if len(addr) != 0 {
		slot.backend = s.getBackendConn(addr, true)
	}
	if len(from) != 0 {
		slot.migrate = s.getBackendConn(from, true)
	}

	if !locked {
		slot.unblock()
	}

	if slot.migrate != nil {
		log.Warnf("fill slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
			id, slot.backend.Addr(), slot.migrate.Addr(), locked)
	} else {
		log.Warnf("fill slot %04d, backend.addr = %s, locked = %t",
			id, slot.backend.Addr(), locked)
	}
	return nil
}
<|endoftext|>