<commit_before>package renter

import (
	"crypto/rand"
	"errors"
	"io"
	"net"
	"os"
	"sync/atomic"
	"time"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

var (
	downloadAttempts = 5
)

// A Download is a file download that has been queued by the renter.
type Download struct {
	// Implementation note: received is declared first to ensure that it is
	// 64-bit aligned. This is necessary to ensure that atomic operations work
	// correctly on ARM and x86-32.
	received uint64

	complete    bool
	filesize    uint64
	destination string
	nickname    string

	pieces []filePiece
	file   *os.File
}

// Complete returns whether the file is ready to be used.
func (d *Download) Complete() bool {
	return d.complete
}

// Filesize returns the size of the file.
func (d *Download) Filesize() uint64 {
	return d.filesize
}

// Received returns the number of bytes downloaded so far.
func (d *Download) Received() uint64 {
	return d.received
}

// Destination returns the file's location on disk.
func (d *Download) Destination() string {
	return d.destination
}

// Nickname returns the identifier assigned to the file when it was uploaded.
func (d *Download) Nickname() string {
	return d.nickname
}

// Write implements the io.Writer interface. Each write updates the Download's
// received field. This allows download progress to be monitored in real-time.
func (d *Download) Write(b []byte) (int, error) {
	n, err := d.file.Write(b)
	// atomically update d.received
	atomic.AddUint64(&d.received, uint64(n))
	return n, err
}

// downloadPiece attempts to retrieve a file piece from a host.
func (d *Download) downloadPiece(piece filePiece) error {
	conn, err := net.DialTimeout("tcp", string(piece.HostIP), 10e9)
	if err != nil {
		return err
	}
	defer conn.Close()
	err = encoding.WriteObject(conn, [8]byte{'R', 'e', 't', 'r', 'i', 'e', 'v', 'e'})
	if err != nil {
		return err
	}

	// Send the ID of the contract for the file piece we're requesting.
	if err := encoding.WriteObject(conn, piece.ContractID); err != nil {
		return err
	}

	// Simultaneously download, decrypt, and calculate the Merkle root of the file.
	tee := io.TeeReader(
		// Use a LimitedReader to ensure we don't read indefinitely.
		io.LimitReader(conn, int64(piece.Contract.FileSize)),
		// Write the decrypted bytes to the file.
		piece.EncryptionKey.NewWriter(d),
	)
	merkleRoot, err := crypto.ReaderMerkleRoot(tee)
	if err != nil {
		return err
	}

	if merkleRoot != piece.Contract.FileMerkleRoot {
		return errors.New("host provided a file that's invalid")
	}

	return nil
}

// start initiates the download of a File.
func (d *Download) start() {
	// We only need one piece, so iterate through the hosts until a download
	// succeeds.
	for i := 0; i < downloadAttempts; i++ {
		for _, piece := range d.pieces {
			downloadErr := d.downloadPiece(piece)
			if downloadErr == nil {
				d.complete = true
				d.file.Close()
				return
			}
			// Reset seek, since the file may have been partially written. The
			// next attempt will overwrite these bytes.
			d.file.Seek(0, 0)
			atomic.SwapUint64(&d.received, 0)
		}

		// This iteration failed, no hosts returned the piece. Try again
		// after waiting a random amount of time.
		randSource := make([]byte, 1)
		rand.Read(randSource)
		time.Sleep(time.Second * time.Duration(i*i) * time.Duration(randSource[0]))
	}

	// File could not be downloaded; delete the copy on disk.
	d.file.Close()
	os.Remove(d.destination)
}

// newDownload initializes a new Download object.
func newDownload(file *file, destination string) (*Download, error) {
	// Create the download destination file.
	handle, err := os.Create(destination)
	if err != nil {
		return nil, err
	}

	// Filter out the inactive pieces.
	var activePieces []filePiece
	for _, piece := range file.Pieces {
		if piece.Active {
			activePieces = append(activePieces, piece)
		}
	}
	if len(activePieces) == 0 {
		return nil, errors.New("no active pieces")
	}

	return &Download{
		complete: false,
		// for now, all the pieces are equivalent
		filesize:    file.Pieces[0].Contract.FileSize,
		received:    0,
		destination: destination,
		nickname:    file.Name,

		pieces: activePieces,
		file:   handle,
	}, nil
}

// Download downloads a file, identified by its nickname, to the destination
// specified.
func (r *Renter) Download(nickname, destination string) error {
	lockID := r.mu.Lock()
	defer r.mu.Unlock(lockID)

	// Lookup the File associated with the nickname.
	file, exists := r.files[nickname]
	if !exists {
		return errors.New("no file of that nickname")
	}

	// Create the download object and spawn the download process.
	d, err := newDownload(file, destination)
	if err != nil {
		return err
	}
	go d.start()

	// Add the download to the download queue.
	r.downloadQueue = append(r.downloadQueue, d)
	return nil
}

// DownloadQueue returns the list of downloads in the queue.
func (r *Renter) DownloadQueue() []modules.DownloadInfo {
	lockID := r.mu.RLock()
	defer r.mu.RUnlock(lockID)

	downloads := make([]modules.DownloadInfo, len(r.downloadQueue))
	for i := range r.downloadQueue {
		downloads[i] = r.downloadQueue[i]
	}
	return downloads
}
<commit_msg>synchronous downloads<commit_after>package renter

import (
	"crypto/rand"
	"errors"
	"io"
	"net"
	"os"
	"sync/atomic"
	"time"

	"github.com/NebulousLabs/Sia/crypto"
	"github.com/NebulousLabs/Sia/encoding"
	"github.com/NebulousLabs/Sia/modules"
)

var (
	downloadAttempts = 5
)

// A Download is a file download that has been queued by the renter. It
// implements the modules.DownloadInfo interface.
type Download struct {
	// Implementation note: received is declared first to ensure that it is
	// 64-bit aligned. This is necessary to ensure that atomic operations work
	// correctly on ARM and x86-32.
	received uint64

	complete    bool
	filesize    uint64
	destination string
	nickname    string

	pieces []filePiece
	file   *os.File
}

// Complete returns whether the file is ready to be used.
func (d *Download) Complete() bool {
	return d.complete
}

// Filesize returns the size of the file.
func (d *Download) Filesize() uint64 {
	return d.filesize
}

// Received returns the number of bytes downloaded so far.
func (d *Download) Received() uint64 {
	return d.received
}

// Destination returns the file's location on disk.
func (d *Download) Destination() string {
	return d.destination
}

// Nickname returns the identifier assigned to the file when it was uploaded.
func (d *Download) Nickname() string {
	return d.nickname
}

// Write implements the io.Writer interface. Each write updates the Download's
// received field. This allows download progress to be monitored in real-time.
func (d *Download) Write(b []byte) (int, error) {
	n, err := d.file.Write(b)
	// atomically update d.received
	// TODO: atomic operations may not be necessary
	atomic.AddUint64(&d.received, uint64(n))
	return n, err
}

// downloadPiece attempts to retrieve a file piece from a host.
func (d *Download) downloadPiece(piece filePiece) error {
	conn, err := net.DialTimeout("tcp", string(piece.HostIP), 10e9)
	if err != nil {
		return err
	}
	defer conn.Close()
	err = encoding.WriteObject(conn, [8]byte{'R', 'e', 't', 'r', 'i', 'e', 'v', 'e'})
	if err != nil {
		return err
	}

	// Send the ID of the contract for the file piece we're requesting.
	if err := encoding.WriteObject(conn, piece.ContractID); err != nil {
		return err
	}

	// Simultaneously download, decrypt, and calculate the Merkle root of the file.
	tee := io.TeeReader(
		// Use a LimitedReader to ensure we don't read indefinitely.
		io.LimitReader(conn, int64(piece.Contract.FileSize)),
		// Write the decrypted bytes to the file.
		piece.EncryptionKey.NewWriter(d),
	)
	merkleRoot, err := crypto.ReaderMerkleRoot(tee)
	if err != nil {
		return err
	}

	if merkleRoot != piece.Contract.FileMerkleRoot {
		return errors.New("host provided a file that's invalid")
	}

	return nil
}

// newDownload initializes a new Download object.
func newDownload(file *file, destination string) (*Download, error) {
	// Create the download destination file.
	handle, err := os.Create(destination)
	if err != nil {
		return nil, err
	}

	// Filter out the inactive pieces.
	var activePieces []filePiece
	for _, piece := range file.Pieces {
		if piece.Active {
			activePieces = append(activePieces, piece)
		}
	}
	if len(activePieces) == 0 {
		return nil, errors.New("no active pieces")
	}

	return &Download{
		complete: false,
		// for now, all the pieces are equivalent
		filesize:    file.Pieces[0].Contract.FileSize,
		received:    0,
		destination: destination,
		nickname:    file.Name,

		pieces: activePieces,
		file:   handle,
	}, nil
}

// Download downloads a file, identified by its nickname, to the destination
// specified.
func (r *Renter) Download(nickname, destination string) error {
	lockID := r.mu.Lock()
	defer r.mu.Unlock(lockID)

	// Lookup the file associated with the nickname.
	file, exists := r.files[nickname]
	if !exists {
		return errors.New("no file of that nickname")
	}

	// Create the download object and spawn the download process.
	d, err := newDownload(file, destination)
	if err != nil {
		return err
	}

	// Add the download to the download queue.
	r.downloadQueue = append(r.downloadQueue, d)

	// Download the file. We only need one piece, so iterate through the hosts
	// until a download succeeds.
	for i := 0; i < downloadAttempts; i++ {
		for _, piece := range d.pieces {
			downloadErr := d.downloadPiece(piece)
			if downloadErr == nil {
				// done
				d.complete = true
				d.file.Close()
				return nil
			}
			// Reset seek, since the file may have been partially written. The
			// next attempt will overwrite these bytes.
			d.file.Seek(0, 0)
			atomic.SwapUint64(&d.received, 0)
		}

		// This iteration failed, no hosts returned the piece. Try again
		// after waiting a random amount of time.
		randSource := make([]byte, 1)
		rand.Read(randSource)
		time.Sleep(time.Second * time.Duration(i*i) * time.Duration(randSource[0]))
	}

	// File could not be downloaded; delete the copy on disk.
	d.file.Close()
	os.Remove(destination)

	return errors.New("could not download any file pieces")
}

// DownloadQueue returns the list of downloads in the queue.
func (r *Renter) DownloadQueue() []modules.DownloadInfo {
	lockID := r.mu.RLock()
	defer r.mu.RUnlock(lockID)

	downloads := make([]modules.DownloadInfo, len(r.downloadQueue))
	for i := range r.downloadQueue {
		downloads[i] = r.downloadQueue[i]
	}
	return downloads
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Ceph-CSI 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 cephfs

import (
	"context"
	"path"
	"strconv"
	"strings"

	"github.com/ceph/ceph-csi/internal/util"

	klog "k8s.io/klog/v2"
)

var (
	// clusterAdditionalInfo contains information regarding if resize is
	// supported in the particular cluster and subvolumegroup is
	// created or not.
	// Subvolumegroup creation and volume resize decisions are
	// taken through this additional cluster information.
	clusterAdditionalInfo = make(map[string]*localClusterState)
)

const (
	cephEntityClientPrefix = "client."
)

// Subvolume holds subvolume information.
type Subvolume struct {
	BytesQuota    int      `json:"bytes_quota"`
	DataPool      string   `json:"data_pool"`
	GID           int      `json:"gid"`
	Mode          int      `json:"mode"`
	MonAddrs      []string `json:"mon_addrs"`
	Path          string   `json:"path"`
	PoolNamespace string   `json:"pool_namespace"`
	Type          string   `json:"type"`
	UID           int      `json:"uid"`
}

func getVolumeRootPathCephDeprecated(volID volumeID) string {
	return path.Join("/", "csi-volumes", string(volID))
}

func getVolumeRootPathCeph(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID) (string, error) {
	stdout, stderr, err := util.ExecCommand(
		ctx,
		"ceph",
		"fs",
		"subvolume",
		"getpath",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix+cr.ID,
		"--keyfile="+cr.KeyFile)

	if err != nil {
		util.ErrorLog(ctx, "failed to get the rootpath for the vol %s(%s) stdError %s", string(volID), err, stderr)
		if strings.Contains(stderr, ErrVolumeNotFound.Error()) {
			return "", util.JoinErrors(ErrVolumeNotFound, err)
		}

		return "", err
	}
	return strings.TrimSuffix(stdout, "\n"), nil
}

func getSubVolumeInfo(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID) (Subvolume, error) {
	info := Subvolume{}
	err := execCommandJSON(
		ctx,
		&info,
		"ceph",
		"fs",
		"subvolume",
		"info",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix+cr.ID,
		"--keyfile="+cr.KeyFile)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to get subvolume info for the vol %s(%s)"), string(volID), err)
		if strings.HasPrefix(err.Error(), ErrVolumeNotFound.Error()) {
			return info, ErrVolumeNotFound
		}
		// Incase the error is other than invalid command return error to the caller.
		if !strings.Contains(err.Error(), ErrInvalidCommand.Error()) {
			return info, ErrInvalidCommand
		}

		return info, err
	}
	return info, nil
}

type localClusterState struct {
	// set true if cluster supports resize functionality.
	resizeSupported bool
	// set true once a subvolumegroup is created
	// for corresponding cluster.
	subVolumeGroupCreated bool
}

func createVolume(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID, bytesQuota int64) error {
	// verify if corresponding ClusterID key is present in the map,
	// and if not, initialize with default values(false).
	if _, keyPresent := clusterAdditionalInfo[volOptions.ClusterID]; !keyPresent {
		clusterAdditionalInfo[volOptions.ClusterID] = &localClusterState{}
	}

	// create subvolumegroup if not already created for the cluster.
	if !clusterAdditionalInfo[volOptions.ClusterID].subVolumeGroupCreated {
		err := execCommandErr(
			ctx,
			"ceph",
			"fs",
			"subvolumegroup",
			"create",
			volOptions.FsName,
			volOptions.SubvolumeGroup,
			"-m", volOptions.Monitors,
			"-c", util.CephConfigPath,
			"-n", cephEntityClientPrefix+cr.ID,
			"--keyfile="+cr.KeyFile)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to create subvolume group %s, for the vol %s(%s)"), volOptions.SubvolumeGroup, string(volID), err)
			return err
		}
		util.DebugLog(ctx, "cephfs: created subvolume group %s", volOptions.SubvolumeGroup)
		clusterAdditionalInfo[volOptions.ClusterID].subVolumeGroupCreated = true
	}

	args := []string{
		"fs",
		"subvolume",
		"create",
		volOptions.FsName,
		string(volID),
		strconv.FormatInt(bytesQuota, 10),
		"--group_name",
		volOptions.SubvolumeGroup,
		"--mode", "777",
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix + cr.ID,
		"--keyfile=" + cr.KeyFile,
	}

	if volOptions.Pool != "" {
		args = append(args, "--pool_layout", volOptions.Pool)
	}

	err := execCommandErr(
		ctx,
		"ceph",
		args[:]...)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to create subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
		return err
	}

	return nil
}

// resizeVolume will try to use ceph fs subvolume resize command to resize the
// subvolume. If the command is not available as a fallback it will use
// CreateVolume to resize the subvolume.
func resizeVolume(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID, bytesQuota int64) error {
	// keyPresent checks whether corresponding clusterID key is present in clusterAdditionalInfo
	var keyPresent bool
	// verify if corresponding ClusterID key is present in the map,
	// and if not, initialize with default values(false).
	if _, keyPresent = clusterAdditionalInfo[volOptions.ClusterID]; !keyPresent {
		clusterAdditionalInfo[volOptions.ClusterID] = &localClusterState{}
	}
	// resize subvolume when either it's supported, or when corresponding
	// clusterID key was not present.
	if clusterAdditionalInfo[volOptions.ClusterID].resizeSupported || !keyPresent {
		args := []string{
			"fs",
			"subvolume",
			"resize",
			volOptions.FsName,
			string(volID),
			strconv.FormatInt(bytesQuota, 10),
			"--group_name",
			volOptions.SubvolumeGroup,
			"-m", volOptions.Monitors,
			"-c", util.CephConfigPath,
			"-n", cephEntityClientPrefix + cr.ID,
			"--keyfile=" + cr.KeyFile,
		}

		err := execCommandErr(
			ctx,
			"ceph",
			args[:]...)

		if err == nil {
			clusterAdditionalInfo[volOptions.ClusterID].resizeSupported = true
			return nil
		}
		// Incase the error is other than invalid command return error to the caller.
		if !strings.Contains(err.Error(), ErrInvalidCommand.Error()) {
			klog.Errorf(util.Log(ctx, "failed to resize subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
			return err
		}
	}
	clusterAdditionalInfo[volOptions.ClusterID].resizeSupported = false
	return createVolume(ctx, volOptions, cr, volID, bytesQuota)
}

func purgeVolume(ctx context.Context, volID volumeID, cr *util.Credentials, volOptions *volumeOptions, force bool) error {
	arg := []string{
		"fs",
		"subvolume",
		"rm",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix + cr.ID,
		"--keyfile=" + cr.KeyFile,
	}
	if force {
		arg = append(arg, "--force")
	}

	err := execCommandErr(ctx, "ceph", arg...)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to purge subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
		if strings.HasPrefix(err.Error(), ErrVolumeNotFound.Error()) {
			return util.JoinErrors(ErrVolumeNotFound, err)
		}
		return err
	}

	return nil
}
<commit_msg>cephfs: replace Errorf with ErrorLog in getSubVolumeInfo<commit_after>/*
Copyright 2018 The Ceph-CSI 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 cephfs

import (
	"context"
	"path"
	"strconv"
	"strings"

	"github.com/ceph/ceph-csi/internal/util"

	klog "k8s.io/klog/v2"
)

var (
	// clusterAdditionalInfo contains information regarding if resize is
	// supported in the particular cluster and subvolumegroup is
	// created or not.
	// Subvolumegroup creation and volume resize decisions are
	// taken through this additional cluster information.
	clusterAdditionalInfo = make(map[string]*localClusterState)
)

const (
	cephEntityClientPrefix = "client."
)

// Subvolume holds subvolume information.
type Subvolume struct {
	BytesQuota    int      `json:"bytes_quota"`
	DataPool      string   `json:"data_pool"`
	GID           int      `json:"gid"`
	Mode          int      `json:"mode"`
	MonAddrs      []string `json:"mon_addrs"`
	Path          string   `json:"path"`
	PoolNamespace string   `json:"pool_namespace"`
	Type          string   `json:"type"`
	UID           int      `json:"uid"`
}

func getVolumeRootPathCephDeprecated(volID volumeID) string {
	return path.Join("/", "csi-volumes", string(volID))
}

func getVolumeRootPathCeph(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID) (string, error) {
	stdout, stderr, err := util.ExecCommand(
		ctx,
		"ceph",
		"fs",
		"subvolume",
		"getpath",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix+cr.ID,
		"--keyfile="+cr.KeyFile)

	if err != nil {
		util.ErrorLog(ctx, "failed to get the rootpath for the vol %s(%s) stdError %s", string(volID), err, stderr)
		if strings.Contains(stderr, ErrVolumeNotFound.Error()) {
			return "", util.JoinErrors(ErrVolumeNotFound, err)
		}

		return "", err
	}
	return strings.TrimSuffix(stdout, "\n"), nil
}

func getSubVolumeInfo(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID) (Subvolume, error) {
	info := Subvolume{}
	err := execCommandJSON(
		ctx,
		&info,
		"ceph",
		"fs",
		"subvolume",
		"info",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix+cr.ID,
		"--keyfile="+cr.KeyFile)
	if err != nil {
		util.ErrorLog(ctx, "failed to get subvolume info for the vol %s(%s)", string(volID), err)
		if strings.HasPrefix(err.Error(), ErrVolumeNotFound.Error()) {
			return info, ErrVolumeNotFound
		}
		// Incase the error is other than invalid command return error to the caller.
		if !strings.Contains(err.Error(), ErrInvalidCommand.Error()) {
			return info, ErrInvalidCommand
		}

		return info, err
	}
	return info, nil
}

type localClusterState struct {
	// set true if cluster supports resize functionality.
	resizeSupported bool
	// set true once a subvolumegroup is created
	// for corresponding cluster.
	subVolumeGroupCreated bool
}

func createVolume(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID, bytesQuota int64) error {
	// verify if corresponding ClusterID key is present in the map,
	// and if not, initialize with default values(false).
	if _, keyPresent := clusterAdditionalInfo[volOptions.ClusterID]; !keyPresent {
		clusterAdditionalInfo[volOptions.ClusterID] = &localClusterState{}
	}

	// create subvolumegroup if not already created for the cluster.
	if !clusterAdditionalInfo[volOptions.ClusterID].subVolumeGroupCreated {
		err := execCommandErr(
			ctx,
			"ceph",
			"fs",
			"subvolumegroup",
			"create",
			volOptions.FsName,
			volOptions.SubvolumeGroup,
			"-m", volOptions.Monitors,
			"-c", util.CephConfigPath,
			"-n", cephEntityClientPrefix+cr.ID,
			"--keyfile="+cr.KeyFile)
		if err != nil {
			klog.Errorf(util.Log(ctx, "failed to create subvolume group %s, for the vol %s(%s)"), volOptions.SubvolumeGroup, string(volID), err)
			return err
		}
		util.DebugLog(ctx, "cephfs: created subvolume group %s", volOptions.SubvolumeGroup)
		clusterAdditionalInfo[volOptions.ClusterID].subVolumeGroupCreated = true
	}

	args := []string{
		"fs",
		"subvolume",
		"create",
		volOptions.FsName,
		string(volID),
		strconv.FormatInt(bytesQuota, 10),
		"--group_name",
		volOptions.SubvolumeGroup,
		"--mode", "777",
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix + cr.ID,
		"--keyfile=" + cr.KeyFile,
	}

	if volOptions.Pool != "" {
		args = append(args, "--pool_layout", volOptions.Pool)
	}

	err := execCommandErr(
		ctx,
		"ceph",
		args[:]...)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to create subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
		return err
	}

	return nil
}

// resizeVolume will try to use ceph fs subvolume resize command to resize the
// subvolume. If the command is not available as a fallback it will use
// CreateVolume to resize the subvolume.
func resizeVolume(ctx context.Context, volOptions *volumeOptions, cr *util.Credentials, volID volumeID, bytesQuota int64) error {
	// keyPresent checks whether corresponding clusterID key is present in clusterAdditionalInfo
	var keyPresent bool
	// verify if corresponding ClusterID key is present in the map,
	// and if not, initialize with default values(false).
	if _, keyPresent = clusterAdditionalInfo[volOptions.ClusterID]; !keyPresent {
		clusterAdditionalInfo[volOptions.ClusterID] = &localClusterState{}
	}
	// resize subvolume when either it's supported, or when corresponding
	// clusterID key was not present.
	if clusterAdditionalInfo[volOptions.ClusterID].resizeSupported || !keyPresent {
		args := []string{
			"fs",
			"subvolume",
			"resize",
			volOptions.FsName,
			string(volID),
			strconv.FormatInt(bytesQuota, 10),
			"--group_name",
			volOptions.SubvolumeGroup,
			"-m", volOptions.Monitors,
			"-c", util.CephConfigPath,
			"-n", cephEntityClientPrefix + cr.ID,
			"--keyfile=" + cr.KeyFile,
		}

		err := execCommandErr(
			ctx,
			"ceph",
			args[:]...)

		if err == nil {
			clusterAdditionalInfo[volOptions.ClusterID].resizeSupported = true
			return nil
		}
		// Incase the error is other than invalid command return error to the caller.
		if !strings.Contains(err.Error(), ErrInvalidCommand.Error()) {
			klog.Errorf(util.Log(ctx, "failed to resize subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
			return err
		}
	}
	clusterAdditionalInfo[volOptions.ClusterID].resizeSupported = false
	return createVolume(ctx, volOptions, cr, volID, bytesQuota)
}

func purgeVolume(ctx context.Context, volID volumeID, cr *util.Credentials, volOptions *volumeOptions, force bool) error {
	arg := []string{
		"fs",
		"subvolume",
		"rm",
		volOptions.FsName,
		string(volID),
		"--group_name",
		volOptions.SubvolumeGroup,
		"-m", volOptions.Monitors,
		"-c", util.CephConfigPath,
		"-n", cephEntityClientPrefix + cr.ID,
		"--keyfile=" + cr.KeyFile,
	}
	if force {
		arg = append(arg, "--force")
	}

	err := execCommandErr(ctx, "ceph", arg...)
	if err != nil {
		klog.Errorf(util.Log(ctx, "failed to purge subvolume %s(%s) in fs %s"), string(volID), err, volOptions.FsName)
		if strings.HasPrefix(err.Error(), ErrVolumeNotFound.Error()) {
			return util.JoinErrors(ErrVolumeNotFound, err)
		}
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package dump

import (
	"archive/tar"
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/restic/restic/internal/archiver"
	"github.com/restic/restic/internal/fs"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
	rtest "github.com/restic/restic/internal/test"
)

func prepareTempdirRepoSrc(t testing.TB, src archiver.TestDir) (tempdir string, repo restic.Repository, cleanup func()) {
	tempdir, removeTempdir := rtest.TempDir(t)
	repo, removeRepository := repository.TestRepository(t)

	archiver.TestCreateFiles(t, tempdir, src)

	cleanup = func() {
		removeRepository()
		removeTempdir()
	}

	return tempdir, repo, cleanup
}

func TestWriteTar(t *testing.T) {
	tests := []struct {
		name   string
		args   archiver.TestDir
		target string
	}{
		{
			name: "single file in root",
			args: archiver.TestDir{
				"file": archiver.TestFile{Content: "string"},
			},
			target: "/",
		},
		{
			name: "multiple files in root",
			args: archiver.TestDir{
				"file1": archiver.TestFile{Content: "string"},
				"file2": archiver.TestFile{Content: "string"},
			},
			target: "/",
		},
		{
			name: "multiple files and folders in root",
			args: archiver.TestDir{
				"file1": archiver.TestFile{Content: "string"},
				"file2": archiver.TestFile{Content: "string"},
				"firstDir": archiver.TestDir{
					"another": archiver.TestFile{Content: "string"},
				},
				"secondDir": archiver.TestDir{
					"another2": archiver.TestFile{Content: "string"},
				},
			},
			target: "/",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			tmpdir, repo, cleanup := prepareTempdirRepoSrc(t, tt.args)
			defer cleanup()

			arch := archiver.New(repo, fs.Track{FS: fs.Local{}}, archiver.Options{})

			back := rtest.Chdir(t, tmpdir)
			defer back()

			sn, _, err := arch.Snapshot(ctx, []string{"."}, archiver.SnapshotOptions{})
			rtest.OK(t, err)

			tree, err := repo.LoadTree(ctx, *sn.Tree)
			rtest.OK(t, err)

			dst := &bytes.Buffer{}
			if err := WriteTar(ctx, repo, tree, tt.target, dst); err != nil {
				t.Fatalf("WriteTar() error = %v", err)
			}
			if err := checkTar(t, tmpdir, dst); err != nil {
				t.Errorf("WriteTar() = tar does not match: %v", err)
			}
		})
	}
}

func checkTar(t *testing.T, testDir string, srcTar *bytes.Buffer) error {
	tr := tar.NewReader(srcTar)

	fileNumber := 0
	tarFiles := 0

	err := filepath.Walk(testDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Name() != filepath.Base(testDir) {
			fileNumber++
		}
		return nil
	})
	if err != nil {
		return err
	}

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		matchPath := filepath.Join(testDir, hdr.Name)
		match, err := os.Stat(matchPath)
		if err != nil {
			return err
		}

		// check metadata, tar header contains time rounded to seconds
		fileTime := match.ModTime().Round(time.Second)
		tarTime := hdr.ModTime
		if !fileTime.Equal(tarTime) {
			return fmt.Errorf("modTime does not match, got: %s, want: %s", fileTime, tarTime)
		}

		if hdr.Typeflag == tar.TypeDir {
			// this is a folder
			if hdr.Name == "." {
				// we don't need to check the root folder
				continue
			}

			filebase := filepath.ToSlash(match.Name())
			if filepath.Base(hdr.Name) != filebase {
				return fmt.Errorf("foldernames don't match got %v want %v", filepath.Base(hdr.Name), filebase)
			}

		} else {
			if match.Size() != hdr.Size {
				return fmt.Errorf("size does not match got %v want %v", hdr.Size, match.Size())
			}
			contentsFile, err := ioutil.ReadFile(matchPath)
			if err != nil {
				t.Fatal(err)
			}
			contentsTar := &bytes.Buffer{}
			_, err = io.Copy(contentsTar, tr)
			if err != nil {
				t.Fatal(err)
			}
			if contentsTar.String() != string(contentsFile) {
				return fmt.Errorf("contents does not match, got %s want %s", contentsTar, contentsFile)
			}
		}
		tarFiles++
	}

	if tarFiles != fileNumber {
		return fmt.Errorf("not the same amount of files got %v want %v", tarFiles, fileNumber)
	}

	return nil
}
<commit_msg>dump: test proper permissions and directory name<commit_after>package dump

import (
	"archive/tar"
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/restic/restic/internal/archiver"
	"github.com/restic/restic/internal/fs"
	"github.com/restic/restic/internal/repository"
	"github.com/restic/restic/internal/restic"
	rtest "github.com/restic/restic/internal/test"
)

func prepareTempdirRepoSrc(t testing.TB, src archiver.TestDir) (tempdir string, repo restic.Repository, cleanup func()) {
	tempdir, removeTempdir := rtest.TempDir(t)
	repo, removeRepository := repository.TestRepository(t)

	archiver.TestCreateFiles(t, tempdir, src)

	cleanup = func() {
		removeRepository()
		removeTempdir()
	}

	return tempdir, repo, cleanup
}

func TestWriteTar(t *testing.T) {
	tests := []struct {
		name   string
		args   archiver.TestDir
		target string
	}{
		{
			name: "single file in root",
			args: archiver.TestDir{
				"file": archiver.TestFile{Content: "string"},
			},
			target: "/",
		},
		{
			name: "multiple files in root",
			args: archiver.TestDir{
				"file1": archiver.TestFile{Content: "string"},
				"file2": archiver.TestFile{Content: "string"},
			},
			target: "/",
		},
		{
			name: "multiple files and folders in root",
			args: archiver.TestDir{
				"file1": archiver.TestFile{Content: "string"},
				"file2": archiver.TestFile{Content: "string"},
				"firstDir": archiver.TestDir{
					"another": archiver.TestFile{Content: "string"},
				},
				"secondDir": archiver.TestDir{
					"another2": archiver.TestFile{Content: "string"},
				},
			},
			target: "/",
		},
		{
			name: "file and symlink in root",
			args: archiver.TestDir{
				"file1": archiver.TestFile{Content: "string"},
				"file2": archiver.TestSymlink{Target: "file1"},
			},
			target: "/",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			tmpdir, repo, cleanup := prepareTempdirRepoSrc(t, tt.args)
			defer cleanup()

			arch := archiver.New(repo, fs.Track{FS: fs.Local{}}, archiver.Options{})

			back := rtest.Chdir(t, tmpdir)
			defer back()

			sn, _, err := arch.Snapshot(ctx, []string{"."}, archiver.SnapshotOptions{})
			rtest.OK(t, err)

			tree, err := repo.LoadTree(ctx, *sn.Tree)
			rtest.OK(t, err)

			dst := &bytes.Buffer{}
			if err := WriteTar(ctx, repo, tree, tt.target, dst); err != nil {
				t.Fatalf("WriteTar() error = %v", err)
			}
			if err := checkTar(t, tmpdir, dst); err != nil {
				t.Errorf("WriteTar() = tar does not match: %v", err)
			}
		})
	}
}

func checkTar(t *testing.T, testDir string, srcTar *bytes.Buffer) error {
	tr := tar.NewReader(srcTar)

	fileNumber := 0
	tarFiles := 0

	err := filepath.Walk(testDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.Name() != filepath.Base(testDir) {
			fileNumber++
		}
		return nil
	})
	if err != nil {
		return err
	}

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}

		matchPath := filepath.Join(testDir, hdr.Name)
		match, err := os.Lstat(matchPath)
		if err != nil {
			return err
		}

		// check metadata, tar header contains time rounded to seconds
		fileTime := match.ModTime().Round(time.Second)
		tarTime := hdr.ModTime
		if !fileTime.Equal(tarTime) {
			return fmt.Errorf("modTime does not match, got: %s, want: %s", fileTime, tarTime)
		}

		if os.FileMode(hdr.Mode).Perm() != match.Mode().Perm() || os.FileMode(hdr.Mode)&^os.ModePerm != 0 {
			return fmt.Errorf("mode does not match, got: %v, want: %v", os.FileMode(hdr.Mode), match.Mode())
		}

		if hdr.Typeflag == tar.TypeDir {
			// this is a folder
			if hdr.Name == "." {
				// we don't need to check the root folder
				continue
			}

			filebase := filepath.ToSlash(match.Name())
			if filepath.Base(hdr.Name) != filebase {
				return fmt.Errorf("foldernames don't match got %v want %v", filepath.Base(hdr.Name), filebase)
			}
			if !strings.HasSuffix(hdr.Name, "/") {
				return fmt.Errorf("foldernames must end with separator got %v", hdr.Name)
			}
		} else if hdr.Typeflag == tar.TypeSymlink {
			target, err := fs.Readlink(matchPath)
			if err != nil {
				return err
			}
			if target != hdr.Linkname {
				return fmt.Errorf("symlink target does not match, got %s want %s", target, hdr.Linkname)
			}
		} else {
			if match.Size() != hdr.Size {
				return fmt.Errorf("size does not match got %v want %v", hdr.Size, match.Size())
			}
			contentsFile, err := ioutil.ReadFile(matchPath)
			if err != nil {
				t.Fatal(err)
			}
			contentsTar := &bytes.Buffer{}
			_, err = io.Copy(contentsTar, tr)
			if err != nil {
				t.Fatal(err)
			}
			if contentsTar.String() != string(contentsFile) {
				return fmt.Errorf("contents does not match, got %s want %s", contentsTar, contentsFile)
			}
		}
		tarFiles++
	}

	if tarFiles != fileNumber {
		return fmt.Errorf("not the same amount of files got %v want %v", tarFiles, fileNumber)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Ebiten 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 pixels

import (
	"image"
	"image/color"

	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
)

type drawImageHistoryItem struct {
	image    *graphics.Image
	vertices []int16
	geom     graphics.Matrix
	colorm   graphics.Matrix
	mode     opengl.CompositeMode
}

// Pixels represents pixels of an image for restoring when GL context is lost.
type Pixels struct {
	image *graphics.Image

	// basePixels and baseColor are exclusive.
	basePixels       []uint8
	baseColor        color.Color
	drawImageHistory []*drawImageHistoryItem
	stale            bool
}

func NewPixels(image *graphics.Image) *Pixels {
	return &Pixels{
		image: image,
	}
}

func (p *Pixels) MakeStale() {
	p.basePixels = nil
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = true
}

func (p *Pixels) Clear() {
	p.basePixels = nil
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) Fill(clr color.Color) {
	p.basePixels = nil
	p.baseColor = clr
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) ReplacePixels(pixels []uint8) {
	if p.basePixels == nil {
		p.basePixels = make([]uint8, len(pixels))
	}
	copy(p.basePixels, pixels)
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) AppendDrawImageHistory(image *graphics.Image, vertices []int16, geom graphics.Matrix, colorm graphics.Matrix, mode opengl.CompositeMode) {
	if p.stale {
		return
	}
	item := &drawImageHistoryItem{
		image:    image,
		vertices: vertices,
		geom:     geom,
		colorm:   colorm,
		mode:     mode,
	}
	p.drawImageHistory = append(p.drawImageHistory, item)
}

// At returns a color value at idx.
//
// Note that this must not be called until context is available.
// This means Pixels members must match with acutal state in GPU.
func (p *Pixels) At(idx int, context *opengl.Context) (color.Color, error) {
	if p.basePixels == nil || p.drawImageHistory != nil || p.stale {
		if err := p.Reset(context); err != nil {
			return nil, err
		}
	}
	r, g, b, a := p.basePixels[idx], p.basePixels[idx+1], p.basePixels[idx+2], p.basePixels[idx+3]
	return color.RGBA{r, g, b, a}, nil
}

func (p *Pixels) DependsOn(target *graphics.Image) bool {
	if p.stale {
		return false
	}
	for _, c := range p.drawImageHistory {
		if c.image == target {
			return true
		}
	}
	return false
}

func (p *Pixels) Reset(context *opengl.Context) error {
	var err error
	p.basePixels, err = p.image.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Pixels) ResetIfStale(context *opengl.Context) error {
	if !p.stale {
		return nil
	}
	return p.Reset(context)
}

func (p *Pixels) HasDependency() bool {
	if p.stale {
		return false
	}
	return p.drawImageHistory != nil
}

// Restore restores the pixels using its history.
//
// Restore is the only function that the pixel data is not present on GPU when this is called.
func (p *Pixels) Restore(context *opengl.Context, width, height int, filter opengl.Filter) (*graphics.Image, error) {
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	if p.basePixels != nil {
		for j := 0; j < height; j++ {
			copy(img.Pix[j*img.Stride:], p.basePixels[j*width*4:(j+1)*width*4])
		}
	}
	gimg, err := graphics.NewImageFromImage(img, filter)
	if err != nil {
		return nil, err
	}
	if p.baseColor != nil {
		if p.basePixels != nil {
			panic("not reach")
		}
		if err := gimg.Fill(p.baseColor); err != nil {
			return nil, err
		}
	}
	for _, c := range p.drawImageHistory {
		// c.image.impl must be already restored.
		/*if c.image.impl.hasHistory() {
			panic("not reach")
		}*/
		if err := gimg.DrawImage(c.image, c.vertices, c.geom, c.colorm, c.mode); err != nil {
			return nil, err
		}
	}
	p.basePixels, err = gimg.Pixels(context)
	if err != nil {
		return nil, err
	}
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
	return gimg, nil
}
<commit_msg>pixels: Ensure pixels is not stale when restoring<commit_after>// Copyright 2016 The Ebiten 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 pixels

import (
	"errors"
	"image"
	"image/color"

	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/graphics/opengl"
)

type drawImageHistoryItem struct {
	image    *graphics.Image
	vertices []int16
	geom     graphics.Matrix
	colorm   graphics.Matrix
	mode     opengl.CompositeMode
}

// Pixels represents pixels of an image for restoring when GL context is lost.
type Pixels struct {
	image *graphics.Image

	// basePixels and baseColor are exclusive.
	basePixels       []uint8
	baseColor        color.Color
	drawImageHistory []*drawImageHistoryItem
	stale            bool
}

func NewPixels(image *graphics.Image) *Pixels {
	return &Pixels{
		image: image,
	}
}

func (p *Pixels) MakeStale() {
	p.basePixels = nil
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = true
}

func (p *Pixels) Clear() {
	p.basePixels = nil
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) Fill(clr color.Color) {
	p.basePixels = nil
	p.baseColor = clr
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) ReplacePixels(pixels []uint8) {
	if p.basePixels == nil {
		p.basePixels = make([]uint8, len(pixels))
	}
	copy(p.basePixels, pixels)
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Pixels) AppendDrawImageHistory(image *graphics.Image, vertices []int16, geom graphics.Matrix, colorm graphics.Matrix, mode opengl.CompositeMode) {
	if p.stale {
		return
	}
	item := &drawImageHistoryItem{
		image:    image,
		vertices: vertices,
		geom:     geom,
		colorm:   colorm,
		mode:     mode,
	}
	p.drawImageHistory = append(p.drawImageHistory, item)
}

// At returns a color value at idx.
//
// Note that this must not be called until context is available.
// This means Pixels members must match with acutal state in GPU.
func (p *Pixels) At(idx int, context *opengl.Context) (color.Color, error) {
	if p.basePixels == nil || p.drawImageHistory != nil || p.stale {
		if err := p.Reset(context); err != nil {
			return nil, err
		}
	}
	r, g, b, a := p.basePixels[idx], p.basePixels[idx+1], p.basePixels[idx+2], p.basePixels[idx+3]
	return color.RGBA{r, g, b, a}, nil
}

func (p *Pixels) DependsOn(target *graphics.Image) bool {
	if p.stale {
		return false
	}
	for _, c := range p.drawImageHistory {
		if c.image == target {
			return true
		}
	}
	return false
}

func (p *Pixels) Reset(context *opengl.Context) error {
	var err error
	p.basePixels, err = p.image.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Pixels) ResetIfStale(context *opengl.Context) error {
	if !p.stale {
		return nil
	}
	return p.Reset(context)
}

func (p *Pixels) HasDependency() bool {
	if p.stale {
		return false
	}
	return p.drawImageHistory != nil
}

// Restore restores the pixels using its history.
//
// Restore is the only function that the pixel data is not present on GPU when this is called.
func (p *Pixels) Restore(context *opengl.Context, width, height int, filter opengl.Filter) (*graphics.Image, error) {
	if p.stale {
		return nil, errors.New("pixels: pixels must not be stale when restoring")
	}
	img := image.NewRGBA(image.Rect(0, 0, width, height))
	if p.basePixels != nil {
		for j := 0; j < height; j++ {
			copy(img.Pix[j*img.Stride:], p.basePixels[j*width*4:(j+1)*width*4])
		}
	}
	gimg, err := graphics.NewImageFromImage(img, filter)
	if err != nil {
		return nil, err
	}
	if p.baseColor != nil {
		if p.basePixels != nil {
			panic("not reach")
		}
		if err := gimg.Fill(p.baseColor); err != nil {
			return nil, err
		}
	}
	for _, c := range p.drawImageHistory {
		// c.image.impl must be already restored.
		/*if c.image.impl.hasHistory() {
			panic("not reach")
		}*/
		if err := gimg.DrawImage(c.image, c.vertices, c.geom, c.colorm, c.mode); err != nil {
			return nil, err
		}
	}
	p.basePixels, err = gimg.Pixels(context)
	if err != nil {
		return nil, err
	}
	p.baseColor = nil
	p.drawImageHistory = nil
	p.stale = false
	return gimg, nil
}
<|endoftext|>


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

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

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

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

package store

import (
	"context"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"

	"k8s.io/kube-state-metrics/v2/pkg/metric"
	generator "k8s.io/kube-state-metrics/v2/pkg/metric_generator"
)

var (
	descServiceAnnotationsName     = "kube_service_annotations"
	descServiceAnnotationsHelp     = "Kubernetes annotations converted to Prometheus labels."
	descServiceLabelsName          = "kube_service_labels"
	descServiceLabelsHelp          = "Kubernetes labels converted to Prometheus labels."
	descServiceLabelsDefaultLabels = []string{"namespace", "service"}
)

func serviceMetricFamilies(allowAnnotationsList, allowLabelsList []string) []generator.FamilyGenerator {
	return []generator.FamilyGenerator{
		*generator.NewFamilyGenerator(
			"kube_service_info",
			"Information about service.",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				m := metric.Metric{
					LabelKeys:   []string{"cluster_ip", "external_name", "load_balancer_ip"},
					LabelValues: []string{s.Spec.ClusterIP, s.Spec.ExternalName, s.Spec.LoadBalancerIP},
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_created",
			"Unix creation timestamp",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if !s.CreationTimestamp.IsZero() {
					m := metric.Metric{
						LabelKeys:   nil,
						LabelValues: nil,
						Value:       float64(s.CreationTimestamp.Unix()),
					}
					return &metric.Family{Metrics: []*metric.Metric{&m}}
				}
				return &metric.Family{Metrics: []*metric.Metric{}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_spec_type",
			"Type about service.",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				m := metric.Metric{

					LabelKeys:   []string{"type"},
					LabelValues: []string{string(s.Spec.Type)},
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			descServiceAnnotationsName,
			descServiceAnnotationsHelp,
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				annotationKeys, annotationValues := createPrometheusLabelKeysValues("annotation", s.Annotations, allowAnnotationsList)
				m := metric.Metric{
					LabelKeys:   annotationKeys,
					LabelValues: annotationValues,
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			descServiceLabelsName,
			descServiceLabelsHelp,
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				labelKeys, labelValues := createPrometheusLabelKeysValues("label", s.Labels, allowLabelsList)
				m := metric.Metric{
					LabelKeys:   labelKeys,
					LabelValues: labelValues,
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_spec_external_ip",
			"Service external ips. One series for each ip",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if len(s.Spec.ExternalIPs) == 0 {
					return &metric.Family{
						Metrics: []*metric.Metric{},
					}
				}

				ms := make([]*metric.Metric, len(s.Spec.ExternalIPs))

				for i, externalIP := range s.Spec.ExternalIPs {
					ms[i] = &metric.Metric{
						LabelKeys:   []string{"external_ip"},
						LabelValues: []string{externalIP},
						Value:       1,
					}
				}

				return &metric.Family{
					Metrics: ms,
				}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_status_load_balancer_ingress",
			"Service load balancer ingress status",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if len(s.Status.LoadBalancer.Ingress) == 0 {
					return &metric.Family{
						Metrics: []*metric.Metric{},
					}
				}

				ms := make([]*metric.Metric, len(s.Status.LoadBalancer.Ingress))

				for i, ingress := range s.Status.LoadBalancer.Ingress {
					ms[i] = &metric.Metric{
						LabelKeys:   []string{"ip", "hostname"},
						LabelValues: []string{ingress.IP, ingress.Hostname},
						Value:       1,
					}
				}

				return &metric.Family{
					Metrics: ms,
				}
			}),
		),
	}
}

func wrapSvcFunc(f func(*v1.Service) *metric.Family) func(interface{}) *metric.Family {
	return func(obj interface{}) *metric.Family {
		svc := obj.(*v1.Service)

		metricFamily := f(svc)

		for _, m := range metricFamily.Metrics {
			m.LabelKeys, m.LabelValues = mergeKeyValues(descServiceLabelsDefaultLabels, []string{svc.Namespace, svc.Name}, m.LabelKeys, m.LabelValues)
		}

		return metricFamily
	}
}

func createServiceListWatch(kubeClient clientset.Interface, ns string, fieldSelector string) cache.ListerWatcher {
	return &cache.ListWatch{
		ListFunc: func(opts metav1.ListOptions) (runtime.Object, error) {
			opts.FieldSelector = fieldSelector
			return kubeClient.CoreV1().Services(ns).List(context.TODO(), opts)
		},
		WatchFunc: func(opts metav1.ListOptions) (watch.Interface, error) {
			opts.FieldSelector = fieldSelector
			return kubeClient.CoreV1().Services(ns).Watch(context.TODO(), opts)
		},
	}
}
<commit_msg>Add default label of uid for service<commit_after>/*
Copyright 2017 The Kubernetes Authors All rights reserved.

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

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

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

package store

import (
	"context"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	clientset "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"

	"k8s.io/kube-state-metrics/v2/pkg/metric"
	generator "k8s.io/kube-state-metrics/v2/pkg/metric_generator"
)

var (
	descServiceAnnotationsName     = "kube_service_annotations"
	descServiceAnnotationsHelp     = "Kubernetes annotations converted to Prometheus labels."
	descServiceLabelsName          = "kube_service_labels"
	descServiceLabelsHelp          = "Kubernetes labels converted to Prometheus labels."
	descServiceLabelsDefaultLabels = []string{"namespace", "service", "uid"}
)

func serviceMetricFamilies(allowAnnotationsList, allowLabelsList []string) []generator.FamilyGenerator {
	return []generator.FamilyGenerator{
		*generator.NewFamilyGenerator(
			"kube_service_info",
			"Information about service.",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				m := metric.Metric{
					LabelKeys:   []string{"cluster_ip", "external_name", "load_balancer_ip"},
					LabelValues: []string{s.Spec.ClusterIP, s.Spec.ExternalName, s.Spec.LoadBalancerIP},
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_created",
			"Unix creation timestamp",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if !s.CreationTimestamp.IsZero() {
					m := metric.Metric{
						LabelKeys:   nil,
						LabelValues: nil,
						Value:       float64(s.CreationTimestamp.Unix()),
					}
					return &metric.Family{Metrics: []*metric.Metric{&m}}
				}
				return &metric.Family{Metrics: []*metric.Metric{}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_spec_type",
			"Type about service.",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				m := metric.Metric{

					LabelKeys:   []string{"type"},
					LabelValues: []string{string(s.Spec.Type)},
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			descServiceAnnotationsName,
			descServiceAnnotationsHelp,
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				annotationKeys, annotationValues := createPrometheusLabelKeysValues("annotation", s.Annotations, allowAnnotationsList)
				m := metric.Metric{
					LabelKeys:   annotationKeys,
					LabelValues: annotationValues,
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			descServiceLabelsName,
			descServiceLabelsHelp,
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				labelKeys, labelValues := createPrometheusLabelKeysValues("label", s.Labels, allowLabelsList)
				m := metric.Metric{
					LabelKeys:   labelKeys,
					LabelValues: labelValues,
					Value:       1,
				}
				return &metric.Family{Metrics: []*metric.Metric{&m}}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_spec_external_ip",
			"Service external ips. One series for each ip",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if len(s.Spec.ExternalIPs) == 0 {
					return &metric.Family{
						Metrics: []*metric.Metric{},
					}
				}

				ms := make([]*metric.Metric, len(s.Spec.ExternalIPs))

				for i, externalIP := range s.Spec.ExternalIPs {
					ms[i] = &metric.Metric{
						LabelKeys:   []string{"external_ip"},
						LabelValues: []string{externalIP},
						Value:       1,
					}
				}

				return &metric.Family{
					Metrics: ms,
				}
			}),
		),
		*generator.NewFamilyGenerator(
			"kube_service_status_load_balancer_ingress",
			"Service load balancer ingress status",
			metric.Gauge,
			"",
			wrapSvcFunc(func(s *v1.Service) *metric.Family {
				if len(s.Status.LoadBalancer.Ingress) == 0 {
					return &metric.Family{
						Metrics: []*metric.Metric{},
					}
				}

				ms := make([]*metric.Metric, len(s.Status.LoadBalancer.Ingress))

				for i, ingress := range s.Status.LoadBalancer.Ingress {
					ms[i] = &metric.Metric{
						LabelKeys:   []string{"ip", "hostname"},
						LabelValues: []string{ingress.IP, ingress.Hostname},
						Value:       1,
					}
				}

				return &metric.Family{
					Metrics: ms,
				}
			}),
		),
	}
}

func wrapSvcFunc(f func(*v1.Service) *metric.Family) func(interface{}) *metric.Family {
	return func(obj interface{}) *metric.Family {
		svc := obj.(*v1.Service)

		metricFamily := f(svc)

		for _, m := range metricFamily.Metrics {
			m.LabelKeys, m.LabelValues = mergeKeyValues(descServiceLabelsDefaultLabels, []string{svc.Namespace, svc.Name, string(svc.UID)}, m.LabelKeys, m.LabelValues)
		}

		return metricFamily
	}
}

func createServiceListWatch(kubeClient clientset.Interface, ns string, fieldSelector string) cache.ListerWatcher {
	return &cache.ListWatch{
		ListFunc: func(opts metav1.ListOptions) (runtime.Object, error) {
			opts.FieldSelector = fieldSelector
			return kubeClient.CoreV1().Services(ns).List(context.TODO(), opts)
		},
		WatchFunc: func(opts metav1.ListOptions) (watch.Interface, error) {
			opts.FieldSelector = fieldSelector
			return kubeClient.CoreV1().Services(ns).Watch(context.TODO(), opts)
		},
	}
}
<|endoftext|>


<commit_before>// Package styxfile provides helper routines and interfaces
// for serving 9P files from Go types.
package styxfile

import (
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"aqwari.net/net/styx/internal/sys"
	"aqwari.net/net/styx/styxproto"
)

// 9P read/write requests contain an offset. This makes them
// well-suited to the io.ReaderAt and io.WriterAt interfaces.
// However to keep our APIs flexible, we do the heavy lifting,
// accepting broader interfaces and forcing them into what we
// need.

// ErrNotSupported is returned when a given type does not
// implement the necessary functionality to complete a given
// read/write operation.
var ErrNotSupported = errors.New("not supported")

// Interface describes the methods a type must implement to
// be used as a file by a 9P file server. The New function converts
// types that implement some, but not all of these methods into
// types that do.
type Interface interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
}

type interfaceWithoutClose interface {
	io.ReaderAt
	io.WriterAt
}

type nopCloser struct {
	interfaceWithoutClose
}

func (nopCloser) Close() error { return nil }

// New creates a new Interface that reads and writes to and from
// rwc. The type of rwc determines the implementation selected
// by New; if rwc already implements Interface, it is used as-is. If
// some methods are missing, wrapper types are used to implement
// missing functionality. If the provided type cannot be adapted into
// an Interface, New returns a non-nil error.
func New(rwc interface{}) (Interface, error) {
	switch rwc := rwc.(type) {
	case Interface:
		return rwc, nil
	case interfaceWithoutClose:
		return nopCloser{rwc}, nil
	case io.Seeker:
		return &seekerAt{rwc: rwc}, nil
	case io.ReadWriter:
		return &dumbPipe{rwc: rwc}, nil
	case io.Reader:
		return &dumbPipe{rwc: rwc}, nil
	case io.Writer:
		return &dumbPipe{rwc: rwc}, nil
	default:
		return nil, fmt.Errorf("Cannot convert type %T into a styxfile.Interface", rwc)
	}
}

// SetDeadline sets read/write deadlines for a file, if the type supports it.
func SetDeadline(file Interface, t time.Time) error {
	var real interface{}
	switch v := file.(type) {
	case *seekerAt:
		real = v.rwc
	case *dumbPipe:
		real = v.rwc
	case *dirReader:
		real = v.Directory
	default:
		real = v
	}
	type deadline interface {
		SetDeadline(time.Time) error
	}
	if v, ok := real.(deadline); ok {
		return v.SetDeadline(t)
	}
	return ErrNotSupported
}

// Stat produces a styxproto.Stat from an open file. If the value
// provides a Stat method matching that of os.File, that is used.
// Otherwise, the styxfile package determines the file's attributes
// based on other characteristics.
func Stat(buf []byte, file Interface, name string, qid styxproto.Qid) (styxproto.Stat, error) {
	var (
		fi  os.FileInfo
		err error
	)
	type hasStat interface {
		Stat() (os.FileInfo, error)
	}
	if v, ok := file.(hasStat); ok {
		fi, err = v.Stat()
		if err != nil {
			return nil, err
		}
	} else {
		fi = statGuess{file, name, qid.Type()}
	}
	uid, gid, muid := sys.FileOwner(fi)
	stat, _, err := styxproto.NewStat(buf, fi.Name(), uid, gid, muid)
	if err != nil {
		return nil, err
	}
	stat.SetLength(fi.Size())
	stat.SetMode(Mode9P(fi.Mode()))
	stat.SetAtime(uint32(fi.ModTime().Unix()))
	stat.SetMtime(uint32(fi.ModTime().Unix()))
	stat.SetQid(qid)
	return stat, nil
}

type statGuess struct {
	file  Interface
	name  string
	qtype uint8
}

func (sg statGuess) Name() string {
	type hasName interface {
		Name() string
	}
	if v, ok := sg.file.(hasName); ok {
		return v.Name()
	}
	return sg.name
}

func (sg statGuess) Size() int64 {
	type hasSize interface {
		Size() int64
	}
	if v, ok := sg.file.(hasSize); ok {
		return v.Size()
	}
	return -1
}

func (sg statGuess) Mode() os.FileMode {
	type hasMode interface {
		Mode() os.FileMode
	}
	if v, ok := sg.file.(hasMode); ok {
		return v.Mode()
	}
	return ModeOS(uint32(sg.qtype)<<24) | 0777
}

func (sg statGuess) IsDir() bool {
	type hasDir interface {
		IsDir() bool
	}
	if v, ok := sg.file.(hasDir); ok {
		return v.IsDir()
	}
	return sg.Mode().IsDir()
}

func (sg statGuess) ModTime() time.Time {
	type hasTime interface {
		ModTime() time.Time
	}
	if v, ok := sg.file.(hasTime); ok {
		return v.ModTime()
	}
	return time.Time{}
}

func (sg statGuess) Sys() interface{} {
	return sg.file
}
<commit_msg>styxfile.Stat: Avoid creating bad messages through the result of fallback `statGuess` guesses We make sure the fallback statGuess name is the base file name, not the absolute path. Otherwise in scenarios where statGuess does not manage to find a Name method on the underlying object, it will return an absolute path, which will generate bad message errors, due to slashes not being allowed in file names.<commit_after>// Package styxfile provides helper routines and interfaces
// for serving 9P files from Go types.
package styxfile

import (
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"aqwari.net/net/styx/internal/sys"
	"aqwari.net/net/styx/styxproto"
)

// 9P read/write requests contain an offset. This makes them
// well-suited to the io.ReaderAt and io.WriterAt interfaces.
// However to keep our APIs flexible, we do the heavy lifting,
// accepting broader interfaces and forcing them into what we
// need.

// ErrNotSupported is returned when a given type does not
// implement the necessary functionality to complete a given
// read/write operation.
var ErrNotSupported = errors.New("not supported")

// Interface describes the methods a type must implement to
// be used as a file by a 9P file server. The New function converts
// types that implement some, but not all of these methods into
// types that do.
type Interface interface {
	io.ReaderAt
	io.WriterAt
	io.Closer
}

type interfaceWithoutClose interface {
	io.ReaderAt
	io.WriterAt
}

type nopCloser struct {
	interfaceWithoutClose
}

func (nopCloser) Close() error { return nil }

// New creates a new Interface that reads and writes to and from
// rwc. The type of rwc determines the implementation selected
// by New; if rwc already implements Interface, it is used as-is. If
// some methods are missing, wrapper types are used to implement
// missing functionality. If the provided type cannot be adapted into
// an Interface, New returns a non-nil error.
func New(rwc interface{}) (Interface, error) {
	switch rwc := rwc.(type) {
	case Interface:
		return rwc, nil
	case interfaceWithoutClose:
		return nopCloser{rwc}, nil
	case io.Seeker:
		return &seekerAt{rwc: rwc}, nil
	case io.ReadWriter:
		return &dumbPipe{rwc: rwc}, nil
	case io.Reader:
		return &dumbPipe{rwc: rwc}, nil
	case io.Writer:
		return &dumbPipe{rwc: rwc}, nil
	default:
		return nil, fmt.Errorf("Cannot convert type %T into a styxfile.Interface", rwc)
	}
}

// SetDeadline sets read/write deadlines for a file, if the type supports it.
func SetDeadline(file Interface, t time.Time) error {
	var real interface{}
	switch v := file.(type) {
	case *seekerAt:
		real = v.rwc
	case *dumbPipe:
		real = v.rwc
	case *dirReader:
		real = v.Directory
	default:
		real = v
	}
	type deadline interface {
		SetDeadline(time.Time) error
	}
	if v, ok := real.(deadline); ok {
		return v.SetDeadline(t)
	}
	return ErrNotSupported
}

// Stat produces a styxproto.Stat from an open file. If the value
// provides a Stat method matching that of os.File, that is used.
// Otherwise, the styxfile package determines the file's attributes
// based on other characteristics.
func Stat(buf []byte, file Interface, name string, qid styxproto.Qid) (styxproto.Stat, error) {
	var (
		fi  os.FileInfo
		err error
	)
	type hasStat interface {
		Stat() (os.FileInfo, error)
	}
	if v, ok := file.(hasStat); ok {
		fi, err = v.Stat()
		if err != nil {
			return nil, err
		}
	} else {
		// name is an absolute path, make sure we don't pass an absolute path to statGuess,
		// otherwise we may get back an absolute path if the file does not have a Name() method,
		// which would be incorrect since stat names cannot contain slashes.
		name := filepath.Base(name)
		fi = statGuess{file, name, qid.Type()}
	}
	uid, gid, muid := sys.FileOwner(fi)
	stat, _, err := styxproto.NewStat(buf, fi.Name(), uid, gid, muid)
	if err != nil {
		return nil, err
	}
	stat.SetLength(fi.Size())
	stat.SetMode(Mode9P(fi.Mode()))
	stat.SetAtime(uint32(fi.ModTime().Unix()))
	stat.SetMtime(uint32(fi.ModTime().Unix()))
	stat.SetQid(qid)
	return stat, nil
}

type statGuess struct {
	file  Interface
	name  string
	qtype uint8
}

func (sg statGuess) Name() string {
	type hasName interface {
		Name() string
	}
	if v, ok := sg.file.(hasName); ok {
		return v.Name()
	}
	return sg.name
}

func (sg statGuess) Size() int64 {
	type hasSize interface {
		Size() int64
	}
	if v, ok := sg.file.(hasSize); ok {
		return v.Size()
	}
	return -1
}

func (sg statGuess) Mode() os.FileMode {
	type hasMode interface {
		Mode() os.FileMode
	}
	if v, ok := sg.file.(hasMode); ok {
		return v.Mode()
	}
	return ModeOS(uint32(sg.qtype)<<24) | 0777
}

func (sg statGuess) IsDir() bool {
	type hasDir interface {
		IsDir() bool
	}
	if v, ok := sg.file.(hasDir); ok {
		return v.IsDir()
	}
	return sg.Mode().IsDir()
}

func (sg statGuess) ModTime() time.Time {
	type hasTime interface {
		ModTime() time.Time
	}
	if v, ok := sg.file.(hasTime); ok {
		return v.ModTime()
	}
	return time.Time{}
}

func (sg statGuess) Sys() interface{} {
	return sg.file
}
<|endoftext|>


<commit_before>package leveldb

import (
    "os"
    "fmt"
    "log"
    "errors"
    "strconv"
    "github.com/Lupino/periodic/driver"
    "github.com/syndtr/goleveldb/leveldb"
    "github.com/syndtr/goleveldb/leveldb/util"
    "github.com/syndtr/goleveldb/leveldb/iterator"
    "github.com/golang/groupcache/lru"
)

const PRE_JOB = "job:"
const PRE_JOB_FUNC = "func:"
const PRE_SEQUENCE = "sequence:"

type LevelDBDriver struct {
    db *leveldb.DB
    cache *lru.Cache
}

func NewLevelDBDriver(dbpath string) LevelDBDriver {
    var db *leveldb.DB
    var err error
    var cache *lru.Cache

    _, err = os.Stat(dbpath)

    if err == nil || os.IsExist(err) {
        db, err = leveldb.RecoverFile(dbpath, nil)
    } else {
        db, err = leveldb.OpenFile(dbpath, nil)
    }
    if err != nil {
        log.Fatal(err)
    }
    cache = lru.New(1000)
    return LevelDBDriver{
        db: db,
        cache: cache,
    }
}

func (l LevelDBDriver) Save(job *driver.Job) (err error) {
    batch := new(leveldb.Batch)
    if job.Id > 0 {
        old, e := l.Get(job.Id)
        if e != nil || old.Id == 0 {
            err = errors.New(fmt.Sprintf("Update Job %d fail, the old job is not exists.", job.Id))
            return
        }
        l.cache.Remove(PRE_JOB + strconv.FormatInt(job.Id, 10))
        if old.Name != job.Name {
            batch.Delete([]byte(PRE_JOB_FUNC + job.Func + ":" + old.Name))
            batch.Put([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name), []byte(strconv.FormatInt(job.Id, 10)))
        }
    } else {
        last_id, e := l.db.Get([]byte(PRE_SEQUENCE + "JOB"), nil)
        if e != nil || last_id == nil {
            job.Id = 1
        } else {
            id, _ := strconv.ParseInt(string(last_id), 10, 64)
            job.Id = id + 1
        }
        batch.Put([]byte(PRE_SEQUENCE + "JOB"), []byte(strconv.FormatInt(job.Id, 10)))
        batch.Put([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name), []byte(strconv.FormatInt(job.Id, 10)))
    }
    batch.Put([]byte(PRE_JOB + strconv.FormatInt(job.Id, 10)), job.Bytes())
    err = l.db.Write(batch, nil)
    return
}

func (l LevelDBDriver) Delete(jobId int64) (err error) {
    var job driver.Job
    batch := new(leveldb.Batch)
    job, err = l.Get(jobId)
    if err != nil {
        return
    }
    batch.Delete([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name))
    batch.Delete([]byte(PRE_JOB + strconv.FormatInt(job.Id, 10)))
    err = l.db.Write(batch, nil)
    l.cache.Remove(PRE_JOB + strconv.FormatInt(jobId, 10))
    return
}

func (l LevelDBDriver) Get(jobId int64) (job driver.Job, err error) {
    var data []byte
    var key = PRE_JOB + strconv.FormatInt(jobId, 10)
    if val, hit := l.cache.Get(key); hit {
        return val.(driver.Job), nil
    }
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    job, err = driver.NewJob(data)
    if err == nil {
        l.cache.Add(key, job)
    }
    return
}

func (l LevelDBDriver) GetOne(Func, name string) (job driver.Job, err error) {
    var data []byte
    var key = PRE_JOB_FUNC + Func + ":" + name
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    key = PRE_JOB + string(data)
    if val, hit := l.cache.Get(key); hit {
        return val.(driver.Job), nil
    }
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    job, err = driver.NewJob(data)
    if err == nil {
        l.cache.Add(key, job)
    }
    return
}

func (l LevelDBDriver) NewIterator(Func []byte) driver.JobIterator {
    var prefix []byte
    if Func == nil {
        prefix = []byte(PRE_JOB)
    } else {
        prefix = []byte(PRE_JOB_FUNC + string(Func))
    }
    iter := l.db.NewIterator(util.BytesPrefix(prefix), nil)
    return &LevelDBIterator{
        l: l,
        iter: iter,
        Func: Func,
    }
}

func (l LevelDBDriver) Close() error {
    err := l.db.Close()
    return err
}

type LevelDBIterator struct {
    l LevelDBDriver
    iter iterator.Iterator
    Func []byte
}

func (iter *LevelDBIterator) Next() bool {
    return iter.iter.Next()
}

func (iter *LevelDBIterator) Value() (job driver.Job) {
    data := iter.iter.Value()
    if iter.Func == nil {
        job, _ = driver.NewJob(data)
        return
    }
    jobId, _ := strconv.ParseInt(string(data), 10, 64)
    job, _ = iter.l.Get(jobId)
    return
}

func (iter *LevelDBIterator) Error() error {
    return iter.iter.Error()
}

func (iter *LevelDBIterator) Close() {
    iter.iter.Release()
}
<commit_msg>less FormatInt<commit_after>package leveldb

import (
    "os"
    "fmt"
    "log"
    "errors"
    "strconv"
    "github.com/Lupino/periodic/driver"
    "github.com/syndtr/goleveldb/leveldb"
    "github.com/syndtr/goleveldb/leveldb/util"
    "github.com/syndtr/goleveldb/leveldb/iterator"
    "github.com/golang/groupcache/lru"
)

const PRE_JOB = "job:"
const PRE_JOB_FUNC = "func:"
const PRE_SEQUENCE = "sequence:"

type LevelDBDriver struct {
    db *leveldb.DB
    cache *lru.Cache
}

func NewLevelDBDriver(dbpath string) LevelDBDriver {
    var db *leveldb.DB
    var err error
    var cache *lru.Cache

    _, err = os.Stat(dbpath)

    if err == nil || os.IsExist(err) {
        db, err = leveldb.RecoverFile(dbpath, nil)
    } else {
        db, err = leveldb.OpenFile(dbpath, nil)
    }
    if err != nil {
        log.Fatal(err)
    }
    cache = lru.New(1000)
    return LevelDBDriver{
        db: db,
        cache: cache,
    }
}

func (l LevelDBDriver) Save(job *driver.Job) (err error) {
    batch := new(leveldb.Batch)
    var isNew = true
    if job.Id > 0 {
        isNew = false
    } else {
        last_id, e := l.db.Get([]byte(PRE_SEQUENCE + "JOB"), nil)
        if e != nil || last_id == nil {
            job.Id = 1
        } else {
            id, _ := strconv.ParseInt(string(last_id), 10, 64)
            job.Id = id + 1
        }
    }
    var strId = strconv.FormatInt(job.Id, 10)
    if isNew {
        batch.Put([]byte(PRE_SEQUENCE + "JOB"), []byte(strId))
        batch.Put([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name), []byte(strId))
    } else {
        old, e := l.Get(job.Id)
        if e != nil || old.Id == 0 {
            err = errors.New(fmt.Sprintf("Update Job %d fail, the old job is not exists.", job.Id))
            return
        }
        l.cache.Remove(PRE_JOB + strId)
        if old.Name != job.Name {
            batch.Delete([]byte(PRE_JOB_FUNC + job.Func + ":" + old.Name))
            batch.Put([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name), []byte(strId))
        }
    }
    batch.Put([]byte(PRE_JOB + strId), job.Bytes())
    err = l.db.Write(batch, nil)
    return
}

func (l LevelDBDriver) Delete(jobId int64) (err error) {
    var job driver.Job
    batch := new(leveldb.Batch)
    job, err = l.Get(jobId)
    if err != nil {
        return
    }
    var strId = strconv.FormatInt(job.Id, 10)
    batch.Delete([]byte(PRE_JOB_FUNC + job.Func + ":" + job.Name))
    batch.Delete([]byte(PRE_JOB + strId))
    err = l.db.Write(batch, nil)
    l.cache.Remove(PRE_JOB + strId)
    return
}

func (l LevelDBDriver) Get(jobId int64) (job driver.Job, err error) {
    var data []byte
    var key = PRE_JOB + strconv.FormatInt(jobId, 10)
    if val, hit := l.cache.Get(key); hit {
        return val.(driver.Job), nil
    }
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    job, err = driver.NewJob(data)
    if err == nil {
        l.cache.Add(key, job)
    }
    return
}

func (l LevelDBDriver) GetOne(Func, name string) (job driver.Job, err error) {
    var data []byte
    var key = PRE_JOB_FUNC + Func + ":" + name
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    key = PRE_JOB + string(data)
    if val, hit := l.cache.Get(key); hit {
        return val.(driver.Job), nil
    }
    data, err = l.db.Get([]byte(key), nil)
    if err != nil {
        return
    }
    job, err = driver.NewJob(data)
    if err == nil {
        l.cache.Add(key, job)
    }
    return
}

func (l LevelDBDriver) NewIterator(Func []byte) driver.JobIterator {
    var prefix []byte
    if Func == nil {
        prefix = []byte(PRE_JOB)
    } else {
        prefix = []byte(PRE_JOB_FUNC + string(Func))
    }
    iter := l.db.NewIterator(util.BytesPrefix(prefix), nil)
    return &LevelDBIterator{
        l: l,
        iter: iter,
        Func: Func,
    }
}

func (l LevelDBDriver) Close() error {
    err := l.db.Close()
    return err
}

type LevelDBIterator struct {
    l LevelDBDriver
    iter iterator.Iterator
    Func []byte
}

func (iter *LevelDBIterator) Next() bool {
    return iter.iter.Next()
}

func (iter *LevelDBIterator) Value() (job driver.Job) {
    data := iter.iter.Value()
    if iter.Func == nil {
        job, _ = driver.NewJob(data)
        return
    }
    jobId, _ := strconv.ParseInt(string(data), 10, 64)
    job, _ = iter.l.Get(jobId)
    return
}

func (iter *LevelDBIterator) Error() error {
    return iter.iter.Error()
}

func (iter *LevelDBIterator) Close() {
    iter.iter.Release()
}
<|endoftext|>


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

import (
	"errors"
	"fmt"
	"net"
	"strings"
	"testing"
)

type shutdownCodec struct {
	responded chan int
	closed    bool
}

func (c *shutdownCodec) WriteRequest(*Request, interface{}) error { return nil }
func (c *shutdownCodec) ReadResponseBody(interface{}) error       { return nil }
func (c *shutdownCodec) ReadResponseHeader(*Response) error {
	c.responded <- 1
	return errors.New("shutdownCodec ReadResponseHeader")
}
func (c *shutdownCodec) Close() error {
	c.closed = true
	return nil
}

func TestCloseCodec(t *testing.T) {
	codec := &shutdownCodec{responded: make(chan int)}
	client := NewClientWithCodec(codec)
	<-codec.responded
	client.Close()
	if !codec.closed {
		t.Error("client.Close did not close codec")
	}
}

// Test that errors in gob shut down the connection. Issue 7689.

type R struct {
	msg []byte // Not exported, so R does not work with gob.
}

type S struct{}

func (s *S) Recv(nul *struct{}, reply *R) error {
	*reply = R{[]byte("foo")}
	return nil
}

func TestGobError(t *testing.T) {
	if runtime.GOOS == "plan9" {
		t.Skip("skipping test; see http://golang.org/issue/8908")
	}
	defer func() {
		err := recover()
		if err == nil {
			t.Fatal("no error")
		}
		if !strings.Contains("reading body EOF", err.(error).Error()) {
			t.Fatal("expected `reading body EOF', got", err)
		}
	}()
	Register(new(S))

	listen, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	go Accept(listen)

	client, err := Dial("tcp", listen.Addr().String())
	if err != nil {
		panic(err)
	}

	var reply Reply
	err = client.Call("S.Recv", &struct{}{}, &reply)
	if err != nil {
		panic(err)
	}

	fmt.Printf("%#v\n", reply)
	client.Close()

	listen.Close()
}
<commit_msg>net/rpc: fix build<commit_after>// Copyright 2014 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 rpc

import (
	"errors"
	"fmt"
	"net"
	"runtime"
	"strings"
	"testing"
)

type shutdownCodec struct {
	responded chan int
	closed    bool
}

func (c *shutdownCodec) WriteRequest(*Request, interface{}) error { return nil }
func (c *shutdownCodec) ReadResponseBody(interface{}) error       { return nil }
func (c *shutdownCodec) ReadResponseHeader(*Response) error {
	c.responded <- 1
	return errors.New("shutdownCodec ReadResponseHeader")
}
func (c *shutdownCodec) Close() error {
	c.closed = true
	return nil
}

func TestCloseCodec(t *testing.T) {
	codec := &shutdownCodec{responded: make(chan int)}
	client := NewClientWithCodec(codec)
	<-codec.responded
	client.Close()
	if !codec.closed {
		t.Error("client.Close did not close codec")
	}
}

// Test that errors in gob shut down the connection. Issue 7689.

type R struct {
	msg []byte // Not exported, so R does not work with gob.
}

type S struct{}

func (s *S) Recv(nul *struct{}, reply *R) error {
	*reply = R{[]byte("foo")}
	return nil
}

func TestGobError(t *testing.T) {
	if runtime.GOOS == "plan9" {
		t.Skip("skipping test; see http://golang.org/issue/8908")
	}
	defer func() {
		err := recover()
		if err == nil {
			t.Fatal("no error")
		}
		if !strings.Contains("reading body EOF", err.(error).Error()) {
			t.Fatal("expected `reading body EOF', got", err)
		}
	}()
	Register(new(S))

	listen, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	go Accept(listen)

	client, err := Dial("tcp", listen.Addr().String())
	if err != nil {
		panic(err)
	}

	var reply Reply
	err = client.Call("S.Recv", &struct{}{}, &reply)
	if err != nil {
		panic(err)
	}

	fmt.Printf("%#v\n", reply)
	client.Close()

	listen.Close()
}
<|endoftext|>


<commit_before>// Copyright 2015 The LUCI 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 datastore

import (
	"fmt"
	"reflect"

	"go.chromium.org/luci/common/errors"

	"google.golang.org/appengine/datastore"
)

type stopErr struct{}

func (stopErr) Error() string { return "stop iteration" }

// These errors are returned by various datastore.Interface methods.
var (
	ErrNoSuchEntity          = datastore.ErrNoSuchEntity
	ErrConcurrentTransaction = datastore.ErrConcurrentTransaction

	// Stop is understood by various services to stop iterative processes. Examples
	// include datastore.Interface.Run's callback.
	Stop = stopErr{}
)

// MakeErrInvalidKey returns an errors.Annotator instance that wraps an invalid
// key error. Calling IsErrInvalidKey on this Annotator or its derivatives will
// return true.
func MakeErrInvalidKey(reason string, args ...interface{}) *errors.Annotator {
	return errors.Annotate(datastore.ErrInvalidKey, reason, args...)
}

// IsErrInvalidKey tests if a given error is a wrapped datastore.ErrInvalidKey
// error.
func IsErrInvalidKey(err error) bool { return errors.Unwrap(err) == datastore.ErrInvalidKey }

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument
// passed to Get or to Iterator.Next.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("gae: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}
<commit_msg>[datastore] Add IsErrNoSuchEntity in datastore/errors.go<commit_after>// Copyright 2015 The LUCI 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 datastore

import (
	"fmt"
	"reflect"

	"go.chromium.org/luci/common/errors"

	"google.golang.org/appengine/datastore"
)

type stopErr struct{}

func (stopErr) Error() string { return "stop iteration" }

// These errors are returned by various datastore.Interface methods.
var (
	ErrNoSuchEntity          = datastore.ErrNoSuchEntity
	ErrConcurrentTransaction = datastore.ErrConcurrentTransaction

	// Stop is understood by various services to stop iterative processes. Examples
	// include datastore.Interface.Run's callback.
	Stop = stopErr{}
)

// MakeErrInvalidKey returns an errors.Annotator instance that wraps an invalid
// key error. Calling IsErrInvalidKey on this Annotator or its derivatives will
// return true.
func MakeErrInvalidKey(reason string, args ...interface{}) *errors.Annotator {
	return errors.Annotate(datastore.ErrInvalidKey, reason, args...)
}

// IsErrInvalidKey tests if a given error is a wrapped datastore.ErrInvalidKey
// error.
func IsErrInvalidKey(err error) bool { return errors.Unwrap(err) == datastore.ErrInvalidKey }

// IsErrNoSuchEntity tests if an error is ErrNoSuchEntity,
// or is a MultiError that contains ErrNoSuchEntity and no other errors.
func IsErrNoSuchEntity(err error) (found bool) {
	errors.Walk(err, func(err error) bool {
		found = err == ErrNoSuchEntity
		// If we found an ErrNoSuchEntity, continue walking.
		// If we found a different type of error, signal Walk to stop walking by returning false.
		// Walk does not walk nil errors.
		return found
	})
	return
}

// ErrFieldMismatch is returned when a field is to be loaded into a different
// type than the one it was stored from, or when a field is missing or
// unexported in the destination struct.
// StructType is the type of the struct pointed to by the destination argument
// passed to Get or to Iterator.Next.
type ErrFieldMismatch struct {
	StructType reflect.Type
	FieldName  string
	Reason     string
}

func (e *ErrFieldMismatch) Error() string {
	return fmt.Sprintf("gae: cannot load field %q into a %q: %s",
		e.FieldName, e.StructType, e.Reason)
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 android

import (
	"context"
	"fmt"
	"io"
	"os"
	"sync/atomic"

	perfetto_pb "perfetto/config"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android"
	"github.com/google/gapid/core/os/android/adb"
	"github.com/google/gapid/core/os/file"
	"github.com/google/gapid/gapis/service"
)

const (
	gpuRenderStagesDataSourceDescriptorName = "gpu.renderstages"

	// perfettoTraceFile is the location on the device where we'll ask Perfetto
	// to store the trace data while tracing.
	perfettoTraceFile                       = "/data/misc/perfetto-traces/gapis-trace"
	prereleaseDriverProperty                = "ro.gfx.driver.1"
	prereleaseDriverOverrideSettingVariable = "gapid.driver_package_override"
	renderStageVulkanLayerName              = "VkRenderStagesProducer"
)

// Process represents a running Perfetto capture.
type Process struct {
	device   adb.Device
	config   *perfetto_pb.TraceConfig
	deferred bool
}

func hasRenderStageEnabled(perfettoConfig *perfetto_pb.TraceConfig) bool {
	for _, dataSource := range perfettoConfig.GetDataSources() {
		if dataSource.Config.GetName() == gpuRenderStagesDataSourceDescriptorName {
			return true
		}
	}
	return false
}

func setupRenderStagesEnvironment(ctx context.Context, d adb.Device, packageName string, perfettoConfig *perfetto_pb.TraceConfig) (app.Cleanup, error) {
	if !hasRenderStageEnabled(perfettoConfig) {
		return nil, nil
	}
	driverPackageName, err := d.SystemProperty(ctx, prereleaseDriverProperty)
	if err != nil {
		return nil, err
	}
	if driverPackageOverride, err := d.SystemSetting(ctx, "global", prereleaseDriverOverrideSettingVariable); driverPackageOverride != "" && driverPackageOverride != "null" && err == nil {
		driverPackageName = driverPackageOverride
	}
	if driverPackageName == "" {
		return nil, nil
	}
	cleanup, err := android.SetupLayer(ctx, d, packageName, driverPackageName, renderStageVulkanLayerName, true)
	if err != nil {
		return nil, log.Err(ctx, err, "Failed to setup gpu.renderstages environment.")
	}
	return cleanup, nil
}

// Start optional starts an app and sets up a Perfetto trace
func Start(ctx context.Context, d adb.Device, a *android.ActivityAction, opts *service.TraceOptions) (*Process, app.Cleanup, error) {
	ctx = log.Enter(ctx, "start")

	var cleanup app.Cleanup
	if a != nil {
		ctx = log.V{
			"package":  a.Package.Name,
			"activity": a.Activity,
		}.Bind(ctx)
		// Before we start the activity, attempt to setup environment if gpu.renderstages is selected.
		var err error
		cleanup, err = setupRenderStagesEnvironment(ctx, d, a.Package.Name, opts.PerfettoConfig)
		if err != nil {
			return nil, cleanup.Invoke(ctx), err
		}
	}

	log.I(ctx, "Unlocking device screen")
	unlocked, err := d.UnlockScreen(ctx)
	if err != nil {
		log.W(ctx, "Failed to determine lock state: %s", err)
	} else if !unlocked {
		return nil, cleanup.Invoke(ctx), log.Err(ctx, nil, "Please unlock your device screen: GAPID can automatically unlock the screen only when no PIN/password/pattern is needed")
	}

	if a != nil {
		if err := d.StartActivity(ctx, *a); err != nil {
			return nil, cleanup.Invoke(ctx), log.Err(ctx, err, "Starting the activity")
		}
	}

	return &Process{
		device:   d,
		config:   opts.PerfettoConfig,
		deferred: opts.DeferStart,
	}, cleanup, nil
}

// Capture starts the perfetto capture.
func (p *Process) Capture(ctx context.Context, start task.Signal, stop task.Signal, w io.Writer, written *int64) (int64, error) {
	tmp, err := file.Temp()
	if err != nil {
		return 0, log.Err(ctx, err, "Failed to create a temp file")
	}

	// Signal that we are ready to start.
	atomic.StoreInt64(written, 1)

	if p.deferred && !start.Wait(ctx) {
		return 0, log.Err(ctx, nil, "Cancelled")
	}

	if err := p.device.StartPerfettoTrace(ctx, p.config, perfettoTraceFile, stop); err != nil {
		return 0, err
	}

	if err := p.device.Pull(ctx, perfettoTraceFile, tmp.System()); err != nil {
		return 0, err
	}

	if err := p.device.RemoveFile(ctx, perfettoTraceFile); err != nil {
		log.E(ctx, "Failed to delete perfetto trace file %v", err)
	}

	size := tmp.Info().Size()
	atomic.StoreInt64(written, size)
	fh, err := os.Open(tmp.System())
	if err != nil {
		return 0, log.Err(ctx, err, fmt.Sprintf("Failed to open %s", tmp))
	}
	return io.Copy(w, fh)
}
<commit_msg>Don't set up layer when driver apk is not found. (#3274)<commit_after>// Copyright (C) 2019 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 android

import (
	"context"
	"fmt"
	"io"
	"os"
	"sync/atomic"

	perfetto_pb "perfetto/config"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/android"
	"github.com/google/gapid/core/os/android/adb"
	"github.com/google/gapid/core/os/file"
	"github.com/google/gapid/gapis/service"
)

const (
	gpuRenderStagesDataSourceDescriptorName = "gpu.renderstages"

	// perfettoTraceFile is the location on the device where we'll ask Perfetto
	// to store the trace data while tracing.
	perfettoTraceFile                       = "/data/misc/perfetto-traces/gapis-trace"
	prereleaseDriverProperty                = "ro.gfx.driver.1"
	prereleaseDriverOverrideSettingVariable = "gapid.driver_package_override"
	renderStageVulkanLayerName              = "VkRenderStagesProducer"
)

// Process represents a running Perfetto capture.
type Process struct {
	device   adb.Device
	config   *perfetto_pb.TraceConfig
	deferred bool
}

func hasRenderStageEnabled(perfettoConfig *perfetto_pb.TraceConfig) bool {
	for _, dataSource := range perfettoConfig.GetDataSources() {
		if dataSource.Config.GetName() == gpuRenderStagesDataSourceDescriptorName {
			return true
		}
	}
	return false
}

func setupRenderStagesEnvironment(ctx context.Context, d adb.Device, packageName string, perfettoConfig *perfetto_pb.TraceConfig) (app.Cleanup, error) {
	if !hasRenderStageEnabled(perfettoConfig) {
		return nil, nil
	}
	driverPackageName, err := d.SystemProperty(ctx, prereleaseDriverProperty)
	if err != nil {
		return nil, err
	}
	if driverPackageOverride, err := d.SystemSetting(ctx, "global", prereleaseDriverOverrideSettingVariable); driverPackageOverride != "" && driverPackageOverride != "null" && err == nil {
		driverPackageName = driverPackageOverride
	}
	if driverPackageName == "" {
		return nil, nil
	}
	if res, err := d.Shell("pm", "path", driverPackageName).Call(ctx); err != nil || res == "" {
		return nil, log.Err(ctx, err, "No driver package found.")
	}
	cleanup, err := android.SetupLayer(ctx, d, packageName, driverPackageName, renderStageVulkanLayerName, true)
	if err != nil {
		return cleanup.Invoke(ctx), log.Err(ctx, err, "Failed to setup gpu.renderstages environment.")
	}
	return cleanup, nil
}

// Start optional starts an app and sets up a Perfetto trace
func Start(ctx context.Context, d adb.Device, a *android.ActivityAction, opts *service.TraceOptions) (*Process, app.Cleanup, error) {
	ctx = log.Enter(ctx, "start")

	var cleanup app.Cleanup
	if a != nil {
		ctx = log.V{
			"package":  a.Package.Name,
			"activity": a.Activity,
		}.Bind(ctx)
		// Before we start the activity, attempt to setup environment if gpu.renderstages is selected.
		var err error
		cleanup, err = setupRenderStagesEnvironment(ctx, d, a.Package.Name, opts.PerfettoConfig)
		if err != nil {
			return nil, cleanup.Invoke(ctx), err
		}
	}

	log.I(ctx, "Unlocking device screen")
	unlocked, err := d.UnlockScreen(ctx)
	if err != nil {
		log.W(ctx, "Failed to determine lock state: %s", err)
	} else if !unlocked {
		return nil, cleanup.Invoke(ctx), log.Err(ctx, nil, "Please unlock your device screen: GAPID can automatically unlock the screen only when no PIN/password/pattern is needed")
	}

	if a != nil {
		if err := d.StartActivity(ctx, *a); err != nil {
			return nil, cleanup.Invoke(ctx), log.Err(ctx, err, "Starting the activity")
		}
	}

	return &Process{
		device:   d,
		config:   opts.PerfettoConfig,
		deferred: opts.DeferStart,
	}, cleanup, nil
}

// Capture starts the perfetto capture.
func (p *Process) Capture(ctx context.Context, start task.Signal, stop task.Signal, w io.Writer, written *int64) (int64, error) {
	tmp, err := file.Temp()
	if err != nil {
		return 0, log.Err(ctx, err, "Failed to create a temp file")
	}

	// Signal that we are ready to start.
	atomic.StoreInt64(written, 1)

	if p.deferred && !start.Wait(ctx) {
		return 0, log.Err(ctx, nil, "Cancelled")
	}

	if err := p.device.StartPerfettoTrace(ctx, p.config, perfettoTraceFile, stop); err != nil {
		return 0, err
	}

	if err := p.device.Pull(ctx, perfettoTraceFile, tmp.System()); err != nil {
		return 0, err
	}

	if err := p.device.RemoveFile(ctx, perfettoTraceFile); err != nil {
		log.E(ctx, "Failed to delete perfetto trace file %v", err)
	}

	size := tmp.Info().Size()
	atomic.StoreInt64(written, size)
	fh, err := os.Open(tmp.System())
	if err != nil {
		return 0, log.Err(ctx, err, fmt.Sprintf("Failed to open %s", tmp))
	}
	return io.Copy(w, fh)
}
<|endoftext|>