<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cherrypickunapproved

import (
	"reflect"
	"regexp"
	"testing"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/labels"
)

type fakeClient struct {
	// current labels
	labels []string
	// labels that are added
	added []string
	// labels that are removed
	removed []string
	// commentsAdded tracks the comments in the client
	commentsAdded map[int][]string
}

// AddLabel adds a label to the specified PR or issue
func (fc *fakeClient) AddLabel(owner, repo string, number int, label string) error {
	fc.added = append(fc.added, label)
	fc.labels = append(fc.labels, label)
	return nil
}

// RemoveLabel removes the label from the specified PR or issue
func (fc *fakeClient) RemoveLabel(owner, repo string, number int, label string) error {
	fc.removed = append(fc.removed, label)

	// remove from existing labels
	for k, v := range fc.labels {
		if label == v {
			fc.labels = append(fc.labels[:k], fc.labels[k+1:]...)
			break
		}
	}

	return nil
}

// GetIssueLabels gets the current labels on the specified PR or issue
func (fc *fakeClient) GetIssueLabels(owner, repo string, number int) ([]github.Label, error) {
	la := []github.Label{}
	for _, l := range fc.labels {
		la = append(la, github.Label{Name: l})
	}
	return la, nil
}

// CreateComment adds and tracks a comment in the client
func (fc *fakeClient) CreateComment(owner, repo string, number int, comment string) error {
	fc.commentsAdded[number] = append(fc.commentsAdded[number], comment)
	return nil
}

// NumComments counts the number of tracked comments
func (fc *fakeClient) NumComments() int {
	n := 0
	for _, comments := range fc.commentsAdded {
		n += len(comments)
	}
	return n
}

type fakePruner struct{}

func (fp *fakePruner) PruneComments(shouldPrune func(github.IssueComment) bool) {}

func makeFakePullRequestEvent(action github.PullRequestEventAction, branch string) github.PullRequestEvent {
	return github.PullRequestEvent{
		Action: action,
		PullRequest: github.PullRequest{
			Base: github.PullRequestBranch{
				Ref: branch,
			},
		},
	}
}

func TestCherryPickUnapprovedLabel(t *testing.T) {
	var testcases = []struct {
		name          string
		branch        string
		action        github.PullRequestEventAction
		labels        []string
		added         []string
		removed       []string
		expectComment bool
	}{
		{
			name:          "unsupported PR action -> no-op",
			branch:        "release-1.10",
			action:        github.PullRequestActionEdited,
			labels:        []string{},
			added:         []string{},
			removed:       []string{},
			expectComment: false,
		},
		{
			name:          "branch that does match regexp -> no-op",
			branch:        "master",
			action:        github.PullRequestActionOpened,
			labels:        []string{},
			added:         []string{},
			removed:       []string{},
			expectComment: false,
		},
		{
			name:          "has cpUnapproved -> no-op",
			branch:        "release-1.10",
			action:        github.PullRequestActionOpened,
			labels:        []string{labels.CpUnapproved},
			added:         []string{},
			removed:       []string{},
			expectComment: false,
		},
		{
			name:          "has both cpApproved and cpUnapproved -> remove cpUnapproved",
			branch:        "release-1.10",
			action:        github.PullRequestActionOpened,
			labels:        []string{labels.CpApproved, labels.CpUnapproved},
			added:         []string{},
			removed:       []string{labels.CpUnapproved},
			expectComment: false,
		},
		{
			name:          "does not have any labels, PR opened against a release branch -> add cpUnapproved and comment",
			branch:        "release-1.10",
			action:        github.PullRequestActionOpened,
			labels:        []string{},
			added:         []string{labels.CpUnapproved},
			removed:       []string{},
			expectComment: true,
		},
		{
			name:          "does not have any labels, PR reopened against a release branch -> add cpUnapproved and comment",
			branch:        "release-1.10",
			action:        github.PullRequestActionReopened,
			labels:        []string{},
			added:         []string{labels.CpUnapproved},
			removed:       []string{},
			expectComment: true,
		},
	}

	for _, tc := range testcases {
		fc := &fakeClient{
			labels:        tc.labels,
			added:         []string{},
			removed:       []string{},
			commentsAdded: make(map[int][]string, 0),
		}

		event := makeFakePullRequestEvent(tc.action, tc.branch)
		branchRe := regexp.MustCompile(`^release-.*$`)
		comment := "dummy cumment"
		err := handlePR(fc, logrus.WithField("plugin", "fake-cherrypick-unapproved"), &event, &fakePruner{}, branchRe, comment)
		switch {
		case err != nil:
			t.Errorf("%s: unexpected error: %v", tc.name, err)
		case !reflect.DeepEqual(tc.added, fc.added):
			t.Errorf("%s: added %v != actual %v", tc.name, tc.added, fc.added)
		case !reflect.DeepEqual(tc.removed, fc.removed):
			t.Errorf("%s: removed %v != actual %v", tc.name, tc.removed, fc.removed)
		}

		// if we expected a comment, verify that a comment was made
		numComments := fc.NumComments()
		if tc.expectComment && numComments != 1 {
			t.Errorf("%s: expected 1 comment but received %d comments", tc.name, numComments)
		}
		if !tc.expectComment && numComments != 0 {
			t.Errorf("%s: expected no comments but received %d comments", tc.name, numComments)
		}
	}
}
<commit_msg>Support changes field in tests<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cherrypickunapproved

import (
	"encoding/json"
	"fmt"
	"reflect"
	"regexp"
	"testing"

	"github.com/sirupsen/logrus"

	"k8s.io/test-infra/prow/github"
	"k8s.io/test-infra/prow/labels"
)

type fakeClient struct {
	// current labels
	labels []string
	// labels that are added
	added []string
	// labels that are removed
	removed []string
	// commentsAdded tracks the comments in the client
	commentsAdded map[int][]string
}

// AddLabel adds a label to the specified PR or issue
func (fc *fakeClient) AddLabel(owner, repo string, number int, label string) error {
	fc.added = append(fc.added, label)
	fc.labels = append(fc.labels, label)
	return nil
}

// RemoveLabel removes the label from the specified PR or issue
func (fc *fakeClient) RemoveLabel(owner, repo string, number int, label string) error {
	fc.removed = append(fc.removed, label)

	// remove from existing labels
	for k, v := range fc.labels {
		if label == v {
			fc.labels = append(fc.labels[:k], fc.labels[k+1:]...)
			break
		}
	}

	return nil
}

// GetIssueLabels gets the current labels on the specified PR or issue
func (fc *fakeClient) GetIssueLabels(owner, repo string, number int) ([]github.Label, error) {
	var la []github.Label
	for _, l := range fc.labels {
		la = append(la, github.Label{Name: l})
	}
	return la, nil
}

// CreateComment adds and tracks a comment in the client
func (fc *fakeClient) CreateComment(owner, repo string, number int, comment string) error {
	fc.commentsAdded[number] = append(fc.commentsAdded[number], comment)
	return nil
}

// NumComments counts the number of tracked comments
func (fc *fakeClient) NumComments() int {
	n := 0
	for _, comments := range fc.commentsAdded {
		n += len(comments)
	}
	return n
}

type fakePruner struct{}

func (fp *fakePruner) PruneComments(shouldPrune func(github.IssueComment) bool) {}

func makeFakePullRequestEvent(action github.PullRequestEventAction, branch, oldBranch string) github.PullRequestEvent {
	event := github.PullRequestEvent{
		Action: action,
		PullRequest: github.PullRequest{
			Base: github.PullRequestBranch{
				Ref: branch,
			},
		},
	}

	if oldBranch != "" {
		event.Changes = json.RawMessage(fmt.Sprintf(`{"base": {"ref": {"from": "%s"}, "sha": {"from": "sha"}}}`, oldBranch))
	}

	return event
}

func TestCherryPickUnapprovedLabel(t *testing.T) {
	var testcases = []struct {
		name           string
		branch         string
		previousBranch string
		action         github.PullRequestEventAction
		labels         []string
		added          []string
		removed        []string
		expectComment  bool
	}{
		{
			name:           "unsupported PR action -> no-op",
			branch:         "release-1.10",
			previousBranch: "",
			action:         github.PullRequestActionClosed,
			labels:         []string{},
			added:          []string{},
			removed:        []string{},
			expectComment:  false,
		},
		{
			name:           "branch that does match regexp -> no-op",
			branch:         "master",
			previousBranch: "",
			action:         github.PullRequestActionOpened,
			labels:         []string{},
			added:          []string{},
			removed:        []string{},
			expectComment:  false,
		},
		{
			name:           "has cpUnapproved -> no-op",
			branch:         "release-1.10",
			previousBranch: "",
			action:         github.PullRequestActionOpened,
			labels:         []string{labels.CpUnapproved},
			added:          []string{},
			removed:        []string{},
			expectComment:  false,
		},
		{
			name:           "has both cpApproved and cpUnapproved -> remove cpUnapproved",
			branch:         "release-1.10",
			previousBranch: "",
			action:         github.PullRequestActionOpened,
			labels:         []string{labels.CpApproved, labels.CpUnapproved},
			added:          []string{},
			removed:        []string{labels.CpUnapproved},
			expectComment:  false,
		},
		{
			name:           "does not have any labels, PR opened against a release branch -> add cpUnapproved and comment",
			branch:         "release-1.10",
			previousBranch: "",
			action:         github.PullRequestActionOpened,
			labels:         []string{},
			added:          []string{labels.CpUnapproved},
			removed:        []string{},
			expectComment:  true,
		},
		{
			name:           "does not have any labels, PR reopened against a release branch -> add cpUnapproved and comment",
			branch:         "release-1.10",
			previousBranch: "",
			action:         github.PullRequestActionReopened,
			labels:         []string{},
			added:          []string{labels.CpUnapproved},
			removed:        []string{},
			expectComment:  true,
		},
	}

	for _, tc := range testcases {
		fc := &fakeClient{
			labels:        tc.labels,
			added:         []string{},
			removed:       []string{},
			commentsAdded: make(map[int][]string, 0),
		}

		event := makeFakePullRequestEvent(tc.action, tc.branch, tc.previousBranch)
		branchRe := regexp.MustCompile(`^release-.*$`)
		comment := "dummy cumment"
		err := handlePR(fc, logrus.WithField("plugin", "fake-cherrypick-unapproved"), &event, &fakePruner{}, branchRe, comment)
		switch {
		case err != nil:
			t.Errorf("%s: unexpected error: %v", tc.name, err)
		case !reflect.DeepEqual(tc.added, fc.added):
			t.Errorf("%s: added %v != actual %v", tc.name, tc.added, fc.added)
		case !reflect.DeepEqual(tc.removed, fc.removed):
			t.Errorf("%s: removed %v != actual %v", tc.name, tc.removed, fc.removed)
		}

		// if we expected a comment, verify that a comment was made
		numComments := fc.NumComments()
		if tc.expectComment && numComments != 1 {
			t.Errorf("%s: expected 1 comment but received %d comments", tc.name, numComments)
		}
		if !tc.expectComment && numComments != 0 {
			t.Errorf("%s: expected no comments but received %d comments", tc.name, numComments)
		}
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"errors"
	"time"

	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"github.com/docker/swarmkit/protobuf/ptypes"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

const dispatcherRPCTimeout = 5 * time.Second

var (
	errSessionDisconnect = errors.New("agent: session disconnect") // instructed to disconnect
	errSessionClosed     = errors.New("agent: session closed")
)

// session encapsulates one round of registration with the manager. session
// starts the registration and heartbeat control cycle. Any failure will result
// in a complete shutdown of the session and it must be reestablished.
//
// All communication with the master is done through session.  Changes that
// flow into the agent, such as task assignment, are called back into the
// agent through errs, messages and tasks.
type session struct {
	agent     *Agent
	sessionID string
	session   api.Dispatcher_SessionClient
	errs      chan error
	messages  chan *api.SessionMessage
	tasks     chan *api.TasksMessage

	registered chan struct{} // closed registration
	closed     chan struct{}
}

func newSession(ctx context.Context, agent *Agent, delay time.Duration) *session {
	s := &session{
		agent:      agent,
		errs:       make(chan error),
		messages:   make(chan *api.SessionMessage),
		tasks:      make(chan *api.TasksMessage),
		registered: make(chan struct{}),
		closed:     make(chan struct{}),
	}

	go s.run(ctx, delay)
	return s
}

func (s *session) run(ctx context.Context, delay time.Duration) {
	time.Sleep(delay) // delay before registering.

	if err := s.start(ctx); err != nil {
		select {
		case s.errs <- err:
		case <-s.closed:
		case <-ctx.Done():
		}
		return
	}

	ctx = log.WithLogger(ctx, log.G(ctx).WithField("session.id", s.sessionID))

	go runctx(ctx, s.closed, s.errs, s.heartbeat)
	go runctx(ctx, s.closed, s.errs, s.watch)
	go runctx(ctx, s.closed, s.errs, s.listen)

	close(s.registered)
}

// start begins the session and returns the first SessionMessage.
func (s *session) start(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).start")

	client := api.NewDispatcherClient(s.agent.config.Conn)

	description, err := s.agent.config.Executor.Describe(ctx)
	if err != nil {
		log.G(ctx).WithError(err).WithField("executor", s.agent.config.Executor).
			Errorf("node description unavailable")
		return err
	}
	// Override hostname
	if s.agent.config.Hostname != "" {
		description.Hostname = s.agent.config.Hostname
	}

	sessionCtx, cancel := context.WithTimeout(ctx, dispatcherRPCTimeout)
	stream, err := client.Session(sessionCtx, &api.SessionRequest{
		Description: description,
	})
	cancel()
	if err != nil {
		return err
	}

	msg, err := stream.Recv()
	if err != nil {
		return err
	}

	s.sessionID = msg.SessionID
	s.session = stream

	return s.handleSessionMessage(ctx, msg)
}

func (s *session) heartbeat(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).heartbeat")
	client := api.NewDispatcherClient(s.agent.config.Conn)
	heartbeat := time.NewTimer(1) // send out a heartbeat right away
	defer heartbeat.Stop()

	for {
		select {
		case <-heartbeat.C:
			heartbeatCtx, cancel := context.WithTimeout(ctx, dispatcherRPCTimeout)
			resp, err := client.Heartbeat(heartbeatCtx, &api.HeartbeatRequest{
				SessionID: s.sessionID,
			})
			cancel()
			if err != nil {
				if grpc.Code(err) == codes.NotFound {
					err = errNodeNotRegistered
				}

				return err
			}

			period, err := ptypes.Duration(&resp.Period)
			if err != nil {
				return err
			}

			heartbeat.Reset(period)
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) listen(ctx context.Context) error {
	defer s.session.CloseSend()
	log.G(ctx).Debugf("(*session).listen")
	for {
		msg, err := s.session.Recv()
		if err != nil {
			return err
		}

		if err := s.handleSessionMessage(ctx, msg); err != nil {
			return err
		}
	}
}

func (s *session) handleSessionMessage(ctx context.Context, msg *api.SessionMessage) error {
	select {
	case s.messages <- msg:
		return nil
	case <-s.closed:
		return errSessionClosed
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *session) watch(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).watch")
	client := api.NewDispatcherClient(s.agent.config.Conn)
	watch, err := client.Tasks(ctx, &api.TasksRequest{
		SessionID: s.sessionID})
	if err != nil {
		return err
	}

	for {
		resp, err := watch.Recv()
		if err != nil {
			return err
		}

		select {
		case s.tasks <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// sendTaskStatus uses the current session to send the status of a single task.
func (s *session) sendTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {

	client := api.NewDispatcherClient(s.agent.config.Conn)
	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		SessionID: s.sessionID,
		Updates: []*api.UpdateTaskStatusRequest_TaskStatusUpdate{
			{
				TaskID: taskID,
				Status: status,
			},
		},
	}); err != nil {
		// TODO(stevvooe): Dispatcher should not return this error. Status
		// reports for unknown tasks should be ignored.
		if grpc.Code(err) == codes.NotFound {
			return errTaskUnknown
		}

		return err
	}

	return nil
}

func (s *session) sendTaskStatuses(ctx context.Context, updates ...*api.UpdateTaskStatusRequest_TaskStatusUpdate) ([]*api.UpdateTaskStatusRequest_TaskStatusUpdate, error) {
	if len(updates) < 1 {
		return nil, nil
	}

	const batchSize = 1024
	select {
	case <-s.registered:
		select {
		case <-s.closed:
			return updates, ErrClosed
		default:
		}
	case <-s.closed:
		return updates, ErrClosed
	case <-ctx.Done():
		return updates, ctx.Err()
	}

	client := api.NewDispatcherClient(s.agent.config.Conn)
	n := batchSize

	if len(updates) < n {
		n = len(updates)
	}

	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		SessionID: s.sessionID,
		Updates:   updates[:n],
	}); err != nil {
		log.G(ctx).WithError(err).Errorf("failed sending task status batch size of %d", len(updates[:n]))
		return updates, err
	}

	return updates[n:], nil
}

func (s *session) close() error {
	select {
	case <-s.closed:
		return errSessionClosed
	default:
		close(s.closed)
		return nil
	}
}
<commit_msg>agent: Use a goroutine to time out Session RPC<commit_after>package agent

import (
	"errors"
	"time"

	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"github.com/docker/swarmkit/protobuf/ptypes"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

const dispatcherRPCTimeout = 5 * time.Second

var (
	errSessionDisconnect = errors.New("agent: session disconnect") // instructed to disconnect
	errSessionClosed     = errors.New("agent: session closed")
)

// session encapsulates one round of registration with the manager. session
// starts the registration and heartbeat control cycle. Any failure will result
// in a complete shutdown of the session and it must be reestablished.
//
// All communication with the master is done through session.  Changes that
// flow into the agent, such as task assignment, are called back into the
// agent through errs, messages and tasks.
type session struct {
	agent     *Agent
	sessionID string
	session   api.Dispatcher_SessionClient
	errs      chan error
	messages  chan *api.SessionMessage
	tasks     chan *api.TasksMessage

	registered chan struct{} // closed registration
	closed     chan struct{}
}

func newSession(ctx context.Context, agent *Agent, delay time.Duration) *session {
	s := &session{
		agent:      agent,
		errs:       make(chan error),
		messages:   make(chan *api.SessionMessage),
		tasks:      make(chan *api.TasksMessage),
		registered: make(chan struct{}),
		closed:     make(chan struct{}),
	}

	go s.run(ctx, delay)
	return s
}

func (s *session) run(ctx context.Context, delay time.Duration) {
	time.Sleep(delay) // delay before registering.

	if err := s.start(ctx); err != nil {
		select {
		case s.errs <- err:
		case <-s.closed:
		case <-ctx.Done():
		}
		return
	}

	ctx = log.WithLogger(ctx, log.G(ctx).WithField("session.id", s.sessionID))

	go runctx(ctx, s.closed, s.errs, s.heartbeat)
	go runctx(ctx, s.closed, s.errs, s.watch)
	go runctx(ctx, s.closed, s.errs, s.listen)

	close(s.registered)
}

// start begins the session and returns the first SessionMessage.
func (s *session) start(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).start")

	client := api.NewDispatcherClient(s.agent.config.Conn)

	description, err := s.agent.config.Executor.Describe(ctx)
	if err != nil {
		log.G(ctx).WithError(err).WithField("executor", s.agent.config.Executor).
			Errorf("node description unavailable")
		return err
	}
	// Override hostname
	if s.agent.config.Hostname != "" {
		description.Hostname = s.agent.config.Hostname
	}

	errChan := make(chan error, 1)
	var (
		msg    *api.SessionMessage
		stream api.Dispatcher_SessionClient
	)
	// Note: we don't defer cancellation of this context, because the
	// streaming RPC is used after this function returned. We only cancel
	// it in the timeout case to make sure the goroutine completes.
	sessionCtx, cancelSession := context.WithCancel(ctx)

	// Need to run Session in a goroutine since there's no way to set a
	// timeout for an individual Recv call in a stream.
	go func() {
		stream, err = client.Session(sessionCtx, &api.SessionRequest{
			Description: description,
		})
		if err != nil {
			errChan <- err
			return
		}

		msg, err = stream.Recv()
		errChan <- err
	}()

	select {
	case err := <-errChan:
		if err != nil {
			return err
		}
	case <-time.After(dispatcherRPCTimeout):
		cancelSession()
		return errors.New("session initiation timed out")
	}

	s.sessionID = msg.SessionID
	s.session = stream

	return s.handleSessionMessage(ctx, msg)
}

func (s *session) heartbeat(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).heartbeat")
	client := api.NewDispatcherClient(s.agent.config.Conn)
	heartbeat := time.NewTimer(1) // send out a heartbeat right away
	defer heartbeat.Stop()

	for {
		select {
		case <-heartbeat.C:
			heartbeatCtx, cancel := context.WithTimeout(ctx, dispatcherRPCTimeout)
			resp, err := client.Heartbeat(heartbeatCtx, &api.HeartbeatRequest{
				SessionID: s.sessionID,
			})
			cancel()
			if err != nil {
				if grpc.Code(err) == codes.NotFound {
					err = errNodeNotRegistered
				}

				return err
			}

			period, err := ptypes.Duration(&resp.Period)
			if err != nil {
				return err
			}

			heartbeat.Reset(period)
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) listen(ctx context.Context) error {
	defer s.session.CloseSend()
	log.G(ctx).Debugf("(*session).listen")
	for {
		msg, err := s.session.Recv()
		if err != nil {
			return err
		}

		if err := s.handleSessionMessage(ctx, msg); err != nil {
			return err
		}
	}
}

func (s *session) handleSessionMessage(ctx context.Context, msg *api.SessionMessage) error {
	select {
	case s.messages <- msg:
		return nil
	case <-s.closed:
		return errSessionClosed
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *session) watch(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).watch")
	client := api.NewDispatcherClient(s.agent.config.Conn)
	watch, err := client.Tasks(ctx, &api.TasksRequest{
		SessionID: s.sessionID})
	if err != nil {
		return err
	}

	for {
		resp, err := watch.Recv()
		if err != nil {
			return err
		}

		select {
		case s.tasks <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// sendTaskStatus uses the current session to send the status of a single task.
func (s *session) sendTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {

	client := api.NewDispatcherClient(s.agent.config.Conn)
	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		SessionID: s.sessionID,
		Updates: []*api.UpdateTaskStatusRequest_TaskStatusUpdate{
			{
				TaskID: taskID,
				Status: status,
			},
		},
	}); err != nil {
		// TODO(stevvooe): Dispatcher should not return this error. Status
		// reports for unknown tasks should be ignored.
		if grpc.Code(err) == codes.NotFound {
			return errTaskUnknown
		}

		return err
	}

	return nil
}

func (s *session) sendTaskStatuses(ctx context.Context, updates ...*api.UpdateTaskStatusRequest_TaskStatusUpdate) ([]*api.UpdateTaskStatusRequest_TaskStatusUpdate, error) {
	if len(updates) < 1 {
		return nil, nil
	}

	const batchSize = 1024
	select {
	case <-s.registered:
		select {
		case <-s.closed:
			return updates, ErrClosed
		default:
		}
	case <-s.closed:
		return updates, ErrClosed
	case <-ctx.Done():
		return updates, ctx.Err()
	}

	client := api.NewDispatcherClient(s.agent.config.Conn)
	n := batchSize

	if len(updates) < n {
		n = len(updates)
	}

	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		SessionID: s.sessionID,
		Updates:   updates[:n],
	}); err != nil {
		log.G(ctx).WithError(err).Errorf("failed sending task status batch size of %d", len(updates[:n]))
		return updates, err
	}

	return updates[n:], nil
}

func (s *session) close() error {
	select {
	case <-s.closed:
		return errSessionClosed
	default:
		close(s.closed)
		return nil
	}
}
<|endoftext|>


<commit_before>package solr

import (
	"fmt"
	"math"
	"net/url"
)

type Document map[string]interface{}

// Has check if a key exist in document
func (d Document) Has(k string) bool {
	_, ok := d[k]
	return ok
}

// Get returns value of a key
func (d Document) Get(k string) interface{} {
	v, _ := d[k]
	return v
}

// Set add a key/value to document
func (d Document) Set(k string, v interface{}) {
	d[k] = v
}

type Collection struct {
	Docs     []Document
	Start    int
	NumFound int
}

type SolrResult struct {
	Status         int         // status quick access to status
	Results        *Collection // results parsed documents, basically response object
	ResponseHeader map[string]interface{}
	FacetCounts    map[string]interface{}
	Highlighting   map[string]interface{}
	Error          map[string]interface{}
	// grouped for grouping result
	// if grouping Results will be empty
	Grouped map[string]interface{}
}

type SolrInterface struct {
	conn *Connection
}

// Return a new instance of SolrInterface
func NewSolrInterface(solrUrl, core string) (*SolrInterface, error) {
	c, err := NewConnection(solrUrl, core)
	if err != nil {
		return nil, err
	}
	return &SolrInterface{conn: c}, nil
}

// Set to new core, this is just wrapper to Connection.SetCore which mean
// it will affect all places that use this Connection instance
func (si *SolrInterface) SetCore(core string) {
	si.conn.SetCore(core)
}

// SetBasicAuth sets the request's Authorization header to use HTTP Basic Authentication with the provided username and password.
// See http://golang.org/pkg/net/http/#Request.SetBasicAuth
func (si *SolrInterface) SetBasicAuth(username, password string) {
	si.conn.SetBasicAuth(username, password)
}

// Return a new instace of Search, q is optional and one can set it later
func (si *SolrInterface) Search(q *Query) *Search {
	return NewSearch(si.conn, q)
}

// makeAddChunks splits the documents into chunks. If chunk_size is less than one it will be default to 100
func makeAddChunks(docs []Document, chunk_size int) []map[string]interface{} {
	if chunk_size < 1 {
		chunk_size = 100
	}
	docs_len := len(docs)
	num_chunk := int(math.Ceil(float64(docs_len) / float64(chunk_size)))
	doc_counter := 0
	chunks := make([]map[string]interface{}, num_chunk)
	for i := 0; i < num_chunk; i++ {
		add := make([]Document, 0, chunk_size)
		for j := 0; j < chunk_size; j++ {
			if doc_counter >= docs_len {
				break
			}
			add = append(add, docs[doc_counter])
			doc_counter++
		}
		chunks[i] = map[string]interface{}{"add": add}
	}
	return chunks
}

// Add will insert documents in batch of chunk_size. success is false as long as one chunk failed.
// The result in UpdateResponse is summery of response from all chunks
// with key chunk_%d
func (si *SolrInterface) Add(docs []Document, chunk_size int, params *url.Values) (*UpdateResponse, error) {
	result := &UpdateResponse{Success: true}
	responses := map[string]interface{}{}
	chunks := makeAddChunks(docs, chunk_size)

	for i := 0; i < len(chunks); i++ {
		res, err := si.Update(chunks[i], params)
		if err != nil {
			return nil, err
		}
		result.Success = result.Success && res.Success
		responses[fmt.Sprintf("chunk_%d", i+1)] = map[string]interface{}{
			"result":  res.Result,
			"success": res.Success,
			"total":   len(chunks[i]["add"].([]Document))}
	}
	result.Result = responses
	return result, nil
}

// Delete take data of type map and optional params which can use to specify addition parameters such as commit=true .
// Only one delete statement is supported, ie data can be { "id":"ID" } .
// If you want to delete more docs use { "query":"QUERY" } .
// Extra params can specify in params or in data such as { "query":"QUERY", "commitWithin":"500" }
func (si *SolrInterface) Delete(data map[string]interface{}, params *url.Values) (*UpdateResponse, error) {
	message := map[string]interface{}{"delete": data}
	return si.Update(message, params)
}

// DeleteAll will remove all documents and commit
func (si *SolrInterface) DeleteAll() (*UpdateResponse, error) {
	params := &url.Values{}
	params.Add("commit", "true")
	return si.Delete(map[string]interface{}{"query": "*:*"}, params)
}

// Update take data of type map and optional params which can use to specify addition parameters such as commit=true
func (si *SolrInterface) Update(data map[string]interface{}, params *url.Values) (*UpdateResponse, error) {
	if si.conn == nil {
		return nil, fmt.Errorf("No connection found for making request to solr")
	}
	return si.conn.Update(data, params)
}

// Commit the changes since the last commit
func (si *SolrInterface) Commit() (*UpdateResponse, error) {
	params := &url.Values{}
	params.Add("commit", "true")
	return si.Update(map[string]interface{}{}, params)
}

func (si *SolrInterface) Optimize(params *url.Values) (*UpdateResponse, error) {
	if params == nil {
		params = &url.Values{}
	}
	params.Set("optimize", "true")
	return si.Update(map[string]interface{}{}, params)
}

// Rollback rollbacks all add/deletes made to the index since the last commit.
// This should use with caution.
// See https://wiki.apache.org/solr/UpdateXmlMessages#A.22rollback.22
func (si *SolrInterface) Rollback() (*UpdateResponse, error) {
	return si.Update(map[string]interface{}{"rollback": map[string]interface{}{}}, nil)
}

// Return new instance of CoreAdmin with provided solrUrl and basic auth
func (si *SolrInterface) CoreAdmin() (*CoreAdmin, error) {
	ca, err := NewCoreAdmin(si.conn.url.String())
	if err != nil {
		return nil, err
	}
	ca.SetBasicAuth(si.conn.username, si.conn.password)
	return ca, nil
}<commit_msg>#24 Schema() for creating a new Schema instance from SolrInterface<commit_after>package solr

import (
	"fmt"
	"math"
	"net/url"
)

type Document map[string]interface{}

// Has check if a key exist in document
func (d Document) Has(k string) bool {
	_, ok := d[k]
	return ok
}

// Get returns value of a key
func (d Document) Get(k string) interface{} {
	v, _ := d[k]
	return v
}

// Set add a key/value to document
func (d Document) Set(k string, v interface{}) {
	d[k] = v
}

type Collection struct {
	Docs     []Document
	Start    int
	NumFound int
}

type SolrResult struct {
	Status         int         // status quick access to status
	Results        *Collection // results parsed documents, basically response object
	ResponseHeader map[string]interface{}
	FacetCounts    map[string]interface{}
	Highlighting   map[string]interface{}
	Error          map[string]interface{}
	// grouped for grouping result
	// if grouping Results will be empty
	Grouped map[string]interface{}
}

type SolrInterface struct {
	conn *Connection
}

// Return a new instance of SolrInterface
func NewSolrInterface(solrUrl, core string) (*SolrInterface, error) {
	c, err := NewConnection(solrUrl, core)
	if err != nil {
		return nil, err
	}
	return &SolrInterface{conn: c}, nil
}

// Set to new core, this is just wrapper to Connection.SetCore which mean
// it will affect all places that use this Connection instance
func (si *SolrInterface) SetCore(core string) {
	si.conn.SetCore(core)
}

// SetBasicAuth sets the request's Authorization header to use HTTP Basic Authentication with the provided username and password.
// See http://golang.org/pkg/net/http/#Request.SetBasicAuth
func (si *SolrInterface) SetBasicAuth(username, password string) {
	si.conn.SetBasicAuth(username, password)
}

// Return a new instace of Search, q is optional and one can set it later
func (si *SolrInterface) Search(q *Query) *Search {
	return NewSearch(si.conn, q)
}

// makeAddChunks splits the documents into chunks. If chunk_size is less than one it will be default to 100
func makeAddChunks(docs []Document, chunk_size int) []map[string]interface{} {
	if chunk_size < 1 {
		chunk_size = 100
	}
	docs_len := len(docs)
	num_chunk := int(math.Ceil(float64(docs_len) / float64(chunk_size)))
	doc_counter := 0
	chunks := make([]map[string]interface{}, num_chunk)
	for i := 0; i < num_chunk; i++ {
		add := make([]Document, 0, chunk_size)
		for j := 0; j < chunk_size; j++ {
			if doc_counter >= docs_len {
				break
			}
			add = append(add, docs[doc_counter])
			doc_counter++
		}
		chunks[i] = map[string]interface{}{"add": add}
	}
	return chunks
}

// Add will insert documents in batch of chunk_size. success is false as long as one chunk failed.
// The result in UpdateResponse is summery of response from all chunks
// with key chunk_%d
func (si *SolrInterface) Add(docs []Document, chunk_size int, params *url.Values) (*UpdateResponse, error) {
	result := &UpdateResponse{Success: true}
	responses := map[string]interface{}{}
	chunks := makeAddChunks(docs, chunk_size)

	for i := 0; i < len(chunks); i++ {
		res, err := si.Update(chunks[i], params)
		if err != nil {
			return nil, err
		}
		result.Success = result.Success && res.Success
		responses[fmt.Sprintf("chunk_%d", i+1)] = map[string]interface{}{
			"result":  res.Result,
			"success": res.Success,
			"total":   len(chunks[i]["add"].([]Document))}
	}
	result.Result = responses
	return result, nil
}

// Delete take data of type map and optional params which can use to specify addition parameters such as commit=true .
// Only one delete statement is supported, ie data can be { "id":"ID" } .
// If you want to delete more docs use { "query":"QUERY" } .
// Extra params can specify in params or in data such as { "query":"QUERY", "commitWithin":"500" }
func (si *SolrInterface) Delete(data map[string]interface{}, params *url.Values) (*UpdateResponse, error) {
	message := map[string]interface{}{"delete": data}
	return si.Update(message, params)
}

// DeleteAll will remove all documents and commit
func (si *SolrInterface) DeleteAll() (*UpdateResponse, error) {
	params := &url.Values{}
	params.Add("commit", "true")
	return si.Delete(map[string]interface{}{"query": "*:*"}, params)
}

// Update take data of type map and optional params which can use to specify addition parameters such as commit=true
func (si *SolrInterface) Update(data map[string]interface{}, params *url.Values) (*UpdateResponse, error) {
	if si.conn == nil {
		return nil, fmt.Errorf("No connection found for making request to solr")
	}
	return si.conn.Update(data, params)
}

// Commit the changes since the last commit
func (si *SolrInterface) Commit() (*UpdateResponse, error) {
	params := &url.Values{}
	params.Add("commit", "true")
	return si.Update(map[string]interface{}{}, params)
}

func (si *SolrInterface) Optimize(params *url.Values) (*UpdateResponse, error) {
	if params == nil {
		params = &url.Values{}
	}
	params.Set("optimize", "true")
	return si.Update(map[string]interface{}{}, params)
}

// Rollback rollbacks all add/deletes made to the index since the last commit.
// This should use with caution.
// See https://wiki.apache.org/solr/UpdateXmlMessages#A.22rollback.22
func (si *SolrInterface) Rollback() (*UpdateResponse, error) {
	return si.Update(map[string]interface{}{"rollback": map[string]interface{}{}}, nil)
}

// Return new instance of CoreAdmin with provided solrUrl and basic auth
func (si *SolrInterface) CoreAdmin() (*CoreAdmin, error) {
	ca, err := NewCoreAdmin(si.conn.url.String())
	if err != nil {
		return nil, err
	}
	ca.SetBasicAuth(si.conn.username, si.conn.password)
	return ca, nil
}

// Return new instance of Schema with provided solrUrl and basic auth
func (si *SolrInterface) Schema() (*Schema, error) {
	s, err := NewSchema(si.conn.url.String(), si.conn.core)
	if err != nil {
		return nil, err
	}
	s.SetBasicAuth(si.conn.username, si.conn.password)
	return s, nil
}
<|endoftext|>


<commit_before><commit_msg>Allows multiple goroutines to call Build at the same time<commit_after><|endoftext|>


<commit_before>package main

import (
	"github.com/fhs/gompd/mpd"
	"strings"
	"testing"
)

func TestSongString(t *testing.T) {
	e := "foo: bar, hoge: fuga"
	r := songString(mpd.Attrs{"hoge": "fuga", "foo": "bar"})
	if r != e {
		t.Errorf(
			"unexpected return\nexpected: %s\nactual:   %s",
			e, r,
		)
	}
}

func TestFindCover(t *testing.T) {
	t.Run("found", func(t *testing.T) {
		cache := make(map[string]string)
		ret := findCover("./hoge", "song_test.*", cache)
		if ret != "song_test.go" {
			t.Errorf("unexpected result: %s", ret)
		}
	})
	t.Run("not found", func(t *testing.T) {
		cache := make(map[string]string)
		ret := findCover("./hoge", "cover_not_found.*", cache)
		if ret != "" {
			t.Errorf("unexpected result: %s", ret)
		}
	})
}

func TestSongAddReadableData(t *testing.T) {
	candidates := []struct {
		input  mpd.Attrs
		expect mpd.Attrs
	}{
		{
			mpd.Attrs{
				"Title": "foo",
				"file":  "path",
			},
			mpd.Attrs{
				"Title":       "foo",
				"file":        "path",
				"DiscNumber":  "0001",
				"TrackNumber": "0000",
				"Length":      "00:00",
			},
		},
		{
			mpd.Attrs{
				"Title": "foo",
				"file":  "path",
				"Disc":  "2",
				"Track": "1",
				"Time":  "121",
			},
			mpd.Attrs{
				"Title":       "foo",
				"file":        "path",
				"Disc":        "2",
				"DiscNumber":  "0002",
				"Track":       "1",
				"TrackNumber": "0001",
				"Time":        "121",
				"Length":      "02:01",
			},
		},
	}

	for _, c := range candidates {
		r := songAddReadableData(c.input)
		if songString(c.expect) != songString(r) {
			t.Errorf(
				"unexpected return\nexpected: %s\nactual:   %s",
				songString(c.expect),
				songString(r),
			)
		}
	}
}

func TestSongSortKey(t *testing.T) {
	candidates := []struct {
		song    mpd.Attrs
		sortkey []string
		expect  string
	}{
		{
			mpd.Attrs{"Title": "foo", "file": "path"},
			[]string{"TrackNumber", "Title", "file"},
			" foopath",
		},
		{
			mpd.Attrs{"Title": "foo", "file": "path"},
			[]string{"foo"},
			" ",
		},
		{
			mpd.Attrs{},
			[]string{"AlbumSort"},
			" ",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"ArtistSort"},
			"foo",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"AlbumArtist"},
			"foo",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"AlbumArtistSort"},
			"foo",
		},
	}
	for _, c := range candidates {
		r := songSortKey(c.song, c.sortkey)
		if r != c.expect {
			t.Errorf(
				"unexpected output for song: %s sortkey: %s\nexpected: \"%s\"\nactual:   \"%s\"",
				songString(c.song),
				strings.Join(c.sortkey, ","),
				c.expect,
				r,
			)
		}
	}
}
<commit_msg>fix cover test<commit_after>package main

import (
	"github.com/fhs/gompd/mpd"
	"strings"
	"testing"
)

func TestSongString(t *testing.T) {
	e := "foo: bar, hoge: fuga"
	r := songString(mpd.Attrs{"hoge": "fuga", "foo": "bar"})
	if r != e {
		t.Errorf(
			"unexpected return\nexpected: %s\nactual:   %s",
			e, r,
		)
	}
}

func TestFindCover(t *testing.T) {
	t.Run("found", func(t *testing.T) {
		cache := make(map[string]string)
		ret := findCover("./", "hoge", "song_test.*", cache)
		if ret != "song_test.go" {
			t.Errorf("unexpected result: %s", ret)
		}
	})
	t.Run("not found", func(t *testing.T) {
		cache := make(map[string]string)
		ret := findCover("./", "hoge", "cover_not_found.*", cache)
		if ret != "" {
			t.Errorf("unexpected result: %s", ret)
		}
	})
}

func TestSongAddReadableData(t *testing.T) {
	candidates := []struct {
		input  mpd.Attrs
		expect mpd.Attrs
	}{
		{
			mpd.Attrs{
				"Title": "foo",
				"file":  "path",
			},
			mpd.Attrs{
				"Title":       "foo",
				"file":        "path",
				"DiscNumber":  "0001",
				"TrackNumber": "0000",
				"Length":      "00:00",
			},
		},
		{
			mpd.Attrs{
				"Title": "foo",
				"file":  "path",
				"Disc":  "2",
				"Track": "1",
				"Time":  "121",
			},
			mpd.Attrs{
				"Title":       "foo",
				"file":        "path",
				"Disc":        "2",
				"DiscNumber":  "0002",
				"Track":       "1",
				"TrackNumber": "0001",
				"Time":        "121",
				"Length":      "02:01",
			},
		},
	}

	for _, c := range candidates {
		r := songAddReadableData(c.input)
		if songString(c.expect) != songString(r) {
			t.Errorf(
				"unexpected return\nexpected: %s\nactual:   %s",
				songString(c.expect),
				songString(r),
			)
		}
	}
}

func TestSongSortKey(t *testing.T) {
	candidates := []struct {
		song    mpd.Attrs
		sortkey []string
		expect  string
	}{
		{
			mpd.Attrs{"Title": "foo", "file": "path"},
			[]string{"TrackNumber", "Title", "file"},
			" foopath",
		},
		{
			mpd.Attrs{"Title": "foo", "file": "path"},
			[]string{"foo"},
			" ",
		},
		{
			mpd.Attrs{},
			[]string{"AlbumSort"},
			" ",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"ArtistSort"},
			"foo",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"AlbumArtist"},
			"foo",
		},
		{
			mpd.Attrs{"Artist": "foo"},
			[]string{"AlbumArtistSort"},
			"foo",
		},
	}
	for _, c := range candidates {
		r := songSortKey(c.song, c.sortkey)
		if r != c.expect {
			t.Errorf(
				"unexpected output for song: %s sortkey: %s\nexpected: \"%s\"\nactual:   \"%s\"",
				songString(c.song),
				strings.Join(c.sortkey, ","),
				c.expect,
				r,
			)
		}
	}
}
<|endoftext|>


<commit_before>package wikiparse

import (
	"compress/bzip2"
	"encoding/xml"
	"io"
	"log"
	"os"
	"sync"
)

type indexChunk struct {
	offset int64
	count  int
}

type multiStreamParser struct {
	siteInfo SiteInfo

	workerch chan indexChunk
	entries  chan *Page
}

func multiStreamIndexWorker(indexfn string, p *multiStreamParser) {
	defer close(p.workerch)

	r, err := os.Open(indexfn)
	if err != nil {
		log.Fatalf("Error opening %v: %v", indexfn, err)
	}
	defer r.Close()

	bz := bzip2.NewReader(r)

	isr, err := NewIndexSummaryReader(bz)
	if err != nil {
		log.Fatalf("Error creating index summary: %v", err)
	}
	for {
		offset, count, err := isr.Next()
		p.workerch <- indexChunk{offset, count}
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("Error reading stream:  %v", err)
		}
	}
}

func multiStreamWorker(datafn string, wg *sync.WaitGroup,
	p *multiStreamParser) {
	defer wg.Done()

	r, err := os.Open(datafn)
	if err != nil {
		log.Fatalf("Error opening %v: %v", datafn, err)
	}
	defer r.Close()

	for idxChunk := range p.workerch {
		_, err := r.Seek(idxChunk.offset, 0)
		if err != nil {
			log.Fatalf("Error seeking to specified offset: %v", err)
		}
		bz := bzip2.NewReader(r)
		d := xml.NewDecoder(bz)

		for i := 0; i < idxChunk.count && err != io.EOF; i++ {
			newpage := &Page{}
			err = d.Decode(newpage)
			if err == nil {
				p.entries <- newpage
			}
		}
	}
}

// Get a wikipedia dump parser reading from the given reader.
func NewIndexedParser(indexfn, datafn string, numWorkers int) (Parser, error) {
	r, err := os.Open(datafn)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	bz := bzip2.NewReader(r)

	d := xml.NewDecoder(bz)
	_, err = d.Token()
	if err != nil {
		return nil, err
	}

	si := SiteInfo{}
	err = d.Decode(&si)
	if err != nil {
		return nil, err
	}

	rv := &multiStreamParser{
		siteInfo: si,
		workerch: make(chan indexChunk, 1000),
		entries:  make(chan *Page, 1000),
	}

	wg := sync.WaitGroup{}
	wg.Add(numWorkers)
	for i := 0; i < numWorkers; i++ {
		go multiStreamWorker(datafn, &wg, rv)
	}

	go multiStreamIndexWorker(indexfn, rv)

	go func() {
		wg.Wait()
		close(rv.entries)
	}()

	return rv, nil
}

func (p *multiStreamParser) Next() (rv *Page, err error) {
	var ok bool
	rv, ok = <-p.entries
	if !ok {
		return nil, io.EOF
	}
	return
}

func (p *multiStreamParser) SiteInfo() SiteInfo {
	return p.siteInfo
}
<commit_msg>Allow multi-stream parsing from things other than os.File<commit_after>package wikiparse

import (
	"compress/bzip2"
	"encoding/xml"
	"io"
	"log"
	"os"
	"sync"
)

type indexChunk struct {
	offset int64
	count  int
}

type multiStreamParser struct {
	siteInfo SiteInfo

	workerch chan indexChunk
	entries  chan *Page
}

func multiStreamIndexWorker(r io.ReadCloser, p *multiStreamParser) {
	defer close(p.workerch)

	defer r.Close()

	bz := bzip2.NewReader(r)

	isr, err := NewIndexSummaryReader(bz)
	if err != nil {
		log.Fatalf("Error creating index summary: %v", err)
	}
	for {
		offset, count, err := isr.Next()
		p.workerch <- indexChunk{offset, count}
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalf("Error reading stream:  %v", err)
		}
	}
}

func multiStreamWorker(src IndexedParseSource, wg *sync.WaitGroup,
	p *multiStreamParser) {
	defer wg.Done()

	r, err := src.OpenData()
	if err != nil {
		log.Fatalf("Error opening data: %v", err)
	}
	defer r.Close()

	for idxChunk := range p.workerch {
		_, err := r.Seek(idxChunk.offset, 0)
		if err != nil {
			log.Fatalf("Error seeking to specified offset: %v", err)
		}
		bz := bzip2.NewReader(r)
		d := xml.NewDecoder(bz)

		for i := 0; i < idxChunk.count && err != io.EOF; i++ {
			newpage := &Page{}
			err = d.Decode(newpage)
			if err == nil {
				p.entries <- newpage
			}
		}
	}
}

type ReadSeekCloser interface {
	io.ReadSeeker
	io.Closer
}

type IndexedParseSource interface {
	OpenIndex() (io.ReadCloser, error)
	OpenData() (ReadSeekCloser, error)
}

type filesSource struct {
	idxfile, datafile string
}

func (f filesSource) OpenIndex() (io.ReadCloser, error) {
	return os.Open(f.idxfile)
}

func (f filesSource) OpenData() (ReadSeekCloser, error) {
	return os.Open(f.datafile)
}

func NewIndexedParserFromSrc(src IndexedParseSource, numWorkers int) (Parser, error) {
	r, err := src.OpenData()
	if err != nil {
		return nil, err
	}
	defer r.Close()

	bz := bzip2.NewReader(r)

	d := xml.NewDecoder(bz)
	_, err = d.Token()
	if err != nil {
		return nil, err
	}

	si := SiteInfo{}
	err = d.Decode(&si)
	if err != nil {
		return nil, err
	}

	rv := &multiStreamParser{
		siteInfo: si,
		workerch: make(chan indexChunk, 1000),
		entries:  make(chan *Page, 1000),
	}

	wg := sync.WaitGroup{}
	wg.Add(numWorkers)
	for i := 0; i < numWorkers; i++ {
		go multiStreamWorker(src, &wg, rv)
	}

	ridx, err := src.OpenIndex()
	if err != nil {
		log.Fatalf("Error opening index: %v", err)
	}
	defer ridx.Close()

	go multiStreamIndexWorker(ridx, rv)

	go func() {
		wg.Wait()
		close(rv.entries)
	}()

	return rv, nil
}

// Get a wikipedia dump parser reading from the given reader.
func NewIndexedParser(indexfn, datafn string, numWorkers int) (Parser, error) {
	return NewIndexedParserFromSrc(filesSource{indexfn, datafn}, numWorkers)
}

func (p *multiStreamParser) Next() (rv *Page, err error) {
	var ok bool
	rv, ok = <-p.entries
	if !ok {
		return nil, io.EOF
	}
	return
}

func (p *multiStreamParser) SiteInfo() SiteInfo {
	return p.siteInfo
}
<|endoftext|>


<commit_before>package multicorecsv

import (
	"bufio"
	"bytes"
	"encoding/csv"
	"io"
	"runtime"
	"sync"
)

type linein struct {
	data []byte
	num  int
}

type lineout struct {
	data []string
	num  int
}

type MulticoreReader struct {
	reader  io.Reader
	linein  chan linein
	lineout chan lineout
	errChan chan error
	// the following are from encoding/csv package and are copied into the underlying csv.Reader
	Comma            rune
	Comment          rune
	FieldsPerRecord  int // we can't implement this without more overhead/synchronization
	LazyQuotes       bool
	TrailingComma    bool
	TrimLeadingSpace bool
	place            int              // how many lines have been returned so far
	queue            map[int][]string // used to buffer lines that come in out of order
	finalError       error
	cancel           chan struct{} // when this is closed, cancel all operations
	once             sync.Once
}

func NewReader(r io.Reader) *MulticoreReader {
	return &MulticoreReader{
		reader:  r,
		Comma:   ',',
		linein:  make(chan linein),
		lineout: make(chan lineout),
		errChan: make(chan error),
		queue:   make(map[int][]string),
		cancel:  make(chan struct{}),
	}
}

func (mcr *MulticoreReader) Close() error {
	close(mcr.cancel)
	return nil
}

func (mcr *MulticoreReader) ReadAll() ([][]string, error) {
	var all [][]string
	out, errChan := mcr.Stream()
	for line := range out {
		all = append(all, line)
	}
	return all, <-errChan
}

func (mcr *MulticoreReader) Stream() (chan []string, chan error) {
	out := make(chan []string)
	errChan := make(chan error, 1)
	go func() {
		defer close(out)
		defer close(errChan)
		for {
			line, err := mcr.Read()
			if len(line) > 0 {
				out <- line
			}
			if err == nil {
				continue
			}
			if err == io.EOF {
				return
			}
			errChan <- err
			return
		}
	}()
	return out, errChan
}

func (mcr *MulticoreReader) Read() ([]string, error) {
	if mcr.finalError != nil {
		return nil, mcr.finalError
	}
	mcr.start()
	line, ok := mcr.queue[mcr.place]
	if ok {
		delete(mcr.queue, mcr.place)
		mcr.place++
		return line, nil
	}
	for {
		//		log.Printf("mcr.Read() reading from lineout")
		line, ok := <-mcr.lineout
		if !ok {
			mcr.finalError = <-mcr.errChan
			return nil, mcr.finalError
		}
		if line.num == mcr.place {
			mcr.place++
			return line.data, nil
		}
		mcr.queue[line.num] = line.data
		// keep going, fetch the next line since this one was out of order
	}
}

func (mcr *MulticoreReader) startReading(err1 chan error) {
	defer close(mcr.linein)
	linenum := 0
	bytesreader := bufio.NewReader(mcr.reader)
	for {
		line, err := bytesreader.ReadBytes('\n')
		if len(line) > 0 {
			if line[0] == '\n' || rune(line[0]) == mcr.Comment {
				continue // we don't want blanks or comments
			}
			select {
			case mcr.linein <- linein{
				data: line,
				num:  linenum,
			}:
				linenum++
			case <-mcr.cancel:
				err1 <- nil
				return
			}
		}
		if err == nil {
			continue
		}
		if err == io.EOF {
			err = nil
		}
		err1 <- err
		return
	}
}

func (mcr *MulticoreReader) parseCSVLines(err2 chan error) {
	var buf bytes.Buffer
	r := csv.NewReader(&buf)
	r.Comma = mcr.Comma
	r.Comment = mcr.Comment
	r.LazyQuotes = mcr.LazyQuotes
	r.TrailingComma = mcr.TrailingComma
	r.TrimLeadingSpace = mcr.TrimLeadingSpace
	for b := range mcr.linein {
		buf.Reset()
		buf.Write(b.data)
		if line, err := r.Read(); err != nil {
			pe := err.(*csv.ParseError)
			pe.Line = b.num + 1
			err2 <- err
			return
		} else {
			select {
			case mcr.lineout <- lineout{
				data: line,
				num:  b.num,
			}:
			case <-mcr.cancel:
				err2 <- nil
				return
			}
		}
	}
	err2 <- nil
}

func (mcr *MulticoreReader) waitForDone(err1, err2 chan error) {
	//	log.Printf("mcr.waitForDone waiting on read error")
	foundError := <-err1
	//	log.Printf("mcr.waitForDone waiting on parse errors")
	for i := 0; i < runtime.NumCPU(); i++ {
		err := <-err2
		if err != nil && err != io.EOF && foundError == nil {
			foundError = err
		}
	}
	if foundError == nil {
		foundError = io.EOF
	}
	//	log.Printf("mcr.waitForDone closing lineout")
	close(mcr.lineout)
	mcr.errChan <- foundError
	//	log.Printf("mcr.waitForDone ending")
}

func (mcr *MulticoreReader) start() {
	mcr.once.Do(func() {
		err1 := make(chan error, 1)
		err2 := make(chan error, runtime.NumCPU())
		go mcr.startReading(err1)
		for i := 0; i < runtime.NumCPU(); i++ {
			go mcr.parseCSVLines(err2)
		}
		go mcr.waitForDone(err1, err2)
	})
}
<commit_msg>Adding comments/docs Fixed a nasty bug in Read() where additional lines could be left in the queue<commit_after>package multicorecsv

import (
	"bufio"
	"bytes"
	"encoding/csv"
	"io"
	"runtime"
	"sync"
)

type linein struct {
	data []byte
	num  int
}

type lineout struct {
	data []string
	num  int
}

type MulticoreReader struct {
	reader  io.Reader
	linein  chan linein
	lineout chan lineout
	errChan chan error
	// the following are from encoding/csv package and are copied into the underlying csv.Reader
	Comma            rune
	Comment          rune
	FieldsPerRecord  int // we can't implement this without more overhead/synchronization
	LazyQuotes       bool
	TrailingComma    bool
	TrimLeadingSpace bool
	place            int              // how many lines have been returned so far
	queue            map[int][]string // used to buffer lines that come in out of order
	finalError       error
	cancel           chan struct{} // when this is closed, cancel all operations
	once             sync.Once
}

// NewReader returns a new Reader that reads from r.
func NewReader(r io.Reader) *MulticoreReader {
	return &MulticoreReader{
		reader:  r,
		Comma:   ',',
		linein:  make(chan linein),
		lineout: make(chan lineout),
		errChan: make(chan error),
		queue:   make(map[int][]string),
		cancel:  make(chan struct{}),
	}
}

// Close will clean up any goroutines that aren't finished
// It will also close the underlying Reader if it implements io.ReadCloser
func (mcr *MulticoreReader) Close() error {
	close(mcr.cancel)
	if c, ok := mcr.reader.(io.ReadCloser); ok {
		return c.Close()
	}
	return nil
}

// ReadAll reads all the remaining records from r.
// Each record is a slice of fields.
// A successful call returns err == nil, not err == EOF. Because ReadAll is
// defined to read until EOF, it does not treat end of file as an error to be
// reported.
func (mcr *MulticoreReader) ReadAll() ([][]string, error) {
	var all [][]string
	out, errChan := mcr.Stream()
	for line := range out {
		all = append(all, line)
	}
	return all, <-errChan
}

// Stream returns a chan of []string representing a row in the CSV file
// Lines are sent on the channel in order they were in the source file
// The caller must receive all rows and receive the error from the error chan,
// otherwise the caller must call Close to clean up any goroutines
func (mcr *MulticoreReader) Stream() (chan []string, chan error) {
	out := make(chan []string)
	errChan := make(chan error, 1)
	go func() {
		defer close(out)
		defer close(errChan)
		for {
			line, err := mcr.Read()
			if len(line) > 0 {
				out <- line
			}
			if err == nil {
				continue
			}
			if err == io.EOF {
				return
			}
			errChan <- err
			return
		}
	}()
	return out, errChan
}

// Read reads one record from r.  The record is a slice of strings with each
// string representing one field.  In the background, the internal io.Reader
// will be read from ahead of the caller utilizing Read() to pull every row
func (mcr *MulticoreReader) Read() ([]string, error) {
	if mcr.finalError != nil {
		return nil, mcr.finalError
	}
	mcr.start()
	for {
		line, ok := mcr.queue[mcr.place]
		if !ok {
			break // next value isn't in the queue, move on
		}
		delete(mcr.queue, mcr.place)
		mcr.place++
		if len(line) == 0 {
			continue
		}
		return line, nil
	}
	for line := range mcr.lineout {
		if line.num == mcr.place {
			mcr.place++
			if len(line.data) == 0 {
				return mcr.Read()
			}
			return line.data, nil
		}
		mcr.queue[line.num] = line.data
		// keep going, fetch the next line since this one was out of order
	}
	mcr.finalError = <-mcr.errChan
	return nil, mcr.finalError
}

func (mcr *MulticoreReader) startReading(err1 chan error) {
	defer close(mcr.linein)
	linenum := 0
	bytesreader := bufio.NewReader(mcr.reader)
	for {
		line, err := bytesreader.ReadBytes('\n')
		if len(line) > 0 {
			select {
			case mcr.linein <- linein{
				data: line,
				num:  linenum,
			}:
				linenum++
			case <-mcr.cancel:
				err1 <- nil
				return
			}
		}
		if err == nil {
			continue
		}
		if err == io.EOF {
			err = nil
		}
		err1 <- err
		return
	}
}

func (mcr *MulticoreReader) parseCSVLines() error {
	var buf bytes.Buffer
	r := csv.NewReader(&buf)
	r.Comma = mcr.Comma
	r.Comment = mcr.Comment
	r.LazyQuotes = mcr.LazyQuotes
	r.TrailingComma = mcr.TrailingComma
	r.TrimLeadingSpace = mcr.TrimLeadingSpace
	for b := range mcr.linein {
		buf.Reset()
		buf.Write(b.data)
		char, _, err := buf.ReadRune()
		if err != nil {
			return err
		}
		if char == '\n' || char == mcr.Comment {
			select {
			case mcr.lineout <- lineout{
				data: nil,
				num:  b.num,
			}:
			case <-mcr.cancel:
				return nil
			}
			continue
		}
		buf.UnreadRune()
		if line, err := r.Read(); err != nil {
			pe, ok := err.(*csv.ParseError)
			if ok {
				pe.Line = b.num + 1
			}
			return err
		} else {
			select {
			case mcr.lineout <- lineout{
				data: line,
				num:  b.num,
			}:
			case <-mcr.cancel:
				return nil
			}
		}
	}
	return nil
}

func (mcr *MulticoreReader) waitForDone(err1, err2 chan error) {
	foundError := <-err1
	for i := 0; i < runtime.NumCPU(); i++ {
		err := <-err2
		if err != nil && err != io.EOF && foundError == nil {
			foundError = err
		}
	}
	if foundError == nil {
		foundError = io.EOF
	}
	close(mcr.lineout)
	mcr.errChan <- foundError
}

func (mcr *MulticoreReader) start() {
	mcr.once.Do(func() {
		err1 := make(chan error, 1)
		err2 := make(chan error)
		go mcr.startReading(err1)
		for i := 0; i < runtime.NumCPU(); i++ {
			go func() {
				err2 <- mcr.parseCSVLines()
			}()
		}
		go mcr.waitForDone(err1, err2)
	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2013, Suresh Sundriyal. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file.

package murmur3_test

import (
	. "../murmur3"
	"testing"
)

func TestAll(t *testing.T) {
	s := []byte("hello")

	x := []byte(`Lorem Ipsum is simply dummy text of the printing and typesetting industry. Lorem Ipsum has been the industry's standard dummy text ever since the 1500s, when an unknown printer took a galley of type and scrambled it to make a type specimen book. It has survived not only five centuries, but also the leap into electronic typesetting, remaining essentially unchanged. It was popularised in the 1960s with the release of Letraset sheets containing Lorem Ipsum passages, and more recently with desktop publishing software like Aldus PageMaker including versions of Lorem Ipsum.`)

	// Test the x86 32-bit version of Murmur3 by hashing 'hello'
	h32 := New32a(0)
	h32.Write(s)
	h := h32.Sum32()
	if h != 1743816747 {
		t.Error("x86_32: ", s, h)
	}

	// Test the x86 32-bit version of Murmur3 by hashing a longer string.
	h32.Reset()
	h32.Write(x)
	h = h32.Sum32()
	if h != 4193992801 {
		t.Error("x86_33: ", x, h)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing 'hello'
	h128 := New64(0)
	h128.Write(s)
	h1, h2 := h128.Sum128()

	if h1 != 14688674573012802306 || h2 != 6565844092913065241 {
		t.Error("x86_64: ", s, h1, h2)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing 'hello' with a seed.
	h128.Reset()
	h128.SetSeed(12345)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 17440987278262125697 || h2 != 15376406881033980724 {
		t.Error("x86_64(seed): ", s, h1, h2)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing a longer string.
	h128.Reset()
	h128.Write(x)
	h1, h2 = h128.Sum128()

	if h1 != 1706326840306453215 || h2 != 5127165288307402704 {
		t.Error("x86_64: ", x, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing 'hello'.
	h128 = New32(0)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 15821672119091348640 || h2 != 11158567162092401078 {
		t.Error("x86: ", s, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing a longer string.
	h128.Reset()
	h128.Write(x)
	h1, h2 = h128.Sum128()

	if h1 != 223949659430422294 || h2 != 10022274208940483369 {
		t.Error("x86: ", x, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing 'hello' with a seed.
	h128.Reset()
	h128.SetSeed(12345)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 7882561715466346695 || h2 != 11883514271246235972 {
		t.Error("x86 (seed): ", s, h1, h2)
	}

}
<commit_msg>test docstring<commit_after>// Copyright (c) 2013, Suresh Sundriyal. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file.

package murmur3_test

import (
	. "../murmur3"
	"testing"
)

func TestAll(t *testing.T) {
	s := []byte("hello")

	x := []byte(`Lorem Ipsum is simply dummy text of the printing and typesetting industry. Lorem Ipsum has been the industry's standard dummy text ever since the 1500s, when an unknown printer took a galley of type and scrambled it to make a type specimen book. It has survived not only five centuries, but also the leap into electronic typesetting, remaining essentially unchanged. It was popularised in the 1960s with the release of Letraset sheets containing Lorem Ipsum passages, and more recently with desktop publishing software like Aldus PageMaker including versions of Lorem Ipsum.`)

	// Test the x86 32-bit version of Murmur3 by hashing 'hello'.
	h32 := New32a(0)
	h32.Write(s)
	h := h32.Sum32()
	if h != 1743816747 {
		t.Error("x86_32: ", s, h)
	}

	// Test the x86 32-bit version of Murmur3 by hashing a longer string.
	h32.Reset()
	h32.Write(x)
	h = h32.Sum32()
	if h != 4193992801 {
		t.Error("x86_33: ", x, h)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing 'hello'.
	h128 := New64(0)
	h128.Write(s)
	h1, h2 := h128.Sum128()

	if h1 != 14688674573012802306 || h2 != 6565844092913065241 {
		t.Error("x86_64: ", s, h1, h2)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing 'hello' with a seed.
	h128.Reset()
	h128.SetSeed(12345)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 17440987278262125697 || h2 != 15376406881033980724 {
		t.Error("x86_64(seed): ", s, h1, h2)
	}

	// Test the x86_64 128-bit version of Murmur3 by hashing a longer string.
	h128.Reset()
	h128.Write(x)
	h1, h2 = h128.Sum128()

	if h1 != 1706326840306453215 || h2 != 5127165288307402704 {
		t.Error("x86_64: ", x, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing 'hello'.
	h128 = New32(0)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 15821672119091348640 || h2 != 11158567162092401078 {
		t.Error("x86: ", s, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing a longer string.
	h128.Reset()
	h128.Write(x)
	h1, h2 = h128.Sum128()

	if h1 != 223949659430422294 || h2 != 10022274208940483369 {
		t.Error("x86: ", x, h1, h2)
	}

	// Test the x86 128-bit version of Murmur3 by hashing 'hello' with a seed.
	h128.Reset()
	h128.SetSeed(12345)
	h128.Write(s)
	h1, h2 = h128.Sum128()

	if h1 != 7882561715466346695 || h2 != 11883514271246235972 {
		t.Error("x86 (seed): ", s, h1, h2)
	}

}
<|endoftext|>


<commit_before>package files

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

type fakeFile struct {
	fileName  string
	fileSize  int64
	fakeFiles []fakeFile
}

func (f fakeFile) Name() string       { return f.fileName }
func (f fakeFile) Size() int64        { return f.fileSize }
func (f fakeFile) Mode() os.FileMode  { return 0 }
func (f fakeFile) ModTime() time.Time { return time.Now() }
func (f fakeFile) IsDir() bool        { return len(f.fakeFiles) > 0 }
func (f fakeFile) Sys() interface{}   { return nil }

func createReadDir(ff fakeFile) ReadDir {
	return func(path string) ([]os.FileInfo, error) {
		names := strings.Split(path, "/")
		fakeFolder := ff
		var found bool
		for _, name := range names {
			found = false
			for _, testFile := range fakeFolder.fakeFiles {
				if testFile.fileName == name {
					fakeFolder = testFile
					found = true
					break
				}
			}
			if !found {
				return []os.FileInfo{}, fmt.Errorf("file not found")
			}

		}
		result := make([]os.FileInfo, len(fakeFolder.fakeFiles))
		for i, resultFile := range fakeFolder.fakeFiles {
			result[i] = resultFile
		}
		return result, nil
	}
}

func TestFilePath(t *testing.T) {
	root := NewTestFolder("root",
		NewTestFolder("folder1",
			NewTestFile("file1", 0),
		),
	)
	want := filepath.Join("root", "folder1", "file1")
	file1 := FindTestFile(root, "file1")
	assert.Equal(t, want, file1.Path())
}

func TestWalkFolderOnSimpleDir(t *testing.T) {
	testStructure := fakeFile{"a", 0, []fakeFile{
		{"b", 0, []fakeFile{
			{"c", 100, []fakeFile{}},
			{"d", 0, []fakeFile{
				{"e", 50, []fakeFile{}},
				{"f", 30, []fakeFile{}},
				{"g", 70, []fakeFile{ //thisfolder should get ignored
					{"h", 10, []fakeFile{}},
					{"i", 20, []fakeFile{}},
				}},
			}},
		}},
	}}
	dummyIgnoreFunction := func(p string) bool { return p == filepath.Join("b", "d", "g") }
	progress := make(chan int, 3)
	result := WalkFolder("b", createReadDir(testStructure), dummyIgnoreFunction, progress)
	buildExpected := func() *File {
		b := &File{"b", nil, 180, true, []*File{}}
		c := &File{"c", b, 100, false, []*File{}}
		d := &File{"d", b, 80, true, []*File{}}
		b.Files = []*File{c, d}

		e := &File{"e", nil, 50, false, []*File{}}
		e.Parent = d
		f := &File{"f", nil, 30, false, []*File{}}
		g := &File{"g", nil, 0, true, []*File{}}
		f.Parent = d
		g.Parent = d
		d.Files = []*File{e, f, g}

		return b
	}
	expected := buildExpected()
	assert.Equal(t, expected, result)
	resultProgress := 0
	resultProgress += <-progress
	resultProgress += <-progress
	_, more := <-progress
	assert.Equal(t, 2, resultProgress)
	assert.False(t, more, "the progress channel should be closed")
}

func TestWalkFolderHandlesError(t *testing.T) {
	failing := func(path string) ([]os.FileInfo, error) {
		return []os.FileInfo{}, errors.New("Not found")
	}
	progress := make(chan int, 2)
	result := WalkFolder("xyz", failing, func(string) bool { return false }, progress)
	assert.Equal(t, File{}, *result, "WalkFolder didn't return empty file on ReadDir failure")
}
<commit_msg>test: fix windows path test again, this time for real<commit_after>package files

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

type fakeFile struct {
	fileName  string
	fileSize  int64
	fakeFiles []fakeFile
}

func (f fakeFile) Name() string       { return f.fileName }
func (f fakeFile) Size() int64        { return f.fileSize }
func (f fakeFile) Mode() os.FileMode  { return 0 }
func (f fakeFile) ModTime() time.Time { return time.Now() }
func (f fakeFile) IsDir() bool        { return len(f.fakeFiles) > 0 }
func (f fakeFile) Sys() interface{}   { return nil }

func createReadDir(ff fakeFile) ReadDir {
	return func(path string) ([]os.FileInfo, error) {
		names := strings.Split(filepath.ToSlash(path), "/") // oh windows, I'm looking at you you silly bugger
		fakeFolder := ff
		var found bool
		for _, name := range names {
			found = false
			for _, testFile := range fakeFolder.fakeFiles {
				if testFile.fileName == name {
					fakeFolder = testFile
					found = true
					break
				}
			}
			if !found {
				return []os.FileInfo{}, fmt.Errorf("file not found")
			}

		}
		result := make([]os.FileInfo, len(fakeFolder.fakeFiles))
		for i, resultFile := range fakeFolder.fakeFiles {
			result[i] = resultFile
		}
		return result, nil
	}
}

func TestFilePath(t *testing.T) {
	root := NewTestFolder("root",
		NewTestFolder("folder1",
			NewTestFile("file1", 0),
		),
	)
	want := filepath.Join("root", "folder1", "file1")
	file1 := FindTestFile(root, "file1")
	assert.Equal(t, want, file1.Path())
}

func TestWalkFolderOnSimpleDir(t *testing.T) {
	testStructure := fakeFile{"a", 0, []fakeFile{
		{"b", 0, []fakeFile{
			{"c", 100, []fakeFile{}},
			{"d", 0, []fakeFile{
				{"e", 50, []fakeFile{}},
				{"f", 30, []fakeFile{}},
				{"g", 70, []fakeFile{ //thisfolder should get ignored
					{"h", 10, []fakeFile{}},
					{"i", 20, []fakeFile{}},
				}},
			}},
		}},
	}}
	dummyIgnoreFunction := func(p string) bool { return p == filepath.Join("b", "d", "g") }
	progress := make(chan int, 3)
	result := WalkFolder("b", createReadDir(testStructure), dummyIgnoreFunction, progress)
	buildExpected := func() *File {
		b := &File{"b", nil, 180, true, []*File{}}
		c := &File{"c", b, 100, false, []*File{}}
		d := &File{"d", b, 80, true, []*File{}}
		b.Files = []*File{c, d}

		e := &File{"e", nil, 50, false, []*File{}}
		e.Parent = d
		f := &File{"f", nil, 30, false, []*File{}}
		g := &File{"g", nil, 0, true, []*File{}}
		f.Parent = d
		g.Parent = d
		d.Files = []*File{e, f, g}

		return b
	}
	expected := buildExpected()
	assert.Equal(t, expected, result)
	resultProgress := 0
	resultProgress += <-progress
	resultProgress += <-progress
	_, more := <-progress
	assert.Equal(t, 2, resultProgress)
	assert.False(t, more, "the progress channel should be closed")
}

func TestWalkFolderHandlesError(t *testing.T) {
	failing := func(path string) ([]os.FileInfo, error) {
		return []os.FileInfo{}, errors.New("Not found")
	}
	progress := make(chan int, 2)
	result := WalkFolder("xyz", failing, func(string) bool { return false }, progress)
	assert.Equal(t, File{}, *result, "WalkFolder didn't return empty file on ReadDir failure")
}
<|endoftext|>


<commit_before><commit_msg>Make http requests reusable<commit_after><|endoftext|>


<commit_before>// Setup helps initialize applications.
package setup

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Shyp/rickover/Godeps/_workspace/src/github.com/Shyp/go-simple-metrics"
	"github.com/Shyp/rickover/models"
	"github.com/Shyp/rickover/models/archived_jobs"
	"github.com/Shyp/rickover/models/db"
	"github.com/Shyp/rickover/models/jobs"
	"github.com/Shyp/rickover/models/queued_jobs"
)

var mu sync.Mutex

// TODO not sure for the best place for this to live.
var activeQueriesStmt *sql.Stmt

func prepare() (err error) {
	if !db.Connected() {
		return errors.New("No DB connection was established, can't query")
	}

	activeQueriesStmt, err = db.Conn.Prepare(`-- setup.GetActiveQueries
SELECT count(*) FROM pg_stat_activity
WHERE state='active'
	`)
	return
}

func GetActiveQueries() (count int64, err error) {
	err = activeQueriesStmt.QueryRow().Scan(&count)
	return
}

// TODO all of these should use a different database connection than the server
// or the worker, to avoid contention.
func MeasureActiveQueries(interval time.Duration) {
	for _ = range time.Tick(interval) {
		count, err := GetActiveQueries()
		if err == nil {
			go metrics.Measure("active_queries.count", count)
		} else {
			go metrics.Increment("active_queries.error")
		}
	}
}

func MeasureQueueDepth(interval time.Duration) {
	for _ = range time.Tick(interval) {
		allCount, readyCount, err := queued_jobs.CountReadyAndAll()
		if err == nil {
			go metrics.Measure("queue_depth.all", int64(allCount))
			go metrics.Measure("queue_depth.ready", int64(readyCount))
		} else {
			go metrics.Increment("queue_depth.error")
		}
	}
}

func MeasureInProgressJobs(interval time.Duration) {
	for _ = range time.Tick(interval) {
		m, err := queued_jobs.GetCountsByStatus(models.StatusInProgress)
		if err == nil {
			count := int64(0)
			for k, v := range m {
				count += v
				go metrics.Measure(fmt.Sprintf("queued_jobs.%s.in_progress", k), v)
			}
			go metrics.Measure("queued_jobs.in_progress", count)
		} else {
			go metrics.Increment("queued_jobs.in_progress.error")
		}
	}
}

// DB initializes a connection to the database, and prepares queries on all
// models.
func DB(connector db.Connector, dbConns int) error {
	mu.Lock()
	defer mu.Unlock()
	if db.Conn == nil {
		db.Conn = &sql.DB{}
	} else {
		if err := db.Conn.Ping(); err == nil {
			// Already connected.
			return nil
		}
	}
	if err := connector.Connect(db.Conn, dbConns); err != nil {
		return errors.New("Could not establish a database connection: " + err.Error())
	}
	if err := db.Conn.Ping(); err != nil {
		return errors.New("Could not establish a database connection: " + err.Error())
	}
	return PrepareAll()
}

func PrepareAll() error {
	if err := jobs.Setup(); err != nil {
		return err
	}
	if err := queued_jobs.Setup(); err != nil {
		return err
	}
	if err := archived_jobs.Setup(); err != nil {
		return err
	}
	if err := prepare(); err != nil {
		return err
	}
	return nil
}
<commit_msg>Add better doc string<commit_after>// Package setup helps initialize the database and all queries.
package setup

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/Shyp/rickover/Godeps/_workspace/src/github.com/Shyp/go-simple-metrics"
	"github.com/Shyp/rickover/models"
	"github.com/Shyp/rickover/models/archived_jobs"
	"github.com/Shyp/rickover/models/db"
	"github.com/Shyp/rickover/models/jobs"
	"github.com/Shyp/rickover/models/queued_jobs"
)

var mu sync.Mutex

// TODO not sure for the best place for this to live.
var activeQueriesStmt *sql.Stmt

func prepare() (err error) {
	if !db.Connected() {
		return errors.New("No DB connection was established, can't query")
	}

	activeQueriesStmt, err = db.Conn.Prepare(`-- setup.GetActiveQueries
SELECT count(*) FROM pg_stat_activity
WHERE state='active'
	`)
	return
}

func GetActiveQueries() (count int64, err error) {
	err = activeQueriesStmt.QueryRow().Scan(&count)
	return
}

// TODO all of these should use a different database connection than the server
// or the worker, to avoid contention.
func MeasureActiveQueries(interval time.Duration) {
	for _ = range time.Tick(interval) {
		count, err := GetActiveQueries()
		if err == nil {
			go metrics.Measure("active_queries.count", count)
		} else {
			go metrics.Increment("active_queries.error")
		}
	}
}

func MeasureQueueDepth(interval time.Duration) {
	for _ = range time.Tick(interval) {
		allCount, readyCount, err := queued_jobs.CountReadyAndAll()
		if err == nil {
			go metrics.Measure("queue_depth.all", int64(allCount))
			go metrics.Measure("queue_depth.ready", int64(readyCount))
		} else {
			go metrics.Increment("queue_depth.error")
		}
	}
}

func MeasureInProgressJobs(interval time.Duration) {
	for _ = range time.Tick(interval) {
		m, err := queued_jobs.GetCountsByStatus(models.StatusInProgress)
		if err == nil {
			count := int64(0)
			for k, v := range m {
				count += v
				go metrics.Measure(fmt.Sprintf("queued_jobs.%s.in_progress", k), v)
			}
			go metrics.Measure("queued_jobs.in_progress", count)
		} else {
			go metrics.Increment("queued_jobs.in_progress.error")
		}
	}
}

// DB initializes a connection to the database, and prepares queries on all
// models.
func DB(connector db.Connector, dbConns int) error {
	mu.Lock()
	defer mu.Unlock()
	if db.Conn == nil {
		db.Conn = &sql.DB{}
	} else {
		if err := db.Conn.Ping(); err == nil {
			// Already connected.
			return nil
		}
	}
	if err := connector.Connect(db.Conn, dbConns); err != nil {
		return errors.New("Could not establish a database connection: " + err.Error())
	}
	if err := db.Conn.Ping(); err != nil {
		return errors.New("Could not establish a database connection: " + err.Error())
	}
	return PrepareAll()
}

func PrepareAll() error {
	if err := jobs.Setup(); err != nil {
		return err
	}
	if err := queued_jobs.Setup(); err != nil {
		return err
	}
	if err := archived_jobs.Setup(); err != nil {
		return err
	}
	if err := prepare(); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package build

import (
	"bufio"
	"bytes"
	"fmt"
	"regexp"
	"strings"
)

func readDockerfileToTarget(dockerfile string, target string) (string, error) {
	targetPattern := regexp.MustCompile(fmt.Sprintf(`\s+[aA][sS]\s+%s\s*$`, target))
	nextPattern := regexp.MustCompile(`\s+[aA][sS]\s+.+\s*$`)

	scanner := bufio.NewScanner(strings.NewReader(dockerfile))
	buf := bytes.NewBuffer(make([]byte, 0, len(dockerfile)))

	targetFound := false
	for scanner.Scan() {
		line := scanner.Text()
		// Read dockerfile until the end of the target stage
		if targetFound && nextPattern.MatchString(line) {
			break
		} else if targetPattern.MatchString(line) {
			targetFound = true
		}
		buf.WriteString(line)
		buf.WriteString("\n")
	}
	if err := scanner.Err(); err != nil {
		return "", fmt.Errorf("[bug] failed to process generated dockerfile: %v", err)
	}
	if !targetFound {
		return "", fmt.Errorf(`build target "%s" does not exist in Dockerfile - double-check that the target is not misspelled`, target)
	}
	return buf.String(), nil
}
<commit_msg>Copy comments and directives as is<commit_after>package build

import (
	"bufio"
	"bytes"
	"fmt"
	"regexp"
	"strings"
)

func readDockerfileToTarget(dockerfile string, target string) (string, error) {
	targetPattern := regexp.MustCompile(fmt.Sprintf(`\s+[aA][sS]\s+%s\s*$`, target))
	nextPattern := regexp.MustCompile(`\s+[aA][sS]\s+.+\s*$`)
	commentOrDirectivePattern := regexp.MustCompile(`^#.*$`)

	scanner := bufio.NewScanner(strings.NewReader(dockerfile))
	buf := bytes.NewBuffer(make([]byte, 0, len(dockerfile)))

	targetFound := false
	for scanner.Scan() {
		line := scanner.Text()
		// Copy comments / directives as is
		if !commentOrDirectivePattern.MatchString(line) {
			// Read dockerfile until the end of the target stage
			if targetFound && nextPattern.MatchString(line) {
				break
			} else if targetPattern.MatchString(line) {
				targetFound = true
			}
		}
		buf.WriteString(line)
		buf.WriteString("\n")
	}
	if err := scanner.Err(); err != nil {
		return "", fmt.Errorf("[bug] failed to process generated dockerfile: %v", err)
	}
	if !targetFound {
		return "", fmt.Errorf(`build target "%s" does not exist in Dockerfile - double-check that the target is not misspelled`, target)
	}
	return buf.String(), nil
}
<|endoftext|>