<commit_before>// Package relyq provides a reliable queue backed by redis
package relyq

import (
	"fmt"
	"github.com/Rafflecopter/golang-simpleq/simpleq"
	"github.com/garyburd/redigo/redis"
	"github.com/yanatan16/gowaiter"
	"io"
)

// A reliable redis-backed queue
type Queue struct {
	// The underlying simpleqs
	Todo, Doing, Done, Failed *simpleq.Queue
	Storage                   Storage
	Cfg                       *Config
	listener                  *Listener
}

// Configuration for Relyq
type Config struct {
	// Required: Prefix on simpleq.Queue key names for redis
	Prefix string
	// Field in any object which contains a unique identifier
	// Defaults to "id"
	IdField string
	// Redis delimeter. Defaults to ":"
	Delimiter string
	// Clean finish (i.e. no Done queue)
	// Defaults to false
	UseDoneQueue bool
	// Should we keep the task stored after they are done?
	// Defaults to false
	KeepDoneTasks bool
}

// A useful alias for a task
type Ider interface {
	// Ensure an id exists by creating it if necessary. Always return it.
	Id() []byte
}

// Storage interface
type Storage interface {
	// Get a task object
	Get(taskid []byte, task interface{}) error
	// Save a task object
	Set(task interface{}, taskid []byte) error
	// Delete the task object in the storage
	Del(taskid []byte) error
	// End the Storage connection
	io.Closer
}

// Create a reliable queue
func New(pool *redis.Pool, storage Storage, cfg *Config) *Queue {
	cfg.Defaults()

	rq := &Queue{
		Todo:    simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"todo"),
		Doing:   simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"doing"),
		Failed:  simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"failed"),
		Storage: storage,
		Cfg:     cfg,
	}

	if cfg.UseDoneQueue {
		rq.Done = simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"done")
	}

	return rq
}

// Push a task onto the queue
func (q *Queue) Push(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if err := q.Storage.Set(task, id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	go func() {
		if _, err := q.Todo.Push(id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Move the next task to the Doing queue. Will decode into task. Returns ok as false if nothing happened
func (q *Queue) Process(task Ider) (ok bool, err error) {
	id, err := q.Todo.PopPipe(q.Doing)
	if err != nil {
		return false, err
	} else if id == nil {
		return false, nil
	}

	err = q.Storage.Get(id, task)
	return err == nil, err
}

// Block and process the next task.
func (q *Queue) BProcess(timeout_secs int, task Ider) error {
	id, err := q.Todo.BPopPipe(q.Doing, timeout_secs)
	if err != nil {
		return err
	}

	err = q.Storage.Get(id, task)
	return err
}

// Move a task to the Done queue if in use
// If a task is not in use, delete if CleanFinishKeepStorage is false
// Sometimes a task is in the Failed queue already (maybe timeout) so we check there if not in Finish
func (q *Queue) Finish(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if q.Cfg.KeepDoneTasks {
			if err := q.Storage.Set(task, id); err != nil {
				w.Errors <- err
			}
		} else {
			if err := q.Storage.Del(id); err != nil {
				w.Errors <- err
			}
		}
		w.Done <- true
	}()

	go func() {
		if q.Cfg.UseDoneQueue {
			if n, err := q.Doing.SPullPipe(q.Done, id); err != nil {
				w.Errors <- err
			} else if n == 0 {
				if n, err := q.Failed.SPullPipe(q.Done, id); err != nil {
					w.Errors <- err
				} else if n == 0 {
					w.Errors <- fmt.Errorf("Task %s not found in Doing or Failed queues.", id)
				}
			}
		} else {
			if n, err := q.Doing.Pull(id); err != nil {
				w.Errors <- err
			} else if n == 0 {
				if n, err := q.Failed.Pull(id); err != nil {
					w.Errors <- err
				} else if n == 0 {
					w.Errors <- fmt.Errorf("Task %s not found in Doing or Failed queues.", id)
				}
			}
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Move a task to the Failed queue
func (q *Queue) Fail(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if err := q.Storage.Set(task, id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	go func() {
		if n, err := q.Doing.SPullPipe(q.Failed, id); err != nil {
			w.Errors <- err
		} else if n == 0 {
			w.Errors <- fmt.Errorf("Task %s not found in Doing queue.", id)
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Remove a task from a queue
// If dontDelete (single extra arg) is true, then no delete call will be done for the task
func (q *Queue) Remove(subq *simpleq.Queue, task Ider, keepInStorage ...bool) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if len(keepInStorage) > 0 && keepInStorage[0] {
			if err := q.Storage.Set(task, id); err != nil {
				w.Errors <- err
			}
		} else {
			if err := q.Storage.Del(id); err != nil {
				w.Errors <- err
			}
		}
		w.Done <- true
	}()

	go func() {
		if n, err := subq.Pull(id); err != nil {
			w.Errors <- err
		} else if n == 0 {
			w.Errors <- fmt.Errorf("Task %s not found in queue.", id)
		}
		w.Done <- true
	}()

	return w.Wait()
}

// End the queue
func (q *Queue) Close() error {
	w := waiter.New(5)

	w.Close(q.Todo)
	w.Close(q.Doing)
	w.Close(q.Failed)
	w.Close(q.Done)
	w.Close(q.Storage)

	return w.Wait()
}

func (cfg *Config) Defaults() {
	if cfg.Prefix == "" {
		panic("Prefix required for relyq")
	}

	if cfg.IdField == "" {
		cfg.IdField = "id"
	}

	if cfg.Delimiter == "" {
		cfg.Delimiter = ":"
	}
}
<commit_msg>Small fix for nil.Defaults()<commit_after>// Package relyq provides a reliable queue backed by redis
package relyq

import (
	"fmt"
	"github.com/Rafflecopter/golang-simpleq/simpleq"
	"github.com/garyburd/redigo/redis"
	"github.com/yanatan16/gowaiter"
	"io"
)

// A reliable redis-backed queue
type Queue struct {
	// The underlying simpleqs
	Todo, Doing, Done, Failed *simpleq.Queue
	Storage                   Storage
	Cfg                       *Config
	listener                  *Listener
}

// Configuration for Relyq
type Config struct {
	// Required: Prefix on simpleq.Queue key names for redis
	Prefix string
	// Field in any object which contains a unique identifier
	// Defaults to "id"
	IdField string
	// Redis delimeter. Defaults to ":"
	Delimiter string
	// Clean finish (i.e. no Done queue)
	// Defaults to false
	UseDoneQueue bool
	// Should we keep the task stored after they are done?
	// Defaults to false
	KeepDoneTasks bool
}

// A useful alias for a task
type Ider interface {
	// Ensure an id exists by creating it if necessary. Always return it.
	Id() []byte
}

// Storage interface
type Storage interface {
	// Get a task object
	Get(taskid []byte, task interface{}) error
	// Save a task object
	Set(task interface{}, taskid []byte) error
	// Delete the task object in the storage
	Del(taskid []byte) error
	// End the Storage connection
	io.Closer
}

// Create a reliable queue
func New(pool *redis.Pool, storage Storage, cfg *Config) *Queue {
	cfg.Defaults()

	rq := &Queue{
		Todo:    simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"todo"),
		Doing:   simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"doing"),
		Failed:  simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"failed"),
		Storage: storage,
		Cfg:     cfg,
	}

	if cfg.UseDoneQueue {
		rq.Done = simpleq.New(pool, cfg.Prefix+cfg.Delimiter+"done")
	}

	return rq
}

// Push a task onto the queue
func (q *Queue) Push(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if err := q.Storage.Set(task, id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	go func() {
		if _, err := q.Todo.Push(id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Move the next task to the Doing queue. Will decode into task. Returns ok as false if nothing happened
func (q *Queue) Process(task Ider) (ok bool, err error) {
	id, err := q.Todo.PopPipe(q.Doing)
	if err != nil {
		return false, err
	} else if id == nil {
		return false, nil
	}

	err = q.Storage.Get(id, task)
	return err == nil, err
}

// Block and process the next task.
func (q *Queue) BProcess(timeout_secs int, task Ider) error {
	id, err := q.Todo.BPopPipe(q.Doing, timeout_secs)
	if err != nil {
		return err
	}

	err = q.Storage.Get(id, task)
	return err
}

// Move a task to the Done queue if in use
// If a task is not in use, delete if CleanFinishKeepStorage is false
// Sometimes a task is in the Failed queue already (maybe timeout) so we check there if not in Finish
func (q *Queue) Finish(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if q.Cfg.KeepDoneTasks {
			if err := q.Storage.Set(task, id); err != nil {
				w.Errors <- err
			}
		} else {
			if err := q.Storage.Del(id); err != nil {
				w.Errors <- err
			}
		}
		w.Done <- true
	}()

	go func() {
		if q.Cfg.UseDoneQueue {
			if n, err := q.Doing.SPullPipe(q.Done, id); err != nil {
				w.Errors <- err
			} else if n == 0 {
				if n, err := q.Failed.SPullPipe(q.Done, id); err != nil {
					w.Errors <- err
				} else if n == 0 {
					w.Errors <- fmt.Errorf("Task %s not found in Doing or Failed queues.", id)
				}
			}
		} else {
			if n, err := q.Doing.Pull(id); err != nil {
				w.Errors <- err
			} else if n == 0 {
				if n, err := q.Failed.Pull(id); err != nil {
					w.Errors <- err
				} else if n == 0 {
					w.Errors <- fmt.Errorf("Task %s not found in Doing or Failed queues.", id)
				}
			}
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Move a task to the Failed queue
func (q *Queue) Fail(task Ider) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if err := q.Storage.Set(task, id); err != nil {
			w.Errors <- err
		}
		w.Done <- true
	}()

	go func() {
		if n, err := q.Doing.SPullPipe(q.Failed, id); err != nil {
			w.Errors <- err
		} else if n == 0 {
			w.Errors <- fmt.Errorf("Task %s not found in Doing queue.", id)
		}
		w.Done <- true
	}()

	return w.Wait()
}

// Remove a task from a queue
// If dontDelete (single extra arg) is true, then no delete call will be done for the task
func (q *Queue) Remove(subq *simpleq.Queue, task Ider, keepInStorage ...bool) error {
	id := task.Id()
	w := waiter.New(2)

	go func() {
		if len(keepInStorage) > 0 && keepInStorage[0] {
			if err := q.Storage.Set(task, id); err != nil {
				w.Errors <- err
			}
		} else {
			if err := q.Storage.Del(id); err != nil {
				w.Errors <- err
			}
		}
		w.Done <- true
	}()

	go func() {
		if n, err := subq.Pull(id); err != nil {
			w.Errors <- err
		} else if n == 0 {
			w.Errors <- fmt.Errorf("Task %s not found in queue.", id)
		}
		w.Done <- true
	}()

	return w.Wait()
}

// End the queue
func (q *Queue) Close() error {
	w := waiter.New(5)

	w.Close(q.Todo)
	w.Close(q.Doing)
	w.Close(q.Failed)
	w.Close(q.Done)
	w.Close(q.Storage)

	return w.Wait()
}

func (cfg *Config) Defaults() {
	if cfg == nil || cfg.Prefix == "" {
		panic("Prefix required for relyq")
	}

	if cfg.IdField == "" {
		cfg.IdField = "id"
	}

	if cfg.Delimiter == "" {
		cfg.Delimiter = ":"
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

package grpc_retry

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
	"golang.org/x/net/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	AttemptMetadataKey = "x-retry-attempty"
)

// UnaryClientInterceptor returns a new retrying unary client interceptor.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
func UnaryClientInterceptor(optFuncs ...CallOption) grpc.UnaryClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(parentCtx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return invoker(parentCtx, method, req, reply, cc, grpcOpts...)
		}
		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
				return err
			}
			callCtx := perCallContext(parentCtx, callOpts, attempt)
			lastErr = invoker(callCtx, method, req, reply, cc, grpcOpts...)
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			if lastErr == nil {
				return nil
			}
			logTrace(parentCtx, "grpc_retry attempt: %d, got err: %v", attempt, lastErr)
			if isContextError(lastErr) {
				if parentCtx.Err() != nil {
					logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
					// its the parent context deadline or cancellation.
					return lastErr
				} else if callOpts.perCallTimeout != 0 {
					// We have set a perCallTimeout in the retry middleware, which would result in a context error if
					// the deadline was exceeded, in which case try again.
					logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
					continue
				}
			}
			if !isRetriable(lastErr, callOpts) {
				return lastErr
			}
		}
		return lastErr
	}
}

// StreamClientInterceptor returns a new retrying stream client interceptor for server side streaming calls.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
//
// Retry logic is available *only for ServerStreams*, i.e. 1:n streams, as the internal logic needs
// to buffer the messages sent by the client. If retry is enabled on any other streams (ClientStreams,
// BidiStreams), the retry interceptor will fail the call.
func StreamClientInterceptor(optFuncs ...CallOption) grpc.StreamClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(parentCtx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return streamer(parentCtx, desc, cc, method, grpcOpts...)
		}
		if desc.ClientStreams {
			return nil, status.Errorf(codes.Unimplemented, "grpc_retry: cannot retry on ClientStreams, set grpc_retry.Disable()")
		}

		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
				return nil, err
			}
			callCtx := perCallContext(parentCtx, callOpts, 0)

			var newStreamer grpc.ClientStream
			newStreamer, lastErr = streamer(callCtx, desc, cc, method, grpcOpts...)
			if lastErr == nil {
				retryingStreamer := &serverStreamingRetryingStream{
					ClientStream: newStreamer,
					callOpts:     callOpts,
					parentCtx:    parentCtx,
					streamerCall: func(ctx context.Context) (grpc.ClientStream, error) {
						return streamer(ctx, desc, cc, method, grpcOpts...)
					},
				}
				return retryingStreamer, nil
			}

			logTrace(parentCtx, "grpc_retry attempt: %d, got err: %v", attempt, lastErr)
			if isContextError(lastErr) {
				if parentCtx.Err() != nil {
					logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
					// its the parent context deadline or cancellation.
					return nil, lastErr
				} else if callOpts.perCallTimeout != 0 {
					// We have set a perCallTimeout in the retry middleware, which would result in a context error if
					// the deadline was exceeded, in which case try again.
					logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
					continue
				}
			}
			if !isRetriable(lastErr, callOpts) {
				return nil, lastErr
			}
		}
		return nil, lastErr
	}
}

// type serverStreamingRetryingStream is the implementation of grpc.ClientStream that acts as a
// proxy to the underlying call. If any of the RecvMsg() calls fail, it will try to reestablish
// a new ClientStream according to the retry policy.
type serverStreamingRetryingStream struct {
	grpc.ClientStream
	bufferedSends []interface{} // single message that the client can sen
	wasClosedSend bool          // indicates that CloseSend was closed
	parentCtx     context.Context
	callOpts      *options
	streamerCall  func(ctx context.Context) (grpc.ClientStream, error)
	mu            sync.RWMutex
}

func (s *serverStreamingRetryingStream) setStream(clientStream grpc.ClientStream) {
	s.mu.Lock()
	s.ClientStream = clientStream
	s.mu.Unlock()
}

func (s *serverStreamingRetryingStream) getStream() grpc.ClientStream {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.ClientStream
}

func (s *serverStreamingRetryingStream) SendMsg(m interface{}) error {
	s.mu.Lock()
	s.bufferedSends = append(s.bufferedSends, m)
	s.mu.Unlock()
	return s.getStream().SendMsg(m)
}

func (s *serverStreamingRetryingStream) CloseSend() error {
	s.mu.Lock()
	s.wasClosedSend = true
	s.mu.Unlock()
	return s.getStream().CloseSend()
}

func (s *serverStreamingRetryingStream) Header() (metadata.MD, error) {
	return s.getStream().Header()
}

func (s *serverStreamingRetryingStream) Trailer() metadata.MD {
	return s.getStream().Trailer()
}

func (s *serverStreamingRetryingStream) RecvMsg(m interface{}) error {
	attemptRetry, lastErr := s.receiveMsgAndIndicateRetry(m)
	if !attemptRetry {
		return lastErr // success or hard failure
	}
	// We start off from attempt 1, because zeroth was already made on normal SendMsg().
	for attempt := uint(1); attempt < s.callOpts.max; attempt++ {
		if err := waitRetryBackoff(attempt, s.parentCtx, s.callOpts); err != nil {
			return err
		}
		callCtx := perCallContext(s.parentCtx, s.callOpts, attempt)
		newStream, err := s.reestablishStreamAndResendBuffer(callCtx)
		if err != nil {
			// Retry dial and transport errors of establishing stream as grpc doesn't retry.
			if isRetriable(err, s.callOpts) {
				continue
			}
			return err
		}

		s.setStream(newStream)
		attemptRetry, lastErr = s.receiveMsgAndIndicateRetry(m)
		//fmt.Printf("Received message and indicate: %v  %v\n", attemptRetry, lastErr)
		if !attemptRetry {
			return lastErr
		}
	}
	return lastErr
}

func (s *serverStreamingRetryingStream) receiveMsgAndIndicateRetry(m interface{}) (bool, error) {
	err := s.getStream().RecvMsg(m)
	if err == nil || err == io.EOF {
		return false, err
	}
	if isContextError(err) {
		if s.parentCtx.Err() != nil {
			logTrace(s.parentCtx, "grpc_retry parent context error: %v", s.parentCtx.Err())
			return false, err
		} else if s.callOpts.perCallTimeout != 0 {
			// We have set a perCallTimeout in the retry middleware, which would result in a context error if
			// the deadline was exceeded, in which case try again.
			logTrace(s.parentCtx, "grpc_retry context error from retry call")
			return true, err
		}
	}
	return isRetriable(err, s.callOpts), err
}

func (s *serverStreamingRetryingStream) reestablishStreamAndResendBuffer(
	callCtx context.Context,
) (grpc.ClientStream, error) {
	s.mu.RLock()
	bufferedSends := s.bufferedSends
	s.mu.RUnlock()
	newStream, err := s.streamerCall(callCtx)
	if err != nil {
		logTrace(callCtx, "grpc_retry failed redialing new stream: %v", err)
		return nil, err
	}
	for _, msg := range bufferedSends {
		if err := newStream.SendMsg(msg); err != nil {
			logTrace(callCtx, "grpc_retry failed resending message: %v", err)
			return nil, err
		}
	}
	if err := newStream.CloseSend(); err != nil {
		logTrace(callCtx, "grpc_retry failed CloseSend on new stream %v", err)
		return nil, err
	}
	return newStream, nil
}

func waitRetryBackoff(attempt uint, parentCtx context.Context, callOpts *options) error {
	var waitTime time.Duration = 0
	if attempt > 0 {
		waitTime = callOpts.backoffFunc(parentCtx, attempt)
	}
	if waitTime > 0 {
		logTrace(parentCtx, "grpc_retry attempt: %d, backoff for %v", attempt, waitTime)
		timer := time.NewTimer(waitTime)
		select {
		case <-parentCtx.Done():
			timer.Stop()
			return contextErrToGrpcErr(parentCtx.Err())
		case <-timer.C:
		}
	}
	return nil
}

func isRetriable(err error, callOpts *options) bool {
	errCode := status.Code(err)
	if isContextError(err) {
		// context errors are not retriable based on user settings.
		return false
	}
	for _, code := range callOpts.codes {
		if code == errCode {
			return true
		}
	}
	return false
}

func isContextError(err error) bool {
	code := status.Code(err)
	return code == codes.DeadlineExceeded || code == codes.Canceled
}

func perCallContext(parentCtx context.Context, callOpts *options, attempt uint) context.Context {
	ctx := parentCtx
	if callOpts.perCallTimeout != 0 {
		ctx, _ = context.WithTimeout(ctx, callOpts.perCallTimeout)
	}
	if attempt > 0 && callOpts.includeHeader {
		mdClone := metautils.ExtractOutgoing(ctx).Clone().Set(AttemptMetadataKey, fmt.Sprintf("%d", attempt))
		ctx = mdClone.ToOutgoing(ctx)
	}
	return ctx
}

func contextErrToGrpcErr(err error) error {
	switch err {
	case context.DeadlineExceeded:
		return status.Error(codes.DeadlineExceeded, err.Error())
	case context.Canceled:
		return status.Error(codes.Canceled, err.Error())
	default:
		return status.Error(codes.Unknown, err.Error())
	}
}

func logTrace(ctx context.Context, format string, a ...interface{}) {
	tr, ok := trace.FromContext(ctx)
	if !ok {
		return
	}
	tr.LazyPrintf(format, a...)
}
<commit_msg>use strconv.FormatUint instead of fmt.Sprintf (#503)<commit_after>// Copyright 2016 Michal Witkowski. All Rights Reserved.
// See LICENSE for licensing terms.

package grpc_retry

import (
	"context"
	"io"
	"strconv"
	"sync"
	"time"

	"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
	"golang.org/x/net/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

const (
	AttemptMetadataKey = "x-retry-attempty"
)

// UnaryClientInterceptor returns a new retrying unary client interceptor.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
func UnaryClientInterceptor(optFuncs ...CallOption) grpc.UnaryClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(parentCtx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return invoker(parentCtx, method, req, reply, cc, grpcOpts...)
		}
		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
				return err
			}
			callCtx := perCallContext(parentCtx, callOpts, attempt)
			lastErr = invoker(callCtx, method, req, reply, cc, grpcOpts...)
			// TODO(mwitkow): Maybe dial and transport errors should be retriable?
			if lastErr == nil {
				return nil
			}
			logTrace(parentCtx, "grpc_retry attempt: %d, got err: %v", attempt, lastErr)
			if isContextError(lastErr) {
				if parentCtx.Err() != nil {
					logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
					// its the parent context deadline or cancellation.
					return lastErr
				} else if callOpts.perCallTimeout != 0 {
					// We have set a perCallTimeout in the retry middleware, which would result in a context error if
					// the deadline was exceeded, in which case try again.
					logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
					continue
				}
			}
			if !isRetriable(lastErr, callOpts) {
				return lastErr
			}
		}
		return lastErr
	}
}

// StreamClientInterceptor returns a new retrying stream client interceptor for server side streaming calls.
//
// The default configuration of the interceptor is to not retry *at all*. This behaviour can be
// changed through options (e.g. WithMax) on creation of the interceptor or on call (through grpc.CallOptions).
//
// Retry logic is available *only for ServerStreams*, i.e. 1:n streams, as the internal logic needs
// to buffer the messages sent by the client. If retry is enabled on any other streams (ClientStreams,
// BidiStreams), the retry interceptor will fail the call.
func StreamClientInterceptor(optFuncs ...CallOption) grpc.StreamClientInterceptor {
	intOpts := reuseOrNewWithCallOptions(defaultOptions, optFuncs)
	return func(parentCtx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		grpcOpts, retryOpts := filterCallOptions(opts)
		callOpts := reuseOrNewWithCallOptions(intOpts, retryOpts)
		// short circuit for simplicity, and avoiding allocations.
		if callOpts.max == 0 {
			return streamer(parentCtx, desc, cc, method, grpcOpts...)
		}
		if desc.ClientStreams {
			return nil, status.Errorf(codes.Unimplemented, "grpc_retry: cannot retry on ClientStreams, set grpc_retry.Disable()")
		}

		var lastErr error
		for attempt := uint(0); attempt < callOpts.max; attempt++ {
			if err := waitRetryBackoff(attempt, parentCtx, callOpts); err != nil {
				return nil, err
			}
			callCtx := perCallContext(parentCtx, callOpts, 0)

			var newStreamer grpc.ClientStream
			newStreamer, lastErr = streamer(callCtx, desc, cc, method, grpcOpts...)
			if lastErr == nil {
				retryingStreamer := &serverStreamingRetryingStream{
					ClientStream: newStreamer,
					callOpts:     callOpts,
					parentCtx:    parentCtx,
					streamerCall: func(ctx context.Context) (grpc.ClientStream, error) {
						return streamer(ctx, desc, cc, method, grpcOpts...)
					},
				}
				return retryingStreamer, nil
			}

			logTrace(parentCtx, "grpc_retry attempt: %d, got err: %v", attempt, lastErr)
			if isContextError(lastErr) {
				if parentCtx.Err() != nil {
					logTrace(parentCtx, "grpc_retry attempt: %d, parent context error: %v", attempt, parentCtx.Err())
					// its the parent context deadline or cancellation.
					return nil, lastErr
				} else if callOpts.perCallTimeout != 0 {
					// We have set a perCallTimeout in the retry middleware, which would result in a context error if
					// the deadline was exceeded, in which case try again.
					logTrace(parentCtx, "grpc_retry attempt: %d, context error from retry call", attempt)
					continue
				}
			}
			if !isRetriable(lastErr, callOpts) {
				return nil, lastErr
			}
		}
		return nil, lastErr
	}
}

// type serverStreamingRetryingStream is the implementation of grpc.ClientStream that acts as a
// proxy to the underlying call. If any of the RecvMsg() calls fail, it will try to reestablish
// a new ClientStream according to the retry policy.
type serverStreamingRetryingStream struct {
	grpc.ClientStream
	bufferedSends []interface{} // single message that the client can sen
	wasClosedSend bool          // indicates that CloseSend was closed
	parentCtx     context.Context
	callOpts      *options
	streamerCall  func(ctx context.Context) (grpc.ClientStream, error)
	mu            sync.RWMutex
}

func (s *serverStreamingRetryingStream) setStream(clientStream grpc.ClientStream) {
	s.mu.Lock()
	s.ClientStream = clientStream
	s.mu.Unlock()
}

func (s *serverStreamingRetryingStream) getStream() grpc.ClientStream {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.ClientStream
}

func (s *serverStreamingRetryingStream) SendMsg(m interface{}) error {
	s.mu.Lock()
	s.bufferedSends = append(s.bufferedSends, m)
	s.mu.Unlock()
	return s.getStream().SendMsg(m)
}

func (s *serverStreamingRetryingStream) CloseSend() error {
	s.mu.Lock()
	s.wasClosedSend = true
	s.mu.Unlock()
	return s.getStream().CloseSend()
}

func (s *serverStreamingRetryingStream) Header() (metadata.MD, error) {
	return s.getStream().Header()
}

func (s *serverStreamingRetryingStream) Trailer() metadata.MD {
	return s.getStream().Trailer()
}

func (s *serverStreamingRetryingStream) RecvMsg(m interface{}) error {
	attemptRetry, lastErr := s.receiveMsgAndIndicateRetry(m)
	if !attemptRetry {
		return lastErr // success or hard failure
	}
	// We start off from attempt 1, because zeroth was already made on normal SendMsg().
	for attempt := uint(1); attempt < s.callOpts.max; attempt++ {
		if err := waitRetryBackoff(attempt, s.parentCtx, s.callOpts); err != nil {
			return err
		}
		callCtx := perCallContext(s.parentCtx, s.callOpts, attempt)
		newStream, err := s.reestablishStreamAndResendBuffer(callCtx)
		if err != nil {
			// Retry dial and transport errors of establishing stream as grpc doesn't retry.
			if isRetriable(err, s.callOpts) {
				continue
			}
			return err
		}

		s.setStream(newStream)
		attemptRetry, lastErr = s.receiveMsgAndIndicateRetry(m)
		//fmt.Printf("Received message and indicate: %v  %v\n", attemptRetry, lastErr)
		if !attemptRetry {
			return lastErr
		}
	}
	return lastErr
}

func (s *serverStreamingRetryingStream) receiveMsgAndIndicateRetry(m interface{}) (bool, error) {
	err := s.getStream().RecvMsg(m)
	if err == nil || err == io.EOF {
		return false, err
	}
	if isContextError(err) {
		if s.parentCtx.Err() != nil {
			logTrace(s.parentCtx, "grpc_retry parent context error: %v", s.parentCtx.Err())
			return false, err
		} else if s.callOpts.perCallTimeout != 0 {
			// We have set a perCallTimeout in the retry middleware, which would result in a context error if
			// the deadline was exceeded, in which case try again.
			logTrace(s.parentCtx, "grpc_retry context error from retry call")
			return true, err
		}
	}
	return isRetriable(err, s.callOpts), err
}

func (s *serverStreamingRetryingStream) reestablishStreamAndResendBuffer(
	callCtx context.Context,
) (grpc.ClientStream, error) {
	s.mu.RLock()
	bufferedSends := s.bufferedSends
	s.mu.RUnlock()
	newStream, err := s.streamerCall(callCtx)
	if err != nil {
		logTrace(callCtx, "grpc_retry failed redialing new stream: %v", err)
		return nil, err
	}
	for _, msg := range bufferedSends {
		if err := newStream.SendMsg(msg); err != nil {
			logTrace(callCtx, "grpc_retry failed resending message: %v", err)
			return nil, err
		}
	}
	if err := newStream.CloseSend(); err != nil {
		logTrace(callCtx, "grpc_retry failed CloseSend on new stream %v", err)
		return nil, err
	}
	return newStream, nil
}

func waitRetryBackoff(attempt uint, parentCtx context.Context, callOpts *options) error {
	var waitTime time.Duration = 0
	if attempt > 0 {
		waitTime = callOpts.backoffFunc(parentCtx, attempt)
	}
	if waitTime > 0 {
		logTrace(parentCtx, "grpc_retry attempt: %d, backoff for %v", attempt, waitTime)
		timer := time.NewTimer(waitTime)
		select {
		case <-parentCtx.Done():
			timer.Stop()
			return contextErrToGrpcErr(parentCtx.Err())
		case <-timer.C:
		}
	}
	return nil
}

func isRetriable(err error, callOpts *options) bool {
	errCode := status.Code(err)
	if isContextError(err) {
		// context errors are not retriable based on user settings.
		return false
	}
	for _, code := range callOpts.codes {
		if code == errCode {
			return true
		}
	}
	return false
}

func isContextError(err error) bool {
	code := status.Code(err)
	return code == codes.DeadlineExceeded || code == codes.Canceled
}

func perCallContext(parentCtx context.Context, callOpts *options, attempt uint) context.Context {
	ctx := parentCtx
	if callOpts.perCallTimeout != 0 {
		ctx, _ = context.WithTimeout(ctx, callOpts.perCallTimeout)
	}
	if attempt > 0 && callOpts.includeHeader {
		mdClone := metautils.ExtractOutgoing(ctx).Clone().Set(AttemptMetadataKey, strconv.FormatUint(uint64(attempt), 10))
		ctx = mdClone.ToOutgoing(ctx)
	}
	return ctx
}

func contextErrToGrpcErr(err error) error {
	switch err {
	case context.DeadlineExceeded:
		return status.Error(codes.DeadlineExceeded, err.Error())
	case context.Canceled:
		return status.Error(codes.Canceled, err.Error())
	default:
		return status.Error(codes.Unknown, err.Error())
	}
}

func logTrace(ctx context.Context, format string, a ...interface{}) {
	tr, ok := trace.FromContext(ctx)
	if !ok {
		return
	}
	tr.LazyPrintf(format, a...)
}
<|endoftext|>


<commit_before>package terminal

import "github.com/nsf/termbox-go"
import "strings"
import "github.com/wx13/sith/syntaxcolor"

type Screen struct {
	row, col int
	fg, bg   termbox.Attribute
	colors   map[string]termbox.Attribute
}

func NewScreen() *Screen {
	screen := Screen{
		row: 0,
		col: 0,
		bg:  termbox.ColorBlack,
		fg:  termbox.ColorWhite,
	}
	termbox.Init()
	return &screen
}

func (screen *Screen) Close() {
	screen.Clear()
	termbox.Flush()
	termbox.Close()
}

func (screen *Screen) Open() {
	termbox.Init()
}

func (screen *Screen) Flush() {
	termbox.Flush()
}

func (screen *Screen) SetCursor(r, c int) {
	screen.row = r
	screen.col = c
	termbox.SetCursor(c, r)
}

func (screen *Screen) Clear() {
	termbox.Clear(screen.fg, screen.bg)
	cols, rows := termbox.Size()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(" ", cols))
	}
}

func (screen *Screen) ReallyClear() {
	cols, rows := termbox.Size()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(".", cols))
	}
	screen.Flush()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(" ", cols))
	}
	screen.Flush()
}

func (screen *Screen) DecorateStatusLine() {
	cells := termbox.CellBuffer()
	cols, rows := termbox.Size()
	for col := 0; col < cols; col++ {
		j := (rows-1)*cols + col
		cells[j].Fg = termbox.ColorBlue
	}
}

func (screen *Screen) WriteString(row, col int, s string) {
	screen.WriteStringColor(row, col, s, screen.fg, screen.bg)
}

func (screen *Screen) Colorize(row int, colors []syntaxcolor.LineColor, offset int) {
	cells := termbox.CellBuffer()
	cols, _ := termbox.Size()
	for _, lc := range colors {
		for col := lc.Start; col < lc.End; col++ {
			j := row*cols + (col-offset)
			if j < 0 || j >= len(cells) {
				continue
			}
			cells[j].Bg, cells[j].Fg = lc.Bg, lc.Fg
		}
	}
}

func (screen *Screen) WriteStringColor(row, col int, s string, fg, bg termbox.Attribute) {
	for k, c := range s {
		termbox.SetCell(col+k, row, c, fg, bg)
	}
}

func (screen *Screen) WriteMessage(msg string) {
	_, rows := termbox.Size()
	screen.WriteString(rows-1, 0, msg)
}

func (screen *Screen) AskYesNo(question string) (bool, error) {
	prompt := MakePrompt(screen)
	return prompt.AskYesNo(question)
}

func (screen *Screen) Ask(question string, history []string) (string, error) {
	prompt := MakePrompt(screen)
	return prompt.Ask(question, history)
}

func (screen *Screen) Highlight(row, col int) {
	cells := termbox.CellBuffer()
	cols, _ := termbox.Size()
	j := row*cols + col
	cell := cells[j]
	cells[j].Bg, cells[j].Fg = cell.Fg, cell.Bg
}

<commit_msg>Add a little padding to prompt<commit_after>package terminal

import "github.com/nsf/termbox-go"
import "strings"
import "github.com/wx13/sith/syntaxcolor"

type Screen struct {
	row, col int
	fg, bg   termbox.Attribute
	colors   map[string]termbox.Attribute
}

func NewScreen() *Screen {
	screen := Screen{
		row: 0,
		col: 0,
		bg:  termbox.ColorBlack,
		fg:  termbox.ColorWhite,
	}
	termbox.Init()
	return &screen
}

func (screen *Screen) Close() {
	screen.Clear()
	termbox.Flush()
	termbox.Close()
}

func (screen *Screen) Open() {
	termbox.Init()
}

func (screen *Screen) Flush() {
	termbox.Flush()
}

func (screen *Screen) SetCursor(r, c int) {
	screen.row = r
	screen.col = c
	termbox.SetCursor(c, r)
}

func (screen *Screen) Clear() {
	termbox.Clear(screen.fg, screen.bg)
	cols, rows := termbox.Size()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(" ", cols))
	}
}

func (screen *Screen) ReallyClear() {
	cols, rows := termbox.Size()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(".", cols))
	}
	screen.Flush()
	for row := 0; row < rows; row++ {
		screen.WriteString(row, 0, strings.Repeat(" ", cols))
	}
	screen.Flush()
}

func (screen *Screen) DecorateStatusLine() {
	cells := termbox.CellBuffer()
	cols, rows := termbox.Size()
	for col := 0; col < cols; col++ {
		j := (rows-1)*cols + col
		cells[j].Fg = termbox.ColorBlue
	}
}

func (screen *Screen) WriteString(row, col int, s string) {
	screen.WriteStringColor(row, col, s, screen.fg, screen.bg)
}

func (screen *Screen) Colorize(row int, colors []syntaxcolor.LineColor, offset int) {
	cells := termbox.CellBuffer()
	cols, _ := termbox.Size()
	for _, lc := range colors {
		for col := lc.Start; col < lc.End; col++ {
			j := row*cols + (col-offset)
			if j < 0 || j >= len(cells) {
				continue
			}
			cells[j].Bg, cells[j].Fg = lc.Bg, lc.Fg
		}
	}
}

func (screen *Screen) WriteStringColor(row, col int, s string, fg, bg termbox.Attribute) {
	for k, c := range s {
		termbox.SetCell(col+k, row, c, fg, bg)
	}
}

func (screen *Screen) WriteMessage(msg string) {
	_, rows := termbox.Size()
	screen.WriteString(rows-1, 0, msg + "  ")
}

func (screen *Screen) AskYesNo(question string) (bool, error) {
	prompt := MakePrompt(screen)
	return prompt.AskYesNo(question)
}

func (screen *Screen) Ask(question string, history []string) (string, error) {
	prompt := MakePrompt(screen)
	return prompt.Ask(question, history)
}

func (screen *Screen) Highlight(row, col int) {
	cells := termbox.CellBuffer()
	cols, _ := termbox.Size()
	j := row*cols + col
	cell := cells[j]
	cells[j].Bg, cells[j].Fg = cell.Fg, cell.Bg
}

<|endoftext|>


<commit_before>package ace

import (
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"testing"
)

var testHandler = func(c *C) { c.Next() }

func TestHTTPMethod(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.GET("/test", func(c *C) {
		c.String(200, "Test")
	})

	a.POST("/test", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.PUT("/", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.PATCH("/", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.DELETE("/", func(c *C) {
		c.String(200, "deleted")
	})

	a.OPTIONS("/", func(c *C) {
		c.String(200, "options")
	})

	a.HEAD("/test", func(c *C) {
		c.String(200, "head")
	})

	r, _ := http.NewRequest("GET", "/test", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("Test", w.Body.String())

	r, _ = http.NewRequest("POST", "/test", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("PUT", "/", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("PATCH", "/", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("DELETE", "/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("deleted", w.Body.String())

	r, _ = http.NewRequest("OPTIONS", "/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("options", w.Body.String())

	r, _ = http.NewRequest("HEAD", "/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("head", w.Body.String())

	//tailing slash
	r, _ = http.NewRequest("GET", "/test/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(301, w.Code)

	r, _ = http.NewRequest("POST", "/test/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(307, w.Code)
}

func TestGroupRoute(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	g1 := a.Group("/g1", testHandler)
	g2 := a.Group("/g2", testHandler)

	g1.GET("/", func(c *C) {
		c.String(200, "g1")
	})

	g1.GET("/test", func(c *C) {
		c.String(200, "g1/test")
	})

	g2.POST("/", func(c *C) {
		c.String(200, "g2")
	})

	g2.POST("/test", func(c *C) {
		c.String(200, "g2/test")
	})

	r, _ := http.NewRequest("GET", "/g1/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g1", w.Body.String())

	r, _ = http.NewRequest("GET", "/g1/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g1/test", w.Body.String())

	r, _ = http.NewRequest("POST", "/g2/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g2", w.Body.String())

	r, _ = http.NewRequest("POST", "/g2/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g2/test", w.Body.String())
}

func TestServeStatic(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.Static("/assets", "./", testHandler)

	r, _ := http.NewRequest("GET", "/assets/README.md", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)

	r, _ = http.NewRequest("GET", "/assets/test.text", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(404, w.Code)
}

func TestConvertHandlerFunc(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.GET("/", a.HTTPHandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Write([]byte("test"))
	}))

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("test", w.Body.String())
}

func TestRouteNotFound(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.RouteNotFound(func(c *C) {
		c.String(404, "test not found")
	})

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(404, w.Code)
	assert.Equal("test not found", w.Body.String())
}

func TestPanic(t *testing.T) {
	assert := assert.New(t)
	mid := ""

	a := New()
	a.Use(Recovery())
	a.Use(func(c *C) {
		mid = "before panic"
		c.Next()
	})

	a.GET("/", func(c *C) {
		panic("panic test")
	})

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(500, w.Code)
	assert.Equal("before panic", mid)
}

func TestStaticPath(t *testing.T) {
	assert := assert.New(t)

	a := New()
	path := a.Router.staticPath("/")
	assert.Equal("/*filepath", path)

	path = a.Router.staticPath("/public")
	assert.Equal("/public/*filepath", path)
}
<commit_msg>add test<commit_after>package ace

import (
	"github.com/plimble/utils/errors2"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"testing"
)

var testHandler = func(c *C) { c.Next() }

func TestHTTPMethod(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.GET("/test", func(c *C) {
		c.String(200, "Test")
	})

	a.POST("/test", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.PUT("/", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.PATCH("/", func(c *C) {
		c.String(200, c.Request.FormValue("test"))
	})

	a.DELETE("/", func(c *C) {
		c.String(200, "deleted")
	})

	a.OPTIONS("/", func(c *C) {
		c.String(200, "options")
	})

	a.HEAD("/test", func(c *C) {
		c.String(200, "head")
	})

	r, _ := http.NewRequest("GET", "/test", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("Test", w.Body.String())

	r, _ = http.NewRequest("POST", "/test", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("PUT", "/", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("PATCH", "/", nil)
	r.ParseForm()
	r.Form.Add("test", "hello")
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("hello", w.Body.String())

	r, _ = http.NewRequest("DELETE", "/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("deleted", w.Body.String())

	r, _ = http.NewRequest("OPTIONS", "/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("options", w.Body.String())

	r, _ = http.NewRequest("HEAD", "/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("head", w.Body.String())

	//tailing slash
	r, _ = http.NewRequest("GET", "/test/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(301, w.Code)

	r, _ = http.NewRequest("POST", "/test/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(307, w.Code)
}

func TestGroupRoute(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	g1 := a.Group("/g1", testHandler)
	g2 := a.Group("/g2", testHandler)

	g1.GET("/", func(c *C) {
		c.String(200, "g1")
	})

	g1.GET("/test", func(c *C) {
		c.String(200, "g1/test")
	})

	g2.POST("/", func(c *C) {
		c.String(200, "g2")
	})

	g2.POST("/test", func(c *C) {
		c.String(200, "g2/test")
	})

	r, _ := http.NewRequest("GET", "/g1/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g1", w.Body.String())

	r, _ = http.NewRequest("GET", "/g1/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g1/test", w.Body.String())

	r, _ = http.NewRequest("POST", "/g2/", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g2", w.Body.String())

	r, _ = http.NewRequest("POST", "/g2/test", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("g2/test", w.Body.String())
}

func TestServeStatic(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.Static("/assets", "./", testHandler)

	r, _ := http.NewRequest("GET", "/assets/README.md", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)

	r, _ = http.NewRequest("GET", "/assets/test.text", nil)
	w = httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(404, w.Code)
}

func TestConvertHandlerFunc(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.GET("/", a.HTTPHandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.Write([]byte("test"))
	}))

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(200, w.Code)
	assert.Equal("test", w.Body.String())
}

func TestRouteNotFound(t *testing.T) {
	assert := assert.New(t)

	a := Default()
	a.RouteNotFound(func(c *C) {
		c.String(404, "test not found")
	})

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(404, w.Code)
	assert.Equal("test not found", w.Body.String())
}

func TestPanic2(t *testing.T) {
	assert := assert.New(t)
	a := New()

	a.Panic(func(c *C, rcv interface{}) {
		err := rcv.(errors2.Error)
		c.JSON(err.HttpStatus(), err)
	})

	a.GET("/", func(c *C) {
		c.Panic(errors2.NewNotFound("not found"))
		c.String(200, "123")
	})

	r, _ := http.NewRequest("GET", "/", nil)
	w := httptest.NewRecorder()
	a.ServeHTTP(w, r)
	assert.Equal(404, w.Code)
	assert.Equal("{\"message\":\"not found\"}\n", w.Body.String())
}

func TestStaticPath(t *testing.T) {
	assert := assert.New(t)

	a := New()
	path := a.Router.staticPath("/")
	assert.Equal("/*filepath", path)

	path = a.Router.staticPath("/public")
	assert.Equal("/public/*filepath", path)
}
<|endoftext|>


<commit_before>package main

import (
  "log"
  "errors"
  "os/exec"
  "strings"
)

type DependencyItem struct {
  targetRelativePath string
  originalPath string
}

type AppDeployer struct {
  processedLibs map[string]bool

  libsChannel chan string
  copyChannel chan DependencyItem
  stripChannel chan string
  rpathChannel chan string
  qtChannel chan string

  additionalLibPaths []string
  destinationPath string
}

func (ad *AppDeployer) DeployApp(exePath string) {
  go func() { ad.libsChannel <- exePath }()
  ad.processLibs()
}

func (ad *AppDeployer) processLibs() {
  for filepath := range ad.libsChannel {
    if _, ok := ad.processedLibs[filepath]; !ok {
      dependencies, err := findLddDependencies(filepath)
      if (err != nil) {
        log.Println(err)
        continue
      }

      ad.processedLibs[filepath] = true
      //go func() { ad.copyChannel <- &{DependencyItem{originalPath: filepath} }()

      for _, dependPath := range dependencies {
        if _, ok := ad.processedLibs[dependPath]; !ok {
          go func() { ad.libsChannel <- dependPath }()
        }
      }
    }
  }
}

func findLddDependencies(filepath string) ([]string, error) {
  log.Printf("Inspecting %v", filepath)

  out, err := exec.Command("ldd", filepath).Output()
  if err != nil { return nil, err }

  dependencies := make([]string, 10)

  output := string(out)
  lines := strings.Split(output, "\n")
  for _, line := range lines {
    line = strings.TrimSpace(line)
    libpath, err := parseLddOutputLine(line)

    if err == nil {
      log.Printf("Found dependency %v", libpath)
      dependencies = append(dependencies, libpath)
    } else {
      log.Printf("Cannot parse ldd line: %v", line)
    }
  }

  return dependencies, nil
}

func parseLddOutputLine(line string) (string, error) {
  if len(line) == 0 { return "", errors.New("Empty") }

  var libpath string

  if strings.Contains(line, " => ") {
    parts := strings.Split(line, " => ")

    if len(parts) != 2 {
      return "", errors.New("Wrong format")
    }

    shortpath := strings.TrimSpace(parts[0])

    if parts[1] == "not found" { return parts[0], nil }
    if len(strings.TrimSpace(parts[1])) == 0 { return "", errors.New("vdso") }

    lastUseful := strings.LastIndex(parts[1], "(0x")
    if lastUseful != -1 {
      libpath = strings.TrimSpace(parts[1][:lastUseful])
    } else {
      libpath = shortpath
    }
  } else {
    log.Printf("Skipping ldd line: %v", line)
  }

  return libpath, nil
}
<commit_msg>Verbosity<commit_after>package main

import (
  "log"
  "errors"
  "os/exec"
  "strings"
)

type DependencyItem struct {
  targetRelativePath string
  originalPath string
}

type AppDeployer struct {
  processedLibs map[string]bool

  libsChannel chan string
  copyChannel chan DependencyItem
  stripChannel chan string
  rpathChannel chan string
  qtChannel chan string

  additionalLibPaths []string
  destinationPath string
}

func (ad *AppDeployer) DeployApp(exePath string) {
  go func() { ad.libsChannel <- exePath }()
  ad.processLibs()
}

func (ad *AppDeployer) processLibs() {
  for filepath := range ad.libsChannel {
    if _, ok := ad.processedLibs[filepath]; !ok {
      dependencies, err := findLddDependencies(filepath)
      if (err != nil) {
        log.Println(err)
        continue
      }

      ad.processedLibs[filepath] = true
      //go func() { ad.copyChannel <- &{DependencyItem{originalPath: filepath} }()

      for _, dependPath := range dependencies {
        if _, ok := ad.processedLibs[dependPath]; !ok {
          go func() { ad.libsChannel <- dependPath }()
        }
      }
    }
  }
}

func findLddDependencies(filepath string) ([]string, error) {
  log.Printf("Inspecting %v", filepath)

  out, err := exec.Command("ldd", filepath).Output()
  if err != nil { return nil, err }

  dependencies := make([]string, 10)

  output := string(out)
  lines := strings.Split(output, "\n")
  for _, line := range lines {
    line = strings.TrimSpace(line)
    libpath, err := parseLddOutputLine(line)

    if err == nil {
      log.Printf("Found dependency %v for line [%v]", libpath, line)
      dependencies = append(dependencies, libpath)
    } else {
      log.Printf("Cannot parse ldd line: %v", line)
    }
  }

  return dependencies, nil
}

func parseLddOutputLine(line string) (string, error) {
  if len(line) == 0 { return "", errors.New("Empty") }

  var libpath string

  if strings.Contains(line, " => ") {
    parts := strings.Split(line, " => ")

    if len(parts) != 2 {
      return "", errors.New("Wrong format")
    }

    shortpath := strings.TrimSpace(parts[0])

    if parts[1] == "not found" { return parts[0], nil }
    if len(strings.TrimSpace(parts[1])) == 0 { return "", errors.New("vdso") }

    lastUseful := strings.LastIndex(parts[1], "(0x")
    if lastUseful != -1 {
      libpath = strings.TrimSpace(parts[1][:lastUseful])
    } else {
      libpath = shortpath
    }
  } else {
    log.Printf("Skipping ldd line: %v", line)
  }

  return libpath, nil
}
<|endoftext|>


<commit_before>package terminal_test

import (
	"testing"

	"github.com/opsidian/parsley/ast"
	"github.com/opsidian/parsley/parser"
	"github.com/opsidian/parsley/text"
	"github.com/opsidian/parsley/text/terminal"
	"github.com/stretchr/testify/assert"
)

func assertCursor(t *testing.T, pos int, line int, col int, r *text.Reader) {
	assert.Equal(t, text.NewPosition(pos, line, col), r.Cursor())
}

func TestRuneShouldMatchCharacter(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('a', "A")(parser.EmptyLeftRecCtx(), r)
	expectedNode := ast.NewTerminalNode("A", text.NewPosition(0, 1, 1), 'a')
	assert.Equal(t, parser.NewResult(expectedNode, r).AsSet(), res)
	assertCursor(t, 1, 1, 2, r)
}

func TestRuneShouldNotUseSpecialChars(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('.', ".")(parser.EmptyLeftRecCtx(), r)
	assert.Nil(t, res)
	assertCursor(t, 0, 1, 1, r)
}

func TestRuneShouldMatchUnicodeCharacter(t *testing.T) {
	r := text.NewReader([]byte("🍕"), true)
	_, res := terminal.Rune('🍕', "PIZZA")(parser.EmptyLeftRecCtx(), r)
	expectedNode := ast.NewTerminalNode("PIZZA", text.NewPosition(0, 1, 1), '🍕')
	assert.Equal(t, parser.NewResult(expectedNode, r).AsSet(), res)
	assertCursor(t, 4, 1, 2, r)
}

func TestRuneShouldNotMatchCharacter(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('b', "B")(parser.EmptyLeftRecCtx(), r)
	assert.Nil(t, res)
	assertCursor(t, 0, 1, 1, r)
}
<commit_msg>Clean up rune test<commit_after>package terminal_test

import (
	"testing"

	"github.com/opsidian/parsley/ast"
	"github.com/opsidian/parsley/parser"
	"github.com/opsidian/parsley/text"
	"github.com/opsidian/parsley/text/terminal"
	"github.com/stretchr/testify/assert"
)

func TestRuneShouldMatchCharacter(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('a', "A")(parser.EmptyLeftRecCtx(), r)
	expectedNode := ast.NewTerminalNode("A", text.NewPosition(0, 1, 1), 'a')
	assert.Equal(t, parser.NewResult(expectedNode, r).AsSet(), res)
	assert.Equal(t, text.NewPosition(1, 1, 2), r.Cursor())
}

func TestRuneShouldNotUseSpecialChars(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('.', ".")(parser.EmptyLeftRecCtx(), r)
	assert.Nil(t, res)
	assert.Equal(t, text.NewPosition(0, 1, 1), r.Cursor())
}

func TestRuneShouldMatchUnicodeCharacter(t *testing.T) {
	r := text.NewReader([]byte("🍕"), true)
	_, res := terminal.Rune('🍕', "PIZZA")(parser.EmptyLeftRecCtx(), r)
	expectedNode := ast.NewTerminalNode("PIZZA", text.NewPosition(0, 1, 1), '🍕')
	assert.Equal(t, parser.NewResult(expectedNode, r).AsSet(), res)
	assert.Equal(t, text.NewPosition(4, 1, 2), r.Cursor())
}

func TestRuneShouldNotMatchCharacter(t *testing.T) {
	r := text.NewReader([]byte("a"), true)
	_, res := terminal.Rune('b', "B")(parser.EmptyLeftRecCtx(), r)
	assert.Nil(t, res)
	assert.Equal(t, text.NewPosition(0, 1, 1), r.Cursor())
}
<|endoftext|>


<commit_before>package arangodb

import (
	"context"
	"crypto/md5"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"time"

	"github.com/arangodb/go-driver"
	"github.com/arangodb/go-driver/http"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	filer.Stores = append(filer.Stores, &ArangodbStore{})
}

type ArangodbStore struct {
	connect    driver.Connection
	client     driver.Client
	database   driver.Database
	collection driver.Collection
}

type Model struct {
	Key       string   `json:"_key"`
	Directory string   `json:"directory"`
	Name      string   `json:"name"`
	Meta      []uint64 `json:"meta"`
}

func (store *ArangodbStore) GetName() string {
	return "arangodb"
}

func (store *ArangodbStore) Initialize(configuration util.Configuration, prefix string) (err error) {
	return store.connection(configuration.GetStringSlice(prefix+"arango_host"),
		configuration.GetString(prefix+"arango_user"),
		configuration.GetString(prefix+"arango_pass"),
	)
}

func (store *ArangodbStore) connection(uris []string, user string, pass string) (err error) {
	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)

	store.connect, err = http.NewConnection(http.ConnectionConfig{
		Endpoints: uris,
		TLSConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	})
	if err != nil {
		return err
	}
	store.client, err = driver.NewClient(driver.ClientConfig{
		Connection:     store.connect,
		Authentication: driver.BasicAuthentication(user, pass),
	})
	if err != nil {
		return err
	}
	db_name := "seaweed-filer"
	ok, err := store.client.DatabaseExists(ctx, db_name)
	if err != nil {
		return err
	}
	if ok {
		store.database, err = store.client.Database(ctx, db_name)
	} else {
		store.database, err = store.client.CreateDatabase(ctx, db_name, &driver.CreateDatabaseOptions{})
	}
	if err != nil {
		return err
	}

	coll_name := "files"
	ok, err = store.database.CollectionExists(ctx, coll_name)
	if err != nil {
		return err
	}
	if ok {
		store.collection, err = store.database.Collection(ctx, coll_name)
	} else {
		store.collection, err = store.database.CreateCollection(ctx, coll_name, &driver.CreateCollectionOptions{})
	}
	if err != nil {
		return err
	}

	// ensure indices

	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"directory", "name"}, &driver.EnsurePersistentIndexOptions{
		Name:   "directory_name_multi",
		Unique: true,
	}); err != nil {
		return err
	}
	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"directory"},
		&driver.EnsurePersistentIndexOptions{Name: "IDX_directory"}); err != nil {
		return err
	}
	// fulltext index not required since no prefix search
	// user should just make one themselves if they intend on using it
	//	if _, _, err = store.collection.EnsureFullTextIndex(ctx, []string{"directory"},
	//		&driver.EnsureFullTextIndexOptions{Name: "IDX_FULLTEXT_directory", MinLength: 1}); err != nil {
	//		return err
	//	}

	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"name"}, &driver.EnsurePersistentIndexOptions{
		Name: "IDX_name",
	}); err != nil {
		return err
	}

	return err
}

type key int

const (
	transactionKey key = 0
)

func (store *ArangodbStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	txn, err := store.database.BeginTransaction(ctx, driver.TransactionCollections{
		Exclusive: []string{"files"},
	}, &driver.BeginTransactionOptions{})
	if err != nil {
		return nil, err
	}

	return context.WithValue(ctx, transactionKey, txn), nil
}

func (store *ArangodbStore) CommitTransaction(ctx context.Context) error {
	val := ctx.Value(transactionKey)
	cast, ok := val.(driver.TransactionID)
	if !ok {
		return fmt.Errorf("txn cast fail %s:", val)
	}
	err := store.database.CommitTransaction(ctx, cast, &driver.CommitTransactionOptions{})
	if err != nil {
		return err
	}
	return nil
}

func (store *ArangodbStore) RollbackTransaction(ctx context.Context) error {
	val := ctx.Value(transactionKey)
	cast, ok := val.(driver.TransactionID)
	if !ok {
		return fmt.Errorf("txn cast fail %s:", val)
	}
	err := store.database.AbortTransaction(ctx, cast, &driver.AbortTransactionOptions{})
	if err != nil {
		return err
	}
	return nil
}

func (store *ArangodbStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {
	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	if len(entry.Chunks) > 50 {
		meta = util.MaybeGzipData(meta)
	}
	model := &Model{
		Key:       hashString(string(entry.FullPath)),
		Directory: dir,
		Name:      name,
		Meta:      bytesToArray(meta),
	}
	_, err = store.collection.CreateDocument(ctx, model)

	if err != nil {
		return fmt.Errorf("UpdateEntry %s: %v", entry.FullPath, err)
	}

	return nil

}

func (store *ArangodbStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	if len(entry.Chunks) > 50 {
		meta = util.MaybeGzipData(meta)
	}
	model := &Model{
		Key:       hashString(string(entry.FullPath)),
		Directory: dir,
		Name:      name,
		Meta:      bytesToArray(meta),
	}

	_, err = store.collection.UpdateDocument(ctx, model.Key, model)

	if err != nil {
		return fmt.Errorf("UpdateEntry %s: %v", entry.FullPath, err)
	}

	return nil
}

func (store *ArangodbStore) FindEntry(ctx context.Context, fullpath util.FullPath) (entry *filer.Entry, err error) {
	var data Model
	_, err = store.collection.ReadDocument(ctx, hashString(string(fullpath)), &data)
	if driver.IsNotFound(err) {
		return nil, filer_pb.ErrNotFound
	}
	if err != nil {
		glog.Errorf("find %s: %v", fullpath, err)
		return nil, filer_pb.ErrNotFound
	}
	if len(data.Meta) == 0 {
		return nil, filer_pb.ErrNotFound
	}
	entry = &filer.Entry{
		FullPath: fullpath,
	}
	err = entry.DecodeAttributesAndChunks(util.MaybeDecompressData(arrayToBytes(data.Meta)))
	if err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}

	return entry, nil
}

func (store *ArangodbStore) DeleteEntry(ctx context.Context, fullpath util.FullPath) error {
	_, err := store.collection.RemoveDocument(ctx, hashString(string(fullpath)))
	if err != nil {
		glog.Errorf("find %s: %v", fullpath, err)
		return fmt.Errorf("delete %s : %v", fullpath, err)
	}

	return nil
}

func (store *ArangodbStore) DeleteFolderChildren(ctx context.Context, fullpath util.FullPath) error {
	dir, _ := fullpath.DirAndName()
	cur, err := store.database.Query(ctx, `
for d in files
filter d.directory == @dir
remove d in files`, map[string]interface{}{"dir": dir})
	if err != nil {
		return fmt.Errorf("delete %s : %v", fullpath, err)
	}
	defer cur.Close()
	return nil
}

func (store *ArangodbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
	return lastFileName, filer.ErrUnsupportedListDirectoryPrefixed
}

//TODO: i must be misunderstanding what this function is supposed to do
//so figure it out is the todo, i guess lol - aaaaa
//func (store *ArangodbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
//	eq := ""
//	if includeStartFile {
//		eq = "filter d.name >= \"" + startFileName + "\""
//	} else {
//		eq = "filter d.name > \"" + startFileName + "\""
//	}
//	query := fmt.Sprintf(`
//for d in fulltext(files,"directory","prefix:%s")
//sort d.name desc
//%s
//limit %d
//return d`, string(dirPath), eq, limit)
//	cur, err := store.database.Query(ctx, query, nil)
//	if err != nil {
//		return lastFileName, fmt.Errorf("failed to list directory entries: find error: %w", err)
//	}
//	defer cur.Close()
//	for cur.HasMore() {
//		var data Model
//		_, err = cur.ReadDocument(ctx, &data)
//		if err != nil {
//			break
//		}
//		entry := &filer.Entry{
//			FullPath: util.NewFullPath(data.Directory, data.Name),
//		}
//		lastFileName = data.Name
//		converted := arrayToBytes(data.Meta)
//		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(converted)); decodeErr != nil {
//			err = decodeErr
//			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
//			break
//		}
//
//		if !eachEntryFunc(entry) {
//			break
//		}
//
//	}
//	return lastFileName, err
//}

func (store *ArangodbStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
	eq := ""
	if includeStartFile {
		eq = "filter d.name >= \"" + startFileName + "\""
	} else {
		eq = "filter d.name > \"" + startFileName + "\""
	}
	query := fmt.Sprintf(`
for d in files
filter d.directory == "%s"
sort d.name desc
%s
limit %d
return d`, string(dirPath), eq, limit)
	cur, err := store.database.Query(ctx, query, nil)
	if err != nil {
		return lastFileName, fmt.Errorf("failed to list directory entries: find error: %w", err)
	}
	defer cur.Close()
	for cur.HasMore() {
		var data Model
		_, err = cur.ReadDocument(ctx, &data)
		if err != nil {
			break
		}
		entry := &filer.Entry{
			FullPath: util.NewFullPath(string(dirPath), data.Name),
		}
		lastFileName = data.Name
		converted := arrayToBytes(data.Meta)
		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(converted)); decodeErr != nil {
			err = decodeErr
			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
			break
		}

		if !eachEntryFunc(entry) {
			break
		}

	}
	return lastFileName, err
}

func (store *ArangodbStore) Shutdown() {
}

//convert a string into arango-key safe hex bytes hash
func hashString(dir string) string {
	h := md5.New()
	io.WriteString(h, dir)
	b := h.Sum(nil)
	return hex.EncodeToString(b)
}

func bytesToArray(bs []byte) []uint64 {
	out := make([]uint64, 0, 2+len(bs)/8)
	out = append(out, uint64(len(bs)))
	for len(bs)%8 != 0 {
		bs = append(bs, 0)
	}
	for i := 0; i < len(bs); i = i + 8 {
		out = append(out, binary.BigEndian.Uint64(bs[i:]))
	}
	return out
}

func arrayToBytes(xs []uint64) []byte {
	if len(xs) < 2 {
		return []byte{}
	}
	first := xs[0]
	out := make([]byte, len(xs)*8)
	for i := 1; i < len(xs); i = i + 1 {
		binary.BigEndian.PutUint64(out[((i-1)*8):], xs[i])
	}
	return out[:first]
}
<commit_msg>revise comment<commit_after>package arangodb

import (
	"context"
	"crypto/md5"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"time"

	"github.com/arangodb/go-driver"
	"github.com/arangodb/go-driver/http"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func init() {
	filer.Stores = append(filer.Stores, &ArangodbStore{})
}

type ArangodbStore struct {
	connect    driver.Connection
	client     driver.Client
	database   driver.Database
	collection driver.Collection
}

type Model struct {
	Key       string   `json:"_key"`
	Directory string   `json:"directory"`
	Name      string   `json:"name"`
	Meta      []uint64 `json:"meta"`
}

func (store *ArangodbStore) GetName() string {
	return "arangodb"
}

func (store *ArangodbStore) Initialize(configuration util.Configuration, prefix string) (err error) {
	return store.connection(configuration.GetStringSlice(prefix+"arango_host"),
		configuration.GetString(prefix+"arango_user"),
		configuration.GetString(prefix+"arango_pass"),
	)
}

func (store *ArangodbStore) connection(uris []string, user string, pass string) (err error) {
	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)

	store.connect, err = http.NewConnection(http.ConnectionConfig{
		Endpoints: uris,
		TLSConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	})
	if err != nil {
		return err
	}
	store.client, err = driver.NewClient(driver.ClientConfig{
		Connection:     store.connect,
		Authentication: driver.BasicAuthentication(user, pass),
	})
	if err != nil {
		return err
	}
	db_name := "seaweed-filer"
	ok, err := store.client.DatabaseExists(ctx, db_name)
	if err != nil {
		return err
	}
	if ok {
		store.database, err = store.client.Database(ctx, db_name)
	} else {
		store.database, err = store.client.CreateDatabase(ctx, db_name, &driver.CreateDatabaseOptions{})
	}
	if err != nil {
		return err
	}

	coll_name := "files"
	ok, err = store.database.CollectionExists(ctx, coll_name)
	if err != nil {
		return err
	}
	if ok {
		store.collection, err = store.database.Collection(ctx, coll_name)
	} else {
		store.collection, err = store.database.CreateCollection(ctx, coll_name, &driver.CreateCollectionOptions{})
	}
	if err != nil {
		return err
	}

	// ensure indices

	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"directory", "name"}, &driver.EnsurePersistentIndexOptions{
		Name:   "directory_name_multi",
		Unique: true,
	}); err != nil {
		return err
	}
	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"directory"},
		&driver.EnsurePersistentIndexOptions{Name: "IDX_directory"}); err != nil {
		return err
	}
	//  fulltext index not required since no prefix search
	//  might change
	//	if _, _, err = store.collection.EnsureFullTextIndex(ctx, []string{"directory"},
	//		&driver.EnsureFullTextIndexOptions{Name: "IDX_FULLTEXT_directory", MinLength: 1}); err != nil {
	//		return err
	//	}

	if _, _, err = store.collection.EnsurePersistentIndex(ctx, []string{"name"}, &driver.EnsurePersistentIndexOptions{
		Name: "IDX_name",
	}); err != nil {
		return err
	}

	return err
}

type key int

const (
	transactionKey key = 0
)

func (store *ArangodbStore) BeginTransaction(ctx context.Context) (context.Context, error) {
	txn, err := store.database.BeginTransaction(ctx, driver.TransactionCollections{
		Exclusive: []string{"files"},
	}, &driver.BeginTransactionOptions{})
	if err != nil {
		return nil, err
	}

	return context.WithValue(ctx, transactionKey, txn), nil
}

func (store *ArangodbStore) CommitTransaction(ctx context.Context) error {
	val := ctx.Value(transactionKey)
	cast, ok := val.(driver.TransactionID)
	if !ok {
		return fmt.Errorf("txn cast fail %s:", val)
	}
	err := store.database.CommitTransaction(ctx, cast, &driver.CommitTransactionOptions{})
	if err != nil {
		return err
	}
	return nil
}

func (store *ArangodbStore) RollbackTransaction(ctx context.Context) error {
	val := ctx.Value(transactionKey)
	cast, ok := val.(driver.TransactionID)
	if !ok {
		return fmt.Errorf("txn cast fail %s:", val)
	}
	err := store.database.AbortTransaction(ctx, cast, &driver.AbortTransactionOptions{})
	if err != nil {
		return err
	}
	return nil
}

func (store *ArangodbStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {
	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	if len(entry.Chunks) > 50 {
		meta = util.MaybeGzipData(meta)
	}
	model := &Model{
		Key:       hashString(string(entry.FullPath)),
		Directory: dir,
		Name:      name,
		Meta:      bytesToArray(meta),
	}
	_, err = store.collection.CreateDocument(ctx, model)

	if err != nil {
		return fmt.Errorf("UpdateEntry %s: %v", entry.FullPath, err)
	}

	return nil

}

func (store *ArangodbStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
	dir, name := entry.FullPath.DirAndName()
	meta, err := entry.EncodeAttributesAndChunks()
	if err != nil {
		return fmt.Errorf("encode %s: %s", entry.FullPath, err)
	}

	if len(entry.Chunks) > 50 {
		meta = util.MaybeGzipData(meta)
	}
	model := &Model{
		Key:       hashString(string(entry.FullPath)),
		Directory: dir,
		Name:      name,
		Meta:      bytesToArray(meta),
	}

	_, err = store.collection.UpdateDocument(ctx, model.Key, model)

	if err != nil {
		return fmt.Errorf("UpdateEntry %s: %v", entry.FullPath, err)
	}

	return nil
}

func (store *ArangodbStore) FindEntry(ctx context.Context, fullpath util.FullPath) (entry *filer.Entry, err error) {
	var data Model
	_, err = store.collection.ReadDocument(ctx, hashString(string(fullpath)), &data)
	if driver.IsNotFound(err) {
		return nil, filer_pb.ErrNotFound
	}
	if err != nil {
		glog.Errorf("find %s: %v", fullpath, err)
		return nil, filer_pb.ErrNotFound
	}
	if len(data.Meta) == 0 {
		return nil, filer_pb.ErrNotFound
	}
	entry = &filer.Entry{
		FullPath: fullpath,
	}
	err = entry.DecodeAttributesAndChunks(util.MaybeDecompressData(arrayToBytes(data.Meta)))
	if err != nil {
		return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
	}

	return entry, nil
}

func (store *ArangodbStore) DeleteEntry(ctx context.Context, fullpath util.FullPath) error {
	_, err := store.collection.RemoveDocument(ctx, hashString(string(fullpath)))
	if err != nil {
		glog.Errorf("find %s: %v", fullpath, err)
		return fmt.Errorf("delete %s : %v", fullpath, err)
	}

	return nil
}

func (store *ArangodbStore) DeleteFolderChildren(ctx context.Context, fullpath util.FullPath) error {
	dir, _ := fullpath.DirAndName()
	cur, err := store.database.Query(ctx, `
for d in files
filter d.directory == @dir
remove d in files`, map[string]interface{}{"dir": dir})
	if err != nil {
		return fmt.Errorf("delete %s : %v", fullpath, err)
	}
	defer cur.Close()
	return nil
}

func (store *ArangodbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
	return lastFileName, filer.ErrUnsupportedListDirectoryPrefixed
}

//TODO: i must be misunderstanding what this function is supposed to do
//so figure it out is the todo, i guess lol - aaaaa
//func (store *ArangodbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
//	eq := ""
//	if includeStartFile {
//		eq = "filter d.name >= \"" + startFileName + "\""
//	} else {
//		eq = "filter d.name > \"" + startFileName + "\""
//	}
//	query := fmt.Sprintf(`
//for d in fulltext(files,"directory","prefix:%s")
//sort d.name desc
//%s
//limit %d
//return d`, string(dirPath), eq, limit)
//	cur, err := store.database.Query(ctx, query, nil)
//	if err != nil {
//		return lastFileName, fmt.Errorf("failed to list directory entries: find error: %w", err)
//	}
//	defer cur.Close()
//	for cur.HasMore() {
//		var data Model
//		_, err = cur.ReadDocument(ctx, &data)
//		if err != nil {
//			break
//		}
//		entry := &filer.Entry{
//			FullPath: util.NewFullPath(data.Directory, data.Name),
//		}
//		lastFileName = data.Name
//		converted := arrayToBytes(data.Meta)
//		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(converted)); decodeErr != nil {
//			err = decodeErr
//			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
//			break
//		}
//
//		if !eachEntryFunc(entry) {
//			break
//		}
//
//	}
//	return lastFileName, err
//}

func (store *ArangodbStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
	eq := ""
	if includeStartFile {
		eq = "filter d.name >= \"" + startFileName + "\""
	} else {
		eq = "filter d.name > \"" + startFileName + "\""
	}
	query := fmt.Sprintf(`
for d in files
filter d.directory == "%s"
sort d.name desc
%s
limit %d
return d`, string(dirPath), eq, limit)
	cur, err := store.database.Query(ctx, query, nil)
	if err != nil {
		return lastFileName, fmt.Errorf("failed to list directory entries: find error: %w", err)
	}
	defer cur.Close()
	for cur.HasMore() {
		var data Model
		_, err = cur.ReadDocument(ctx, &data)
		if err != nil {
			break
		}
		entry := &filer.Entry{
			FullPath: util.NewFullPath(string(dirPath), data.Name),
		}
		lastFileName = data.Name
		converted := arrayToBytes(data.Meta)
		if decodeErr := entry.DecodeAttributesAndChunks(util.MaybeDecompressData(converted)); decodeErr != nil {
			err = decodeErr
			glog.V(0).Infof("list %s : %v", entry.FullPath, err)
			break
		}

		if !eachEntryFunc(entry) {
			break
		}

	}
	return lastFileName, err
}

func (store *ArangodbStore) Shutdown() {
}

//convert a string into arango-key safe hex bytes hash
func hashString(dir string) string {
	h := md5.New()
	io.WriteString(h, dir)
	b := h.Sum(nil)
	return hex.EncodeToString(b)
}

func bytesToArray(bs []byte) []uint64 {
	out := make([]uint64, 0, 2+len(bs)/8)
	out = append(out, uint64(len(bs)))
	for len(bs)%8 != 0 {
		bs = append(bs, 0)
	}
	for i := 0; i < len(bs); i = i + 8 {
		out = append(out, binary.BigEndian.Uint64(bs[i:]))
	}
	return out
}

func arrayToBytes(xs []uint64) []byte {
	if len(xs) < 2 {
		return []byte{}
	}
	first := xs[0]
	out := make([]byte, len(xs)*8)
	for i := 1; i < len(xs); i = i + 1 {
		binary.BigEndian.PutUint64(out[((i-1)*8):], xs[i])
	}
	return out[:first]
}
<|endoftext|>


<commit_before>package runner

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"testing"
)

func TestProgressChunks(t *testing.T) {
	flag.Set("log_chunk_size", "3")

	in := []byte("aaa\naaa\naaa\n")
	ch := make(chan LogChunk)

	go func() {
		processChunks(ch, bytes.NewReader(in), "test")
		close(ch)
	}()

	cnt := 0
	for _ = range ch {
		cnt++
	}

	if cnt != 3 {
		t.Fail()
	}
}

type FormData struct {
	params map[string]string
	files  map[string]string
    path   string
}

func TestCompleteFlow(t *testing.T) {
	var err error
	var formData []FormData
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("OK"))

		if r.Method != "POST" {
			err = fmt.Errorf("Non-POST request received: %s", r.Method)
			return
		}

        r.ParseMultipartForm(1 << 20)
        f := FormData{params: make(map[string]string), path: r.URL.Path}

        for k, v := range r.MultipartForm.Value {
            if k == "date" {
                continue
            }
            if len(v) != 1 {
                err = fmt.Errorf("Multiple values for form field: %s", k)
                return
            }

            f.params[k] = v[0]
        }

        if len(r.MultipartForm.File) > 0 {
            f.files = make(map[string]string)

            files := r.MultipartForm.File
            if len(files) != 1 {
                err = fmt.Errorf("Invalid number of artifacts found")
                return
            }

            for filename, fileHeaders := range files {
                if len(fileHeaders) != 1 {
                    err = fmt.Errorf("Multiple file headers found")
                    return
                }

                file, err := fileHeaders[0].Open()
                if err != nil {
                    return
                }
                fileContents, err := ioutil.ReadAll(file)
                if err != nil {
                    return
                }

                f.files[filename] = string(fileContents)
            }
        }

        formData = append(formData, f)
        return

		err = fmt.Errorf("Unexpected path: %s", r.URL.Path)
	}))
	defer ts.Close()

	// Current running program is definitely an artifact which will be present in the pogram
	required_artifact := os.Args[0]

	template := `
	{
		"commands": [
			{
				"id": "cmd_1",
				"script": "#!/bin/bash\necho -n $VAR",
				"env": {"VAR": "hello world"},
				"cwd": "/tmp",
                "artifacts": ["%s"]
			},
			{
				"id": "cmd_2",
				"script": "#!/bin/bash\necho test",
				"cwd": "/tmp"
			}
		]
	}
	`

	config := &Config{}
    config.Server = ts.URL
    config.JobID = "job_1"
	if json.Unmarshal([]byte(fmt.Sprintf(template, required_artifact)), config) != nil {
		t.Errorf("Failed to parse build config")
	}

	reporter := NewReporter(config.Server)
	RunCmds(reporter, config)
	reporter.Shutdown()

	if err != nil {
		t.Errorf(err.Error())
	}

	expectedFileContents, _ := ioutil.ReadFile(os.Args[0])
	expected := []FormData{
		FormData{
            path: "/jobsteps/job_1/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/commands/cmd_1/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/jobsteps/job_1/logappend/",
			params: map[string]string{
				"text":   "hello world",
				"source": "stdout",
				"offset": "0",
			},
		},
		FormData{
            path: "/commands/cmd_1/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "return_code": "0",
			},
		},
		FormData{
            path: "/jobsteps/job_1/artifacts/",
            params: map[string]string{
                "name": os.Args[0],
            },
			files: map[string]string{
                "file": string(expectedFileContents),
			},
		},
		FormData{
            path: "/commands/cmd_2/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/jobsteps/job_1/logappend/",
			params: map[string]string{
				"text":   "test\n",
				"source": "stdout",
				"offset": "11",
			},
		},
		FormData{
            path: "/commands/cmd_2/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "return_code": "0",
			},
		},
		FormData{
            path: "/jobsteps/job_1/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "result": "passed",
			},
		},
	}

    for i, v := range formData {
        if !reflect.DeepEqual(v, expected[i]) {
            fmt.Println("A", i, v.params, expected[i].params)
            t.Fail()
        }
    }

	if !reflect.DeepEqual(formData, expected) {
		t.Errorf("Form data does not match")
	}
}
<commit_msg>Fix test-cases<commit_after>package runner

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"testing"
)

func TestProgressChunks(t *testing.T) {
	flag.Set("log_chunk_size", "3")

	in := []byte("aaa\naaa\naaa\n")
	ch := make(chan LogChunk)

	go func() {
		processChunks(ch, bytes.NewReader(in), "test")
		close(ch)
	}()

	cnt := 0
	for _ = range ch {
		cnt++
	}

	if cnt != 3 {
		t.Fail()
	}
}

type FormData struct {
	params map[string]string
	files  map[string]string
    path   string
}

func TestCompleteFlow(t *testing.T) {
	var err error
	var formData []FormData
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("OK"))

		if r.Method != "POST" {
			err = fmt.Errorf("Non-POST request received: %s", r.Method)
			return
		}

        r.ParseMultipartForm(1 << 20)
        f := FormData{params: make(map[string]string), path: r.URL.Path}

        for k, v := range r.MultipartForm.Value {
            if k == "date" {
                continue
            }
            if len(v) != 1 {
                err = fmt.Errorf("Multiple values for form field: %s", k)
                return
            }

            f.params[k] = v[0]
        }

        if len(r.MultipartForm.File) > 0 {
            f.files = make(map[string]string)

            files := r.MultipartForm.File
            if len(files) != 1 {
                err = fmt.Errorf("Invalid number of artifacts found")
                return
            }

            for filename, fileHeaders := range files {
                if len(fileHeaders) != 1 {
                    err = fmt.Errorf("Multiple file headers found")
                    return
                }

                file, err := fileHeaders[0].Open()
                if err != nil {
                    return
                }
                fileContents, err := ioutil.ReadAll(file)
                if err != nil {
                    return
                }

                f.files[filename] = string(fileContents)
            }
        }

        formData = append(formData, f)
        return

		err = fmt.Errorf("Unexpected path: %s", r.URL.Path)
	}))
	defer ts.Close()

	// Current running program is definitely an artifact which will be present in the pogram
	required_artifact := os.Args[0]

	template := `
	{
		"commands": [
			{
				"id": "cmd_1",
				"script": "#!/bin/bash\necho -n $VAR",
				"env": {"VAR": "hello world"},
				"cwd": "/tmp",
                "artifacts": ["%s"]
			},
			{
				"id": "cmd_2",
				"script": "#!/bin/bash\necho test",
				"cwd": "/tmp"
			}
		]
	}
	`

	config := &Config{}
    config.Server = ts.URL
    config.JobstepID = "job_1"
	if json.Unmarshal([]byte(fmt.Sprintf(template, required_artifact)), config) != nil {
		t.Errorf("Failed to parse build config")
	}

	reporter := NewReporter(config.Server)
	RunCmds(reporter, config)
	reporter.Shutdown()

	if err != nil {
		t.Errorf(err.Error())
	}

	expectedFileContents, _ := ioutil.ReadFile(os.Args[0])
	expected := []FormData{
		FormData{
            path: "/jobsteps/job_1/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/commands/cmd_1/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/jobsteps/job_1/logappend/",
			params: map[string]string{
				"text":   "hello world",
				"source": "stdout",
				"offset": "0",
			},
		},
		FormData{
            path: "/commands/cmd_1/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "return_code": "0",
			},
		},
		FormData{
            path: "/jobsteps/job_1/artifacts/",
            params: map[string]string{
                "name": os.Args[0],
            },
			files: map[string]string{
                "file": string(expectedFileContents),
			},
		},
		FormData{
            path: "/commands/cmd_2/",
			params: map[string]string{
				"status": STATUS_IN_PROGRESS,
			},
		},
		FormData{
            path: "/jobsteps/job_1/logappend/",
			params: map[string]string{
				"text":   "test\n",
				"source": "stdout",
				"offset": "11",
			},
		},
		FormData{
            path: "/commands/cmd_2/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "return_code": "0",
			},
		},
		FormData{
            path: "/jobsteps/job_1/",
			params: map[string]string{
				"status": STATUS_FINISHED,
                "result": "passed",
			},
		},
	}

    for i, v := range formData {
        if !reflect.DeepEqual(v, expected[i]) {
            fmt.Println("A", i, v.params, expected[i].params)
            t.Fail()
        }
    }

	if !reflect.DeepEqual(formData, expected) {
		t.Errorf("Form data does not match")
	}
}
<|endoftext|>


<commit_before>package policy

/*
#include <sys/eventfd.h>
*/
import "C"

import (
	"fmt"
	"io/ioutil"
	"math"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
)

type Evictor struct {
	limit     int
	eventfd   int
	usagePath string
}

func NewEvictor(pkgfile, rootCID string, kb_limit int) (*Evictor, error) {
	basePath := fmt.Sprintf("/sys/fs/cgroup/memory/docker/%s/", rootCID)
	byte_limit := 1024 * kb_limit

	eventfd, err := C.eventfd(0, C.EFD_CLOEXEC)
	if err != nil {
		return nil, err
	}

	usagePath := filepath.Join(basePath, "memory.usage_in_bytes")
	usagefd, err := syscall.Open(usagePath, syscall.O_RDONLY, 0777)
	if err != nil {
		return nil, err
	}

	eventPath := filepath.Join(basePath, "cgroup.event_control")

	eventStr := fmt.Sprintf("'%d %d %d'", eventfd, usagefd, byte_limit)
	echo := exec.Command("echo", eventStr, ">", eventPath)
	if err = echo.Run(); err != nil {
		return nil, err
	}

	e := &Evictor{
		limit:     byte_limit,
		eventfd:   int(eventfd),
		usagePath: usagePath,
	}

	return e, nil
}

func (e *Evictor) CheckUsage(servers []*ForkServer, mutex *sync.Mutex) []*ForkServer {
	usage := e.usage()
	if usage > e.limit {
		mutex.Lock()
		defer mutex.Unlock()
		return e.evict(servers)
	}

	return servers
}

func (e *Evictor) usage() (usage int) {
	buf, err := ioutil.ReadFile(e.usagePath)
	if err != nil {
		return 0
	}

	str := strings.TrimSpace(string(buf[:]))
	usage, err = strconv.Atoi(str)
	if err != nil {
		panic(fmt.Sprintf("atoi failed: %v", err))
	}

	return usage
}

func (e *Evictor) evict(servers []*ForkServer) []*ForkServer {
	idx := -1
	worst := float64(math.Inf(+1))

	for k := 1; k < len(servers); k++ {
		if servers[k].Children == 0 && servers[k].Runners {
			if ratio := servers[k].Hits / servers[k].Size; ratio < worst {
				idx = k
				worst = ratio
			}
		}
	}

	if idx != -1 {
		err := servers[idx].Kill()
		if err != nil {
			panic(fmt.Sprintf("failed to evict with: %v", err))
		}

		return append(servers[:idx], servers[idx+1:]...)
	}

	return servers
}
<commit_msg>eviction log message<commit_after>package policy

/*
#include <sys/eventfd.h>
*/
import "C"

import (
	"fmt"
	"log"
	"io/ioutil"
	"math"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
)

type Evictor struct {
	limit     int
	eventfd   int
	usagePath string
}

func NewEvictor(pkgfile, rootCID string, kb_limit int) (*Evictor, error) {
	basePath := fmt.Sprintf("/sys/fs/cgroup/memory/docker/%s/", rootCID)
	byte_limit := 1024 * kb_limit

	eventfd, err := C.eventfd(0, C.EFD_CLOEXEC)
	if err != nil {
		return nil, err
	}

	usagePath := filepath.Join(basePath, "memory.usage_in_bytes")
	usagefd, err := syscall.Open(usagePath, syscall.O_RDONLY, 0777)
	if err != nil {
		return nil, err
	}

	eventPath := filepath.Join(basePath, "cgroup.event_control")

	eventStr := fmt.Sprintf("'%d %d %d'", eventfd, usagefd, byte_limit)
	echo := exec.Command("echo", eventStr, ">", eventPath)
	if err = echo.Run(); err != nil {
		return nil, err
	}

	e := &Evictor{
		limit:     byte_limit,
		eventfd:   int(eventfd),
		usagePath: usagePath,
	}

	return e, nil
}

func (e *Evictor) CheckUsage(servers []*ForkServer, mutex *sync.Mutex) []*ForkServer {
	usage := e.usage()
	if usage > e.limit {
		mutex.Lock()
		defer mutex.Unlock()
		log.Printf("EVICT IMPORT: %v usage / %v limit", usage, e.limit)
		return e.evict(servers)
	}

	return servers
}

func (e *Evictor) usage() (usage int) {
	buf, err := ioutil.ReadFile(e.usagePath)
	if err != nil {
		return 0
	}

	str := strings.TrimSpace(string(buf[:]))
	usage, err = strconv.Atoi(str)
	if err != nil {
		panic(fmt.Sprintf("atoi failed: %v", err))
	}

	return usage
}

func (e *Evictor) evict(servers []*ForkServer) []*ForkServer {
	idx := -1
	worst := float64(math.Inf(+1))

	for k := 1; k < len(servers); k++ {
		if servers[k].Children == 0 && servers[k].Runners {
			if ratio := servers[k].Hits / servers[k].Size; ratio < worst {
				idx = k
				worst = ratio
			}
		}
	}

	if idx != -1 {
		err := servers[idx].Kill()
		if err != nil {
			panic(fmt.Sprintf("failed to evict with: %v", err))
		}

		return append(servers[:idx], servers[idx+1:]...)
	}

	return servers
}
<|endoftext|>