<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sync

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/loggo"

	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/simplestreams"
	"launchpad.net/juju-core/environs/storage"
	envtools "launchpad.net/juju-core/environs/tools"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

var logger = loggo.GetLogger("juju.environs.sync")

// SyncContext describes the context for tool synchronization.
type SyncContext struct {
	// Target holds the destination for the tool synchronization
	Target storage.Storage

	// AllVersions controls the copy of all versions, not only the latest.
	AllVersions bool

	// Copy tools with major version, if MajorVersion > 0.
	MajorVersion int

	// Copy tools with minor version, if MinorVersion > 0.
	MinorVersion int

	// DryRun controls that nothing is copied. Instead it's logged
	// what would be coppied.
	DryRun bool

	// Dev controls the copy of development versions as well as released ones.
	Dev bool

	// Tools are being synced for a public cloud so include mirrors information.
	Public bool

	// Source, if non-empty, specifies a directory in the local file system
	// to use as a source.
	Source string
}

// SyncTools copies the Juju tools tarball from the official bucket
// or a specified source directory into the user's environment.
func SyncTools(syncContext *SyncContext) error {
	sourceDataSource, err := selectSourceDatasource(syncContext)
	if err != nil {
		return err
	}

	logger.Infof("listing available tools")
	if syncContext.MajorVersion == 0 && syncContext.MinorVersion == 0 {
		syncContext.MajorVersion = version.Current.Major
		syncContext.MinorVersion = -1
		if !syncContext.AllVersions {
			syncContext.MinorVersion = version.Current.Minor
		}
	} else if !syncContext.Dev && syncContext.MinorVersion != -1 {
		// If a major.minor version is specified, we allow dev versions.
		// If Dev is already true, leave it alone.
		syncContext.Dev = true
	}

	released := !syncContext.Dev && !version.Current.IsDev()
	sourceTools, err := envtools.FindToolsForCloud(
		[]simplestreams.DataSource{sourceDataSource}, simplestreams.CloudSpec{},
		syncContext.MajorVersion, syncContext.MinorVersion, coretools.Filter{Released: released})
	if err != nil {
		return err
	}

	logger.Infof("found %d tools", len(sourceTools))
	if !syncContext.AllVersions {
		var latest version.Number
		latest, sourceTools = sourceTools.Newest()
		logger.Infof("found %d recent tools (version %s)", len(sourceTools), latest)
	}
	for _, tool := range sourceTools {
		logger.Debugf("found source tool: %v", tool)
	}

	logger.Infof("listing target tools storage")
	targetStorage := syncContext.Target
	targetTools, err := envtools.ReadList(targetStorage, syncContext.MajorVersion, -1)
	switch err {
	case nil, coretools.ErrNoMatches, envtools.ErrNoTools:
	default:
		return err
	}
	for _, tool := range targetTools {
		logger.Debugf("found target tool: %v", tool)
	}

	missing := sourceTools.Exclude(targetTools)
	logger.Infof("found %d tools in target; %d tools to be copied", len(targetTools), len(missing))
	err = copyTools(missing, syncContext, targetStorage)
	if err != nil {
		return err
	}
	logger.Infof("copied %d tools", len(missing))

	logger.Infof("generating tools metadata")
	if !syncContext.DryRun {
		targetTools = append(targetTools, missing...)
		writeMirrors := envtools.DoNotWriteMirrors
		if syncContext.Public {
			writeMirrors = envtools.WriteMirrors
		}
		err = envtools.MergeAndWriteMetadata(targetStorage, targetTools, writeMirrors)
		if err != nil {
			return err
		}
	}
	logger.Infof("tools metadata written")
	return nil
}

// selectSourceStorage returns a storage reader based on the source setting.
func selectSourceDatasource(syncContext *SyncContext) (simplestreams.DataSource, error) {
	source := syncContext.Source
	if source == "" {
		source = envtools.DefaultBaseURL
	}
	sourceURL, err := envtools.ToolsURL(source)
	if err != nil {
		return nil, err
	}
	logger.Infof("using sync tools source: %v", sourceURL)
	return simplestreams.NewURLDataSource("sync tools source", sourceURL, utils.VerifySSLHostnames), nil
}

// copyTools copies a set of tools from the source to the target.
func copyTools(tools []*coretools.Tools, syncContext *SyncContext, dest storage.Storage) error {
	for _, tool := range tools {
		logger.Infof("copying %s from %s", tool.Version, tool.URL)
		if syncContext.DryRun {
			continue
		}
		if err := copyOneToolsPackage(tool, dest); err != nil {
			return err
		}
	}
	return nil
}

// copyOneToolsPackage copies one tool from the source to the target.
func copyOneToolsPackage(tool *coretools.Tools, dest storage.Storage) error {
	toolsName := envtools.StorageName(tool.Version)
	logger.Infof("copying %v", toolsName)
	resp, err := utils.GetValidatingHTTPClient().Get(tool.URL)
	if err != nil {
		return err
	}
	buf := &bytes.Buffer{}
	srcFile := resp.Body
	defer srcFile.Close()
	tool.SHA256, tool.Size, err = utils.ReadSHA256(io.TeeReader(srcFile, buf))
	if err != nil {
		return err
	}
	sizeInKB := (tool.Size + 512) / 1024
	logger.Infof("downloaded %v (%dkB), uploading", toolsName, sizeInKB)
	logger.Infof("download %dkB, uploading", sizeInKB)
	return dest.Put(toolsName, buf, tool.Size)
}

// UploadFunc is the type of Upload, which may be
// reassigned to control the behaviour of tools
// uploading.
type UploadFunc func(stor storage.Storage, forceVersion *version.Number, series ...string) (*coretools.Tools, error)

// Upload builds whatever version of launchpad.net/juju-core is in $GOPATH,
// uploads it to the given storage, and returns a Tools instance describing
// them. If forceVersion is not nil, the uploaded tools bundle will report
// the given version number; if any fakeSeries are supplied, additional copies
// of the built tools will be uploaded for use by machines of those series.
// Juju tools built for one series do not necessarily run on another, but this
// func exists only for development use cases.
var Upload UploadFunc = upload

func upload(stor storage.Storage, forceVersion *version.Number, fakeSeries ...string) (*coretools.Tools, error) {
	builtTools, err := BuildToolsTarball(forceVersion)
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(builtTools.Dir)
	logger.Debugf("Uploading tools for %v", fakeSeries)
	return SyncBuiltTools(stor, builtTools, fakeSeries...)
}

// cloneToolsForSeries copies the built tools tarball into a tarball for the specified
// series and generates corresponding metadata.
func cloneToolsForSeries(toolsInfo *BuiltTools, series ...string) error {
	// Copy the tools to the target storage, recording a Tools struct for each one.
	var targetTools coretools.List
	targetTools = append(targetTools, &coretools.Tools{
		Version: toolsInfo.Version,
		Size:    toolsInfo.Size,
		SHA256:  toolsInfo.Sha256Hash,
	})
	putTools := func(vers version.Binary) (string, error) {
		name := envtools.StorageName(vers)
		src := filepath.Join(toolsInfo.Dir, toolsInfo.StorageName)
		dest := filepath.Join(toolsInfo.Dir, name)
		err := utils.CopyFile(dest, src)
		if err != nil {
			return "", err
		}
		// Append to targetTools the attributes required to write out tools metadata.
		targetTools = append(targetTools, &coretools.Tools{
			Version: vers,
			Size:    toolsInfo.Size,
			SHA256:  toolsInfo.Sha256Hash,
		})
		return name, nil
	}
	logger.Debugf("generating tarballs for %v", series)
	for _, series := range series {
		_, err := simplestreams.SeriesVersion(series)
		if err != nil {
			return err
		}
		if series != toolsInfo.Version.Series {
			fakeVersion := toolsInfo.Version
			fakeVersion.Series = series
			if _, err := putTools(fakeVersion); err != nil {
				return err
			}
		}
	}
	// The tools have been copied to a temp location from which they will be uploaded,
	// now write out the matching simplestreams metadata so that SyncTools can find them.
	metadataStore, err := filestorage.NewFileStorageWriter(toolsInfo.Dir)
	if err != nil {
		return err
	}
	logger.Debugf("generating tools metadata")
	return envtools.MergeAndWriteMetadata(metadataStore, targetTools, false)
}

// BuiltTools contains metadata for a tools tarball resulting from
// a call to BundleTools.
type BuiltTools struct {
	Version     version.Binary
	Dir         string
	StorageName string
	Sha256Hash  string
	Size        int64
}

// BuildToolsTarballFunc is a function which can build a tools tarball.
type BuildToolsTarballFunc func(forceVersion *version.Number) (*BuiltTools, error)

// Override for testing.
var BuildToolsTarball BuildToolsTarballFunc = buildToolsTarball

// buildToolsTarball bundles a tools tarball and places it in a temp directory in
// the expected tools path.
func buildToolsTarball(forceVersion *version.Number) (builtTools *BuiltTools, err error) {
	// TODO(rog) find binaries from $PATH when not using a development
	// version of juju within a $GOPATH.

	logger.Debugf("Building tools")
	// We create the entire archive before asking the environment to
	// start uploading so that we can be sure we have archived
	// correctly.
	f, err := ioutil.TempFile("", "juju-tgz")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	toolsVersion, sha256Hash, err := envtools.BundleTools(f, forceVersion)
	if err != nil {
		return nil, err
	}
	fileInfo, err := f.Stat()
	if err != nil {
		return nil, fmt.Errorf("cannot stat newly made tools archive: %v", err)
	}
	size := fileInfo.Size()
	logger.Infof("built tools %v (%dkB)", toolsVersion, (size+512)/1024)
	baseToolsDir, err := ioutil.TempDir("", "")
	if err != nil {
		return nil, err
	}

	// If we exit with an error, clean up the built tools directory.
	defer func() {
		if err != nil {
			os.RemoveAll(baseToolsDir)
		}
	}()

	err = os.MkdirAll(filepath.Join(baseToolsDir, storage.BaseToolsPath, "releases"), 0755)
	if err != nil {
		return nil, err
	}
	storageName := envtools.StorageName(toolsVersion)
	err = utils.CopyFile(filepath.Join(baseToolsDir, storageName), f.Name())
	if err != nil {
		return nil, err
	}
	return &BuiltTools{
		Version:     toolsVersion,
		Dir:         baseToolsDir,
		StorageName: storageName,
		Size:        size,
		Sha256Hash:  sha256Hash,
	}, nil
}

// SyncBuiltTools copies to storage a tools tarball and cloned copies for each series.
func SyncBuiltTools(stor storage.Storage, builtTools *BuiltTools, fakeSeries ...string) (*coretools.Tools, error) {
	if err := cloneToolsForSeries(builtTools, fakeSeries...); err != nil {
		return nil, err
	}
	syncContext := &SyncContext{
		Source:       builtTools.Dir,
		Target:       stor,
		AllVersions:  true,
		Dev:          builtTools.Version.IsDev(),
		MajorVersion: builtTools.Version.Major,
		MinorVersion: -1,
	}
	logger.Debugf("uploading tools to cloud storage")
	err := SyncTools(syncContext)
	if err != nil {
		return nil, err
	}
	url, err := stor.URL(builtTools.StorageName)
	if err != nil {
		return nil, err
	}
	return &coretools.Tools{
		Version: builtTools.Version,
		URL:     url,
		Size:    builtTools.Size,
		SHA256:  builtTools.Sha256Hash,
	}, nil
}
<commit_msg>use temp dir prefix<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sync

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/juju/loggo"

	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/simplestreams"
	"launchpad.net/juju-core/environs/storage"
	envtools "launchpad.net/juju-core/environs/tools"
	coretools "launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

var logger = loggo.GetLogger("juju.environs.sync")

// SyncContext describes the context for tool synchronization.
type SyncContext struct {
	// Target holds the destination for the tool synchronization
	Target storage.Storage

	// AllVersions controls the copy of all versions, not only the latest.
	AllVersions bool

	// Copy tools with major version, if MajorVersion > 0.
	MajorVersion int

	// Copy tools with minor version, if MinorVersion > 0.
	MinorVersion int

	// DryRun controls that nothing is copied. Instead it's logged
	// what would be coppied.
	DryRun bool

	// Dev controls the copy of development versions as well as released ones.
	Dev bool

	// Tools are being synced for a public cloud so include mirrors information.
	Public bool

	// Source, if non-empty, specifies a directory in the local file system
	// to use as a source.
	Source string
}

// SyncTools copies the Juju tools tarball from the official bucket
// or a specified source directory into the user's environment.
func SyncTools(syncContext *SyncContext) error {
	sourceDataSource, err := selectSourceDatasource(syncContext)
	if err != nil {
		return err
	}

	logger.Infof("listing available tools")
	if syncContext.MajorVersion == 0 && syncContext.MinorVersion == 0 {
		syncContext.MajorVersion = version.Current.Major
		syncContext.MinorVersion = -1
		if !syncContext.AllVersions {
			syncContext.MinorVersion = version.Current.Minor
		}
	} else if !syncContext.Dev && syncContext.MinorVersion != -1 {
		// If a major.minor version is specified, we allow dev versions.
		// If Dev is already true, leave it alone.
		syncContext.Dev = true
	}

	released := !syncContext.Dev && !version.Current.IsDev()
	sourceTools, err := envtools.FindToolsForCloud(
		[]simplestreams.DataSource{sourceDataSource}, simplestreams.CloudSpec{},
		syncContext.MajorVersion, syncContext.MinorVersion, coretools.Filter{Released: released})
	if err != nil {
		return err
	}

	logger.Infof("found %d tools", len(sourceTools))
	if !syncContext.AllVersions {
		var latest version.Number
		latest, sourceTools = sourceTools.Newest()
		logger.Infof("found %d recent tools (version %s)", len(sourceTools), latest)
	}
	for _, tool := range sourceTools {
		logger.Debugf("found source tool: %v", tool)
	}

	logger.Infof("listing target tools storage")
	targetStorage := syncContext.Target
	targetTools, err := envtools.ReadList(targetStorage, syncContext.MajorVersion, -1)
	switch err {
	case nil, coretools.ErrNoMatches, envtools.ErrNoTools:
	default:
		return err
	}
	for _, tool := range targetTools {
		logger.Debugf("found target tool: %v", tool)
	}

	missing := sourceTools.Exclude(targetTools)
	logger.Infof("found %d tools in target; %d tools to be copied", len(targetTools), len(missing))
	err = copyTools(missing, syncContext, targetStorage)
	if err != nil {
		return err
	}
	logger.Infof("copied %d tools", len(missing))

	logger.Infof("generating tools metadata")
	if !syncContext.DryRun {
		targetTools = append(targetTools, missing...)
		writeMirrors := envtools.DoNotWriteMirrors
		if syncContext.Public {
			writeMirrors = envtools.WriteMirrors
		}
		err = envtools.MergeAndWriteMetadata(targetStorage, targetTools, writeMirrors)
		if err != nil {
			return err
		}
	}
	logger.Infof("tools metadata written")
	return nil
}

// selectSourceStorage returns a storage reader based on the source setting.
func selectSourceDatasource(syncContext *SyncContext) (simplestreams.DataSource, error) {
	source := syncContext.Source
	if source == "" {
		source = envtools.DefaultBaseURL
	}
	sourceURL, err := envtools.ToolsURL(source)
	if err != nil {
		return nil, err
	}
	logger.Infof("using sync tools source: %v", sourceURL)
	return simplestreams.NewURLDataSource("sync tools source", sourceURL, utils.VerifySSLHostnames), nil
}

// copyTools copies a set of tools from the source to the target.
func copyTools(tools []*coretools.Tools, syncContext *SyncContext, dest storage.Storage) error {
	for _, tool := range tools {
		logger.Infof("copying %s from %s", tool.Version, tool.URL)
		if syncContext.DryRun {
			continue
		}
		if err := copyOneToolsPackage(tool, dest); err != nil {
			return err
		}
	}
	return nil
}

// copyOneToolsPackage copies one tool from the source to the target.
func copyOneToolsPackage(tool *coretools.Tools, dest storage.Storage) error {
	toolsName := envtools.StorageName(tool.Version)
	logger.Infof("copying %v", toolsName)
	resp, err := utils.GetValidatingHTTPClient().Get(tool.URL)
	if err != nil {
		return err
	}
	buf := &bytes.Buffer{}
	srcFile := resp.Body
	defer srcFile.Close()
	tool.SHA256, tool.Size, err = utils.ReadSHA256(io.TeeReader(srcFile, buf))
	if err != nil {
		return err
	}
	sizeInKB := (tool.Size + 512) / 1024
	logger.Infof("downloaded %v (%dkB), uploading", toolsName, sizeInKB)
	logger.Infof("download %dkB, uploading", sizeInKB)
	return dest.Put(toolsName, buf, tool.Size)
}

// UploadFunc is the type of Upload, which may be
// reassigned to control the behaviour of tools
// uploading.
type UploadFunc func(stor storage.Storage, forceVersion *version.Number, series ...string) (*coretools.Tools, error)

// Upload builds whatever version of launchpad.net/juju-core is in $GOPATH,
// uploads it to the given storage, and returns a Tools instance describing
// them. If forceVersion is not nil, the uploaded tools bundle will report
// the given version number; if any fakeSeries are supplied, additional copies
// of the built tools will be uploaded for use by machines of those series.
// Juju tools built for one series do not necessarily run on another, but this
// func exists only for development use cases.
var Upload UploadFunc = upload

func upload(stor storage.Storage, forceVersion *version.Number, fakeSeries ...string) (*coretools.Tools, error) {
	builtTools, err := BuildToolsTarball(forceVersion)
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(builtTools.Dir)
	logger.Debugf("Uploading tools for %v", fakeSeries)
	return SyncBuiltTools(stor, builtTools, fakeSeries...)
}

// cloneToolsForSeries copies the built tools tarball into a tarball for the specified
// series and generates corresponding metadata.
func cloneToolsForSeries(toolsInfo *BuiltTools, series ...string) error {
	// Copy the tools to the target storage, recording a Tools struct for each one.
	var targetTools coretools.List
	targetTools = append(targetTools, &coretools.Tools{
		Version: toolsInfo.Version,
		Size:    toolsInfo.Size,
		SHA256:  toolsInfo.Sha256Hash,
	})
	putTools := func(vers version.Binary) (string, error) {
		name := envtools.StorageName(vers)
		src := filepath.Join(toolsInfo.Dir, toolsInfo.StorageName)
		dest := filepath.Join(toolsInfo.Dir, name)
		err := utils.CopyFile(dest, src)
		if err != nil {
			return "", err
		}
		// Append to targetTools the attributes required to write out tools metadata.
		targetTools = append(targetTools, &coretools.Tools{
			Version: vers,
			Size:    toolsInfo.Size,
			SHA256:  toolsInfo.Sha256Hash,
		})
		return name, nil
	}
	logger.Debugf("generating tarballs for %v", series)
	for _, series := range series {
		_, err := simplestreams.SeriesVersion(series)
		if err != nil {
			return err
		}
		if series != toolsInfo.Version.Series {
			fakeVersion := toolsInfo.Version
			fakeVersion.Series = series
			if _, err := putTools(fakeVersion); err != nil {
				return err
			}
		}
	}
	// The tools have been copied to a temp location from which they will be uploaded,
	// now write out the matching simplestreams metadata so that SyncTools can find them.
	metadataStore, err := filestorage.NewFileStorageWriter(toolsInfo.Dir)
	if err != nil {
		return err
	}
	logger.Debugf("generating tools metadata")
	return envtools.MergeAndWriteMetadata(metadataStore, targetTools, false)
}

// BuiltTools contains metadata for a tools tarball resulting from
// a call to BundleTools.
type BuiltTools struct {
	Version     version.Binary
	Dir         string
	StorageName string
	Sha256Hash  string
	Size        int64
}

// BuildToolsTarballFunc is a function which can build a tools tarball.
type BuildToolsTarballFunc func(forceVersion *version.Number) (*BuiltTools, error)

// Override for testing.
var BuildToolsTarball BuildToolsTarballFunc = buildToolsTarball

// buildToolsTarball bundles a tools tarball and places it in a temp directory in
// the expected tools path.
func buildToolsTarball(forceVersion *version.Number) (builtTools *BuiltTools, err error) {
	// TODO(rog) find binaries from $PATH when not using a development
	// version of juju within a $GOPATH.

	logger.Debugf("Building tools")
	// We create the entire archive before asking the environment to
	// start uploading so that we can be sure we have archived
	// correctly.
	f, err := ioutil.TempFile("", "juju-tgz")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	defer os.Remove(f.Name())
	toolsVersion, sha256Hash, err := envtools.BundleTools(f, forceVersion)
	if err != nil {
		return nil, err
	}
	fileInfo, err := f.Stat()
	if err != nil {
		return nil, fmt.Errorf("cannot stat newly made tools archive: %v", err)
	}
	size := fileInfo.Size()
	logger.Infof("built tools %v (%dkB)", toolsVersion, (size+512)/1024)
	baseToolsDir, err := ioutil.TempDir("", "juju-tools")
	if err != nil {
		return nil, err
	}

	// If we exit with an error, clean up the built tools directory.
	defer func() {
		if err != nil {
			os.RemoveAll(baseToolsDir)
		}
	}()

	err = os.MkdirAll(filepath.Join(baseToolsDir, storage.BaseToolsPath, "releases"), 0755)
	if err != nil {
		return nil, err
	}
	storageName := envtools.StorageName(toolsVersion)
	err = utils.CopyFile(filepath.Join(baseToolsDir, storageName), f.Name())
	if err != nil {
		return nil, err
	}
	return &BuiltTools{
		Version:     toolsVersion,
		Dir:         baseToolsDir,
		StorageName: storageName,
		Size:        size,
		Sha256Hash:  sha256Hash,
	}, nil
}

// SyncBuiltTools copies to storage a tools tarball and cloned copies for each series.
func SyncBuiltTools(stor storage.Storage, builtTools *BuiltTools, fakeSeries ...string) (*coretools.Tools, error) {
	if err := cloneToolsForSeries(builtTools, fakeSeries...); err != nil {
		return nil, err
	}
	syncContext := &SyncContext{
		Source:       builtTools.Dir,
		Target:       stor,
		AllVersions:  true,
		Dev:          builtTools.Version.IsDev(),
		MajorVersion: builtTools.Version.Major,
		MinorVersion: -1,
	}
	logger.Debugf("uploading tools to cloud storage")
	err := SyncTools(syncContext)
	if err != nil {
		return nil, err
	}
	url, err := stor.URL(builtTools.StorageName)
	if err != nil {
		return nil, err
	}
	return &coretools.Tools{
		Version: builtTools.Version,
		URL:     url,
		Size:    builtTools.Size,
		SHA256:  builtTools.Sha256Hash,
	}, nil
}
<|endoftext|>


<commit_before>package eagain

import (
	"io"
	"syscall"

	"github.com/lxc/lxd/shared"
)

// Reader represents an io.Reader that handles EAGAIN
type Reader struct {
	Reader io.Reader
}

// Read behaves like io.Reader.Read but will retry on EAGAIN
func (er Reader) Read(p []byte) (int, error) {
again:
	n, err := er.Reader.Read(p)
	if err == nil {
		return n, nil
	}

	// keep retrying on EAGAIN
	errno, ok := shared.GetErrno(err)
	if ok && errno == syscall.EAGAIN {
		goto again
	}

	return n, err
}

// Writer represents an io.Writer that handles EAGAIN
type Writer struct {
	Writer io.Writer
}

// Write behaves like io.Writer.Write but will retry on EAGAIN
func (ew Writer) Write(p []byte) (int, error) {
again:
	n, err := ew.Writer.Write(p)
	if err == nil {
		return n, nil
	}

	// keep retrying on EAGAIN
	errno, ok := shared.GetErrno(err)
	if ok && errno == syscall.EAGAIN {
		goto again
	}

	return n, err
}
<commit_msg>reader: Handle EINTR<commit_after>package eagain

import (
	"io"
	"syscall"

	"github.com/lxc/lxd/shared"
)

// Reader represents an io.Reader that handles EAGAIN
type Reader struct {
	Reader io.Reader
}

// Read behaves like io.Reader.Read but will retry on EAGAIN
func (er Reader) Read(p []byte) (int, error) {
again:
	n, err := er.Reader.Read(p)
	if err == nil {
		return n, nil
	}

	// keep retrying on EAGAIN
	errno, ok := shared.GetErrno(err)
	if ok && (errno == syscall.EAGAIN || errno == syscall.EINTR) {
		goto again
	}

	return n, err
}

// Writer represents an io.Writer that handles EAGAIN
type Writer struct {
	Writer io.Writer
}

// Write behaves like io.Writer.Write but will retry on EAGAIN
func (ew Writer) Write(p []byte) (int, error) {
again:
	n, err := ew.Writer.Write(p)
	if err == nil {
		return n, nil
	}

	// keep retrying on EAGAIN
	errno, ok := shared.GetErrno(err)
	if ok && (errno == syscall.EAGAIN || errno == syscall.EINTR) {
		goto again
	}

	return n, err
}
<|endoftext|>


<commit_before>package cli

import (
	"context"

	"github.com/cosmos/cosmos-sdk/client"
	"github.com/cosmos/cosmos-sdk/client/flags"
	"github.com/spf13/cobra"
	"github.com/zigbee-alliance/distributed-compliance-ledger/x/dclauth/types"
)

func CmdListRevokedAccount() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "list-revoked-account",
		Short: "list all RevokedAccount",
		RunE: func(cmd *cobra.Command, args []string) error {
			clientCtx := client.GetClientContextFromCmd(cmd)

			pageReq, err := client.ReadPageRequest(cmd.Flags())
			if err != nil {
				return err
			}

			queryClient := types.NewQueryClient(clientCtx)

			params := &types.QueryAllRevokedAccountRequest{
				Pagination: pageReq,
			}

			res, err := queryClient.RevokedAccountAll(context.Background(), params)
			if err != nil {
				return err
			}

			return clientCtx.PrintProto(res)
		},
	}

	flags.AddPaginationFlagsToCmd(cmd, cmd.Use)
	flags.AddQueryFlagsToCmd(cmd)

	return cmd
}

func CmdShowRevokedAccount() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "revoked-account [address]",
		Short: "shows a RevokedAccount",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) (err error) {
			clientCtx := client.GetClientContextFromCmd(cmd)

			queryClient := types.NewQueryClient(clientCtx)

			argAddress := args[0]

			params := &types.QueryGetRevokedAccountRequest{
				Address: argAddress,
			}

			res, err := queryClient.RevokedAccount(context.Background(), params)
			if err != nil {
				return err
			}

			return clientCtx.PrintProto(res)
		},
	}

	cmd.Flags().String(FlagAddress, "", "Bech32 encoded account address")
	flags.AddQueryFlagsToCmd(cmd)

	_ = cmd.MarkFlagRequired(FlagAddress)

	return cmd
}
<commit_msg>Change revoked account query<commit_after>package cli

import (
	"context"

	"github.com/cosmos/cosmos-sdk/client"
	"github.com/cosmos/cosmos-sdk/client/flags"
	sdk "github.com/cosmos/cosmos-sdk/types"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/zigbee-alliance/distributed-compliance-ledger/utils/cli"
	"github.com/zigbee-alliance/distributed-compliance-ledger/x/dclauth/types"
)

func CmdListRevokedAccount() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "list-revoked-account",
		Short: "list all RevokedAccount",
		RunE: func(cmd *cobra.Command, args []string) error {
			clientCtx := client.GetClientContextFromCmd(cmd)

			pageReq, err := client.ReadPageRequest(cmd.Flags())
			if err != nil {
				return err
			}

			queryClient := types.NewQueryClient(clientCtx)

			params := &types.QueryAllRevokedAccountRequest{
				Pagination: pageReq,
			}

			res, err := queryClient.RevokedAccountAll(context.Background(), params)
			if err != nil {
				return err
			}

			return clientCtx.PrintProto(res)
		},
	}

	flags.AddPaginationFlagsToCmd(cmd, cmd.Use)
	flags.AddQueryFlagsToCmd(cmd)

	return cmd
}

func CmdShowRevokedAccount() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "revoked-account",
		Short: "shows a RevokedAccount",
		Args:  cobra.ExactArgs(0),
		RunE: func(cmd *cobra.Command, args []string) (err error) {
			clientCtx := client.GetClientContextFromCmd(cmd)

			argAddress, err := sdk.AccAddressFromBech32(viper.GetString(FlagAddress))
			if err != nil {
				return err
			}

			var res types.RevokedAccount
			return cli.QueryWithProof(
				clientCtx,
				types.StoreKey,
				types.PendingAccountRevocationKeyPrefix,
				types.PendingAccountRevocationKey(argAddress),
				&res,
			)
		},
	}

	cmd.Flags().String(FlagAddress, "", "Bech32 encoded account address")
	flags.AddQueryFlagsToCmd(cmd)

	_ = cmd.MarkFlagRequired(FlagAddress)

	return cmd
}
<|endoftext|>


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

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	ErrorString string
}

func (err *Error) String() string { return err.ErrorString }

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.  Even if Read
// returns n < len(p), it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available instead of waiting for more.
//
// When Read encounters an error or end-of-file condition after
// successfully reading n > 0 bytes, it returns the number of
// bytes read.  It may return the (non-nil) error from the same call
// or return the error (and n == 0) from a subsequent call.
// An instance of this general case is that a Reader returning
// a non-zero number of bytes at the end of the input stream may
// return either err == os.EOF or err == nil.  The next Read should
// return 0, os.EOF regardless.
//
// Callers should always process the n > 0 bytes returned before
// considering the error err.  Doing so correctly handles I/O errors
// that happen after reading some bytes and also both of the
// allowed EOF behaviors.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying input source.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// When ReadAt returns n < len(p), it returns a non-nil error
// explaining why more bytes were not returned.  In this respect,
// ReadAt is stricter than Read.
//
// Even if ReadAt returns n < len(p), it may use all of p as scratch
// space during the call.  If some data is available but not len(p) bytes,
// ReadAt blocks until either all the data is available or an error occurs.
// In this respect ReadAt is different from Read.
//
// If the n = len(p) bytes returned by ReadAt are at the end of the
// input source, ReadAt may return either err == os.EOF or err == nil.
//
// If ReadAt is reading from an input source with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// ByteScanner is the interface that adds the UnreadByte method to the
// basic ReadByte method.
//
// UnreadByte causes the next call to ReadByte to return the same byte
// as the previous call to ReadByte.
// It may be an error to call UnreadByte twice without an intervening
// call to ReadByte.
type ByteScanner interface {
	ByteReader
	UnreadByte() os.Error
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// RuneScanner is the interface that adds the UnreadRune method to the
// basic ReadRune method.
//
// UnreadRune causes the next call to ReadRune to return the same rune
// as the previous call to ReadRune.
// It may be an error to call UnreadRune twice without an intervening
// call to ReadRune.
type RuneScanner interface {
	RuneReader
	UnreadRune() os.Error
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the earliest
// error encountered while copying.  Because Read can
// return the full amount requested as well as an error
// (including os.EOF), so can Copyn.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the first error encountered while copying, if any.
//
// A successful Copy returns err == nil, not err == os.EOF.
// Because Copy is defined to read from src until EOF, it does
// not treat an EOF from Read as an error to be reported.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<commit_msg>io.WriteString: if the object has a WriteString method, use it<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package io provides basic interfaces to I/O primitives.
// Its primary job is to wrap existing implementations of such primitives,
// such as those in package os, into shared public interfaces that
// abstract the functionality, plus some other related primitives.
package io

import "os"

// Error represents an unexpected I/O behavior.
type Error struct {
	ErrorString string
}

func (err *Error) String() string { return err.ErrorString }

// ErrShortWrite means that a write accepted fewer bytes than requested
// but failed to return an explicit error.
var ErrShortWrite os.Error = &Error{"short write"}

// ErrShortBuffer means that a read required a longer buffer than was provided.
var ErrShortBuffer os.Error = &Error{"short buffer"}

// ErrUnexpectedEOF means that os.EOF was encountered in the
// middle of reading a fixed-size block or data structure.
var ErrUnexpectedEOF os.Error = &Error{"unexpected EOF"}

// Reader is the interface that wraps the basic Read method.
//
// Read reads up to len(p) bytes into p.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.  Even if Read
// returns n < len(p), it may use all of p as scratch space during the call.
// If some data is available but not len(p) bytes, Read conventionally
// returns what is available instead of waiting for more.
//
// When Read encounters an error or end-of-file condition after
// successfully reading n > 0 bytes, it returns the number of
// bytes read.  It may return the (non-nil) error from the same call
// or return the error (and n == 0) from a subsequent call.
// An instance of this general case is that a Reader returning
// a non-zero number of bytes at the end of the input stream may
// return either err == os.EOF or err == nil.  The next Read should
// return 0, os.EOF regardless.
//
// Callers should always process the n > 0 bytes returned before
// considering the error err.  Doing so correctly handles I/O errors
// that happen after reading some bytes and also both of the
// allowed EOF behaviors.
type Reader interface {
	Read(p []byte) (n int, err os.Error)
}

// Writer is the interface that wraps the basic Write method.
//
// Write writes len(p) bytes from p to the underlying data stream.
// It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// Write must return a non-nil error if it returns n < len(p).
type Writer interface {
	Write(p []byte) (n int, err os.Error)
}

// Closer is the interface that wraps the basic Close method.
type Closer interface {
	Close() os.Error
}

// Seeker is the interface that wraps the basic Seek method.
//
// Seek sets the offset for the next Read or Write to offset,
// interpreted according to whence: 0 means relative to the origin of
// the file, 1 means relative to the current offset, and 2 means
// relative to the end.  Seek returns the new offset and an Error, if
// any.
type Seeker interface {
	Seek(offset int64, whence int) (ret int64, err os.Error)
}

// ReadWriter is the interface that groups the basic Read and Write methods.
type ReadWriter interface {
	Reader
	Writer
}

// ReadCloser is the interface that groups the basic Read and Close methods.
type ReadCloser interface {
	Reader
	Closer
}

// WriteCloser is the interface that groups the basic Write and Close methods.
type WriteCloser interface {
	Writer
	Closer
}

// ReadWriteCloser is the interface that groups the basic Read, Write and Close methods.
type ReadWriteCloser interface {
	Reader
	Writer
	Closer
}

// ReadSeeker is the interface that groups the basic Read and Seek methods.
type ReadSeeker interface {
	Reader
	Seeker
}

// WriteSeeker is the interface that groups the basic Write and Seek methods.
type WriteSeeker interface {
	Writer
	Seeker
}

// ReadWriteSeeker is the interface that groups the basic Read, Write and Seek methods.
type ReadWriteSeeker interface {
	Reader
	Writer
	Seeker
}

// ReaderFrom is the interface that wraps the ReadFrom method.
type ReaderFrom interface {
	ReadFrom(r Reader) (n int64, err os.Error)
}

// WriterTo is the interface that wraps the WriteTo method.
type WriterTo interface {
	WriteTo(w Writer) (n int64, err os.Error)
}

// ReaderAt is the interface that wraps the basic ReadAt method.
//
// ReadAt reads len(p) bytes into p starting at offset off in the
// underlying input source.  It returns the number of bytes
// read (0 <= n <= len(p)) and any error encountered.
//
// When ReadAt returns n < len(p), it returns a non-nil error
// explaining why more bytes were not returned.  In this respect,
// ReadAt is stricter than Read.
//
// Even if ReadAt returns n < len(p), it may use all of p as scratch
// space during the call.  If some data is available but not len(p) bytes,
// ReadAt blocks until either all the data is available or an error occurs.
// In this respect ReadAt is different from Read.
//
// If the n = len(p) bytes returned by ReadAt are at the end of the
// input source, ReadAt may return either err == os.EOF or err == nil.
//
// If ReadAt is reading from an input source with a seek offset,
// ReadAt should not affect nor be affected by the underlying
// seek offset.
type ReaderAt interface {
	ReadAt(p []byte, off int64) (n int, err os.Error)
}

// WriterAt is the interface that wraps the basic WriteAt method.
//
// WriteAt writes len(p) bytes from p to the underlying data stream
// at offset off.  It returns the number of bytes written from p (0 <= n <= len(p))
// and any error encountered that caused the write to stop early.
// WriteAt must return a non-nil error if it returns n < len(p).
type WriterAt interface {
	WriteAt(p []byte, off int64) (n int, err os.Error)
}

// ByteReader is the interface that wraps the ReadByte method.
//
// ReadByte reads and returns the next byte from the input.
// If no byte is available, err will be set.
type ByteReader interface {
	ReadByte() (c byte, err os.Error)
}

// ByteScanner is the interface that adds the UnreadByte method to the
// basic ReadByte method.
//
// UnreadByte causes the next call to ReadByte to return the same byte
// as the previous call to ReadByte.
// It may be an error to call UnreadByte twice without an intervening
// call to ReadByte.
type ByteScanner interface {
	ByteReader
	UnreadByte() os.Error
}

// RuneReader is the interface that wraps the ReadRune method.
//
// ReadRune reads a single UTF-8 encoded Unicode character
// and returns the rune and its size in bytes. If no character is
// available, err will be set.
type RuneReader interface {
	ReadRune() (rune int, size int, err os.Error)
}

// RuneScanner is the interface that adds the UnreadRune method to the
// basic ReadRune method.
//
// UnreadRune causes the next call to ReadRune to return the same rune
// as the previous call to ReadRune.
// It may be an error to call UnreadRune twice without an intervening
// call to ReadRune.
type RuneScanner interface {
	RuneReader
	UnreadRune() os.Error
}

// stringWriter is the interface that wraps the WriteString method.
type stringWriter interface {
	WriteString(s string) (n int, err os.Error)
}

// WriteString writes the contents of the string s to w, which accepts an array of bytes.
func WriteString(w Writer, s string) (n int, err os.Error) {
	if sw, ok := w.(stringWriter); ok {
		return sw.WriteString(s)
	}
	return w.Write([]byte(s))
}

// ReadAtLeast reads from r into buf until it has read at least min bytes.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading fewer than min bytes,
// ReadAtLeast returns ErrUnexpectedEOF.
// If min is greater than the length of buf, ReadAtLeast returns ErrShortBuffer.
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err os.Error) {
	if len(buf) < min {
		return 0, ErrShortBuffer
	}
	for n < min && err == nil {
		var nn int
		nn, err = r.Read(buf[n:])
		n += nn
	}
	if err == os.EOF {
		if n >= min {
			err = nil
		} else if n > 0 {
			err = ErrUnexpectedEOF
		}
	}
	return
}

// ReadFull reads exactly len(buf) bytes from r into buf.
// It returns the number of bytes copied and an error if fewer bytes were read.
// The error is os.EOF only if no bytes were read.
// If an EOF happens after reading some but not all the bytes,
// ReadFull returns ErrUnexpectedEOF.
func ReadFull(r Reader, buf []byte) (n int, err os.Error) {
	return ReadAtLeast(r, buf, len(buf))
}

// Copyn copies n bytes (or until an error) from src to dst.
// It returns the number of bytes copied and the earliest
// error encountered while copying.  Because Read can
// return the full amount requested as well as an error
// (including os.EOF), so can Copyn.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
func Copyn(dst Writer, src Reader, n int64) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids a buffer allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		written, err = rt.ReadFrom(LimitReader(src, n))
		if written < n && err == nil {
			// rt stopped early; must have been EOF.
			err = os.EOF
		}
		return
	}
	buf := make([]byte, 32*1024)
	for written < n {
		l := len(buf)
		if d := n - written; d < int64(l) {
			l = int(d)
		}
		nr, er := src.Read(buf[0:l])
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs.  It returns the number of bytes
// copied and the first error encountered while copying, if any.
//
// A successful Copy returns err == nil, not err == os.EOF.
// Because Copy is defined to read from src until EOF, it does
// not treat an EOF from Read as an error to be reported.
//
// If dst implements the ReaderFrom interface,
// the copy is implemented by calling dst.ReadFrom(src).
// Otherwise, if src implements the WriterTo interface,
// the copy is implemented by calling src.WriteTo(dst).
func Copy(dst Writer, src Reader) (written int64, err os.Error) {
	// If the writer has a ReadFrom method, use it to do the copy.
	// Avoids an allocation and a copy.
	if rt, ok := dst.(ReaderFrom); ok {
		return rt.ReadFrom(src)
	}
	// Similarly, if the reader has a WriteTo method, use it to do the copy.
	if wt, ok := src.(WriterTo); ok {
		return wt.WriteTo(dst)
	}
	buf := make([]byte, 32*1024)
	for {
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = ErrShortWrite
				break
			}
		}
		if er == os.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return written, err
}

// LimitReader returns a Reader that reads from r
// but stops with os.EOF after n bytes.
// The underlying implementation is a *LimitedReader.
func LimitReader(r Reader, n int64) Reader { return &LimitedReader{r, n} }

// A LimitedReader reads from R but limits the amount of
// data returned to just N bytes. Each call to Read
// updates N to reflect the new amount remaining.
type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err os.Error) {
	if l.N <= 0 {
		return 0, os.EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

// NewSectionReader returns a SectionReader that reads from r
// starting at offset off and stops with os.EOF after n bytes.
func NewSectionReader(r ReaderAt, off int64, n int64) *SectionReader {
	return &SectionReader{r, off, off, off + n}
}

// SectionReader implements Read, Seek, and ReadAt on a section
// of an underlying ReaderAt.
type SectionReader struct {
	r     ReaderAt
	base  int64
	off   int64
	limit int64
}

func (s *SectionReader) Read(p []byte) (n int, err os.Error) {
	if s.off >= s.limit {
		return 0, os.EOF
	}
	if max := s.limit - s.off; int64(len(p)) > max {
		p = p[0:max]
	}
	n, err = s.r.ReadAt(p, s.off)
	s.off += int64(n)
	return
}

func (s *SectionReader) Seek(offset int64, whence int) (ret int64, err os.Error) {
	switch whence {
	default:
		return 0, os.EINVAL
	case 0:
		offset += s.base
	case 1:
		offset += s.off
	case 2:
		offset += s.limit
	}
	if offset < s.base || offset > s.limit {
		return 0, os.EINVAL
	}
	s.off = offset
	return offset - s.base, nil
}

func (s *SectionReader) ReadAt(p []byte, off int64) (n int, err os.Error) {
	if off < 0 || off >= s.limit-s.base {
		return 0, os.EOF
	}
	off += s.base
	if max := s.limit - off; int64(len(p)) > max {
		p = p[0:max]
	}
	return s.r.ReadAt(p, off)
}

// Size returns the size of the section in bytes.
func (s *SectionReader) Size() int64 { return s.limit - s.base }
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package apiserver

import (
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

// Operation represents an ongoing action which the server is performing.
type Operation struct {
	ID       string
	result   interface{}
	awaiting <-chan interface{}
	finished *time.Time
	lock     sync.Mutex
	notify   chan bool
}

// Operations tracks all the ongoing operations.
type Operations struct {
	// Access only using functions from atomic.
	lastID int64

	// 'lock' guards the ops map.
	lock sync.Mutex
	ops  map[string]*Operation
}

// NewOperations returns a new Operations repository.
func NewOperations() *Operations {
	ops := &Operations{
		ops: map[string]*Operation{},
	}
	go util.Forever(func() { ops.expire(10 * time.Minute) }, 5*time.Minute)
	return ops
}

// NewOperation adds a new operation. It is lock-free.
func (ops *Operations) NewOperation(from <-chan interface{}) *Operation {
	id := atomic.AddInt64(&ops.lastID, 1)
	op := &Operation{
		ID:       strconv.FormatInt(id, 10),
		awaiting: from,
		notify:   make(chan bool, 1),
	}
	go op.wait()
	go ops.insert(op)
	return op
}

// Inserts op into the ops map.
func (ops *Operations) insert(op *Operation) {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	ops.ops[op.ID] = op
}

// List operations for an API client.
func (ops *Operations) List() api.ServerOpList {
	ops.lock.Lock()
	defer ops.lock.Unlock()

	ids := []string{}
	for id := range ops.ops {
		ids = append(ids, id)
	}
	sort.StringSlice(ids).Sort()
	ol := api.ServerOpList{}
	for _, id := range ids {
		ol.Items = append(ol.Items, api.ServerOp{JSONBase: api.JSONBase{ID: id}})
	}
	return ol
}

// Get returns the operation with the given ID, or nil
func (ops *Operations) Get(id string) *Operation {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	return ops.ops[id]
}

// Garbage collect operations that have finished longer than maxAge ago.
func (ops *Operations) expire(maxAge time.Duration) {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	keep := map[string]*Operation{}
	limitTime := time.Now().Add(-maxAge)
	for id, op := range ops.ops {
		if !op.expired(limitTime) {
			keep[id] = op
		}
	}
	ops.ops = keep
}

// Waits forever for the operation to complete; call via go when
// the operation is created. Sets op.finished when the operation
// does complete, and sends on the notify channel, in case there
// are any WaitFor() calls in progress.
// Does not keep op locked while waiting.
func (op *Operation) wait() {
	defer util.HandleCrash()
	result := <-op.awaiting

	op.lock.Lock()
	defer op.lock.Unlock()
	op.result = result
	finished := time.Now()
	op.finished = &finished
	op.notify <- true
}

// WaitFor waits for the specified duration, or until the operation finishes,
// whichever happens first.
func (op *Operation) WaitFor(timeout time.Duration) {
	select {
	case <-time.After(timeout):
	case <-op.notify:
		// Re-send on this channel in case there are others
		// waiting for notification.
		op.notify <- true
	}
}

// Returns true if this operation finished before limitTime.
func (op *Operation) expired(limitTime time.Time) bool {
	op.lock.Lock()
	defer op.lock.Unlock()
	if op.finished == nil {
		return false
	}
	return op.finished.Before(limitTime)
}

// StatusOrResult returns status information or the result of the operation if it is complete,
// with a bool indicating true in the latter case.
func (op *Operation) StatusOrResult() (description interface{}, finished bool) {
	op.lock.Lock()
	defer op.lock.Unlock()

	if op.finished == nil {
		return api.Status{
			Status:  api.StatusWorking,
			Details: op.ID,
		}, false
	}
	return op.result, true
}
<commit_msg>Changed op.notify<-true to close(op.notify).<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package apiserver

import (
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

// Operation represents an ongoing action which the server is performing.
type Operation struct {
	ID       string
	result   interface{}
	awaiting <-chan interface{}
	finished *time.Time
	lock     sync.Mutex
	notify   chan struct{}
}

// Operations tracks all the ongoing operations.
type Operations struct {
	// Access only using functions from atomic.
	lastID int64

	// 'lock' guards the ops map.
	lock sync.Mutex
	ops  map[string]*Operation
}

// NewOperations returns a new Operations repository.
func NewOperations() *Operations {
	ops := &Operations{
		ops: map[string]*Operation{},
	}
	go util.Forever(func() { ops.expire(10 * time.Minute) }, 5*time.Minute)
	return ops
}

// NewOperation adds a new operation. It is lock-free.
func (ops *Operations) NewOperation(from <-chan interface{}) *Operation {
	id := atomic.AddInt64(&ops.lastID, 1)
	op := &Operation{
		ID:       strconv.FormatInt(id, 10),
		awaiting: from,
		notify:   make(chan struct{}),
	}
	go op.wait()
	go ops.insert(op)
	return op
}

// Inserts op into the ops map.
func (ops *Operations) insert(op *Operation) {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	ops.ops[op.ID] = op
}

// List operations for an API client.
func (ops *Operations) List() api.ServerOpList {
	ops.lock.Lock()
	defer ops.lock.Unlock()

	ids := []string{}
	for id := range ops.ops {
		ids = append(ids, id)
	}
	sort.StringSlice(ids).Sort()
	ol := api.ServerOpList{}
	for _, id := range ids {
		ol.Items = append(ol.Items, api.ServerOp{JSONBase: api.JSONBase{ID: id}})
	}
	return ol
}

// Get returns the operation with the given ID, or nil
func (ops *Operations) Get(id string) *Operation {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	return ops.ops[id]
}

// Garbage collect operations that have finished longer than maxAge ago.
func (ops *Operations) expire(maxAge time.Duration) {
	ops.lock.Lock()
	defer ops.lock.Unlock()
	keep := map[string]*Operation{}
	limitTime := time.Now().Add(-maxAge)
	for id, op := range ops.ops {
		if !op.expired(limitTime) {
			keep[id] = op
		}
	}
	ops.ops = keep
}

// Waits forever for the operation to complete; call via go when
// the operation is created. Sets op.finished when the operation
// does complete, and sends on the notify channel, in case there
// are any WaitFor() calls in progress.
// Does not keep op locked while waiting.
func (op *Operation) wait() {
	defer util.HandleCrash()
	result := <-op.awaiting

	op.lock.Lock()
	defer op.lock.Unlock()
	op.result = result
	finished := time.Now()
	op.finished = &finished
	close(op.notify)
}

// WaitFor waits for the specified duration, or until the operation finishes,
// whichever happens first.
func (op *Operation) WaitFor(timeout time.Duration) {
	select {
	case <-time.After(timeout):
	case <-op.notify:
	}
}

// Returns true if this operation finished before limitTime.
func (op *Operation) expired(limitTime time.Time) bool {
	op.lock.Lock()
	defer op.lock.Unlock()
	if op.finished == nil {
		return false
	}
	return op.finished.Before(limitTime)
}

// StatusOrResult returns status information or the result of the operation if it is complete,
// with a bool indicating true in the latter case.
func (op *Operation) StatusOrResult() (description interface{}, finished bool) {
	op.lock.Lock()
	defer op.lock.Unlock()

	if op.finished == nil {
		return api.Status{
			Status:  api.StatusWorking,
			Details: op.ID,
		}, false
	}
	return op.result, true
}
<|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 filters

import (
	"context"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"

	flowcontrol "k8s.io/api/flowcontrol/v1beta2"
	apitypes "k8s.io/apimachinery/pkg/types"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/httplog"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	flowcontrolrequest "k8s.io/apiserver/pkg/util/flowcontrol/request"
	"k8s.io/klog/v2"
)

// PriorityAndFairnessClassification identifies the results of
// classification for API Priority and Fairness
type PriorityAndFairnessClassification struct {
	FlowSchemaName    string
	FlowSchemaUID     apitypes.UID
	PriorityLevelName string
	PriorityLevelUID  apitypes.UID
}

// waitingMark tracks requests waiting rather than being executed
var waitingMark = &requestWatermark{
	phase:            epmetrics.WaitingPhase,
	readOnlyObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.ReadOnlyKind}).RequestsWaiting,
	mutatingObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.MutatingKind}).RequestsWaiting,
}

var atomicMutatingExecuting, atomicReadOnlyExecuting int32
var atomicMutatingWaiting, atomicReadOnlyWaiting int32

// newInitializationSignal is defined for testing purposes.
var newInitializationSignal = utilflowcontrol.NewInitializationSignal

func truncateLogField(s string) string {
	const maxFieldLogLength = 64

	if len(s) > maxFieldLogLength {
		s = s[0:maxFieldLogLength]
	}
	return s
}

// WithPriorityAndFairness limits the number of in-flight
// requests in a fine-grained way.
func WithPriorityAndFairness(
	handler http.Handler,
	longRunningRequestCheck apirequest.LongRunningRequestCheck,
	fcIfc utilflowcontrol.Interface,
	workEstimator flowcontrolrequest.WorkEstimatorFunc,
) http.Handler {
	if fcIfc == nil {
		klog.Warningf("priority and fairness support not found, skipping")
		return handler
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		requestInfo, ok := apirequest.RequestInfoFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no RequestInfo found in context"))
			return
		}
		user, ok := apirequest.UserFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no User found in context"))
			return
		}

		isWatchRequest := watchVerbs.Has(requestInfo.Verb)

		// Skip tracking long running non-watch requests.
		if longRunningRequestCheck != nil && longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
			klog.V(6).Infof("Serving RequestInfo=%#+v, user.Info=%#+v as longrunning\n", requestInfo, user)
			handler.ServeHTTP(w, r)
			return
		}

		var classification *PriorityAndFairnessClassification
		note := func(fs *flowcontrol.FlowSchema, pl *flowcontrol.PriorityLevelConfiguration, flowDistinguisher string) {
			classification = &PriorityAndFairnessClassification{
				FlowSchemaName:    fs.Name,
				FlowSchemaUID:     fs.UID,
				PriorityLevelName: pl.Name,
				PriorityLevelUID:  pl.UID}

			httplog.AddKeyValue(ctx, "apf_pl", truncateLogField(pl.Name))
			httplog.AddKeyValue(ctx, "apf_fs", truncateLogField(fs.Name))
			httplog.AddKeyValue(ctx, "apf_fd", truncateLogField(flowDistinguisher))
		}

		var served bool
		isMutatingRequest := !nonMutatingRequestVerbs.Has(requestInfo.Verb)
		noteExecutingDelta := func(delta int32) {
			if isMutatingRequest {
				watermark.recordMutating(int(atomic.AddInt32(&atomicMutatingExecuting, delta)))
			} else {
				watermark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyExecuting, delta)))
			}
		}
		noteWaitingDelta := func(delta int32) {
			if isMutatingRequest {
				waitingMark.recordMutating(int(atomic.AddInt32(&atomicMutatingWaiting, delta)))
			} else {
				waitingMark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyWaiting, delta)))
			}
		}
		queueNote := func(inQueue bool) {
			if inQueue {
				noteWaitingDelta(1)
			} else {
				noteWaitingDelta(-1)
			}
		}

		// find the estimated amount of work of the request
		// TODO: Estimate cost should also take fcIfc.GetWatchCount(requestInfo) as a parameter.
		workEstimate := workEstimator.EstimateWork(r)
		digest := utilflowcontrol.RequestDigest{
			RequestInfo:  requestInfo,
			User:         user,
			WorkEstimate: workEstimate,
		}

		if isWatchRequest {
			// This channel blocks calling handler.ServeHTTP() until closed, and is closed inside execute().
			// If APF rejects the request, it is never closed.
			shouldStartWatchCh := make(chan struct{})

			watchInitializationSignal := newInitializationSignal()
			// This wraps the request passed to handler.ServeHTTP(),
			// setting a context that plumbs watchInitializationSignal to storage
			var watchReq *http.Request
			// This is set inside execute(), prior to closing shouldStartWatchCh.
			// If the request is rejected by APF it is left nil.
			var forgetWatch utilflowcontrol.ForgetWatchFunc

			defer func() {
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				if watchInitializationSignal != nil {
					watchInitializationSignal.Signal()
				}
				// Forget the watcher if it was registered.
				//
				// // This is race-free because by this point, one of the following occurred:
				// case <-shouldStartWatchCh: execute() completed the assignment to forgetWatch
				// case <-resultCh: Handle() completed, and Handle() does not return
				//   while execute() is running
				if forgetWatch != nil {
					forgetWatch()
				}
			}()

			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				forgetWatch = fcIfc.RegisterWatch(requestInfo)

				// Notify the main thread that we're ready to start the watch.
				close(shouldStartWatchCh)

				// Wait until the request is finished from the APF point of view
				// (which is when its initialization is done).
				watchInitializationSignal.Wait()
			}

			// Ensure that an item can be put to resultCh asynchronously.
			resultCh := make(chan interface{}, 1)

			// Call Handle in a separate goroutine.
			// The reason for it is that from APF point of view, the request processing
			// finishes as soon as watch is initialized (which is generally orders of
			// magnitude faster then the watch request itself). This means that Handle()
			// call finishes much faster and for performance reasons we want to reduce
			// the number of running goroutines - so we run the shorter thing in a
			// dedicated goroutine and the actual watch handler in the main one.
			go func() {
				defer func() {
					err := recover()
					// do not wrap the sentinel ErrAbortHandler panic value
					if err != nil && err != http.ErrAbortHandler {
						// Same as stdlib http server code. Manually allocate stack
						// trace buffer size to prevent excessively large logs
						const size = 64 << 10
						buf := make([]byte, size)
						buf = buf[:runtime.Stack(buf, false)]
						err = fmt.Sprintf("%v\n%s", err, buf)
					}

					// Ensure that the result is put into resultCh independently of the panic.
					resultCh <- err
				}()

				// We create handleCtx with explicit cancelation function.
				// The reason for it is that Handle() underneath may start additional goroutine
				// that is blocked on context cancellation. However, from APF point of view,
				// we don't want to wait until the whole watch request is processed (which is
				// when it context is actually cancelled) - we want to unblock the goroutine as
				// soon as the request is processed from the APF point of view.
				//
				// Note that we explicitly do NOT call the actuall handler using that context
				// to avoid cancelling request too early.
				handleCtx, handleCtxCancel := context.WithCancel(ctx)
				defer handleCtxCancel()

				// Note that Handle will return irrespective of whether the request
				// executes or is rejected. In the latter case, the function will return
				// without calling the passed `execute` function.
				fcIfc.Handle(handleCtx, digest, note, queueNote, execute)
			}()

			select {
			case <-shouldStartWatchCh:
				watchCtx := utilflowcontrol.WithInitializationSignal(ctx, watchInitializationSignal)
				watchReq = r.WithContext(watchCtx)
				handler.ServeHTTP(w, watchReq)
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				// It has to happen before waiting on the resultCh below.
				watchInitializationSignal.Signal()
				// TODO: Consider finishing the request as soon as Handle call panics.
				if err := <-resultCh; err != nil {
					panic(err)
				}
			case err := <-resultCh:
				if err != nil {
					panic(err)
				}
			}
		} else {
			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				handler.ServeHTTP(w, r)
			}

			fcIfc.Handle(ctx, digest, note, queueNote, execute)
		}

		if !served {
			setResponseHeaders(classification, w)

			if isMutatingRequest {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.MutatingKind).Inc()
			} else {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.ReadOnlyKind).Inc()
			}
			epmetrics.RecordRequestTermination(r, requestInfo, epmetrics.APIServerComponent, http.StatusTooManyRequests)
			tooManyRequests(r, w)
		}
	})
}

// StartPriorityAndFairnessWatermarkMaintenance starts the goroutines to observe and maintain watermarks for
// priority-and-fairness requests.
func StartPriorityAndFairnessWatermarkMaintenance(stopCh <-chan struct{}) {
	startWatermarkMaintenance(watermark, stopCh)
	startWatermarkMaintenance(waitingMark, stopCh)
}

func setResponseHeaders(classification *PriorityAndFairnessClassification, w http.ResponseWriter) {
	if classification == nil {
		return
	}

	// We intentionally set the UID of the flow-schema and priority-level instead of name. This is so that
	// the names that cluster-admins choose for categorization and priority levels are not exposed, also
	// the names might make it obvious to the users that they are rejected due to classification with low priority.
	w.Header().Set(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, string(classification.PriorityLevelUID))
	w.Header().Set(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, string(classification.FlowSchemaUID))
}
<commit_msg>apf: print watch init latency in httplog<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 filters

import (
	"context"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"
	"time"

	flowcontrol "k8s.io/api/flowcontrol/v1beta2"
	apitypes "k8s.io/apimachinery/pkg/types"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/httplog"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	flowcontrolrequest "k8s.io/apiserver/pkg/util/flowcontrol/request"
	"k8s.io/klog/v2"
)

// PriorityAndFairnessClassification identifies the results of
// classification for API Priority and Fairness
type PriorityAndFairnessClassification struct {
	FlowSchemaName    string
	FlowSchemaUID     apitypes.UID
	PriorityLevelName string
	PriorityLevelUID  apitypes.UID
}

// waitingMark tracks requests waiting rather than being executed
var waitingMark = &requestWatermark{
	phase:            epmetrics.WaitingPhase,
	readOnlyObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.ReadOnlyKind}).RequestsWaiting,
	mutatingObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.MutatingKind}).RequestsWaiting,
}

var atomicMutatingExecuting, atomicReadOnlyExecuting int32
var atomicMutatingWaiting, atomicReadOnlyWaiting int32

// newInitializationSignal is defined for testing purposes.
var newInitializationSignal = utilflowcontrol.NewInitializationSignal

func truncateLogField(s string) string {
	const maxFieldLogLength = 64

	if len(s) > maxFieldLogLength {
		s = s[0:maxFieldLogLength]
	}
	return s
}

// WithPriorityAndFairness limits the number of in-flight
// requests in a fine-grained way.
func WithPriorityAndFairness(
	handler http.Handler,
	longRunningRequestCheck apirequest.LongRunningRequestCheck,
	fcIfc utilflowcontrol.Interface,
	workEstimator flowcontrolrequest.WorkEstimatorFunc,
) http.Handler {
	if fcIfc == nil {
		klog.Warningf("priority and fairness support not found, skipping")
		return handler
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		requestInfo, ok := apirequest.RequestInfoFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no RequestInfo found in context"))
			return
		}
		user, ok := apirequest.UserFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no User found in context"))
			return
		}

		isWatchRequest := watchVerbs.Has(requestInfo.Verb)

		// Skip tracking long running non-watch requests.
		if longRunningRequestCheck != nil && longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
			klog.V(6).Infof("Serving RequestInfo=%#+v, user.Info=%#+v as longrunning\n", requestInfo, user)
			handler.ServeHTTP(w, r)
			return
		}

		var classification *PriorityAndFairnessClassification
		note := func(fs *flowcontrol.FlowSchema, pl *flowcontrol.PriorityLevelConfiguration, flowDistinguisher string) {
			classification = &PriorityAndFairnessClassification{
				FlowSchemaName:    fs.Name,
				FlowSchemaUID:     fs.UID,
				PriorityLevelName: pl.Name,
				PriorityLevelUID:  pl.UID}

			httplog.AddKeyValue(ctx, "apf_pl", truncateLogField(pl.Name))
			httplog.AddKeyValue(ctx, "apf_fs", truncateLogField(fs.Name))
			httplog.AddKeyValue(ctx, "apf_fd", truncateLogField(flowDistinguisher))
		}

		var served bool
		isMutatingRequest := !nonMutatingRequestVerbs.Has(requestInfo.Verb)
		noteExecutingDelta := func(delta int32) {
			if isMutatingRequest {
				watermark.recordMutating(int(atomic.AddInt32(&atomicMutatingExecuting, delta)))
			} else {
				watermark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyExecuting, delta)))
			}
		}
		noteWaitingDelta := func(delta int32) {
			if isMutatingRequest {
				waitingMark.recordMutating(int(atomic.AddInt32(&atomicMutatingWaiting, delta)))
			} else {
				waitingMark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyWaiting, delta)))
			}
		}
		queueNote := func(inQueue bool) {
			if inQueue {
				noteWaitingDelta(1)
			} else {
				noteWaitingDelta(-1)
			}
		}

		// find the estimated amount of work of the request
		// TODO: Estimate cost should also take fcIfc.GetWatchCount(requestInfo) as a parameter.
		workEstimate := workEstimator.EstimateWork(r)
		digest := utilflowcontrol.RequestDigest{
			RequestInfo:  requestInfo,
			User:         user,
			WorkEstimate: workEstimate,
		}

		if isWatchRequest {
			// This channel blocks calling handler.ServeHTTP() until closed, and is closed inside execute().
			// If APF rejects the request, it is never closed.
			shouldStartWatchCh := make(chan struct{})

			watchInitializationSignal := newInitializationSignal()
			// This wraps the request passed to handler.ServeHTTP(),
			// setting a context that plumbs watchInitializationSignal to storage
			var watchReq *http.Request
			// This is set inside execute(), prior to closing shouldStartWatchCh.
			// If the request is rejected by APF it is left nil.
			var forgetWatch utilflowcontrol.ForgetWatchFunc

			defer func() {
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				if watchInitializationSignal != nil {
					watchInitializationSignal.Signal()
				}
				// Forget the watcher if it was registered.
				//
				// // This is race-free because by this point, one of the following occurred:
				// case <-shouldStartWatchCh: execute() completed the assignment to forgetWatch
				// case <-resultCh: Handle() completed, and Handle() does not return
				//   while execute() is running
				if forgetWatch != nil {
					forgetWatch()
				}
			}()

			execute := func() {
				startedAt := time.Now()
				defer func() {
					httplog.AddKeyValue(ctx, "apf_init_latency", time.Now().Sub(startedAt))
				}()
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				forgetWatch = fcIfc.RegisterWatch(requestInfo)

				// Notify the main thread that we're ready to start the watch.
				close(shouldStartWatchCh)

				// Wait until the request is finished from the APF point of view
				// (which is when its initialization is done).
				watchInitializationSignal.Wait()
			}

			// Ensure that an item can be put to resultCh asynchronously.
			resultCh := make(chan interface{}, 1)

			// Call Handle in a separate goroutine.
			// The reason for it is that from APF point of view, the request processing
			// finishes as soon as watch is initialized (which is generally orders of
			// magnitude faster then the watch request itself). This means that Handle()
			// call finishes much faster and for performance reasons we want to reduce
			// the number of running goroutines - so we run the shorter thing in a
			// dedicated goroutine and the actual watch handler in the main one.
			go func() {
				defer func() {
					err := recover()
					// do not wrap the sentinel ErrAbortHandler panic value
					if err != nil && err != http.ErrAbortHandler {
						// Same as stdlib http server code. Manually allocate stack
						// trace buffer size to prevent excessively large logs
						const size = 64 << 10
						buf := make([]byte, size)
						buf = buf[:runtime.Stack(buf, false)]
						err = fmt.Sprintf("%v\n%s", err, buf)
					}

					// Ensure that the result is put into resultCh independently of the panic.
					resultCh <- err
				}()

				// We create handleCtx with explicit cancelation function.
				// The reason for it is that Handle() underneath may start additional goroutine
				// that is blocked on context cancellation. However, from APF point of view,
				// we don't want to wait until the whole watch request is processed (which is
				// when it context is actually cancelled) - we want to unblock the goroutine as
				// soon as the request is processed from the APF point of view.
				//
				// Note that we explicitly do NOT call the actuall handler using that context
				// to avoid cancelling request too early.
				handleCtx, handleCtxCancel := context.WithCancel(ctx)
				defer handleCtxCancel()

				// Note that Handle will return irrespective of whether the request
				// executes or is rejected. In the latter case, the function will return
				// without calling the passed `execute` function.
				fcIfc.Handle(handleCtx, digest, note, queueNote, execute)
			}()

			select {
			case <-shouldStartWatchCh:
				watchCtx := utilflowcontrol.WithInitializationSignal(ctx, watchInitializationSignal)
				watchReq = r.WithContext(watchCtx)
				handler.ServeHTTP(w, watchReq)
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				// It has to happen before waiting on the resultCh below.
				watchInitializationSignal.Signal()
				// TODO: Consider finishing the request as soon as Handle call panics.
				if err := <-resultCh; err != nil {
					panic(err)
				}
			case err := <-resultCh:
				if err != nil {
					panic(err)
				}
			}
		} else {
			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				handler.ServeHTTP(w, r)
			}

			fcIfc.Handle(ctx, digest, note, queueNote, execute)
		}

		if !served {
			setResponseHeaders(classification, w)

			if isMutatingRequest {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.MutatingKind).Inc()
			} else {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.ReadOnlyKind).Inc()
			}
			epmetrics.RecordRequestTermination(r, requestInfo, epmetrics.APIServerComponent, http.StatusTooManyRequests)
			tooManyRequests(r, w)
		}
	})
}

// StartPriorityAndFairnessWatermarkMaintenance starts the goroutines to observe and maintain watermarks for
// priority-and-fairness requests.
func StartPriorityAndFairnessWatermarkMaintenance(stopCh <-chan struct{}) {
	startWatermarkMaintenance(watermark, stopCh)
	startWatermarkMaintenance(waitingMark, stopCh)
}

func setResponseHeaders(classification *PriorityAndFairnessClassification, w http.ResponseWriter) {
	if classification == nil {
		return
	}

	// We intentionally set the UID of the flow-schema and priority-level instead of name. This is so that
	// the names that cluster-admins choose for categorization and priority levels are not exposed, also
	// the names might make it obvious to the users that they are rejected due to classification with low priority.
	w.Header().Set(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, string(classification.PriorityLevelUID))
	w.Header().Set(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, string(classification.FlowSchemaUID))
}
<|endoftext|>