<commit_before>package layout

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/containers/image/v5/directory/explicitfilepath"
	"github.com/containers/image/v5/docker/reference"
	"github.com/containers/image/v5/internal/image"
	"github.com/containers/image/v5/oci/internal"
	"github.com/containers/image/v5/transports"
	"github.com/containers/image/v5/types"
	"github.com/opencontainers/go-digest"
	imgspecv1 "github.com/opencontainers/image-spec/specs-go/v1"
)

func init() {
	transports.Register(Transport)
}

var (
	// Transport is an ImageTransport for OCI directories.
	Transport = ociTransport{}

	// ErrMoreThanOneImage is an error returned when the manifest includes
	// more than one image and the user should choose which one to use.
	ErrMoreThanOneImage = errors.New("more than one image in oci, choose an image")
)

type ociTransport struct{}

func (t ociTransport) Name() string {
	return "oci"
}

// ParseReference converts a string, which should not start with the ImageTransport.Name prefix, into an ImageReference.
func (t ociTransport) ParseReference(reference string) (types.ImageReference, error) {
	return ParseReference(reference)
}

// ValidatePolicyConfigurationScope checks that scope is a valid name for a signature.PolicyTransportScopes keys
// (i.e. a valid PolicyConfigurationIdentity() or PolicyConfigurationNamespaces() return value).
// It is acceptable to allow an invalid value which will never be matched, it can "only" cause user confusion.
// scope passed to this function will not be "", that value is always allowed.
func (t ociTransport) ValidatePolicyConfigurationScope(scope string) error {
	return internal.ValidateScope(scope)
}

// ociReference is an ImageReference for OCI directory paths.
type ociReference struct {
	// Note that the interpretation of paths below depends on the underlying filesystem state, which may change under us at any time!
	// Either of the paths may point to a different, or no, inode over time.  resolvedDir may contain symbolic links, and so on.

	// Generally we follow the intent of the user, and use the "dir" member for filesystem operations (e.g. the user can use a relative path to avoid
	// being exposed to symlinks and renames in the parent directories to the working directory).
	// (But in general, we make no attempt to be completely safe against concurrent hostile filesystem modifications.)
	dir         string // As specified by the user. May be relative, contain symlinks, etc.
	resolvedDir string // Absolute path with no symlinks, at least at the time of its creation. Primarily used for policy namespaces.
	// If image=="", it means the "only image" in the index.json is used in the case it is a source
	// for destinations, the image name annotation "image.ref.name" is not added to the index.json
	image string
}

// ParseReference converts a string, which should not start with the ImageTransport.Name prefix, into an OCI ImageReference.
func ParseReference(reference string) (types.ImageReference, error) {
	dir, image := internal.SplitPathAndImage(reference)
	return NewReference(dir, image)
}

// NewReference returns an OCI reference for a directory and a image.
//
// We do not expose an API supplying the resolvedDir; we could, but recomputing it
// is generally cheap enough that we prefer being confident about the properties of resolvedDir.
func NewReference(dir, image string) (types.ImageReference, error) {
	resolved, err := explicitfilepath.ResolvePathToFullyExplicit(dir)
	if err != nil {
		return nil, err
	}

	if err := internal.ValidateOCIPath(dir); err != nil {
		return nil, err
	}

	if err = internal.ValidateImageName(image); err != nil {
		return nil, err
	}

	return ociReference{dir: dir, resolvedDir: resolved, image: image}, nil
}

func (ref ociReference) Transport() types.ImageTransport {
	return Transport
}

// StringWithinTransport returns a string representation of the reference, which MUST be such that
// reference.Transport().ParseReference(reference.StringWithinTransport()) returns an equivalent reference.
// NOTE: The returned string is not promised to be equal to the original input to ParseReference;
// e.g. default attribute values omitted by the user may be filled in in the return value, or vice versa.
// WARNING: Do not use the return value in the UI to describe an image, it does not contain the Transport().Name() prefix.
func (ref ociReference) StringWithinTransport() string {
	return fmt.Sprintf("%s:%s", ref.dir, ref.image)
}

// DockerReference returns a Docker reference associated with this reference
// (fully explicit, i.e. !reference.IsNameOnly, but reflecting user intent,
// not e.g. after redirect or alias processing), or nil if unknown/not applicable.
func (ref ociReference) DockerReference() reference.Named {
	return nil
}

// PolicyConfigurationIdentity returns a string representation of the reference, suitable for policy lookup.
// This MUST reflect user intent, not e.g. after processing of third-party redirects or aliases;
// The value SHOULD be fully explicit about its semantics, with no hidden defaults, AND canonical
// (i.e. various references with exactly the same semantics should return the same configuration identity)
// It is fine for the return value to be equal to StringWithinTransport(), and it is desirable but
// not required/guaranteed that it will be a valid input to Transport().ParseReference().
// Returns "" if configuration identities for these references are not supported.
func (ref ociReference) PolicyConfigurationIdentity() string {
	// NOTE: ref.image is not a part of the image identity, because "$dir:$someimage" and "$dir:" may mean the
	// same image and the two can’t be statically disambiguated.  Using at least the repository directory is
	// less granular but hopefully still useful.
	return ref.resolvedDir
}

// PolicyConfigurationNamespaces returns a list of other policy configuration namespaces to search
// for if explicit configuration for PolicyConfigurationIdentity() is not set.  The list will be processed
// in order, terminating on first match, and an implicit "" is always checked at the end.
// It is STRONGLY recommended for the first element, if any, to be a prefix of PolicyConfigurationIdentity(),
// and each following element to be a prefix of the element preceding it.
func (ref ociReference) PolicyConfigurationNamespaces() []string {
	res := []string{}
	path := ref.resolvedDir
	for {
		lastSlash := strings.LastIndex(path, "/")
		// Note that we do not include "/"; it is redundant with the default "" global default,
		// and rejected by ociTransport.ValidatePolicyConfigurationScope above.
		if lastSlash == -1 || path == "/" {
			break
		}
		res = append(res, path)
		path = path[:lastSlash]
	}
	return res
}

// NewImage returns a types.ImageCloser for this reference, possibly specialized for this ImageTransport.
// The caller must call .Close() on the returned ImageCloser.
// NOTE: If any kind of signature verification should happen, build an UnparsedImage from the value returned by NewImageSource,
// verify that UnparsedImage, and convert it into a real Image via image.FromUnparsedImage.
// WARNING: This may not do the right thing for a manifest list, see image.FromSource for details.
func (ref ociReference) NewImage(ctx context.Context, sys *types.SystemContext) (types.ImageCloser, error) {
	return image.FromReference(ctx, sys, ref)
}

// getIndex returns a pointer to the index references by this ociReference. If an error occurs opening an index nil is returned together
// with an error.
func (ref ociReference) getIndex() (*imgspecv1.Index, error) {
	indexJSON, err := os.Open(ref.indexPath())
	if err != nil {
		return nil, err
	}
	defer indexJSON.Close()

	index := &imgspecv1.Index{}
	if err := json.NewDecoder(indexJSON).Decode(index); err != nil {
		return nil, err
	}
	return index, nil
}

func (ref ociReference) getManifestDescriptor() (imgspecv1.Descriptor, error) {
	index, err := ref.getIndex()
	if err != nil {
		return imgspecv1.Descriptor{}, err
	}

	var d *imgspecv1.Descriptor
	if ref.image == "" {
		// return manifest if only one image is in the oci directory
		if len(index.Manifests) == 1 {
			d = &index.Manifests[0]
		} else {
			// ask user to choose image when more than one image in the oci directory
			return imgspecv1.Descriptor{}, ErrMoreThanOneImage
		}
	} else {
		// if image specified, look through all manifests for a match
		for _, md := range index.Manifests {
			if md.MediaType != imgspecv1.MediaTypeImageManifest && md.MediaType != imgspecv1.MediaTypeImageIndex {
				continue
			}
			refName, ok := md.Annotations[imgspecv1.AnnotationRefName]
			if !ok {
				continue
			}
			if refName == ref.image {
				d = &md
				break
			}
		}
	}
	if d == nil {
		return imgspecv1.Descriptor{}, ImageNotFoundError{ref}
	}
	return *d, nil
}

// LoadManifestDescriptor loads the manifest descriptor to be used to retrieve the image name
// when pulling an image
func LoadManifestDescriptor(imgRef types.ImageReference) (imgspecv1.Descriptor, error) {
	ociRef, ok := imgRef.(ociReference)
	if !ok {
		return imgspecv1.Descriptor{}, errors.New("error typecasting, need type ociRef")
	}
	return ociRef.getManifestDescriptor()
}

// NewImageSource returns a types.ImageSource for this reference.
// The caller must call .Close() on the returned ImageSource.
func (ref ociReference) NewImageSource(ctx context.Context, sys *types.SystemContext) (types.ImageSource, error) {
	return newImageSource(sys, ref)
}

// NewImageDestination returns a types.ImageDestination for this reference.
// The caller must call .Close() on the returned ImageDestination.
func (ref ociReference) NewImageDestination(ctx context.Context, sys *types.SystemContext) (types.ImageDestination, error) {
	return newImageDestination(sys, ref)
}

// DeleteImage deletes the named image from the registry, if supported.
func (ref ociReference) DeleteImage(ctx context.Context, sys *types.SystemContext) error {
	return errors.New("Deleting images not implemented for oci: images")
}

// ociLayoutPath returns a path for the oci-layout within a directory using OCI conventions.
func (ref ociReference) ociLayoutPath() string {
	return filepath.Join(ref.dir, "oci-layout")
}

// indexPath returns a path for the index.json within a directory using OCI conventions.
func (ref ociReference) indexPath() string {
	return filepath.Join(ref.dir, "index.json")
}

// blobPath returns a path for a blob within a directory using OCI image-layout conventions.
func (ref ociReference) blobPath(digest digest.Digest, sharedBlobDir string) (string, error) {
	if err := digest.Validate(); err != nil {
		return "", fmt.Errorf("unexpected digest reference %s: %w", digest, err)
	}
	blobDir := filepath.Join(ref.dir, "blobs")
	if sharedBlobDir != "" {
		blobDir = sharedBlobDir
	}
	return filepath.Join(blobDir, digest.Algorithm().String(), digest.Hex()), nil
}
<commit_msg>Simplify ociReference.getManifestDescriptor<commit_after>package layout

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/containers/image/v5/directory/explicitfilepath"
	"github.com/containers/image/v5/docker/reference"
	"github.com/containers/image/v5/internal/image"
	"github.com/containers/image/v5/oci/internal"
	"github.com/containers/image/v5/transports"
	"github.com/containers/image/v5/types"
	"github.com/opencontainers/go-digest"
	imgspecv1 "github.com/opencontainers/image-spec/specs-go/v1"
)

func init() {
	transports.Register(Transport)
}

var (
	// Transport is an ImageTransport for OCI directories.
	Transport = ociTransport{}

	// ErrMoreThanOneImage is an error returned when the manifest includes
	// more than one image and the user should choose which one to use.
	ErrMoreThanOneImage = errors.New("more than one image in oci, choose an image")
)

type ociTransport struct{}

func (t ociTransport) Name() string {
	return "oci"
}

// ParseReference converts a string, which should not start with the ImageTransport.Name prefix, into an ImageReference.
func (t ociTransport) ParseReference(reference string) (types.ImageReference, error) {
	return ParseReference(reference)
}

// ValidatePolicyConfigurationScope checks that scope is a valid name for a signature.PolicyTransportScopes keys
// (i.e. a valid PolicyConfigurationIdentity() or PolicyConfigurationNamespaces() return value).
// It is acceptable to allow an invalid value which will never be matched, it can "only" cause user confusion.
// scope passed to this function will not be "", that value is always allowed.
func (t ociTransport) ValidatePolicyConfigurationScope(scope string) error {
	return internal.ValidateScope(scope)
}

// ociReference is an ImageReference for OCI directory paths.
type ociReference struct {
	// Note that the interpretation of paths below depends on the underlying filesystem state, which may change under us at any time!
	// Either of the paths may point to a different, or no, inode over time.  resolvedDir may contain symbolic links, and so on.

	// Generally we follow the intent of the user, and use the "dir" member for filesystem operations (e.g. the user can use a relative path to avoid
	// being exposed to symlinks and renames in the parent directories to the working directory).
	// (But in general, we make no attempt to be completely safe against concurrent hostile filesystem modifications.)
	dir         string // As specified by the user. May be relative, contain symlinks, etc.
	resolvedDir string // Absolute path with no symlinks, at least at the time of its creation. Primarily used for policy namespaces.
	// If image=="", it means the "only image" in the index.json is used in the case it is a source
	// for destinations, the image name annotation "image.ref.name" is not added to the index.json
	image string
}

// ParseReference converts a string, which should not start with the ImageTransport.Name prefix, into an OCI ImageReference.
func ParseReference(reference string) (types.ImageReference, error) {
	dir, image := internal.SplitPathAndImage(reference)
	return NewReference(dir, image)
}

// NewReference returns an OCI reference for a directory and a image.
//
// We do not expose an API supplying the resolvedDir; we could, but recomputing it
// is generally cheap enough that we prefer being confident about the properties of resolvedDir.
func NewReference(dir, image string) (types.ImageReference, error) {
	resolved, err := explicitfilepath.ResolvePathToFullyExplicit(dir)
	if err != nil {
		return nil, err
	}

	if err := internal.ValidateOCIPath(dir); err != nil {
		return nil, err
	}

	if err = internal.ValidateImageName(image); err != nil {
		return nil, err
	}

	return ociReference{dir: dir, resolvedDir: resolved, image: image}, nil
}

func (ref ociReference) Transport() types.ImageTransport {
	return Transport
}

// StringWithinTransport returns a string representation of the reference, which MUST be such that
// reference.Transport().ParseReference(reference.StringWithinTransport()) returns an equivalent reference.
// NOTE: The returned string is not promised to be equal to the original input to ParseReference;
// e.g. default attribute values omitted by the user may be filled in in the return value, or vice versa.
// WARNING: Do not use the return value in the UI to describe an image, it does not contain the Transport().Name() prefix.
func (ref ociReference) StringWithinTransport() string {
	return fmt.Sprintf("%s:%s", ref.dir, ref.image)
}

// DockerReference returns a Docker reference associated with this reference
// (fully explicit, i.e. !reference.IsNameOnly, but reflecting user intent,
// not e.g. after redirect or alias processing), or nil if unknown/not applicable.
func (ref ociReference) DockerReference() reference.Named {
	return nil
}

// PolicyConfigurationIdentity returns a string representation of the reference, suitable for policy lookup.
// This MUST reflect user intent, not e.g. after processing of third-party redirects or aliases;
// The value SHOULD be fully explicit about its semantics, with no hidden defaults, AND canonical
// (i.e. various references with exactly the same semantics should return the same configuration identity)
// It is fine for the return value to be equal to StringWithinTransport(), and it is desirable but
// not required/guaranteed that it will be a valid input to Transport().ParseReference().
// Returns "" if configuration identities for these references are not supported.
func (ref ociReference) PolicyConfigurationIdentity() string {
	// NOTE: ref.image is not a part of the image identity, because "$dir:$someimage" and "$dir:" may mean the
	// same image and the two can’t be statically disambiguated.  Using at least the repository directory is
	// less granular but hopefully still useful.
	return ref.resolvedDir
}

// PolicyConfigurationNamespaces returns a list of other policy configuration namespaces to search
// for if explicit configuration for PolicyConfigurationIdentity() is not set.  The list will be processed
// in order, terminating on first match, and an implicit "" is always checked at the end.
// It is STRONGLY recommended for the first element, if any, to be a prefix of PolicyConfigurationIdentity(),
// and each following element to be a prefix of the element preceding it.
func (ref ociReference) PolicyConfigurationNamespaces() []string {
	res := []string{}
	path := ref.resolvedDir
	for {
		lastSlash := strings.LastIndex(path, "/")
		// Note that we do not include "/"; it is redundant with the default "" global default,
		// and rejected by ociTransport.ValidatePolicyConfigurationScope above.
		if lastSlash == -1 || path == "/" {
			break
		}
		res = append(res, path)
		path = path[:lastSlash]
	}
	return res
}

// NewImage returns a types.ImageCloser for this reference, possibly specialized for this ImageTransport.
// The caller must call .Close() on the returned ImageCloser.
// NOTE: If any kind of signature verification should happen, build an UnparsedImage from the value returned by NewImageSource,
// verify that UnparsedImage, and convert it into a real Image via image.FromUnparsedImage.
// WARNING: This may not do the right thing for a manifest list, see image.FromSource for details.
func (ref ociReference) NewImage(ctx context.Context, sys *types.SystemContext) (types.ImageCloser, error) {
	return image.FromReference(ctx, sys, ref)
}

// getIndex returns a pointer to the index references by this ociReference. If an error occurs opening an index nil is returned together
// with an error.
func (ref ociReference) getIndex() (*imgspecv1.Index, error) {
	indexJSON, err := os.Open(ref.indexPath())
	if err != nil {
		return nil, err
	}
	defer indexJSON.Close()

	index := &imgspecv1.Index{}
	if err := json.NewDecoder(indexJSON).Decode(index); err != nil {
		return nil, err
	}
	return index, nil
}

func (ref ociReference) getManifestDescriptor() (imgspecv1.Descriptor, error) {
	index, err := ref.getIndex()
	if err != nil {
		return imgspecv1.Descriptor{}, err
	}

	if ref.image == "" {
		// return manifest if only one image is in the oci directory
		if len(index.Manifests) == 1 {
			return index.Manifests[0], nil
		} else {
			// ask user to choose image when more than one image in the oci directory
			return imgspecv1.Descriptor{}, ErrMoreThanOneImage
		}
	} else {
		// if image specified, look through all manifests for a match
		for _, md := range index.Manifests {
			if md.MediaType != imgspecv1.MediaTypeImageManifest && md.MediaType != imgspecv1.MediaTypeImageIndex {
				continue
			}
			refName, ok := md.Annotations[imgspecv1.AnnotationRefName]
			if !ok {
				continue
			}
			if refName == ref.image {
				return md, nil
			}
		}
	}
	return imgspecv1.Descriptor{}, ImageNotFoundError{ref}
}

// LoadManifestDescriptor loads the manifest descriptor to be used to retrieve the image name
// when pulling an image
func LoadManifestDescriptor(imgRef types.ImageReference) (imgspecv1.Descriptor, error) {
	ociRef, ok := imgRef.(ociReference)
	if !ok {
		return imgspecv1.Descriptor{}, errors.New("error typecasting, need type ociRef")
	}
	return ociRef.getManifestDescriptor()
}

// NewImageSource returns a types.ImageSource for this reference.
// The caller must call .Close() on the returned ImageSource.
func (ref ociReference) NewImageSource(ctx context.Context, sys *types.SystemContext) (types.ImageSource, error) {
	return newImageSource(sys, ref)
}

// NewImageDestination returns a types.ImageDestination for this reference.
// The caller must call .Close() on the returned ImageDestination.
func (ref ociReference) NewImageDestination(ctx context.Context, sys *types.SystemContext) (types.ImageDestination, error) {
	return newImageDestination(sys, ref)
}

// DeleteImage deletes the named image from the registry, if supported.
func (ref ociReference) DeleteImage(ctx context.Context, sys *types.SystemContext) error {
	return errors.New("Deleting images not implemented for oci: images")
}

// ociLayoutPath returns a path for the oci-layout within a directory using OCI conventions.
func (ref ociReference) ociLayoutPath() string {
	return filepath.Join(ref.dir, "oci-layout")
}

// indexPath returns a path for the index.json within a directory using OCI conventions.
func (ref ociReference) indexPath() string {
	return filepath.Join(ref.dir, "index.json")
}

// blobPath returns a path for a blob within a directory using OCI image-layout conventions.
func (ref ociReference) blobPath(digest digest.Digest, sharedBlobDir string) (string, error) {
	if err := digest.Validate(); err != nil {
		return "", fmt.Errorf("unexpected digest reference %s: %w", digest, err)
	}
	blobDir := filepath.Join(ref.dir, "blobs")
	if sharedBlobDir != "" {
		blobDir = sharedBlobDir
	}
	return filepath.Join(blobDir, digest.Algorithm().String(), digest.Hex()), nil
}
<|endoftext|>


<commit_before>package main

import (
  "flag"
  "github.com/lestrrat/go-cpan-maybe-faster"
)

func main() {
  notest            := flag.Bool("notest", false, "Do not run unit tests")
  verbose           := flag.Bool("verbose", false, "be verbose")
  localLib          := flag.String("local-lib", "", "Specify the install base to install modules")
  localLibContained := flag.String("local-lib-contained", "", "Specify the install base to install all non-core modules")
  flag.Parse()

  c := cpan.NewClient()
  c.SetVerbose(*verbose)
  c.SetNotest(*notest)
  c.SetLocalLib(*localLib)
  c.SetLocalLibContained(*localLibContained)

  args := flag.Args()
  if len(args) > 0 {
    c.Install(args[0])
  }
}<commit_msg>attempt to install all modules passed<commit_after>package main

import (
  "flag"
  "github.com/lestrrat/go-cpan-maybe-faster"
)

func main() {
  notest            := flag.Bool("notest", false, "Do not run unit tests")
  verbose           := flag.Bool("verbose", false, "be verbose")
  localLib          := flag.String("local-lib", "", "Specify the install base to install modules")
  localLibContained := flag.String("local-lib-contained", "", "Specify the install base to install all non-core modules")
  flag.Parse()

  c := cpan.NewClient()
  c.SetVerbose(*verbose)
  c.SetNotest(*notest)
  c.SetLocalLib(*localLib)
  c.SetLocalLibContained(*localLibContained)

  args := flag.Args()
  for i := 0; i < len(args); i++ {
    c.Install(args[i])
  }
}<|endoftext|>


<commit_before>package mindcli

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/user"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type MindCliConfig struct {
	Image             string
	ContainerSkillDir string
	ServeMPKPort      int
	ServeRemotePort   int
}

type MindCli struct {
	config       *MindCliConfig
	userConfig   *UserConfig
	RobotScanner *RobotScanner
}

func NewMindCli(robotScanner *RobotScanner, userConfig *UserConfig, config *MindCliConfig) *MindCli {
	if userConfig.DockerImage != "" {
		config.Image = userConfig.DockerImage
	}
	mindcli := &MindCli{
		RobotScanner: robotScanner,
		userConfig:   userConfig,
		config:       config,
	}
	return mindcli
}

func (mindcli *MindCli) skillMountPoint() string {
	curdir, err := os.Getwd()
	if runtime.GOOS == "windows" {
		re := regexp.MustCompile(`^(?P<disk>[A-z]):(?P<path>.*)$`)
		n1 := re.SubexpNames()
		r2 := re.FindAllStringSubmatch(curdir, -1)[0]
		md := map[string]string{}
		for i, n := range r2 {
			md[n1[i]] = n
		}
		return fmt.Sprintf(
			"/%s%s:%s",
			strings.ToLower(md["disk"]),
			strings.Replace(md["path"], "\\", "/", -1),
			mindcli.config.ContainerSkillDir)
	}
	if err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
	return fmt.Sprintf("%s:%s", curdir, mindcli.config.ContainerSkillDir)
}

func (mindcli *MindCli) dockerUpgradeImageArgs() []string {
	return []string{
		"pull",
		fmt.Sprintf("%s", mindcli.config.Image),
	}
}

func mindUser() string {
	if runtime.GOOS == "linux" {
		user, err := user.Current()
		if err != nil {
			return ""
		}
		return user.Uid
	}
	return ""
}

func (mindcli *MindCli) dockerXArgs(args ...string) []string {
	c := []string{
		"run",
		"-it",
		"--rm",
		"-p", fmt.Sprintf("%d:%d", mindcli.config.ServeRemotePort, mindcli.config.ServeRemotePort),
		"-p", fmt.Sprintf("%d:%d", mindcli.config.ServeMPKPort, mindcli.config.ServeMPKPort),
		"-e", fmt.Sprintf("MIND_USER=%s", mindUser()),
		"-e", fmt.Sprintf("SERVE_REMOTE_PORT=%d", mindcli.config.ServeRemotePort),
		"-e", fmt.Sprintf("SERVE_MPK_PORT=%d", mindcli.config.ServeMPKPort),
		"-e", fmt.Sprintf("USER_HASH=%s", mindcli.userConfig.userHash),
		"-e", fmt.Sprintf("SKILLDIR=%s", mindcli.config.ContainerSkillDir),
		"-v", mindcli.skillMountPoint(),
		fmt.Sprintf("%s", mindcli.config.Image),
	}
	return append(c, args...)
}

func (mindcli *MindCli) execDocker(args []string) {
	cmd := exec.Command("docker", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if err != nil {
		fmt.Println(err)
	}
}

func (mindcli *MindCli) execDockerOutput(args []string) string {
	var outb bytes.Buffer
	cmd := exec.Command("docker", args...)
	cmd.Stdout = &outb
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if err != nil {
		fmt.Println(err)
	}
	return outb.String()
}

func (mindcli *MindCli) UpgradeImage() {
	mindcli.execDocker(mindcli.dockerUpgradeImageArgs())
}

func (mindcli *MindCli) X(args ...string) {
	mindcli.execDocker(mindcli.dockerXArgs(args...))
}

func (mindcli *MindCli) XOutput(args ...string) string {
	return mindcli.execDockerOutput(mindcli.dockerXArgs(args...))
}

func (mindcli *MindCli) RunSkill(noInstall bool, robotIp string, args ...string) {
	var robot *Robot
	if len(args) < 1 {
		robot = mindcli.RobotByName(mindcli.DefaultRobotName())
	} else {
		robot = mindcli.RobotByName(args[0])
	}
	if robot == nil {
		fmt.Println("Could not find robot with that name\nPlease run `mind scan`")
		os.Exit(-1)
	}
	robots, _ := mindcli.RobotScanner.ScanIP(robot.IP, time.Duration(2))
	if len(robots) < 1 {
		fmt.Println("Could not find any robots\nPlease run `mind scan`")
		os.Exit(-1)
	}
	if robots[0].Name != robot.Name {
		fmt.Println("Looks like robot has changed IP\nPlease run `mind scan`")
		os.Exit(-1)
	}
	ip, _ := GetLocalIPByNeighbourIP(robot.IP)
	xParams := []string{"mindcli-run"}
	if noInstall {
		xParams = append(xParams, "-n")
	}
	if robotIp == "" {
		robotIp = robot.IP
	}
	if ret := net.ParseIP(robotIp); ret == nil {
		fmt.Println("The IP address you just typed does NOT meet the standard ipv4 format.")
		return
	}
	xParams = append(xParams,
		"skill.mpk",
		fmt.Sprintf("http://%s:%d", ip.String(), mindcli.config.ServeMPKPort),
		fmt.Sprintf("%v", robotIp),
	)
	mindcli.X(append(xParams, args...)...)

}

func (mindcli *MindCli) Login(email, pass string) error {
	userhash := strings.Trim(mindcli.XOutput(append([]string{"mindcli-login"}, email, pass)...), "\r\n")
	if len(userhash) == 0 {
		return errors.New("Login failed")
	}
	mindcli.userConfig.userHash = userhash
	return mindcli.userConfig.Write()
}

func (mindcli *MindCli) Scan(waitDuration int, args ...string) ([]Robot, error) {
	var robots []Robot
	var err error
	if len(args) < 1 {
		robots, err = mindcli.RobotScanner.BroadcastToNetwork(time.Duration(waitDuration))
		if err != nil {
			return nil, err
		}
	} else {
		robots, err = mindcli.RobotScanner.ScanIP(args[0], time.Duration(waitDuration))
		if err != nil {
			return nil, err
		}
	}
	mindcli.userConfig.Robots = robots
	return robots, mindcli.userConfig.Write()
}

func (mindcli *MindCli) RobotByName(robotName string) *Robot {
	for _, robot := range mindcli.userConfig.Robots {
		if robot.Name == robotName {
			return &robot
		}
	}
	return nil
}

func (mindcli *MindCli) SetDefaultRobotName(robotName string) error {
	for _, robot := range mindcli.userConfig.Robots {
		if robot.Name == robotName {
			mindcli.userConfig.DefaultRobotName = robotName
			mindcli.userConfig.Write()
			return nil
		}
	}
	return errors.New("Could not find robot with that name")
}

func (mindcli *MindCli) DefaultRobotName() string {
	return mindcli.userConfig.DefaultRobotName
}

func (mindcli *MindCli) DefaultRobotIP() string {
	robot := mindcli.RobotByName(mindcli.userConfig.DefaultRobotName)
	if robot == nil {
		return ""
	}
	return robot.IP
}

func (mindcli *MindCli) RunFlightTest(args ...string) {
	ips, err := GetLocalIPs()
	if err != nil || len(ips) < 1 {
		fmt.Println("Cannot find a local IP, check your network settings.")
		return
	}
	mindcli.X(append([]string{
		"mindcli-flighttest",
		"skill.mpk",
		fmt.Sprintf("%v", ips[0]),
		fmt.Sprintf("%v", mindcli.config.ServeMPKPort),
	}, args...)...)
}
<commit_msg>Use --ip with no robot info in .mind.json (#22)<commit_after>package mindcli

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/user"
	"regexp"
	"runtime"
	"strings"
	"time"
)

type MindCliConfig struct {
	Image             string
	ContainerSkillDir string
	ServeMPKPort      int
	ServeRemotePort   int
}

type MindCli struct {
	config       *MindCliConfig
	userConfig   *UserConfig
	RobotScanner *RobotScanner
}

func NewMindCli(robotScanner *RobotScanner, userConfig *UserConfig, config *MindCliConfig) *MindCli {
	if userConfig.DockerImage != "" {
		config.Image = userConfig.DockerImage
	}
	mindcli := &MindCli{
		RobotScanner: robotScanner,
		userConfig:   userConfig,
		config:       config,
	}
	return mindcli
}

func (mindcli *MindCli) skillMountPoint() string {
	curdir, err := os.Getwd()
	if runtime.GOOS == "windows" {
		re := regexp.MustCompile(`^(?P<disk>[A-z]):(?P<path>.*)$`)
		n1 := re.SubexpNames()
		r2 := re.FindAllStringSubmatch(curdir, -1)[0]
		md := map[string]string{}
		for i, n := range r2 {
			md[n1[i]] = n
		}
		return fmt.Sprintf(
			"/%s%s:%s",
			strings.ToLower(md["disk"]),
			strings.Replace(md["path"], "\\", "/", -1),
			mindcli.config.ContainerSkillDir)
	}
	if err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
	return fmt.Sprintf("%s:%s", curdir, mindcli.config.ContainerSkillDir)
}

func (mindcli *MindCli) dockerUpgradeImageArgs() []string {
	return []string{
		"pull",
		fmt.Sprintf("%s", mindcli.config.Image),
	}
}

func mindUser() string {
	if runtime.GOOS == "linux" {
		user, err := user.Current()
		if err != nil {
			return ""
		}
		return user.Uid
	}
	return ""
}

func (mindcli *MindCli) dockerXArgs(args ...string) []string {
	c := []string{
		"run",
		"-it",
		"--rm",
		"-p", fmt.Sprintf("%d:%d", mindcli.config.ServeRemotePort, mindcli.config.ServeRemotePort),
		"-p", fmt.Sprintf("%d:%d", mindcli.config.ServeMPKPort, mindcli.config.ServeMPKPort),
		"-e", fmt.Sprintf("MIND_USER=%s", mindUser()),
		"-e", fmt.Sprintf("SERVE_REMOTE_PORT=%d", mindcli.config.ServeRemotePort),
		"-e", fmt.Sprintf("SERVE_MPK_PORT=%d", mindcli.config.ServeMPKPort),
		"-e", fmt.Sprintf("USER_HASH=%s", mindcli.userConfig.userHash),
		"-e", fmt.Sprintf("SKILLDIR=%s", mindcli.config.ContainerSkillDir),
		"-v", mindcli.skillMountPoint(),
		fmt.Sprintf("%s", mindcli.config.Image),
	}
	return append(c, args...)
}

func (mindcli *MindCli) execDocker(args []string) {
	cmd := exec.Command("docker", args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if err != nil {
		fmt.Println(err)
	}
}

func (mindcli *MindCli) execDockerOutput(args []string) string {
	var outb bytes.Buffer
	cmd := exec.Command("docker", args...)
	cmd.Stdout = &outb
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	err := cmd.Run()
	if err != nil {
		fmt.Println(err)
	}
	return outb.String()
}

func (mindcli *MindCli) UpgradeImage() {
	mindcli.execDocker(mindcli.dockerUpgradeImageArgs())
}

func (mindcli *MindCli) X(args ...string) {
	mindcli.execDocker(mindcli.dockerXArgs(args...))
}

func (mindcli *MindCli) XOutput(args ...string) string {
	return mindcli.execDockerOutput(mindcli.dockerXArgs(args...))
}

func (mindcli *MindCli) RunSkill(noInstall bool, robotIp string, args ...string) {
	if robotIp == "" {
		var robot *Robot
		if len(args) < 1 {
			robot = mindcli.RobotByName(mindcli.DefaultRobotName())
		} else {
			robot = mindcli.RobotByName(args[0])
		}
		if robot == nil {
			fmt.Println("Could not find robot with that name\nPlease run `mind scan`")
			os.Exit(-1)
		}
		robots, _ := mindcli.RobotScanner.ScanIP(robot.IP, time.Duration(2))
		if len(robots) < 1 {
			fmt.Println("Could not find any robots\nPlease run `mind scan`")
			os.Exit(-1)
		}
		if robots[0].Name != robot.Name {
			fmt.Println("Looks like robot has changed IP\nPlease run `mind scan`")
			os.Exit(-1)
		}

		robotIp = robot.IP
	}
	ip, _ := GetLocalIPByNeighbourIP(robotIp)
	xParams := []string{"mindcli-run"}
	if noInstall {
		xParams = append(xParams, "-n")
	}
	if ret := net.ParseIP(robotIp); ret == nil {
		fmt.Println("The IP address you just typed does NOT meet the standard ipv4 format.")
		return
	}
	xParams = append(xParams,
		"skill.mpk",
		fmt.Sprintf("http://%s:%d", ip.String(), mindcli.config.ServeMPKPort),
		fmt.Sprintf("%v", robotIp),
	)
	mindcli.X(append(xParams, args...)...)

}

func (mindcli *MindCli) Login(email, pass string) error {
	userhash := strings.Trim(mindcli.XOutput(append([]string{"mindcli-login"}, email, pass)...), "\r\n")
	if len(userhash) == 0 {
		return errors.New("Login failed")
	}
	mindcli.userConfig.userHash = userhash
	return mindcli.userConfig.Write()
}

func (mindcli *MindCli) Scan(waitDuration int, args ...string) ([]Robot, error) {
	var robots []Robot
	var err error
	if len(args) < 1 {
		robots, err = mindcli.RobotScanner.BroadcastToNetwork(time.Duration(waitDuration))
		if err != nil {
			return nil, err
		}
	} else {
		robots, err = mindcli.RobotScanner.ScanIP(args[0], time.Duration(waitDuration))
		if err != nil {
			return nil, err
		}
	}
	mindcli.userConfig.Robots = robots
	return robots, mindcli.userConfig.Write()
}

func (mindcli *MindCli) RobotByName(robotName string) *Robot {
	for _, robot := range mindcli.userConfig.Robots {
		if robot.Name == robotName {
			return &robot
		}
	}
	return nil
}

func (mindcli *MindCli) SetDefaultRobotName(robotName string) error {
	for _, robot := range mindcli.userConfig.Robots {
		if robot.Name == robotName {
			mindcli.userConfig.DefaultRobotName = robotName
			mindcli.userConfig.Write()
			return nil
		}
	}
	return errors.New("Could not find robot with that name")
}

func (mindcli *MindCli) DefaultRobotName() string {
	return mindcli.userConfig.DefaultRobotName
}

func (mindcli *MindCli) DefaultRobotIP() string {
	robot := mindcli.RobotByName(mindcli.userConfig.DefaultRobotName)
	if robot == nil {
		return ""
	}
	return robot.IP
}

func (mindcli *MindCli) RunFlightTest(args ...string) {
	ips, err := GetLocalIPs()
	if err != nil || len(ips) < 1 {
		fmt.Println("Cannot find a local IP, check your network settings.")
		return
	}
	mindcli.X(append([]string{
		"mindcli-flighttest",
		"skill.mpk",
		fmt.Sprintf("%v", ips[0]),
		fmt.Sprintf("%v", mindcli.config.ServeMPKPort),
	}, args...)...)
}
<|endoftext|>


<commit_before>package artifice

import (
	"time"

	"github.com/antihax/evedata/internal/redisqueue"
)

func init() {
	registerTrigger("mutatedItem", resolveItems, time.NewTicker(time.Second*30))
}

func resolveItems(s *Artifice) error {
	entities, err := s.db.Query(
		`SELECT I.itemID, I.typeID FROM evedata.contractItems I
			INNER JOIN eve.invTypes T ON T.typeID = I.typeID
			LEFT OUTER JOIN evedata.mutations M ON M.itemID = I.itemID
			WHERE typeName LIKE "%Abyssal%" AND M.itemID IS NULL`)
	if err != nil {
		return err
	}
	defer entities.Close()

	work := []redisqueue.Work{}

	// Loop the entities
	for entities.Next() {
		var (
			itemID int64
			typeID int32
		)

		err = entities.Scan(&itemID, &typeID)
		if err != nil {
			return err
		}

		work = append(work, redisqueue.Work{
			Operation: "mutatedItem",
			Parameter: []interface{}{
				itemID, typeID,
			}})
	}

	return s.QueueWork(work)
}
<commit_msg>Needs a priority<commit_after>package artifice

import (
	"time"

	"github.com/antihax/evedata/internal/redisqueue"
)

func init() {
	registerTrigger("mutatedItem", resolveItems, time.NewTicker(time.Second*30))
}

func resolveItems(s *Artifice) error {
	entities, err := s.db.Query(
		`SELECT I.itemID, I.typeID FROM evedata.contractItems I
			INNER JOIN eve.invTypes T ON T.typeID = I.typeID
			LEFT OUTER JOIN evedata.mutations M ON M.itemID = I.itemID
			WHERE typeName LIKE "%Abyssal%" AND M.itemID IS NULL`)
	if err != nil {
		return err
	}
	defer entities.Close()

	work := []redisqueue.Work{}

	// Loop the entities
	for entities.Next() {
		var (
			itemID int64
			typeID int32
		)

		err = entities.Scan(&itemID, &typeID)
		if err != nil {
			return err
		}

		work = append(work, redisqueue.Work{
			Operation: "mutatedItem",
			Parameter: []interface{}{
				itemID, typeID,
			}})
	}

	return s.QueueWork(work, redisqueue.Priority_High)
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package lib

import (
	"bytes"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
	"time"
)

var (
	// The directory where all the views are being stored.
	ViewsDir = "views"
)

const (
	// The path to the layout file.
	layout = "application/layout"

	// The extension of views.
	viewsExt = "tpl"
)

// This struct holds all the data that can be passed to a view.
type ViewData struct {
	// The id of the current user.
	Id string

	// Set to true if the current user is logged in.
	LoggedIn bool

	// The name of the javascript file to be used.
	JS string

	// Set to true if an error has happenned.
	Error bool
}

// Returns the path to be used to open the view with the given name.
func view(name string) string {
	return path.Join(ViewsDir, name+"."+viewsExt)
}

// Render the view with the given name after evaluating the passed data. The
// rendered view will be written to the given writer.
func Render(res http.ResponseWriter, name string, data interface{}) {
	b, e := ioutil.ReadFile(view(layout))
	if e != nil {
		panic("Could not read layout file!")
	}
	t, e := template.New("l").Funcs(layoutHelpers(name, data)).Parse(string(b))
	if e != nil {
		panic("Could not parse layout file!")
	}
	t.Execute(res, data)
}

// Returns all the helpers used by the layout template. Right now only the
// "yield" helpers has been implemented.
func layoutHelpers(name string, data interface{}) template.FuncMap {
	return template.FuncMap{
		"yield": func() template.HTML {
			var buffer bytes.Buffer

			b, e := ioutil.ReadFile(view(name))
			if e != nil {
				r := fmt.Sprintf("Could not read: %v => %v", name, e)
				panic(r)
			}
			t := template.New(name).Funcs(viewHelpers())
			t, e = t.Parse(string(b))
			if e != nil {
				r := fmt.Sprintf("Could not parse: %v => %v", name, e)
				panic(r)
			}
			t.Execute(&buffer, data)
			return template.HTML(buffer.String())
		},
		"view": func() template.HTML {
			s := strings.SplitN(name, "/", 2)
			if len(s) == 2 {
				return template.HTML(s[0] + "_" + s[1])
			}
			return template.HTML("")
		},
	}
}

// Returns all the helpers available to any view. We have the following
// helpers: fmtDate and inc. The inc helper just increases the given integer
// value by one. The fmtDate helper executes the fmtDate function.
func viewHelpers() template.FuncMap {
	return template.FuncMap{
		"fmtDate": fmtDate,
		"inc": func(n int) int {
			return n + 1
		},
		"noescape": func(str string) template.HTML {
			return template.HTML(str)
		},
	}
}

// Returns a string with the given time formatted as expected by the view.
func fmtDate(t time.Time) string {
	return fmt.Sprintf("%02d/%02d/%04d", t.Day(), t.Month(), t.Year())
}
<commit_msg>lib: we don't need the majority of attributes from the ViewData struct.<commit_after>// Copyright (C) 2014 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package lib

import (
	"bytes"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"path"
	"strings"
	"time"
)

var (
	// The directory where all the views are being stored.
	ViewsDir = "views"
)

const (
	// The path to the layout file.
	layout = "application/layout"

	// The extension of views.
	viewsExt = "tpl"
)

// This struct holds all the data that can be passed to a view.
type ViewData struct {
	// The name of the javascript file to be used.
	JS string

	// The error message.
	Error string
}

// Returns the path to be used to open the view with the given name.
func view(name string) string {
	return path.Join(ViewsDir, name+"."+viewsExt)
}

// Render the view with the given name after evaluating the passed data. The
// rendered view will be written to the given writer.
func Render(res http.ResponseWriter, name string, data interface{}) {
	b, e := ioutil.ReadFile(view(layout))
	if e != nil {
		panic("Could not read layout file!")
	}
	t, e := template.New("l").Funcs(layoutHelpers(name, data)).Parse(string(b))
	if e != nil {
		panic("Could not parse layout file!")
	}
	t.Execute(res, data)
}

// Returns all the helpers used by the layout template. Right now only the
// "yield" helpers has been implemented.
func layoutHelpers(name string, data interface{}) template.FuncMap {
	return template.FuncMap{
		"yield": func() template.HTML {
			var buffer bytes.Buffer

			b, e := ioutil.ReadFile(view(name))
			if e != nil {
				r := fmt.Sprintf("Could not read: %v => %v", name, e)
				panic(r)
			}
			t := template.New(name).Funcs(viewHelpers())
			t, e = t.Parse(string(b))
			if e != nil {
				r := fmt.Sprintf("Could not parse: %v => %v", name, e)
				panic(r)
			}
			t.Execute(&buffer, data)
			return template.HTML(buffer.String())
		},
		"view": func() template.HTML {
			s := strings.SplitN(name, "/", 2)
			if len(s) == 2 {
				return template.HTML(s[0] + "_" + s[1])
			}
			return template.HTML("")
		},
	}
}

// Returns all the helpers available to any view. We have the following
// helpers: fmtDate and inc. The inc helper just increases the given integer
// value by one. The fmtDate helper executes the fmtDate function.
func viewHelpers() template.FuncMap {
	return template.FuncMap{
		"fmtDate": fmtDate,
		"inc": func(n int) int {
			return n + 1
		},
		"noescape": func(str string) template.HTML {
			return template.HTML(str)
		},
	}
}

// Returns a string with the given time formatted as expected by the view.
func fmtDate(t time.Time) string {
	return fmt.Sprintf("%02d/%02d/%04d", t.Day(), t.Month(), t.Year())
}
<|endoftext|>


<commit_before>package cluster

import (
	"net"
	"sync"

	"gopkg.in/fatih/pool.v2"
)

type clientPool struct {
	mu   sync.RWMutex
	pool map[uint64]pool.Pool
}

func newClientPool() *clientPool {
	return &clientPool{
		pool: make(map[uint64]pool.Pool),
	}
}

func (c *clientPool) setPool(nodeID uint64, p pool.Pool) {
	c.mu.Lock()
	c.pool[nodeID] = p
	c.mu.Unlock()
}

func (c *clientPool) getPool(nodeID uint64) (pool.Pool, bool) {
	c.mu.Lock()
	p, ok := c.pool[nodeID]
	c.mu.Unlock()
	return p, ok
}

func (c *clientPool) size() int {
	c.mu.RLock()
	var size int
	for _, p := range c.pool {
		size += p.Len()
	}
	c.mu.RUnlock()
	return size
}

func (c *clientPool) conn(nodeID uint64) (net.Conn, error) {
	c.mu.Lock()
	conn, err := c.pool[nodeID].Get()
	c.mu.Unlock()
	return conn, err
}

func (c *clientPool) close() {
	c.mu.Lock()
	for _, p := range c.pool {
		p.Close()
	}
	c.mu.Unlock()
}
<commit_msg>Use read locks instead of write lock for connection pools checkout<commit_after>package cluster

import (
	"net"
	"sync"

	"gopkg.in/fatih/pool.v2"
)

type clientPool struct {
	mu   sync.RWMutex
	pool map[uint64]pool.Pool
}

func newClientPool() *clientPool {
	return &clientPool{
		pool: make(map[uint64]pool.Pool),
	}
}

func (c *clientPool) setPool(nodeID uint64, p pool.Pool) {
	c.mu.Lock()
	c.pool[nodeID] = p
	c.mu.Unlock()
}

func (c *clientPool) getPool(nodeID uint64) (pool.Pool, bool) {
	c.mu.RLock()
	p, ok := c.pool[nodeID]
	c.mu.RUnlock()
	return p, ok
}

func (c *clientPool) size() int {
	c.mu.RLock()
	var size int
	for _, p := range c.pool {
		size += p.Len()
	}
	c.mu.RUnlock()
	return size
}

func (c *clientPool) conn(nodeID uint64) (net.Conn, error) {
	c.mu.RLock()
	conn, err := c.pool[nodeID].Get()
	c.mu.RUnlock()
	return conn, err
}

func (c *clientPool) close() {
	c.mu.Lock()
	for _, p := range c.pool {
		p.Close()
	}
	c.mu.Unlock()
}
<|endoftext|>


<commit_before>package controller

import (
	"errors"
	"github.com/crawlerclub/x/crawler"
	"github.com/crawlerclub/x/store"
	"github.com/crawlerclub/x/types"
	"github.com/golang/glog"
	"github.com/syndtr/goleveldb/leveldb/util"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ErrNotInited      = errors.New("controller/controller.go call Controller.Init first")
	ErrWorkerCount    = errors.New("controller/controller.go worker count must between 0 and 1000")
	ErrNilCrawlerItem = errors.New("controller/controller.go CrawlerItem is nil")
	ErrNamesNotSame   = errors.New("controller/controller.go CrawlerNames are not the same")
	ErrNameTaken      = errors.New("controller/controller.go CrawlerName already taken")
	ErrNoName         = errors.New("controller/controller.go no CrawlerName")
)

var StoreNames = []string{"crawler", "seed", "running", "crontab"}

type Controller struct {
	Crawlers    map[string]crawler.Crawler
	Schduler    CrawlerScheduler
	Stores      map[string]*store.LevelStore
	WorkerCount int

	workDir  string
	isInited bool
}

func timeStr(t int64) string {
	return time.Unix(t, 0).Format("20060102030405")
}

func (self *Controller) Init(dir string, wc int) error {
	glog.Info("call Init")
	if wc <= 0 || wc > 1000 {
		return ErrWorkerCount
	}
	self.workDir = dir
	self.WorkerCount = wc
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	self.Stores = make(map[string]*store.LevelStore)
	for _, name := range StoreNames {
		self.Stores[name], err = store.NewLevelStore(dir + "/db/" + name)
		if err != nil {
			return err
		}
	}

	err = self.initCrawlersFromDB()
	if err != nil {
		return err
	}
	self.isInited = true
	return nil
}

func (self *Controller) runCrawler(item *types.CrawlerItem) error {
	glog.Info("call runCrawler: ", item.CrawlerName)
	var crawler crawler.Crawler
	crawler.Conf = &item.Conf
	err := crawler.InitEs()
	if err != nil {
		//return err
		glog.Error(err)
	}
	dir := self.workDir + "/queue"
	err = crawler.InitTaskQueue(dir)
	if err != nil {
		return err
	}
	if _, ok := self.Crawlers[item.CrawlerName]; ok {
		self.Schduler.Remove(item.CrawlerName)
	}
	self.Crawlers[item.CrawlerName] = crawler
	var sitem Item
	sitem.CrawlerName = item.CrawlerName
	sitem.Weight = item.Weight
	glog.Info("SortedList add ", sitem)
	err = self.Schduler.Insert(sitem)
	if err != nil {
		return err
	}
	for _, url := range item.Conf.StartUrls {
		task := types.Task{
			CrawlerName: item.Conf.CrawlerName,
			ParserName:  item.Conf.StartParserName,
			IsSeedUrl:   true,
			Url:         url,
		}
		if ret, _ := self.Stores["seed"].Has(task.Id()); !ret {
			// enqueue new start urls
			if _, err = crawler.TaskQueue.EnqueueObject(task); err != nil {
				return err
			}
		}
		value, _ := store.ObjectToBytes(task)
		self.Stores["seed"].Put(task.Id(), value) // update seed
	}
	return nil
}

func (self *Controller) initCrawlersFromDB() error {
	glog.Info("call initCrawlersFromDB")
	self.Crawlers = make(map[string]crawler.Crawler)
	self.Schduler.Init()

	crawlerStore := self.Stores["crawler"]
	count := 0
	err := crawlerStore.ForEach(nil, func(key, value []byte) (bool, error) {
		count += 1
		var item types.CrawlerItem
		e := store.BytesToObject(value, &item)
		if e != nil {
			return false, e
		}
		e = self.runCrawler(&item)
		if e != nil {
			return false, nil
		}
		return true, nil
	})
	if err != nil {
		return err
	}
	glog.Info("loaded ", count, " crawler items")
	return nil
}

func (self *Controller) CloseCrawler(name string) error {
	err := self.Schduler.Remove(name)
	if err != nil && err != ErrNameNotFound {
		return err
	}
	if crawler, ok := self.Crawlers[name]; ok {
		crawler.Close()
	}
	delete(self.Crawlers, name)
	return nil
}

func (self *Controller) DelCrawler(name string) error {
	err := self.CloseCrawler(name)
	if err != nil {
		return err
	}
	return self.Stores["crawler"].Delete(name)
}

func (self *Controller) UpdateCrawler(item *types.CrawlerItem, isNew bool) error {
	if item == nil {
		return ErrNilCrawlerItem
	}
	ok, err := item.Conf.IsValid()
	if !ok {
		return err
	}
	if item.CrawlerName != item.Conf.CrawlerName {
		return ErrNamesNotSame
	}
	has, err := self.Stores["crawler"].Has(item.CrawlerName)
	if err != nil {
		return err
	}
	if isNew && has {
		return ErrNameTaken
	}
	if !isNew && !has {
		return ErrNoName
	}
	value, err := store.ObjectToBytes(item)
	if err != nil {
		return err
	}
	err = self.Stores["crawler"].Put(item.CrawlerName, value)
	if err != nil {
		return err
	}
	if item.Status == "enabled" && item.Weight > 0 {
		err = self.runCrawler(item)
		if err != nil {
			return err
		}
	} else {
	}
	return nil
}

func (self *Controller) Finish() {
	if !self.isInited {
		return
	}
	for _, v := range self.Stores {
		v.Close()
	}
}

func (self *Controller) enqueueTask(wg *sync.WaitGroup, exitCh chan int, name string) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	defer wg.Done()
	glog.Info("start ", name, " worker")
	defer glog.Info("exit ", name, " worker")
	for {
		select {
		case <-exitCh:
			return
		default:
			glog.Info("begin ", name)
			now := time.Now().Format("20060102030405")
			self.Stores[name].ForEach(&util.Range{Limit: []byte(now)},
				func(key, value []byte) (bool, error) {
					var task types.Task
					store.BytesToObject(value, &task)
					if crawler, ok := self.Crawlers[task.CrawlerName]; ok {
						crawler.TaskQueue.EnqueueObject(task)
					}
					self.Stores[name].Delete(string(key))
					return true, nil
				})
			// every 5 seconds
			time.Sleep(5 * time.Second)
		}
	}
}

func (self *Controller) cron(wg *sync.WaitGroup, exitCh chan int) {
	self.enqueueTask(wg, exitCh, "crontab")
}

func (self *Controller) retry(wg *sync.WaitGroup, exitCh chan int) {
	self.enqueueTask(wg, exitCh, "running")
}

func (self *Controller) startWorker(worker int, wg *sync.WaitGroup, exitCh chan int) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	defer wg.Done()
	glog.Info("start crawler worker: ", worker)
	defer glog.Info("exit crawler worker: ", worker)
	for {
		select {
		case <-exitCh:
			return
		default:
			glog.Info("Work on next task! worker: ", worker)
			name, err := self.Schduler.WeightedChoice()
			if err != nil {
				glog.Error(err)
				time.Sleep(10 * time.Second)
				continue
			}
			glog.Info("worker ", worker, " is working on ", name)
			if crawler, ok := self.Crawlers[name]; ok {
				item, err := crawler.TaskQueue.Dequeue()
				if err != nil {
					glog.Error(err)
					continue
				}
				var task types.Task
				err = item.ToObject(&task)
				if err != nil {
					glog.Error(err)
					continue
				}
				now := time.Now().Unix()
				key := timeStr(now+300) + "\t" + task.Id()
				value, _ := store.ObjectToBytes(task)
				self.Stores["running"].Put(key, value)

				glog.Info("process task:", task)
				tasks, items, err := crawler.Process(&task)
				if err != nil {
					glog.Error(err)
					continue
				}
				// remove task from Running
				self.Stores["running"].Delete(key)

				if parseConf, ok := crawler.Conf.ParseConfs[task.ParserName]; ok {
					if parseConf.RevisitInterval > 0 && task.IsSeedUrl {
						// add this task back to crontab
						task.LastAccessTime = now
						task.RevisitInterval = parseConf.RevisitInterval
						key = timeStr(now+task.RevisitInterval) + "\t" + task.Id()
						value, _ = store.ObjectToBytes(task)
						self.Stores["crontab"].Put(key, value)
					}
				}
				for _, t := range tasks {
					glog.Info("enqueue task:", t)
					// add SeedUrl to Seed
					if t.IsSeedUrl {
						value, _ = store.ObjectToBytes(t)
						self.Stores["seed"].Put(t.Id(), value)
					}
					crawler.TaskQueue.EnqueueObject(t)
				}
				for _, item := range items {
					crawler.Save(item)
				}
			} else {
				glog.Error("No crawler named: ", name)
			}
		}
	}
}

func (self *Controller) stop(sigs chan os.Signal, exitCh chan int) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	<-sigs
	glog.Info("receive stop signal")
	close(exitCh)
}

func (self *Controller) Run() {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	exitCh := make(chan int)
	sigs := make(chan os.Signal)
	var wg sync.WaitGroup
	for i := 0; i < self.WorkerCount; i++ {
		wg.Add(1)
		go self.startWorker(i, &wg, exitCh)
	}
	wg.Add(1)
	go self.cron(&wg, exitCh)
	wg.Add(1)
	go self.retry(&wg, exitCh)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go self.stop(sigs, exitCh)
	wg.Wait()
}
<commit_msg>fix bug, should close crawler after changing its status to disabled<commit_after>package controller

import (
	"errors"
	"github.com/crawlerclub/x/crawler"
	"github.com/crawlerclub/x/store"
	"github.com/crawlerclub/x/types"
	"github.com/golang/glog"
	"github.com/syndtr/goleveldb/leveldb/util"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	ErrNotInited      = errors.New("controller/controller.go call Controller.Init first")
	ErrWorkerCount    = errors.New("controller/controller.go worker count must between 0 and 1000")
	ErrNilCrawlerItem = errors.New("controller/controller.go CrawlerItem is nil")
	ErrNamesNotSame   = errors.New("controller/controller.go CrawlerNames are not the same")
	ErrNameTaken      = errors.New("controller/controller.go CrawlerName already taken")
	ErrNoName         = errors.New("controller/controller.go no CrawlerName")
)

var StoreNames = []string{"crawler", "seed", "running", "crontab"}

type Controller struct {
	Crawlers    map[string]crawler.Crawler
	Schduler    CrawlerScheduler
	Stores      map[string]*store.LevelStore
	WorkerCount int

	workDir  string
	isInited bool
}

func timeStr(t int64) string {
	return time.Unix(t, 0).Format("20060102030405")
}

func (self *Controller) Init(dir string, wc int) error {
	glog.Info("call Init")
	if wc <= 0 || wc > 1000 {
		return ErrWorkerCount
	}
	self.workDir = dir
	self.WorkerCount = wc
	_, err := os.Stat(dir)
	if os.IsNotExist(err) {
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			return err
		}
	}
	self.Stores = make(map[string]*store.LevelStore)
	for _, name := range StoreNames {
		self.Stores[name], err = store.NewLevelStore(dir + "/db/" + name)
		if err != nil {
			return err
		}
	}

	err = self.initCrawlersFromDB()
	if err != nil {
		return err
	}
	self.isInited = true
	return nil
}

func (self *Controller) runCrawler(item *types.CrawlerItem) error {
	glog.Info("call runCrawler: ", item.CrawlerName)
	var crawler crawler.Crawler
	crawler.Conf = &item.Conf
	err := crawler.InitEs()
	if err != nil {
		//return err
		glog.Error(err)
	}
	dir := self.workDir + "/queue"
	err = crawler.InitTaskQueue(dir)
	if err != nil {
		return err
	}
	if _, ok := self.Crawlers[item.CrawlerName]; ok {
		self.Schduler.Remove(item.CrawlerName)
	}
	self.Crawlers[item.CrawlerName] = crawler
	var sitem Item
	sitem.CrawlerName = item.CrawlerName
	sitem.Weight = item.Weight
	glog.Info("SortedList add ", sitem)
	err = self.Schduler.Insert(sitem)
	if err != nil {
		return err
	}
	for _, url := range item.Conf.StartUrls {
		task := types.Task{
			CrawlerName: item.Conf.CrawlerName,
			ParserName:  item.Conf.StartParserName,
			IsSeedUrl:   true,
			Url:         url,
		}
		if ret, _ := self.Stores["seed"].Has(task.Id()); !ret {
			// enqueue new start urls
			if _, err = crawler.TaskQueue.EnqueueObject(task); err != nil {
				return err
			}
		}
		value, _ := store.ObjectToBytes(task)
		self.Stores["seed"].Put(task.Id(), value) // update seed
	}
	return nil
}

func (self *Controller) initCrawlersFromDB() error {
	glog.Info("call initCrawlersFromDB")
	self.Crawlers = make(map[string]crawler.Crawler)
	self.Schduler.Init()

	crawlerStore := self.Stores["crawler"]
	count := 0
	err := crawlerStore.ForEach(nil, func(key, value []byte) (bool, error) {
		count += 1
		var item types.CrawlerItem
		e := store.BytesToObject(value, &item)
		if e != nil {
			return false, e
		}
		e = self.runCrawler(&item)
		if e != nil {
			return false, nil
		}
		return true, nil
	})
	if err != nil {
		return err
	}
	glog.Info("loaded ", count, " crawler items")
	return nil
}

func (self *Controller) CloseCrawler(name string) error {
	err := self.Schduler.Remove(name)
	if err != nil && err != ErrNameNotFound {
		return err
	}
	if crawler, ok := self.Crawlers[name]; ok {
		crawler.Close()
	}
	delete(self.Crawlers, name)
	return nil
}

func (self *Controller) DelCrawler(name string) error {
	err := self.CloseCrawler(name)
	if err != nil {
		return err
	}
	return self.Stores["crawler"].Delete(name)
}

func (self *Controller) UpdateCrawler(item *types.CrawlerItem, isNew bool) error {
	if item == nil {
		return ErrNilCrawlerItem
	}
	ok, err := item.Conf.IsValid()
	if !ok {
		return err
	}
	if item.CrawlerName != item.Conf.CrawlerName {
		return ErrNamesNotSame
	}
	has, err := self.Stores["crawler"].Has(item.CrawlerName)
	if err != nil {
		return err
	}
	if isNew && has {
		return ErrNameTaken
	}
	if !isNew && !has {
		return ErrNoName
	}
	now := time.Now().Unix()
	if isNew {
		item.CreateTime = now
	} else {
		item.ModifyTime = now
	}
	value, err := store.ObjectToBytes(item)
	if err != nil {
		return err
	}
	err = self.Stores["crawler"].Put(item.CrawlerName, value)
	if err != nil {
		return err
	}
	if item.Status == "enabled" && item.Weight > 0 {
		err = self.runCrawler(item)
	} else {
		err = self.CloseCrawler(item.CrawlerName)
	}
	return err
}

func (self *Controller) Finish() {
	if !self.isInited {
		return
	}
	for _, v := range self.Stores {
		v.Close()
	}
}

func (self *Controller) enqueueTask(wg *sync.WaitGroup, exitCh chan int, name string) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	defer wg.Done()
	glog.Info("start ", name, " worker")
	defer glog.Info("exit ", name, " worker")
	for {
		select {
		case <-exitCh:
			return
		default:
			glog.Info("begin ", name)
			now := time.Now().Format("20060102030405")
			self.Stores[name].ForEach(&util.Range{Limit: []byte(now)},
				func(key, value []byte) (bool, error) {
					var task types.Task
					store.BytesToObject(value, &task)
					if crawler, ok := self.Crawlers[task.CrawlerName]; ok {
						crawler.TaskQueue.EnqueueObject(task)
					}
					self.Stores[name].Delete(string(key))
					return true, nil
				})
			// every 5 seconds
			time.Sleep(5 * time.Second)
		}
	}
}

func (self *Controller) cron(wg *sync.WaitGroup, exitCh chan int) {
	self.enqueueTask(wg, exitCh, "crontab")
}

func (self *Controller) retry(wg *sync.WaitGroup, exitCh chan int) {
	self.enqueueTask(wg, exitCh, "running")
}

func (self *Controller) startWorker(worker int, wg *sync.WaitGroup, exitCh chan int) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	defer wg.Done()
	glog.Info("start crawler worker: ", worker)
	defer glog.Info("exit crawler worker: ", worker)
	for {
		select {
		case <-exitCh:
			return
		default:
			glog.Info("Work on next task! worker: ", worker)
			name, err := self.Schduler.WeightedChoice()
			if err != nil {
				glog.Error(err)
				time.Sleep(10 * time.Second)
				continue
			}
			glog.Info("worker ", worker, " is working on ", name)
			if crawler, ok := self.Crawlers[name]; ok {
				item, err := crawler.TaskQueue.Dequeue()
				if err != nil {
					glog.Error(err)
					continue
				}
				var task types.Task
				err = item.ToObject(&task)
				if err != nil {
					glog.Error(err)
					continue
				}
				now := time.Now().Unix()
				key := timeStr(now+300) + "\t" + task.Id()
				value, _ := store.ObjectToBytes(task)
				self.Stores["running"].Put(key, value)

				glog.Info("process task:", task)
				tasks, items, err := crawler.Process(&task)
				if err != nil {
					glog.Error(err)
					continue
				}
				// remove task from Running
				self.Stores["running"].Delete(key)

				if parseConf, ok := crawler.Conf.ParseConfs[task.ParserName]; ok {
					if parseConf.RevisitInterval > 0 && task.IsSeedUrl {
						// add this task back to crontab
						task.LastAccessTime = now
						task.RevisitInterval = parseConf.RevisitInterval
						key = timeStr(now+task.RevisitInterval) + "\t" + task.Id()
						value, _ = store.ObjectToBytes(task)
						self.Stores["crontab"].Put(key, value)
					}
				}
				for _, t := range tasks {
					glog.Info("enqueue task:", t)
					// add SeedUrl to Seed
					if t.IsSeedUrl {
						value, _ = store.ObjectToBytes(t)
						self.Stores["seed"].Put(t.Id(), value)
					}
					crawler.TaskQueue.EnqueueObject(t)
				}
				for _, item := range items {
					crawler.Save(item)
				}
			} else {
				glog.Error("No crawler named: ", name)
			}
		}
	}
}

func (self *Controller) stop(sigs chan os.Signal, exitCh chan int) {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	<-sigs
	glog.Info("receive stop signal")
	close(exitCh)
}

func (self *Controller) Run() {
	if !self.isInited {
		glog.Error(ErrNotInited)
		return
	}
	exitCh := make(chan int)
	sigs := make(chan os.Signal)
	var wg sync.WaitGroup
	for i := 0; i < self.WorkerCount; i++ {
		wg.Add(1)
		go self.startWorker(i, &wg, exitCh)
	}
	wg.Add(1)
	go self.cron(&wg, exitCh)
	wg.Add(1)
	go self.retry(&wg, exitCh)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	go self.stop(sigs, exitCh)
	wg.Wait()
}
<|endoftext|>


<commit_before>package controller

import (
	// "fmt"
	"github.com/eaciit/colony-core/v0"
	"github.com/eaciit/colony-manager/helper"
	// "github.com/eaciit/dbox"
	"github.com/eaciit/knot/knot.v1"
	// "github.com/eaciit/sedotan/sedotan.v1"
	"github.com/eaciit/sedotan/sedotan.v1/webapps/modules"
	"github.com/eaciit/toolkit"
	"path"
	f "path/filepath"
	"reflect"
	"strings"
	// "strconv"
	// "time"
)

type WebGrabberController struct {
	App
}

func CreateWebGrabberController(s *knot.Server) *WebGrabberController {
	var controller = new(WebGrabberController)
	controller.Server = s
	return controller
}

func (w *WebGrabberController) PrepareHistoryPath() {
	modules.HistoryPath = AppBasePath + toolkit.PathSeparator + f.Join("config", "webgrabber", "history") + toolkit.PathSeparator
	modules.HistoryRecPath = AppBasePath + toolkit.PathSeparator + f.Join("config", "webgrabber", "historyrec") + toolkit.PathSeparator
}

func (w *WebGrabberController) GetScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	cursor, err := colonycore.Find(new(colonycore.WebGrabber), nil)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	data := []colonycore.WebGrabber{}
	err = cursor.Fetch(&data, 0, false)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}
	defer cursor.Close()

	return helper.CreateResult(true, data, "")
}

func (w *WebGrabberController) SelectScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, payload, "")
}

func (w *WebGrabberController) SaveScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Delete(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Save(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, payload, "")
}

func (w *WebGrabberController) FetchContent(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	param := toolkit.M{} //.Set("formvalues", payload.Parameter)
	res, err := toolkit.HttpCall(payload.URL, payload.CallType, nil, param)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	data := toolkit.HttpContentString(res)
	return helper.CreateResult(true, data, "")
}

func (w *WebGrabberController) StartService(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err, isRun := modules.Process([]interface{}{o})
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(isRun, nil, "")
}

func (w *WebGrabberController) StopService(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err, isRun := modules.StopProcess([]interface{}{o})
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(isRun, nil, "")
}

func (w *WebGrabberController) Stat(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	grabStatus := modules.NewGrabService().CheckStat([]interface{}{o})
	return helper.CreateResult(true, grabStatus, "")
}

func (w *WebGrabberController) GetHistory(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	module := modules.NewHistory(payload.ID)
	history := module.OpenHistory()

	if reflect.ValueOf(history).Kind() == reflect.String {
		if strings.Contains(history.(string), "Cannot Open File") {
			return helper.CreateResult(false, nil, "Cannot Open File")
		}
	}

	return helper.CreateResult(true, history, "")
}

func (w *WebGrabberController) GetFetchedData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := struct {
		Driver     string
		Host       string
		Database   string
		Collection string
		Username   string
		Password   string
	}{}
	err := r.GetPayload(&payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	var data []toolkit.M

	if payload.Driver == "csv" {
		config := toolkit.M{"useheader": true, "delimiter": ","}
		query := helper.Query("csv", payload.Host, "", "", "", config)
		data, err = query.SelectAll("")
	} else {
		query := helper.Query("mongo", payload.Host, payload.Database, payload.Username, payload.Password)
		data, err = query.SelectAll(payload.Collection)
	}

	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, data, "")
}

func (w *WebGrabberController) GetLog(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := struct {
		ID   string `json:"_id"`
		Date string `json:"date"`
	}{}
	err := r.GetPayload(&payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	wg := new(colonycore.WebGrabber)
	err = colonycore.Get(wg, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(wg)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	history := modules.NewHistory(payload.ID)
	logs := history.GetLog([]interface{}{o}, payload.Date)

	return helper.CreateResult(true, logs, "")
}

func (w *WebGrabberController) RemoveGrabber(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Delete(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, nil, "")
}

func (w *WebGrabberController) InsertSampleData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	wg := new(colonycore.WebGrabber)
	wg.CallType = "POST"
	wg.DataSettings = []*colonycore.DataSetting{
		&colonycore.DataSetting{
			ColumnSettings: []*colonycore.ColumnSetting{
				&colonycore.ColumnSetting{Alias: "Contract", Index: 0, Selector: "td:nth-child(1)"},
				&colonycore.ColumnSetting{Alias: "Open", Index: 0, Selector: "td:nth-child(2)"},
				&colonycore.ColumnSetting{Alias: "High", Index: 0, Selector: "td:nth-child(3)"},
				&colonycore.ColumnSetting{Alias: "Low", Index: 0, Selector: "td:nth-child(4)"},
				&colonycore.ColumnSetting{Alias: "Close", Index: 0, Selector: "td:nth-child(5)"},
				&colonycore.ColumnSetting{Alias: "Prev Settle", Index: 0, Selector: "td:nth-child(6)"},
				&colonycore.ColumnSetting{Alias: "Prev Settle", Index: 0, Selector: "td:nth-child(7)"},
				&colonycore.ColumnSetting{Alias: "Settle", Index: 0, Selector: "td:nth-child(8)"},
				&colonycore.ColumnSetting{Alias: "Chg", Index: 0, Selector: "td:nth-child(9)"},
				&colonycore.ColumnSetting{Alias: "Volume", Index: 0, Selector: "td:nth-child(10)"},
				&colonycore.ColumnSetting{Alias: "OI", Index: 0, Selector: "td:nth-child(11)"},
				&colonycore.ColumnSetting{Alias: "OI Chg", Index: 0, Selector: "td:nth-child(12)"},
				&colonycore.ColumnSetting{Alias: "Turnover", Index: 0, Selector: "td:nth-child(13)"},
			},
			ConnectionInfo: &colonycore.ConnectionInfo{
				Collection: "irondcecom",
				Database:   "valegrab",
				Host:       "localhost:27017",
			},
			DestinationType: "mongo",
			Name:            "GoldTab01",
			RowSelector:     "table .table tbody tr",
			RowDeleteCondition: toolkit.M{
				"$or": []toolkit.M{
					toolkit.M{"Contract": "Contract"},
					toolkit.M{"Contract": "Iron Ore Subtotal"},
					toolkit.M{"Contract": "Total"},
				},
			},
		},
	}
	wg.GrabConfiguration = toolkit.M{
		"data": toolkit.M{
			"Pu00231_Input.trade_date": 2.0151214e+07,
			"Pu00231_Input.trade_type": 0,
			"Pu00231_Input.variety":    "i",
			"Submit":                   "Go",
			"action":                   "Pu00231_result",
		},
	}
	wg.GrabInterval = 20
	wg.IntervalType = "seconds"
	wg.LogConfiguration = &colonycore.LogConfiguration{
		FileName:    "LOG-GRABDCE",
		FilePattern: "20060102",
		LogPath:     path.Join(AppBasePath, "config", "webgrabber", "log"),
	}
	wg.ID = "irondcecomcn"
	wg.IDBackup = "irondcecomcn"
	wg.SourceType = "SourceType_Http"
	wg.TimeoutInterval = 5
	wg.URL = "http://www.dce.com.cn/PublicWeb/MainServlet"

	colonycore.Save(wg)

	return helper.CreateResult(true, wg, "")
}
<commit_msg>add function findwebgrabber<commit_after>package controller

import (
	// "fmt"
	"github.com/eaciit/colony-core/v0"
	"github.com/eaciit/colony-manager/helper"
	// "github.com/eaciit/dbox"
	"github.com/eaciit/knot/knot.v1"
	// "github.com/eaciit/sedotan/sedotan.v1"
	"github.com/eaciit/sedotan/sedotan.v1/webapps/modules"
	"github.com/eaciit/toolkit"
	"path"
	f "path/filepath"
	"reflect"
	"strings"
	// "strconv"
	// "time"
)

type WebGrabberController struct {
	App
}

func CreateWebGrabberController(s *knot.Server) *WebGrabberController {
	var controller = new(WebGrabberController)
	controller.Server = s
	return controller
}

func (w *WebGrabberController) PrepareHistoryPath() {
	modules.HistoryPath = AppBasePath + toolkit.PathSeparator + f.Join("config", "webgrabber", "history") + toolkit.PathSeparator
	modules.HistoryRecPath = AppBasePath + toolkit.PathSeparator + f.Join("config", "webgrabber", "historyrec") + toolkit.PathSeparator
}

func (w *WebGrabberController) GetScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	cursor, err := colonycore.Find(new(colonycore.WebGrabber), nil)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	data := []colonycore.WebGrabber{}
	err = cursor.Fetch(&data, 0, false)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}
	defer cursor.Close()

	return helper.CreateResult(true, data, "")
}

func (d *WebGrabberController) FindDataGrabber(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	payload := map[string]string{"inputText":"GRAB_TEST","inputRequest":"","inputType":""}
	//~ payload := map[string]interface{}{}

	//~ err := r.GetPayload(&payload)
	//~ if err != nil {
		//~ return helper.CreateResult(false, nil, err.Error())
	//~ }

	text := payload["inputText"]
	req  := payload["inputRequest"]
	tipe := payload["inputType"]

	// == bug, cant find if autocomplite, just full text can be get result
	var query *dbox.Filter
	if text != "" {
		query = dbox.Or(dbox.Eq("_id",text),dbox.Eq("calltype",text),dbox.Eq("sourcetype",text),dbox.Eq("intervaltype",text),dbox.Eq("grabinterval",text),dbox.Eq("timeoutinterval",text))
	}

	if req != "" {
		query = dbox.And(query, dbox.Eq("calltype",req))
	}

	if tipe != "" {
		query = dbox.And(query, dbox.Eq("sourcetype",tipe))
	}

	data := []colonycore.WebGrabber{}
	cursor, err := colonycore.Find(new(colonycore.WebGrabber), query)
	cursor.Fetch(&data, 0, false)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}
	defer cursor.Close()

	return helper.CreateResult(true, payload, "")
}

func (w *WebGrabberController) SelectScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, payload, "")
}

func (w *WebGrabberController) SaveScrapperData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Delete(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Save(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, payload, "")
}

func (w *WebGrabberController) FetchContent(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	param := toolkit.M{} //.Set("formvalues", payload.Parameter)
	res, err := toolkit.HttpCall(payload.URL, payload.CallType, nil, param)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	data := toolkit.HttpContentString(res)
	return helper.CreateResult(true, data, "")
}

func (w *WebGrabberController) StartService(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err, isRun := modules.Process([]interface{}{o})
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(isRun, nil, "")
}

func (w *WebGrabberController) StopService(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err, isRun := modules.StopProcess([]interface{}{o})
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(isRun, nil, "")
}

func (w *WebGrabberController) Stat(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Get(payload, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	grabStatus := modules.NewGrabService().CheckStat([]interface{}{o})
	return helper.CreateResult(true, grabStatus, "")
}

func (w *WebGrabberController) GetHistory(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	module := modules.NewHistory(payload.ID)
	history := module.OpenHistory()

	if reflect.ValueOf(history).Kind() == reflect.String {
		if strings.Contains(history.(string), "Cannot Open File") {
			return helper.CreateResult(false, nil, "Cannot Open File")
		}
	}

	return helper.CreateResult(true, history, "")
}

func (w *WebGrabberController) GetFetchedData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := struct {
		Driver     string
		Host       string
		Database   string
		Collection string
		Username   string
		Password   string
	}{}
	err := r.GetPayload(&payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	var data []toolkit.M

	if payload.Driver == "csv" {
		config := toolkit.M{"useheader": true, "delimiter": ","}
		query := helper.Query("csv", payload.Host, "", "", "", config)
		data, err = query.SelectAll("")
	} else {
		query := helper.Query("mongo", payload.Host, payload.Database, payload.Username, payload.Password)
		data, err = query.SelectAll(payload.Collection)
	}

	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, data, "")
}

func (w *WebGrabberController) GetLog(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := struct {
		ID   string `json:"_id"`
		Date string `json:"date"`
	}{}
	err := r.GetPayload(&payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	wg := new(colonycore.WebGrabber)
	err = colonycore.Get(wg, payload.ID)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	o, err := toolkit.ToM(wg)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	history := modules.NewHistory(payload.ID)
	logs := history.GetLog([]interface{}{o}, payload.Date)

	return helper.CreateResult(true, logs, "")
}

func (w *WebGrabberController) RemoveGrabber(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	w.PrepareHistoryPath()

	payload := new(colonycore.WebGrabber)
	err := r.GetPayload(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	err = colonycore.Delete(payload)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}

	return helper.CreateResult(true, nil, "")
}

func (w *WebGrabberController) InsertSampleData(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	wg := new(colonycore.WebGrabber)
	wg.CallType = "POST"
	wg.DataSettings = []*colonycore.DataSetting{
		&colonycore.DataSetting{
			ColumnSettings: []*colonycore.ColumnSetting{
				&colonycore.ColumnSetting{Alias: "Contract", Index: 0, Selector: "td:nth-child(1)"},
				&colonycore.ColumnSetting{Alias: "Open", Index: 0, Selector: "td:nth-child(2)"},
				&colonycore.ColumnSetting{Alias: "High", Index: 0, Selector: "td:nth-child(3)"},
				&colonycore.ColumnSetting{Alias: "Low", Index: 0, Selector: "td:nth-child(4)"},
				&colonycore.ColumnSetting{Alias: "Close", Index: 0, Selector: "td:nth-child(5)"},
				&colonycore.ColumnSetting{Alias: "Prev Settle", Index: 0, Selector: "td:nth-child(6)"},
				&colonycore.ColumnSetting{Alias: "Prev Settle", Index: 0, Selector: "td:nth-child(7)"},
				&colonycore.ColumnSetting{Alias: "Settle", Index: 0, Selector: "td:nth-child(8)"},
				&colonycore.ColumnSetting{Alias: "Chg", Index: 0, Selector: "td:nth-child(9)"},
				&colonycore.ColumnSetting{Alias: "Volume", Index: 0, Selector: "td:nth-child(10)"},
				&colonycore.ColumnSetting{Alias: "OI", Index: 0, Selector: "td:nth-child(11)"},
				&colonycore.ColumnSetting{Alias: "OI Chg", Index: 0, Selector: "td:nth-child(12)"},
				&colonycore.ColumnSetting{Alias: "Turnover", Index: 0, Selector: "td:nth-child(13)"},
			},
			ConnectionInfo: &colonycore.ConnectionInfo{
				Collection: "irondcecom",
				Database:   "valegrab",
				Host:       "localhost:27017",
			},
			DestinationType: "mongo",
			Name:            "GoldTab01",
			RowSelector:     "table .table tbody tr",
			RowDeleteCondition: toolkit.M{
				"$or": []toolkit.M{
					toolkit.M{"Contract": "Contract"},
					toolkit.M{"Contract": "Iron Ore Subtotal"},
					toolkit.M{"Contract": "Total"},
				},
			},
		},
	}
	wg.GrabConfiguration = toolkit.M{
		"data": toolkit.M{
			"Pu00231_Input.trade_date": 2.0151214e+07,
			"Pu00231_Input.trade_type": 0,
			"Pu00231_Input.variety":    "i",
			"Submit":                   "Go",
			"action":                   "Pu00231_result",
		},
	}
	wg.GrabInterval = 20
	wg.IntervalType = "seconds"
	wg.LogConfiguration = &colonycore.LogConfiguration{
		FileName:    "LOG-GRABDCE",
		FilePattern: "20060102",
		LogPath:     path.Join(AppBasePath, "config", "webgrabber", "log"),
	}
	wg.ID = "irondcecomcn"
	wg.IDBackup = "irondcecomcn"
	wg.SourceType = "SourceType_Http"
	wg.TimeoutInterval = 5
	wg.URL = "http://www.dce.com.cn/PublicWeb/MainServlet"

	colonycore.Save(wg)

	return helper.CreateResult(true, wg, "")
}
<|endoftext|>


<commit_before>/*
Copyright 2011 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 main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/client"
)

var (
	flagLoop    = flag.Bool("loop", false, "sync in a loop once done; requires --removesrc")
	flagVerbose = flag.Bool("verbose", false, "be verbose")

	flagSrc      = flag.String("src", "", "Source blobserver prefix (generally a mirrored queue partition)")
	flagSrcPass  = flag.String("srcpassword", "", "Source password")
	flagDest     = flag.String("dest", "", "Destination blobserver, or 'stdout' to just enumerate the --src blobs to stdout")
	flagDestPass = flag.String("destpassword", "", "Destination password")

	flagRemoveSource = flag.Bool("removesrc", false,
		"remove each blob from the source after syncing to the destination; for queue processing")
)

type SyncStats struct {
	BlobsCopied int
	BytesCopied int64
	ErrorCount  int
}

func usage(err string) {
	if err != "" {
		fmt.Fprintf(os.Stderr, "Error: %s\n\nUsage:\n", err)
	}
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	flag.Parse()

	if *flagSrc == "" {
		usage("No --src specified.")
	}
	if *flagDest == "" {
		usage("No --dest specified.")
	}
	if *flagLoop && !*flagRemoveSource {
		usage("Can't use --loop without --removesrc")
	}

	// TODO(mpl): adapt to the userpass scheme once it has settled
	sc := client.New(*flagSrc)
	sc.SetupAuth()
	dc := client.New(*flagDest)
	dc.SetupAuth()

	var logger *log.Logger = nil
	if *flagVerbose {
		logger = log.New(os.Stderr, "", 0)
	}
	sc.SetLogger(logger)
	dc.SetLogger(logger)

	passNum := 0
	for {
		passNum++
		stats, err := doPass(sc, dc, passNum)
		if err != nil {
			log.Fatalf("sync failed: %v", err)
		}
		if *flagVerbose {
			log.Printf("sync stats - pass: %d, blobs: %d, bytes %d\n", passNum, stats.BlobsCopied, stats.BytesCopied)
		}
		if !*flagLoop {
			break
		}
	}
}

func doPass(sc, dc *client.Client, passNum int) (stats SyncStats, retErr error) {
	srcBlobs := make(chan blobref.SizedBlobRef, 100)
	destBlobs := make(chan blobref.SizedBlobRef, 100)
	srcErr := make(chan error)
	destErr := make(chan error)

	go func() {
		srcErr <- sc.SimpleEnumerateBlobs(srcBlobs)
	}()
	checkSourceError := func() {
		if err := <-srcErr; err != nil {
			retErr = errors.New(fmt.Sprintf("Enumerate error from source: %v", err))
		}
	}

	if *flagDest == "stdout" {
		for sb := range srcBlobs {
			fmt.Printf("%s %d\n", sb.BlobRef, sb.Size)
		}
		checkSourceError()
		return
	}

	go func() {
		destErr <- dc.SimpleEnumerateBlobs(destBlobs)
	}()
	checkDestError := func() {
		if err := <-destErr; err != nil {
			retErr = errors.New(fmt.Sprintf("Enumerate error from destination: %v", err))
		}
	}

	destNotHaveBlobs := make(chan blobref.SizedBlobRef, 100)
	go client.ListMissingDestinationBlobs(destNotHaveBlobs, srcBlobs, destBlobs)
	for sb := range destNotHaveBlobs {
		fmt.Printf("Destination needs blob: %s\n", sb)

		blobReader, size, err := sc.FetchStreaming(sb.BlobRef)
		if err != nil {
			stats.ErrorCount++
			log.Printf("Error fetching %s: %v", sb.BlobRef, err)
			continue
		}
		if size != sb.Size {
			stats.ErrorCount++
			log.Printf("Source blobserver's enumerate size of %d for blob %s doesn't match its Get size of %d",
				sb.Size, sb.BlobRef, size)
			continue
		}
		uh := &client.UploadHandle{BlobRef: sb.BlobRef, Size: size, Contents: blobReader}
		pr, err := dc.Upload(uh)
		if err != nil {
			stats.ErrorCount++
			log.Printf("Upload of %s to destination blobserver failed: %v", sb.BlobRef, err)
			continue
		}
		if !pr.Skipped {
			stats.BlobsCopied++
			stats.BytesCopied += pr.Size
		}
		if *flagRemoveSource {
			if err = sc.RemoveBlob(sb.BlobRef); err != nil {
				stats.ErrorCount++
				log.Printf("Failed to delete %s from source: %v", sb.BlobRef, err)
			}
		}
	}

	checkSourceError()
	checkDestError()
	if retErr == nil && stats.ErrorCount > 0 {
		retErr = errors.New(fmt.Sprintf("%d errors during sync", stats.ErrorCount))
	}
	return stats, retErr
}
<commit_msg>Make camsync verbose more verbose; echo position every second.<commit_after>/*
Copyright 2011 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 main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"camlistore.org/pkg/blobref"
	"camlistore.org/pkg/client"
)

var (
	flagLoop    = flag.Bool("loop", false, "sync in a loop once done; requires --removesrc")
	flagVerbose = flag.Bool("verbose", false, "be verbose")

	flagSrc      = flag.String("src", "", "Source blobserver prefix (generally a mirrored queue partition)")
	flagSrcPass  = flag.String("srcpassword", "", "Source password")
	flagDest     = flag.String("dest", "", "Destination blobserver, or 'stdout' to just enumerate the --src blobs to stdout")
	flagDestPass = flag.String("destpassword", "", "Destination password")

	flagRemoveSource = flag.Bool("removesrc", false,
		"remove each blob from the source after syncing to the destination; for queue processing")
)

type SyncStats struct {
	BlobsCopied int
	BytesCopied int64
	ErrorCount  int
}

func usage(err string) {
	if err != "" {
		fmt.Fprintf(os.Stderr, "Error: %s\n\nUsage:\n", err)
	}
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	flag.Parse()

	if *flagSrc == "" {
		usage("No --src specified.")
	}
	if *flagDest == "" {
		usage("No --dest specified.")
	}
	if *flagLoop && !*flagRemoveSource {
		usage("Can't use --loop without --removesrc")
	}

	// TODO(mpl): adapt to the userpass scheme once it has settled
	sc := client.New(*flagSrc)
	sc.SetupAuth()
	dc := client.New(*flagDest)
	dc.SetupAuth()

	var logger *log.Logger = nil
	if *flagVerbose {
		logger = log.New(os.Stderr, "", 0)
	}
	sc.SetLogger(logger)
	dc.SetLogger(logger)

	passNum := 0
	for {
		passNum++
		stats, err := doPass(sc, dc, passNum)
		if err != nil {
			log.Fatalf("sync failed: %v", err)
		}
		if *flagVerbose {
			log.Printf("sync stats - pass: %d, blobs: %d, bytes %d\n", passNum, stats.BlobsCopied, stats.BytesCopied)
		}
		if !*flagLoop {
			break
		}
	}
}

func doPass(sc, dc *client.Client, passNum int) (stats SyncStats, retErr error) {
	srcBlobs := make(chan blobref.SizedBlobRef, 100)
	destBlobs := make(chan blobref.SizedBlobRef, 100)
	srcErr := make(chan error)
	destErr := make(chan error)

	go func() {
		srcErr <- sc.SimpleEnumerateBlobs(srcBlobs)
	}()
	checkSourceError := func() {
		if err := <-srcErr; err != nil {
			retErr = fmt.Errorf("Enumerate error from source: %v", err)
		}
	}

	if *flagDest == "stdout" {
		for sb := range srcBlobs {
			fmt.Printf("%s %d\n", sb.BlobRef, sb.Size)
		}
		checkSourceError()
		return
	}

	go func() {
		destErr <- dc.SimpleEnumerateBlobs(destBlobs)
	}()
	checkDestError := func() {
		if err := <-destErr; err != nil {
			retErr = errors.New(fmt.Sprintf("Enumerate error from destination: %v", err))
		}
	}

	destNotHaveBlobs := make(chan blobref.SizedBlobRef, 100)
	readSrcBlobs := srcBlobs
	if *flagVerbose {
		readSrcBlobs = loggingBlobRefChannel(srcBlobs)
	}
	go client.ListMissingDestinationBlobs(destNotHaveBlobs, readSrcBlobs, destBlobs)
	for sb := range destNotHaveBlobs {
		fmt.Printf("Destination needs blob: %s\n", sb)

		blobReader, size, err := sc.FetchStreaming(sb.BlobRef)
		if err != nil {
			stats.ErrorCount++
			log.Printf("Error fetching %s: %v", sb.BlobRef, err)
			continue
		}
		if size != sb.Size {
			stats.ErrorCount++
			log.Printf("Source blobserver's enumerate size of %d for blob %s doesn't match its Get size of %d",
				sb.Size, sb.BlobRef, size)
			continue
		}
		uh := &client.UploadHandle{BlobRef: sb.BlobRef, Size: size, Contents: blobReader}
		pr, err := dc.Upload(uh)
		if err != nil {
			stats.ErrorCount++
			log.Printf("Upload of %s to destination blobserver failed: %v", sb.BlobRef, err)
			continue
		}
		if !pr.Skipped {
			stats.BlobsCopied++
			stats.BytesCopied += pr.Size
		}
		if *flagRemoveSource {
			if err = sc.RemoveBlob(sb.BlobRef); err != nil {
				stats.ErrorCount++
				log.Printf("Failed to delete %s from source: %v", sb.BlobRef, err)
			}
		}
	}

	checkSourceError()
	checkDestError()
	if retErr == nil && stats.ErrorCount > 0 {
		retErr = errors.New(fmt.Sprintf("%d errors during sync", stats.ErrorCount))
	}
	return stats, retErr
}

func loggingBlobRefChannel(ch <-chan blobref.SizedBlobRef) chan blobref.SizedBlobRef {
	ch2 := make(chan blobref.SizedBlobRef)
	go func() {
		defer close(ch2)
		var last time.Time
		for v := range ch {
			ch2 <- v
			now := time.Now()
			if last.IsZero() || now.After(last.Add(1*time.Second)) {
				last = now
				log.Printf("At source blob %v", v)
			}
		}
	}()
	return ch2
}
<|endoftext|>